대규모 언어 모델 (Large Language Models, LLMs)은 파라미터 수가 계속해서 증가하고 있지만, 이를 대규모로 서비스하기 위해서는 단순한 연산 능력 이상의 것이 필요합니다. 모델 프루닝 (Model Pruning)은 처음부터 완전히 다시 학습시키지 않고도 메모리 대역폭 (Memory Bandwidth)과 추론 지연 시간 (Inference Latency)을 줄이기 위해 네트워크에서 중복된 가중치 (Weights) 또는 전체 구조를 제거하는 기술입니다. 프로덕션 AI 인프라를 구축하는 팀에게 있어, 프루닝 기술을 이해하는 것은 정확도와 처리량 (Throughput) 및 비용 사이의 균형을 맞추는 데 필수적입니다.

모델 프루닝 (Model Pruning)이란 무엇인가?

본질적으로 모델 프루닝은 네트워크의 출력에 거의 기여하지 않는 파라미터를 제거하는 과정입니다. 이 기술은 트랜스포머 (Transformers) 이전부터 존재했지만, LLM은 그 깊이, 어텐션 메커니즘 (Attention Mechanisms), 그리고 엄청난 규모로 인해 독특한 과제들을 제시합니다. 프루닝은 개별 가중치부터 전체 어텐션 헤드 (Attention Heads)에 이르기까지 다양한 입도 (Granularities)로 적용될 수 있으며, 학습 중, 학습 후, 또는 양자화 (Quantization) 및 증류 (Distillation)를 포함하는 압축 파이프라인의 일부로 수행될 수 있습니다.

비구조적 희소성 (Unstructured Sparsity) vs. 구조적 희소성 (Structured Sparsity)

비구조적 희소성 (Unstructured Sparsity)은 절대적 크기 (Absolute Magnitude)나 그래디언트 (Gradient) 정보와 같은 중요도 기준에 따라 개별 가중치를 0으로 만듭니다. 0이 아닌 요소들이 규칙적인 패턴을 따르지 않기 때문에, 추론 속도를 높이려면 NVIDIA Ampere의 2:4 구조적 희소성 제약 조건이나 커스텀 CUDA 구현과 같은 희소 행렬 커널 (Sparse Matrix Kernels) 및 하드웨어 지원이 필요합니다. 이러한 지원이 없다면 밀집 GEMM (Dense GEMM) 라이브러리는 이전과 동일한 양의 작업을 수행하게 되며, 이는 메모리 절약이 자동으로 낮은 지연 시간으로 이어지지 않음을 의미합니다.

반면, 구조적 희소성 (Structured Sparsity)은 어텐션 헤드, 피드포워드 차원 (Feed-forward Dimensions), 또는 레이어 전체와 같이 일관된 블록을 제거합니다. 결과적으로 생성된 모델은 하드웨어 스케줄러가 고려하는 모든 차원에서 더 작아지므로, 표준 밀집 BLAS 루틴이 더 작은 행렬에서 더 빠르게 실행됩니다. 트레이드오프(Tradeoff)는 제거된 파라미터당 정확도 하락 폭이 더 크다는 점이며, 이 때문에 구조적 프루닝은 종종 가벼운 복구 미세 조정 (Recovery Fine-tuning) 또는 증류 (Distillation)와 함께 사용됩니다.

트랜스포머 (Transformers)를 위한 프루닝 기술

트랜스포머 (Transformers) 기반의 LLM을 위해 특별히 개발된 몇 가지 방법들이 등장했습니다:

• Magnitude pruning (크기 기반 프루닝). 가장 단순한 베이스라인입니다. 절대값이 가장 낮은 가중치 (weights)를 0으로 설정합니다. 구현이 쉽지만, 가중치와 활성화 값 (activations) 사이의 상호작용을 무시합니다.
• Wanda. 가중치의 크기에 그에 상응하는 입력 활성화 노름 (input activation norms)을 곱하여 가중치를 프루닝하며, 재학습 없이도 강력한 원샷 (one-shot) 결과를 달성합니다.
• SparseGPT. 근사 2차 정보 (approximate second-order information)를 사용하여 프루닝 후 남은 가중치들을 조정하며, 거대한 GPT 스타일 모델의 정확한 원샷 압축을 가능하게 합니다.
• LLM-Pruner. 쿼리-키-값 (query-key-value) 그룹화 및 피드포워드 확장 (feed-forward expansions)과 같은 트랜스포머 의존성을 존중하는 구조적 접근 방식(structured approach)으로, 최소한의 복구 튜닝 (recovery tuning)만으로 태스크 불가지론적 (task-agnostic) 프루닝을 허용합니다.
• Sheared LLaMA. 구조적 프루닝 (structured pruning)과 타겟 데이터 분포에 대한 지속적 사전 학습 (continued pre-training)을 결합하여, 모델 크기를 영구적으로 줄이면서도 정확도를 회복합니다.

현재 대부분의 프로덕션 파이프라인은 프루닝을 INT8 또는 INT4로의 양자화 (quantization) 및 어텐션 레이어 (attention layers)의 런타임 퓨전 (runtime fusion)을 포함할 수 있는 더 넓은 최적화 워크플로의 한 단계로 취급합니다.

코드 예시: PyTorch를 이용한 Magnitude Pruning

다음 스니펫은 선형 레이어 (linear layer)에 대한 비구조적 (unstructured) L1 magnitude pruning과, 이어서 전체 출력 채널을 제거하는 구조적 (structured) 변형을 보여줍니다. 이러한 패턴은 트랜스포머 스택 내부의 다층 퍼셉트론 (multi-layer perceptron) 블록으로 일반화될 수 있습니다.

import torch
import torch.nn.utils.prune as prune

...