Technische Samenvatting: DisagMoE

1. Probleemstelling

Grootschalige Mixture-of-Experts (MoE)-architecturen maken het trainen van Large Language Models (LLM's) met biljoenen parameters mogelijk door een subset van experts spaarzaam te activeren. Het trainen van deze modellen stuit echter op ernstige knelpunten vanwege de all-to-all communicatie die vereist is voor expertparallelisme (EP). Naarmate modellen schalen over meerdere GPU-nodes, verergert de beperkte netwerkbandbreedte tussen nodes (bijv. InfiniBand) de communicatielatentie, die tot 50% van de tijd van een trainingsstap kan in beslag nemen.

Bestaande optimalisatiestrategieën richten zich voornamelijk op het overlappen van communicatie met berekening op het operator-niveau (bijv. het chunken van FFN-berekeningen) of het gebruik van micro-batch pipelining. Deze benaderingen lijden echter onder twee fundamentele beperkingen:

1. Residuele Stallingen: Inherent onbalans in de verhouding tussen berekening en communicatie tussen attention- en FFN-lagen voorkomt een perfecte overlapping. Attention-berekening schaalt kwadratisch met de sequentielengte, terwijl FFN en communicatie lineair schalen. Dit creëert "communicatie-uiteinden" die niet verborgen kunnen worden, vooral naarmate de sequentielengtes toenemen.
2. Resource-inefficiëntie: Eerdere gedesegreerde benaderingen (het scheiden van attention en FFN op verschillende hardware) vereisten vaak dat elke groep de volledige modelparameters in het geheugen hield, wat leidde tot Out-of-Memory (OOM)-fouten bij schaling, of faalden in het aanpassen van de resource-allokatie aan de specifieke reken-communicatiekenmerken van elke component.

2. Methodologie: DisagMoE

DisagMoE is een trainsysteem dat is ontworpen om all-to-all knelpunten te mitigeren door modelplaatsing en planning gezamenlijk te optimaliseren via een gedesegreerde architectuur.

2.1 Gedesegreerde Componentplaatsing

In tegenstelling tot traditioneel Expert Parallelisme, waarbij attention- en FFN-lagen op dezelfde GPU-groep coëxisteren, partitioneert DisagMoE het model op componenttype:

• Attention-groepen (A-Workers): Deze groepen bevatten self-attention-lagen uit meerdere transformerlagen. Ze repliceren dichte componenten en verdelen trainingsdata via Data Parallelism (DP).
• FFN-groepen (F-Workers): Deze groepen bevatten FFN-lagen (experts) uit meerdere transformerlagen. Experts worden over GPU's binnen de groep verspreid via Expert Parallelism (EP).
• Lagetoewijzing: Lagen worden op een verweven manier aan groepen toegewezen (bijv. Groep 0 behandelt lagen 0, $p$, $2p$; Groep 1 behandelt lagen 1, $p+1$, enz.), wat de apparaatbenutting maximaliseert onder geheugenbeperkingen.

2.2 AF-Pipe: Meerfasige Pipeline met Overlapping

DisagMoE introduceert AF-Pipe, een pijplijnplanning die communicatie behandelt als een eerste klas fase in plaats van een neveneffect.

• Veel-naar-Veel (M2N) Communicatie: In plaats van standaard all-to-all uitwisselingen binnen een enkele node-groep, vestigt AF-Pipe unidirectionele, veel-naar-veel token-uitwisselingen tussen de Attention- en FFN-groepen.
• Fase-uitlijning: De pijplijn aligneert grenzen over groepen heen, en voegt de traditionele "dispatch"- en "combine"-operaties samen tot één geünificeerde communicatiefase. Dit vermindert de communicatie- overhead met ongeveer $1/k$ (waarbij $k$ de top-$k$ experts is) en elimineert redundante gegevensoverdrachten.
• Asynchrone Executie: Het systeem gebruikt drie gecoördineerde streams (forward, backward en communicatie) binnen elke worker. Dit stelt de FFN-groep in staat om gradienten te sturen (N2M) terwijl het tegelijkertijd de forward pass berekent op beschikbare data, wat niet-blokkerende uitvoering en bijna continue GPU-benutting garandeert.

2.3 Adaptieve Worker-allokatie

Om het onbalans tussen attention (rekengebonden) en FFN (communicatiegebonden) aan te pakken, maakt DisagMoE gebruik van een adaptieve allocator die wordt geleid door een Reken-Communicatie Roofline Model.

• Modellering: Het systeem modelleert het keerpunt ($\hat{I}$) waar een component overschakelt van netwerkgebonden naar rekengebonden. Attention-lagen bereiken de rekenplafond sneller dan FFN-lagen naarmate de sequentielengte toeneemt.
• Optimalisatie: Het systeem formuleert de partitionering van GPU's en Network Interface Cards (NIC's) als een Mixed Integer Linear Programming (MILP)-probleem. Het alloceert dynamisch meer netwerkbandbreedte aan FFN-groepen (om hun communicatieplafond te verhogen) en meer rekenresources aan Attention-groepen, waardoor de pijplijnfasen in evenwicht worden gebracht om bubbels te minimaliseren.
• Zoekstrategie: Een algoritme gebruikt het roofline-model om een zaadconfiguratie te genereren, die vervolgens wordt verfijnd via lokaal zoeken met gebruik van gemeten tijden per stap om niet-gemodelleerde overheads in aanmerking te nemen.

3. Belangrijkste Bijdragen

1. Gedesegreerde Componentplaatsing: Een nieuwe architectuur die attention- en FFN-lagen scheidt in disjuncte worker-groepen, waardoor onafhankelijke schaling en resource-allokatie mogelijk zijn.
2. AF-Pipe Planning: Een meerfasige pijplijn die conventioneel all-to-all vervangt door inter-groep M2N-uitwisselingen, en systematisch communicatie overlapt met zowel attention- als FFN-berekeningen.
3. Adaptieve Resource-allokatie: Een roofline-geleide mechanisme dat de grootte van GPU- en NIC-resources voor attention- en FFN-groepen gezamenlijk bepaalt, waardoor de reken-communicatieverhouding effectief in evenwicht wordt gebracht over variërende sequentielengtes en modelconfiguraties.

4. Experimentele Resultaten

De auteurs evalueerden DisagMoE op clusters van 16 nodes (128× H800 GPU's) met drie MoE-modellen: DeepSeek-MoE, GPT-OSS-120B en Qwen3, met sequentielengtes variërend van 4K tot 32K.

• Throughputverbetering: DisagMoE presteerde consequent beter dan sterke baselines, waaronder Megatron-LM (1F1B verweven), Tutel, Comet en DualPipe.
  • Bereikte tot 1,81× snelheidswinst ten opzichte van Megatron-1F1B.
  • Bereikte tot 1,34× snelheidswinst ten opzichte van state-of-the-art MoE-overlapsystemen.
• Communicatiereductie: In vergelijking met baselines verkleinde DisagMoE de niet-overlappende communicatietijd met tot 88% (vs. Tutel), 75% (vs. Comet) en 45% (vs. DualPipe).
• Schalbaarheid: Het systeem toonde effectiviteit over verschillende top-$k$-waarden en Expert Parallelism (EP)-groottes.
• Resourcesensitiviteit: Experimenten bevestigden dat de optimale GPU-verhoudingen tussen attention- en FFN-groepen variëren met de sequentielengte (bijv. langere sequenties bevoordelen meer attention-resources), wat de noodzaak van adaptieve allokatie valideert.

5. Betekenis en Claims

Het artikel claimt dat DisagMoE een fundamentele beperking in huidige MoE-training aanpakt: het onvermogen om communicatie perfect te overlappen vanwege het structurele onbalans tussen attention- en FFN-lagen. Door deze componenten te desegregeren en resources adaptief toe te wijzen op basis van een roofline-model, breidt het systeem het effectieve overlappingvenster uit verder dan wat symmetrische bandbreedtevoorziening kan bereiken.

De auteurs benadrukken dat hun benadering bijzonder significant is voor grootschalige pre-training workloads met vaste sequentielengtes. Zij merken op dat hoewel hun huidige ontwerp stabiele vormen aanneemt (vaste sequentielengte en micro-batch), het kerninzicht – dat toekomstige hardware en software moeten mikken op gunstige reken-communicatieverhoudingen over componenten heen in plaats van simpelweg ruwe bandbreedte te schalen – een nieuwe richting biedt voor het efficiënt schalen van MoE-training. Het systeem is geïmplementeerd op Megatron-LM en blijkt een haalbare oplossing te zijn voor het trainen van modellen met honderden miljarden parameters zonder de geheugenbeperkingen die eerdere gedesegreerde trainingspogingen teisterden.