[Paper] TACO: 효율적인 중간 텐서 통신 압축을 통한 확장 가능한 텐서-패럴렐 LLM 훈련

Source: arXiv - 2604.24088v1

개요

오늘날 거대한 언어 모델을 학습할 때는 **텐서 병렬화(TP)**에 크게 의존합니다. 텐서 병렬화는 큰 가중치 행렬을 여러 GPU에 나누어 놓아 계산을 확장하지만, 동시에 거대한 중간 텐서를 자주 교환해야 하므로 통신 병목 현상을 초래합니다. 이 논문에서는 TACO라는 경량 FP8 기반 압축 프레임워크를 소개합니다. TACO는 모델 품질을 손상시키지 않으면서 TP 트래픽 양을 크게 줄이며, GPT‑계열 및 Qwen 모델에서 최대 1.87배 빠른 엔드‑투‑엔드 학습 속도를 제공합니다.

주요 기여

• FP8‑centric compression pipeline은 데이터‑드리븐 재배열 단계와 **Adaptive Scale‑Hadamard Transform (ASHT)**를 결합하여 중간 텐서의 고‑충실도 양자화를 수행합니다.
• Dual‑Scale Quantization (DSQ) 메커니즘은 전체 훈련 과정에서 수치 안정성을 유지하여 일반적으로 저정밀도 방식에서 발생하는 오버플로/언더플로를 방지합니다.
• Highly fused compression operator는 재배열, 스케일링, 양자화를 하나의 GPU 커널로 통합하여 메모리 트래픽과 커널‑런치 오버헤드를 크게 감소시킵니다.
• 기존 데이터‑패럴렐(DP) 및 파이프라인‑패럴렐(PP) 런타임과 Seamless integration을 제공하여 3‑D parallel training stack(DP × PP × TP)를 구현하고, 이를 Megatron‑LM, DeepSpeed와 같은 인기 프레임워크에 바로 적용할 수 있습니다.
• GPT‑2/3‑스케일 모델 및 Qwen 계열에 대한 Extensive empirical validation을 수행했으며, 거의 손실 없는 퍼플렉시티/정확도를 유지하면서 처리량을 최대 **1.87×**까지 향상시켰음을 보여줍니다.

방법론

1. Reshaping & Distribution Awareness – 압축 전에 각 중간 텐서는 짧은 캘리브레이션 실행을 통해 학습된 경험적 값 분포에 따라 재배열됩니다. 이 “데이터‑드리븐 리쉐이핑”은 신호의 대부분을 더 작은 서브‑스페이스에 집중시켜 이후 양자화가 보다 효과적으로 이루어지게 합니다.
2. Adaptive Scale‑Hadamard Transform (ASHT) – 가벼운 직교 변환(Hadamard)이 텐서별 스케일링 팩터와 함께 적용되며, 이 팩터는 학습 중 관찰된 동적 범위에 맞춰 적응합니다. 변환은 데이터를 비상관화하여 분포를 0 주변으로 더욱 압축합니다.
3. FP8 Quantization + Dual‑Scale Quantization – 변환된 텐서는 8‑비트 부동소수점(FP8)으로 양자화됩니다. DSQ는 전방 패스와 역방 패스 각각에 대해 두 개의 스케일링 팩터를 유지함으로써, 동일한 압축 표현을 재사용하더라도 그래디언트와 활성값이 충분한 정밀도를 유지하도록 합니다.
4. Fused Compression Kernel – 모든 단계(reshape → ASHT → scaling → FP8 cast)가 단일 CUDA 커널에서 구현되어 중간 버퍼를 없애고, 커널이 NCCL 통신과 동시에 실행될 수 있게 합니다.
5. 3‑D Parallel Integration – TACO의 압축/압축 해제 훅이 기존 3‑D 병렬 트레이너의 TP 레이어에서 기본 all‑reduce/all‑gather 호출을 대체하며, DP와 PP 로직은 그대로 유지됩니다.

결과 및 발견

• 통신량이 FP8 압축 후 평균 약 45 % 감소했습니다.
• 커널 실행 오버헤드가 결합 연산자 덕분에 약 30 % 감소했습니다.
• 학습 안정성은 전체 정밀도 TP와 비슷하게 유지되었으며, 손실 곡선이 거의 완벽하게 겹쳤습니다.
• 이 접근법은 디코더 전용(GPT) 및 인코더‑디코더(Qwen) 아키텍처 모두에서 작동하여 광범위한 적용 가능성을 보여줍니다.

실용적인 시사점

• 비용 절감: GPU 간 트래픽을 줄임으로써 클라우드 사용자는 동일한 하드웨어 예산으로 더 큰 모델을 학습하거나 학습 사이클을 더 빠르게 마쳐 GPU‑시간 비용을 감소시킬 수 있습니다.
• 확장성: TACO는 일반적인 64‑GPU 한계를 넘어 TP를 적용할 수 있게 하여 네트워크 병목 현상을 피하고, 일반 클러스터에서 실제 페타스케일 LLM을 구현할 수 있는 길을 엽니다.
• 프레임워크 채택: TACO가 TP 통신 프리미티브의 드롭‑인 대체 구현으로 제공되기 때문에 Megatron‑LM, DeepSpeed, 또는 FairScale를 사용하는 개발자는 최소한의 코드 수정으로 이를 활성화할 수 있습니다.
• 엣지‑투‑클라우드 연속성: 동일한 FP8 양자화 파이프라인을 추론 시 텐서 압축(예: 다중 노드 엣지 클러스터에서 모델 병렬 추론)으로 재활용할 수 있어 지연 시간을 절반으로 줄일 수 있습니다.
• 하드웨어 정렬: FP8이 현재 NVIDIA Hopper와 향후 출시될 AMD GPU에서 지원되므로 TACO는 양자화/역양자화 단계에서 네이티브 텐서‑코어 가속을 활용해 성능을 더욱 향상시킬 수 있습니다.

제한 사항 및 향후 연구

• 하드웨어 의존성: 현재 구현은 빠른 FP8 지원이 가능한 GPU를 가정합니다; 구형 하드웨어는 에뮬레이션된 FP8으로 대체되어 이득이 감소합니다.
• 보정 오버헤드: 데이터‑드리븐 재구성은 훈련 시작 시 짧은 보정 단계가 필요합니다; 규모는 크지 않지만 완전 동적 워크로드를 위해 자동화가 필요할 수 있습니다.
• TP를 넘어 확장: 이 논문은 TP 텐서에 초점을 맞추고 있으며, 동일한 압축 아이디어를 DP 그래디언트나 PP 활성화에 적용하는 것은 아직 미해결 과제입니다.
• 극한 스케일링에 대한 견고성: 실험은 256 GPU까지 진행되었습니다; 향후 연구에서는 수천 노드로 확장할 때 네트워크 토폴로지 효과가 더욱 두드러지는 안정성 및 속도 향상을 검증해야 합니다.

전반적으로 TACO는 텐서‑패럴렐 LLM 훈련에서 오랫동안 병목 현상이 되어 온 통신 문제에 대한 실용적이고 높은 영향을 미치는 해결책을 제공하며, 보다 비용 효율적인 대규모 모델 개발의 길을 열어줍니다.

저자

• Man Liu
• Xingchen Liu
• Xingjian Tian
• Bing Lu
• Shengkay Lyu
• Shengquan Yin
• Wenjing Huang
• Zheng Wei
• Hairui Zhao
• Guangming Tan
• Dingwen Tao

논문 정보

• arXiv ID: 2604.24088v1
• 분류: cs.DC, cs.AI
• 출판일: 2026년 4월 27일
• PDF: PDF 다운로드