TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training

다음은 "TACO: 확장 가능한 텐서 병렬 LLM 훈련을 위한 중간 텐서의 효율적인 통신 압축" 논문에 대한 상세한 기술 요약입니다.

1. 문제 제기

대규모 언어 모델 (LLM) 의 확장된 훈련은 3 차원 병렬화(데이터 병렬화, 파이프라인 병렬화, 텐서 병렬화) 에 크게 의존합니다. 데이터 병렬화 (DP) 와 파이프라인 병렬화 (PP) 는 성공적인 통신 압축 기법을 경험해 왔으나, **텐서 병렬화 (TP)**는 여전히 치명적인 병목 현상으로 남아 있습니다.

• 통신 오버헤드: TP 는 모든 레이어와 반복 단계에서 중간 텐서 (활성화 및 기울기) 의 빈번하고 엄격하게 동기화된 통신을 요구합니다. 대규모 훈련에서 TP 통신은 총 통신 시간의 **50~60%**를 차지하며, 중요한 실행 경로에 위치하여 연산과 겹쳐 수행하기 어렵습니다.
• 압축의 도전 과제:
  • 오류 누적: DP/PP 와 달리 TP 는 고빈도 동기화를 수반합니다. 중간 텐서의 미세한 양자화 오차는 어텐션 및 잔차 연결을 통해 전파되어 역전파 과정에서 증폭되고 훈련 발산을 초래합니다.
  • 데이터 분포: TP 중간 텐서는 긴 꼬리를 가진 밀집된 근사 0 분포를 보입니다. 표준 양자화 방법 (예: INT8) 은 균일한 간격을 사용하므로 텐서를 지배하는 작은 크기의 값에 대해 막대한 정보 손실을 초래하여 여기서 실패합니다.
  • 기존 방법의 비효율성: 이전의 추론 지향 압축 (예: SmoothQuant) 이나 DP 중심 방법 (예: SDP4bit) 은 수치적 불안정성으로 인한 치명적인 실패 없이 TP 훈련에 직접 적용할 수 없습니다.

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

저자들은 TP 중간 텐서에 특화된 강력한 FP8 기반 프레임워크인 **TACO(Tensor-parallel Adaptive COmmunication compression)**를 제안합니다.

A. 왜 FP8 인가?

이 논문은 TP 텐서의 통계적 특성을 분석하여 FP8이 INT8 보다 우월하다고 결론지었습니다:

• INT8 의 한계: 균일 양자화는 서로 다른 작은 값들이 동일한 정수로 매핑되는 "충돌"을 유발하여 밀집된 0 중심 영역에서 높은 상대 오차를 초래합니다.
• FP8 의 장점: FP8(특히 E4M3 포맷) 은 지수 스케일을 사용하여 작은 값 (0 근처) 에 대해 높은 정밀도를 제공하고, 이상치에 대해 넓은 동적 범위를 제공합니다. 이는 TP 텐서의 통계적 분포와 완벽하게 일치합니다.

B. 핵심 기법

TACO 는 TP 훈련을 위해 FP8 압축을 실현 가능하게 만드는 세 가지 주요 혁신을 도입합니다:

1. 적응형 스케일–해다마드 (ASH) 변환:

   • 문제: 표준 해다마드 변환은 등거리 (에너지를 보존) 이며 TP 텐서의 밀집된 0 클러스터를 효과적으로 분산시킬 수 없습니다. 이들은 0 근처에 갇혀 FP8 의 정밀도를 충분히 활용하지 못합니다.
   • 해결책: TACO 는 해다마드 변환 전에 블록 단위 적응형 재스케일링을 적용합니다. 각 블록의 2 차 원 모멘트 (에너지) 를 계산하여 에너지를 목표 수준으로 정규화하는 스케일링 인자를 적용합니다.
   • 효과: 이는 밀집된 0 클러스터를 "펼쳐" FP8 표현 범위에 걸쳐 분산시킨 후, 양자화에 적합한 가우스와 유사한 분포를 만들기 위해 직교 해다마드 회전을 수행합니다.

2. 이중 스케일 (DS) 양자화:

   • 문제: ASH 를 사용하더라도 극단적인 값은 오버플로우를, 작은 값은 언더플로우를 유발하여 불안정성을 초래할 수 있습니다.
   • 해결책: TACO 는 블록당 두 개의 서로 다른 스칼라를 사용합니다:
     • 변환 전 스케일 ($\alpha_k$): 변환 전 블록의 에너지를 제어합니다.
     • 변환 후 스케일 ($s_k$): 변환된 값이 FP8 동적 범위 ($Q_{max}$) 내에 엄격하게 들어오도록 보장합니다.
   • 재구성: 압축 해제 시, 이 두 가지 스케일을 사용하여 과정을 역전시켜 고충실도 재구성을 보장하고 오류 증폭을 방지합니다.

3. 시스템 수준 커널 퓨전 및 오버랩:

   • 퓨즈드 연산자: TP 가 중요한 경로에 있으므로 압축의 계산 오버헤드를 최소화하기 위해 TACO 는 분산 계산, ASH 변환, DS 양자화를 단일 GPU 커널로 퓨즈합니다. 이는 불필요한 글로벌 메모리 접근과 커널 런치 오버헤드를 제거합니다.
   • 통신 - 연산 오버랩: TACO 는 COCCL(Compressible Collective Communication Library) 과 통합하여 "두 번의 샷" AllReduce 전략 (ReduceScatter + AllGather) 을 사용합니다. 압축/압축 해제 커널을 통신과 겹쳐 지연 시간을 효과적으로 숨깁니다.

3. 주요 기여

1. 체계적 분석: TP 텐서에서의 "Zero-Collapse(0 붕괴)" 현상을 식별하여 표준 INT8 과 고정 변환이 부적절함을 증명하고, FP8 을 최적 포맷으로 확립했습니다.
2. TACO 프레임워크: 대규모 TP 훈련에서 안정적인 수렴을 달성한 최초의 3 차원 병렬 통신 압축 시스템입니다. TP 텐서의 고유한 분포를 처리하기 위해 ASH 변환과 이중 스케일 양자화를 도입했습니다.
3. 고성능 구현: 메모리 트래픽을 줄이고 통신과의 효율적인 오버랩을 가능하게 하여 빈번한 압축의 오버헤드를 완화하는 고도로 퓨즈된 압축 연산자를 개발했습니다.
4. 완전한 3 차원 통합: 최첨단 DP(SDP4bit) 및 PP(TahQuant) 압축기와 TACO 를 성공적으로 통합하여 완전히 압축이 활성화된 3 차원 병렬 훈련 시스템을 구축했습니다.

4. 실험 결과

저자들은 NVIDIA H100 GPU 를 사용하여 GPT 모델(2.7B, 6.7B, 13B) 및 Qwen2.5-7B에 대해 TACO 를 평가했습니다.

• 정확도 (수렴):

  • TACO 는 거의 손실 없는 정확도를 달성합니다. GPT-350M 의 경우, BF16 기준 대비 검증 손실 저하는 **+0.25%**에 불과한 반면, 기존 방법 (TahQuant) 은 약 2.9% 저하됩니다.
  • 더 큰 모델 (GPT-6.7B, Qwen-2.5-7B) 에서는 저하가 미미하게 (<0.7%) 유지되어 규모 전반에 걸친 안정성을 입증합니다.
  • 제거 실험 (Ablation studies) 은 ASH 나 DS 를 제거할 경우 치명적인 발산이나 심각한 정확도 손실이 발생함을 확인시켜 줍니다.

• 성능 (처리량):

  • TP 압축: TACO 는 압축되지 않은 기준 대비 TP 통신에 대해 최대 1.87 배의 엔드투엔드 처리량 향상을 달성합니다.
  • 완전한 3 차원 병렬화: TP=4, PP=2, DP=2 인 완전한 3 차원 병렬 설정에 적용 시, TACO 는 기준 대비 최대 1.53 배의 속도 향상을 제공합니다.
  • 확장성: TP 통신이 지배적인 병목 현상이 됨에 따라 텐서 병렬화 (TP) 정도가 증가함에 따라 성능 이득도 증가합니다.

5. 의의

• TP 병목 현상 해결: TACO 는 분산 LLM 훈련에서 가장 어려운 측면인 모델 품질을 희생하지 않고 텐서 병렬화의 고빈도, 수치적으로 민감한 중간 텐서를 압축하는 문제를 해결합니다.
• 더 큰 모델 가능: FP8 을 통해 통신 오버헤드를 거의 50% 줄이면서 안정성을 유지함으로써, 기존 하드웨어 클러스터에서 더 큰 모델을 훈련하거나 기존 모델을 훨씬 빠르게 훈련할 수 있게 합니다.
• 하드웨어 독립성: FP8(최신 GPU 인 H100 등 지원) 에 최적화되어 있지만, 네이티브 FP8 지원이 없는 하드웨어에서는 INT8 으로 우아하게 저하되도록 설계되어 광범위한 적용성을 보장합니다.
• 미래 영향: 이 연구는 압축이 데이터 분포에 수학적으로 맞춤화될 경우 훈련의 중요한 경로에서도 공격적인 통신 압축이 가능함을 입증함으로써 엑사스케일 훈련의 길을 닦습니다.