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

1. Probleemstelling

De training van Large Language Models (LLM) op grote schaal is sterk afhankelijk van 3D-parallelisme (Data Parallelism, Pipeline Parallelism en Tensor Parallelism). Hoewel voor Data Parallelism (DP) en Pipeline Parallelism (PP) succesvolle communicatiecompressietechnieken zijn ontwikkeld, blijft Tensor Parallelism (TP) een kritieke knelpunt.

• Communicatie-Overhead: TP vereist frequente, nauwkeurig gesynchroniseerde communicatie van intermediaire tensoren (activaties en gradiënten) bij elke laag en iteratie. Bij training op grote schaal staat TP-communicatie voor 50–60% van de totale communicatietijd en ligt het op het kritieke uitvoeringspad, waardoor het moeilijk te overlappen is met berekening.
• De Compressie-uitdaging:
  • Foutaccumulatie: In tegenstelling tot DP/PP omvat TP frequentie-synchronisatie. Kleine kwantisatiefouten in intermediaire tensoren verspreiden zich via attention- en residual-verbindingen, versterken tijdens backpropagation en veroorzaken divergentie van de training.
  • Data-distributie: TP-intermediaire tensoren vertonen een dichte, bijna-nul distributie met een lange staart. Standaard kwantisatiemethoden (zoals INT8) falen hier omdat ze uniforme tussenruimte gebruiken, wat leidt tot massaal informatieverlies voor de waarden met kleine magnitude die de tensor domineren.
  • Inefficiëntie van bestaande methoden: Eerdere compressiemethoden gericht op inferentie (bijv. SmoothQuant) of op DP gerichte methoden (bijv. SDP4bit) kunnen niet direct worden toegepast op TP-training zonder catastrofale uitval door numerieke instabiliteit.

2. Methodologie: TACO Framework

De auteurs stellen TACO (Tensor-parallel Adaptive COmmunication compression) voor, een robuust op FP8 gebaseerd framework dat specifiek is ontworpen voor TP-intermediaire tensoren.

A. Waarom FP8?

Het artikel analyseert de statistische eigenschappen van TP-tensoren en concludeert dat FP8 superieur is aan INT8:

• INT8-beperkingen: Uniforme kwantisatie veroorzaakt "botsingen" waarbij verschillende kleine waarden worden gemapt naar hetzelfde gehele getal, wat leidt tot een hoge relatieve fout in de dichte, rond nul gecentreerde regio.
• FP8-voordeel: FP8 (specifiek het E4M3-formaat) gebruikt een exponentiële schaal, wat hoge precisie voor kleine waarden (dicht bij nul) biedt en een groot dynamisch bereik voor uitschieters. Dit sluit perfect aan bij de statistische distributie van TP-tensoren.

B. Kerntechnieken

TACO introduceert drie belangrijke innovaties om FP8-compressie haalbaar te maken voor TP-training:

1. Adaptive Scale–Hadamard (ASH) Transform:

   • Probleem: Standaard Hadamard-transformaties zijn isometrisch (behouden energie) en kunnen de dichte nul-clusters van TP-tensoren niet effectief verspreiden; ze blijven gevangen rond nul, waardoor de precisie van FP8 onvoldoende wordt benut.
   • Oplossing: TACO hanteert een bloksgewijze adaptieve herschaling vóór de Hadamard-transformatie. Het berekent de tweede-orde ruwe moment (energie) van elk blok en past een schalingsfactor toe om de energie te normaliseren naar een doelniveau.
   • Effect: Dit "ontvouwt" de dichte nul-clusters, spreidt ze over het representeerbare bereik van FP8, gevolgd door een orthogonale Hadamard-rotatie om een Gauss-achtige distributie te creëren die geschikt is voor kwantisatie.

2. Dual-Scale (DS) Quantization:

   • Probleem: Zelfs met ASH kunnen extreme waarden overloop veroorzaken, of kleine waarden onderloop, wat leidt tot instabiliteit.
   • Oplossing: TACO gebruikt twee verschillende scalairen per blok:
     • Pre-rotatie schaal ($\alpha_k$): Regelt de energie van het blok vóór transformatie.
     • Post-rotatie schaal ($s_k$): Zorgt ervoor dat de getransformeerde waarden strikt binnen het dynamische bereik van FP8 ($Q_{max}$) vallen.
   • Reconstructie: Tijdens decompressie worden deze twee schalen gebruikt om het proces om te draaien, waardoor reconstructie met hoge fideliteit wordt gewaarborgd en foutversterking wordt voorkomen.

3. Systeemniveau Kernel Fusion & Overlap:

   • Gefuseerde Operator: Om de rekenkundige overhead van compressie te minimaliseren (wat kritiek is aangezien TP op het kritieke pad ligt), fuseert TACO variantierekening, ASH-transformatie en DS-kwantisatie tot een enkele GPU-kernel. Dit elimineert redundante toegang tot globaal geheugen en kernel-launch-overhead.
   • Communicatie-Berekening Overlap: TACO integreert met de COCCL (Compressible Collective Communication Library) en gebruikt een "two-shot" AllReduce-strategie (ReduceScatter + AllGather). Het overlapt de compressie/decompressie-kernels met communicatie, waardoor latentie effectief wordt verborgen.

3. Belangrijkste Bijdragen

1. Systematische Analyse: Identificatie van het "Zero-Collapse"-fenomeen in TP-tensoren, wat aantoont dat standaard INT8 en vaste transformaties ontoereikend zijn, en FP8 vestigt als het optimale formaat.
2. TACO Framework: Het eerste 3D-parallelle communicatiecompressiesysteem dat stabiele convergentie bereikt in TP-training op grote schaal. Het introduceert de ASH-transformatie en Dual-Scale kwantisatie om de unieke distributie van TP-tensoren te hanteren.
3. Hoge Prestaties Implementatie: Ontwikkeling van een sterk gefuseerde compressie-operator die het verkeer in het geheugen vermindert en efficiënte overlapping met communicatie mogelijk maakt, waardoor de overhead van frequente compressie wordt beperkt.
4. Volledige 3D-Integratie: Succesvolle integratie van TACO met state-of-the-art DP (SDP4bit) en PP (TahQuant) compressoren om een volledig op compressie gebaseerd 3D-parallel trainingssysteem te creëren.

4. Experimentele Resultaten

De auteurs evalueerden TACO op GPT-modellen (2.7B, 6.7B, 13B) en Qwen2.5-7B met NVIDIA H100 GPU's.

• Nauwkeurigheid (Convergentie):

  • TACO bereikt nagenoeg verliesvrije nauwkeurigheid. Op GPT-350M is de degradatie van de validatieverlies slechts +0,25% vergeleken met de BF16-baseline, terwijl bestaande methoden (TahQuant) degradatie van ongeveer 2,9% vertonen.
  • Bij grotere modellen (GPT-6.7B, Qwen-2.5-7B) blijft de degradatie verwaarloosbaar (<0,7%), wat stabiliteit over schalen bewijst.
  • Ablatiestudies bevestigen dat het verwijderen van ASH of DS leidt tot catastrofale divergentie of aanzienlijk verlies aan nauwkeurigheid.

• Prestaties (Throughput):

  • TP-Compressie: TACO bereikt tot 1,87× verbetering in end-to-end throughput voor TP-communicatie vergeleken met de niet-gecomprimeerde baseline.
  • Volledig 3D-Parallelisme: Wanneer toegepast op een volledig 3D-parallelle setup (TP=4, PP=2, DP=2), levert TACO tot 1,53× snelheidswinst op ten opzichte van de baseline.
  • Schaalbaarheid: De prestatiewinsten nemen toe met de mate van Tensor Parallelism (TP), aangezien TP-communicatie de dominante bottleneck wordt.

5. Betekenis

• Oplossen van de TP-bottleneck: TACO adresseert het moeilijkste aspect van gedistribueerde LLM-training: het comprimeren van de frequente, numeriek gevoelige intermediaire tensoren in Tensor Parallelism zonder in te leveren op modelkwaliteit.
• Aanbieden van grotere modellen: Door de communicatie-Overhead met bijna 50% te verminderen (via FP8) terwijl stabiliteit wordt behouden, stelt TACO het trainen van grotere modellen op bestaande hardware-clusters mogelijk, of het trainen van bestaande modellen aanzienlijk sneller.
• Hardware-onafhankelijk: Hoewel geoptimaliseerd voor FP8 (ondersteund door moderne GPU's zoals de H100), is het framework ontworpen om op hardware zonder native FP8-ondersteuning graceful te degraderen naar INT8, wat brede toepasbaarheid garandeert.
• Toekomstige Impact: Dit werk ebaant de weg voor exascale-training door aan te tonen dat agressieve communicatiecompressie haalbaar is, zelfs op het kritieke pad van training, mits de compressie wiskundig is toegesneden op de data-distributie.