프로덕션 환경에서의 LLM 최적화: 정밀도, 어텐션 및 아키텍처

프로덕션 환경에 대규모 언어 모델(LLM)을 배포하려면 두 가지 주요 병목 현상을 극복해야 합니다. 바로 수십억 개의 파라미터로 인한 막대한 VRAM 요구량과 긴 입력 시퀀스와 관련된 이차 함수적(quadratic) 메모리 증가 문제입니다. 이를 해결하기 위해 Hugging Face는 저정밀도 양자화(quantization), 최적화된 어텐션 알고리즘, 그리고 전략적 아키텍처 선택의 조합을 권장합니다.

저정밀도를 통한 메모리 사용량 감소

수치 정밀도를 낮추면 모델 가중치를 로드하는 데 필요한 VRAM이 줄어들어, 더 큰 모델을 더 작거나 접근성이 좋은 하드웨어에서 실행할 수 있습니다.

정밀도별 VRAM 요구사항

짧은 텍스트 입력(1024 토큰 미만)의 경우 모델 가중치를 로드하는 것이 메모리 비용의 대부분을 차지합니다. VRAM 요구사항에 대한 일반적인 규칙은 다음과 같습니다:

  • float32: X0억 개의 파라미터를 가진 모델의 경우 대략 4 * X GB의 VRAM이 필요합니다.
  • bfloat16/float16: X0억 개의 파라미터를 가진 모델의 경우 대략 2 * X GB의 VRAM이 필요합니다.

예를 들어, Llama-2-70b는 bfloat16에서 약 140 GB의 VRAM이 필요하며, 이는 단일 A100(80GB)의 용량을 초과하므로 텐서 병렬성(tensor parallelism) 또는 파이프라인 병렬성(pipeline parallelism)이 필요합니다.

양자화 (8-bit 및 4-bit)

양자화는 정밀도를 8비트 또는 4비트로 더 낮추어, 텍스트 생성 시 정확도에 미치는 영향을 최소화하면서 메모리 사용량을 크게 줄입니다. 이는 텍스트 생성이 정확한 값보다는 다음 토큰 로짓(logits)의 상대적 분포에 의존하기 때문입니다.

  • 8-bit 양자화: VRAM 사용량을 크게 줄입니다 (예: OctoCoder의 피크 메모리가 ~32GB에서 ~15GB로 감소). 계산 중 동적 역양자화(dynamic dequantization)가 필요하기 때문에 추론 속도가 약간 느려질 수 있습니다.
  • 4-bit 양자화: VRAM을 더욱 줄여줍니다 (예: OctoCoder가 ~9.5GB로 감소), 이를 통해 RTX 3090과 같은 소비자용 GPU에서 모델을 실행할 수 있게 합니다. 하지만 8-bit 양자화보다 정확도 저하가 더 눈에 띄고 추론 속도가 느려질 수 있습니다.

Flash Attention을 이용한 추론 가속화

표준 셀프 어텐션(self-attention)은 시퀀스 길이($N$)에 따라 이차 함수적인 계산 및 메모리 복잡도를 가지므로, 긴 컨텍스트(예: 16,000+ 토큰)에서는 비용이 매우 많이 듭니다.

Flash Attention 알고리즘

Flash Attention은 계산을 더 작은 청크(chunk)로 나누고 여러 소프트맥스(softmax) 단계를 반복함으로써 어텐션 메커니즘을 최적화합니다. 이는 거대한 $QK^T$ 행렬의 생성을 피하므로, 메모리 비용이 $N$에 따라 이차 함수적이 아닌 선형적으로 증가하게 합니다.

성능 이점

Flash Attention은 소프트맥스 정규화 통계치를 재계산해야 하므로 더 많은 FLOPs를 필요로 하지만, 느린 고대역폭 메모리(VRAM)에 대한 액세스를 최소화하고 빠른 온칩 SRAM의 사용을 극대화하기 때문에 실제로는 더 빠릅니다. 또한 기본 셀프 어텐션 알고리즘과 수치적으로 동일한 출력을 생성합니다.

긴 컨텍스트 및 채팅을 위한 아키텍처 최적화

학습 중에 이루어지는 아키텍처 선택은 모델이 긴 시퀀스와 다회차 대화를 얼마나 효율적으로 처리하는지를 결정합니다. 두 가지 핵심 영역은 위치 임베딩(positional embeddings)과 키-값(KV) 캐시입니다.

상대적 위치 임베딩

절대적 위치 임베딩(sinusoidal 또는 학습된 방식)은 긴 텍스트에서 성능이 저하되는 경우가 많으며 학습된 길이 이상의 외삽(extrapolate)에 어려움을 겪습니다. 상대적 위치 임베딩이 더 효과적입니다:

  • Rotary Position Embedding (RoPE): 쿼리-키 쌍을 회전시켜 위치를 인코딩합니다. Falcon, Llama, PaLM에서 사용됩니다.
  • ALiBi: $QK^T$ 행렬에 미리 정의된 값으로 스케일링된 음수 정수를 더합니다. MPT와 BLOOM에서 사용되며, 일반적으로 RoPE보다 긴 시퀀스로의 외삽 능력이 더 뛰어납니다.

키-값 (KV) 캐시 최적화

자기회귀(Auto-regressive) 생성은 모든 이전 토큰에 대한 키-값 벡터를 저장하기 위해 KV 캐시를 사용하여 매 단계마다 이를 재계산할 필요가 없게 합니다. 이는 $QK^T$ 계산을 벡터-행렬 곱셈($\text{query} \times \text{KV cache}$)으로 변환하여 속도를 크게 향상시킵니다.

하지만 KV 캐시는 메모리 병목 현상이 될 수 있습니다. 두 가지 아키텍처가 이 오버헤드를 줄여줍니다:

  • Multi-Query Attention (MQA): 모든 어텐션 헤드에서 공유되는 단일 키-값 프로젝션 헤드를 사용합니다. 이는 캐시 크기를 획기적으로 줄여주며(예: OctoCoder의 16,000 토큰 시퀀스에서 15 GB에서 400 MB 미만으로 감소), 메모리 대역폭 병목을 줄입니다. Falcon, PaLM, MPT, BLOOM에서 사용됩니다.
  • Grouped-Query Attention (GQA): MQA와 표준 멀티 헤드 어텐션의 중간 단계입니다. 적은 수의 KV 프로젝션 헤드(예: 2, 4 또는 8개)를 사용하여 MQA보다 더 많은 모델 용량을 유지하면서도 대부분의 효율성을 유지합니다. Llama-2에서 사용됩니다.

Sources