Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

1. Probleemstelling

Het trainen en infereren van grote taalmodellen (LLM's) staat voor twee primaire geheugenflessenhals:

1. Parametergeheugen: De opslagruimte die nodig is voor modelgewichten, gradiënten en optimizer-toestanden.
2. Activeringsgeheugen: De opslagruimte die nodig is voor intermediaire activeringen, die kwadratisch schaalt met de sequentielengte ($S^2$) in attention-mechanismen.

Bestaande parallelismestrategieën adresseren deze problemen geïsoleerd of met compromissen:

• Tensor Parallelism (TP): Verdeelt modelgewichten over apparaten, wat het parametergeheugen vermindert maar het activeringsgeheugen volledig gekopieerd laat. Het wordt inefficiënt bij lange sequentielengtes door het hoge activeringsgeheugen en het communicatievolume.
• Sequence Parallelism (SP): Verdeelt de inputsequentie over apparaten, wat het activeringsgeheugen vermindert maar de modelgewichten volledig gekopieerd laat. Het is inefficiënt wanneer parametergeheugen de beperkende factor is.
• Conventioneel TP+SP: Combineert beide door ze toe te wijzen aan orthogonale assen in een apparaatmesh (bijvoorbeeld een $T \times \Sigma$-rooster). Hoewel effectief, verbruikt dit een groot aantal apparaten voor modelparallelisme ($T \times \Sigma$), waardoor er minder apparaten overblijven voor dataparallelisme. Bovendien dwingt het parallelle groepen vaak om langzamere inter-node links te overspannen, waardoor hoogwaardige intra-node interconnects (zoals NVLink of AMD Infinity Fabric) onderbenut blijven.

De Kernuitdaging: Hoe kan men zowel parameter- als activeringsgeheugen gelijktijdig verminderen op één apparaatas om geheugenefficiëntie te maximaliseren en apparaten vrij te maken voor dataparallelisme, zonder prohibitieve communicatiekosten te maken.

2. Methodologie: Tensor and Sequence Parallelism (TSP)

De auteurs stellen TSP voor, een strategie die tensor- en sequence-parallelisme "vouwt" op een enkele apparaatas. In plaats van TP en SP toe te wijzen aan orthogonale dimensies, wijst TSP elke rank (GPU) zowel een shard van de modelgewichten toe als een shard van de sequentie.

Belangrijkste Ontwerpelementen:

• Gecentraliseerde Sharding: Elk apparaat bevat $1/D$ van de modelparameters en $1/D$ van de sequentietokens. Dit vermindert zowel het gewichtsproportionele geheugen als het activeringsgeheugen met een factor $D$ op één as.
• Runtime-schema's: Om het feit te hanteren dat gewichten en activeringen verschillend worden geshard, hanteert TSP specifieke communicatieschema's voor verschillende lagen:
  • Attention-lagen: Gebruikt een broadcast-schema.
    • Ranks itereren over gewichtsshards (gepakte $W_Q, W_K, W_V, W_O$).
    • In elke iteratie broadcastt één rank zijn gewichtsshard naar alle anderen.
    • Alle ranks berekenen lokale projecties op hun sequentieshards.
    • Een All-Gather-operatie reconstrueert de volledige Key/Value-context over de sequentiedimensie om FlashAttention uit te voeren.
    • Een Zigzag-partitioneringsschema wordt gebruikt om de causale attention-werklast in evenwicht te brengen (zorgen dat latere tokens niet onevenredig meer werk genereren).
  • Gated MLP-lagen: Gebruikt een ringcommunicatie-schema.
    • Gewichtsshards ($W_1, W_2, W_3$) circuleren punt-tot-punt rond de ring.
    • Elke rank acumuleert lokale deelresultaten.
    • Dit elimineert de noodzaak van een All-Reduce-operatie (standaard bij TP), aangezien de sequentie lokaal blijft en alleen gewichten bewegen.
• Hardware-bewustzijn: Het ontwerp is geoptimaliseerd voor topologieën zoals AMD's MI300X (Infinity Fabric), waarbij de collectieve bandbreedte schaalt met het aantal deelnemende GPU's. TSP zorgt ervoor dat communicatieprimitieven (broadcast en ring) de volledige groep betrekken of worden overdekt met berekening om de efficiëntie te maximaliseren.

3. Belangrijkste Bijdragen

1. Nieuwe Parallelismestrategie: Introductie van TSP, dat twee orthogonale parallelisme-assen samenvoegt tot één, waardoor gelijktijdige vermindering van parameter- en activeringsgeheugen mogelijk is.
2. Theoretische Analyse: Een strikt wiskundig model dat TSP vergelijkt met DP, TP, SP en TP+SP met betrekking tot:
   • Geheugen: TSP is het enige schema dat zowel gewichtsproportioneel als activeringsgeheugen vermindert met $1/D$ zonder $D^2$ apparaten te vereisen.
   • Communicatie: Hoewel TSP communicatie voor gewichtsverplaatsing introduceert (wat TP+SP vermijdt), bewijzen de auteurs dat voor lange sequenties of gemiddelde batchgroottes, de activeringsproportionele communicatie domineert, waardoor TSP concurrerend blijft.
   • FLOPs: TSP behoudt dezelfde per-apparaat rekenefficiëntie als TP en SP.
3. Implementatie & Overlap: Gedetailleerde implementatie van TSP-attention en MLP-blokken met expliciete strategieën voor communicatie/berekening-overlap (pipelining van gewichtsoverdrachten achter GEMM-berekeningen).
4. Topologie-optimalisatie: Aantonen dat TSP toelaat dat de volledige modelreplica binnen een enkele hoogwaardige node (intra-node) past, waardoor de prestatieboetes worden vermeden die vaak worden veroorzaakt door het overspannen van langzamere inter-node links die nodig zijn bij 2D TP+SP-meshes.

4. Resultaten

De auteurs hebben TSP geëvalueerd op een cluster van maximaal 1.024 AMD MI300X GPU's (128 nodes) met een dicht transformer-model van 7B parameters.

• Geheugenefficiëntie:
  • TSP bereikte consistent de laagste piekgeheugengebruik over alle geteste sequentielengtes (16K tot 128K tokens).
  • Bij korte contexten kwam TSP overeen met TP (gedomineerd door parametergeheugen).
  • Bij lange contexten presteerde TSP aanzienlijk beter dan TP en TP+SP (gedomineerd door activeringsgeheugen), waardoor Out-Of-Memory (OOM)-fouten werden voorkomen waar andere strategieën faalden.
• Doorvoer:
  • Ondanks een hoger theoretisch communicatievolume (door gewichtsverplaatsing), behield TSP concurrerende doorvoer met TP+SP.
  • De extra communicatie werd succesvol verborgen achter berekening via pipelining.
  • De doorvoervoordelen groeiden met toenemende parallelismegraden ($D$) en sequentielengtes.
• Schaalbaarheid:
  • TSP maakte grotere micro-batches mogelijk door het verminderde geheugenfootprint, wat de doorvoer verder verbeterde.
  • Het stelde in staat om modelreplica's volledig binnen een enkele node (8 GPU's) te plaatsen, waardoor ranks vrijkwamen voor Dataparallelisme (DP) in grotere clusters.

5. Betekenis

• Hardware-bewuste schaling: TSP biedt een praktische oplossing voor moderne HPC-clusters waar de intra-node bandbreedte hoog is, maar de inter-node bandbreedte een bottleneck vormt. Het maximaliseert het nut van dure intra-node interconnects.
• Haalbaarheid voor lange context: Naarmate modellen bewegen naar langere contextvensters (100k+ tokens), wordt activeringsgeheugen de primaire bottleneck. TSP is uniek gepositioneerd om dit aan te pakken zonder de enorme apparaatoverhead van 2D-parallelisme.
• Flexibiliteit: TSP kan orthogonaal worden samengesteld met andere parallelismeschema's (Pipeline Parallelism, Expert Parallelism, Dataparallelism), wat een flexibele "as" biedt voor systeemontwerpers om te optimaliseren voor specifieke hardwaretopologieën.
• Kostenefficiëntie: Door het aantal apparaten dat nodig is voor modelparallelisme te verminderen, staat TSP toe dat meer apparaten worden gebruikt voor dataparallelisme, wat potentieel de totale clustergrootte die nodig is voor een bepaalde trainingstask verkleint of de effectieve batchgrootte vergroot.

Concluderend vertegenwoordigt TSP een significante evolutie in gedistribueerde trainingsstrategieën, waarbij een beheersbare toename in communicatievolume wordt geruild voor aanzienlijke winsten in geheugenefficiëntie en hardwarebenutting, waardoor het een cruciaal instrument wordt voor het trainen van next-generation lange-context LLM's.