AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs

다음은 "AGoQ: Activation and Gradient Quantization for Memory-Efficient Distributed Training of LLMs"라는 논문에 대한 상세한 기술 요약입니다.

1. 문제 제기

대규모 언어 모델 (LLM) 의 학습에는 모델 파라미터, 옵티마이저 상태, 그래디언트, 그리고 **활성화 (activations)**에 의해 주로 소비되는 막대한 GPU 메모리가 필요합니다.

• 메모리 병목 현상: 학습 중 활성화는 일반적으로 메모리의 가장 큰 부분을 차지하며, 시퀀스 길이와 배치 크기에 비례하여 선형적으로 증가합니다.
• 기존 방법의 한계:
  • 활성화 재계산/오프로딩: 메모리를 줄이지만 재계산으로 인한 상당한 계산 오버헤드 또는 오프로딩으로 인한 I/O 지연을 초래합니다.
  • 표준 양자화: 8 비트 양자화 (INT8/FP8) 가 일반적이지만, 4 비트 활성화나 8 비트 그래디언트에 이를 적용하면 누적 및 All-Reduce 연산 중 발생하는 수치적 오버플로우로 인해 심각한 정확도 손실, 수렴 지연 또는 학습 불안정이 초래되는 경우가 많습니다.
  • 그래디언트 통신: 현재 시스템은 종종 그래디언트를 고정밀도 (FP32) 로 저장하고 통신 시에만 양자화하여, 저장된 그래디언트의 메모리 사용량을 줄이지 못합니다.

2. 방법론: AGoQ 프레임워크

저자들은 모델 수렴을 희생하지 않으면서 거의 4 비트 활성화 저장 및 8 비트 그래디언트 저장/통신을 달성하기 위해 두 가지 새로운 기술을 통합한 AGoQ라는 포괄적인 프레임워크를 제안합니다.

A. 레이어 인식 활성화 양자화 (LAAQ)

모든 레이어를 균일하게 4 비트로 양자화하면 정확도 손실이 발생하므로, LAAQ 는 레이어 유형과 파이프라인 스테이지에 따라 비트 너비를 동적으로 할당합니다.

• 이론적 오차 분석: 저자들은 RMSNorm, SiLU, Attention, GEMM 등 다양한 연산에 대한 오차 한계를 유도했습니다.
  • 결과: Attention 모듈의 활성화 (Q, K, V) 를 양자화하는 것은 다른 모듈에 비해 훨씬 큰 그래디언트 오차를 유발합니다. 따라서 Attention 활성화는 양자화되지 않습니다 (더 높은 정밀도로 유지).
  • 전략: 다른 레이어 (RMSNorm, SiLU, FFN) 의 경우 시스템은 4 비트 블록 단위 양자화를 사용합니다.
  • 재계산 vs 캐싱: 분석 결과, RMSNorm 및 SiLU 와 같은 연산의 경우 양자화된 중간 값을 캐싱하는 것보다 양자화된 입력만 저장하고 역전파 동안 중간 값을 재계산하는 것이 더 엄격한 오차 한계를 제공합니다. 따라서 AGoQ 는 4 비트 입력만 저장하고 중간 값을 재계산합니다.
• 동적 비트 너비 보상 (DBCA-PP):
  • 파이프라인 병렬성 (PP) 에서 서로 다른 파이프라인 스테이지는 서로 다른 수의 활성화 배치를 보유합니다 (불균등한 메모리 사용).
  • AGoQ 는 이를 활용하여 활성 배치가 적은 스테이지에는 더 높은 비트 너비(예: 5, 6 또는 8 비트) 를 할당하고, 배치가 많은 스테이지에는 더 낮은 비트 너비(4 비트) 를 할당합니다. 이는 전체 메모리 사용량을 균형 있게 유지하면서 요소당 평균 약 4 비트를 유지하고, 메모리가 여유 있는 곳에서 정밀도 손실을 보상합니다.

B. 정밀도 보존 그래디언트 양자화 (QuanGrad)

이 기술은 전체 과정에서 8 비트 (FP8) 그래디언트를 사용하여 메모리 사용량과 통신 시간을 모두 줄입니다.

• 로컬 누적: 여러 미니배치에서 그래디언트를 누적할 때 오버플로우를 방지하기 위해 시스템은 FP8 메인 그래디언트를 FP16/BF16 으로 역양자화한 후 새로운 로컬 그래디언트를 더하고, 그 합을 다시 FP8 로 양자화합니다.
• 정밀도 보존 All-Reduce: 표준 FP8 All-Reduce 는 축소 단계에서 오버플로우를 유발합니다. AGoQ 는 All-Reduce 연산을 다음과 같이 분해합니다:
  1. All-to-All: 압축된 FP8 그래디언트를 모든 장치에 전송합니다.
  2. 로컬 축소: 각 장치는 수신된 데이터를 FP32 로 역양자화하여 로컬 합을 수행한 후 다시 FP8 로 양자화합니다.
  3. All-Gather: 축소된 결과를 수집합니다.
  • 이를 통해 통신은 저정밀도 (FP8) 로 유지하면서 누적은 고정밀도 (FP32) 에서 발생하도록 하여 오버플로우를 방지하고 수렴을 유지합니다.

C. 커널 퓨전

양자화/역양자화 연산의 오버헤드를 제거하기 위해 AGoQ 는 이러한 요소별 연산을 인접한 GEMM(일반 행렬 곱셈) 커널과 융합합니다. 이는 텐서 코어가 중량 작업을 수행하도록 하면서 양자화 논리를 동일한 커널 내의 CUDA 코어에서 처리하여 지연 시간을 최소화합니다.

3. 주요 기여

1. 레이어 인식 양자화 전략: 4 비트로 안전하게 압축할 수 있는 레이어 (Attention 제외) 와 오차를 최소화하기 위해 재계산 전략이 필요한 레이어를 식별하는 이론적 및 실용적 접근법.
2. 동적 비트 너비 보상: 각 스테이지의 메모리 가용성에 따라 양자화 정밀도를 동적으로 조정하여 피크 사용량을 증가시키지 않으면서 메모리 효율성을 극대화하는 파이프라인 병렬성을 위한 새로운 알고리즘.
3. 정밀도 보존 그래디언트 통신: 연산을 분할하여 고정밀도 누적을 가능하게 하는 8 비트 All-Reduce 수행 방법으로, 이전에는 저비트 그래디언트 학습을 방해했던 오버플로우 문제를 해결합니다.
4. 시스템 통합: Megatron-LM 에 완전히 통합되어 8 비트 Adam 옵티마이저와 호환되며, NVIDIA GPU 와 Huawei Ascend NPU 를 모두 지원합니다.

4. 실험 결과

LLaMA 모델 (7B, 8B, 13B, 34B) 과 OLMo-1B 를 사용하여 최대 64 개 GPU(NVIDIA A6000/Pro 6000) 클러스터에서 실험이 수행되었습니다.

• 메모리 감소:

  • AGoQ 는 고정밀도 학습 (Megatron-LM) 대비 최대 **52%**의 피크 메모리 사용량을 줄입니다.
  • 최첨단 COAT시스템과 비교할 때, AGoQ 는 활성화 메모리를 추가로 30% 줄이고 그래디언트 메모리를 75% 줄입니다.
  • 기존 시스템에서 Out-Of-Memory(OOM) 오류를 발생시키는 더 긴 시퀀스 길이 (예: 80K 토큰) 의 학습을 가능하게 합니다.

• 학습 속도:

  • Megatron-LM(ZERO 유무) 과 DeepSpeed 대비 종단 간 학습 시간에서 최대 1.34 배의 속도 향상을 달성합니다.
  • 통신 오버헤드와 메모리 제약이 감소함에 따라 더 긴 시퀀스 길이에서 속도 향상 효과가 더 두드러집니다.

• 수렴 및 정확도:

  • 수렴: LLaMA2-7B 와 LLaMA3-8B 에 대한 학습 손실 곡선은 고정밀도 기준선과 밀접하게 일치합니다. 반면, 단순 FP8 All-Reduce 방법은 손실이 현저히 높았습니다.
  • 하류 작업: ARC, HellaSwag, PIQA 등의 데이터셋에 대한 제로샷 정확도는 기준선과 유사하게 유지되며, 일부 작업에서는 약간의 개선이 나타납니다.
  • 안정성: 그래디언트 노름은 학습 내내 안정적으로 유지되어 정밀도 보존 누적의 효과를 확인시켜 줍니다.

5. 의의

• 긴 컨텍스트 학습 가능: 메모리 발자국을 극적으로 줄여 기존 하드웨어에서 대규모 클러스터 없이도 매우 긴 컨텍스트 창 (예: 64K–80K 토큰) 을 가진 LLM 학습을 실현 가능하게 합니다.
• 격차 해소: 모델 품질을 훼손하지 않으면서 활성화에 4 비트, 그래디언트에 8 비트의 양자화를 성공적으로 도입하여, 이전에는 LLM 학습에 너무 불안하다고 간주되던 영역을 개척했습니다.
• 확장성: 이 방법은 다양한 병렬성 전략 (DP, TP, PP) 과 하드웨어 아키텍처 전반에 걸쳐 매우 확장 가능하며, 제한된 자원으로 수조 파라미터 모델을 학습할 수 있는 실용적인 길을 제시합니다.

요약하자면, AGoQ 는 단순한 압축을 넘어 메모리 효율성, 통신 속도, 모델 정확도 사이의 균형을 맞추는 수학적으로 근거 있는 적응형 양자화 프레임워크로, LLM 학습 시스템 분야에서 중요한 발전을 나타냅니다.