Technische Samenvatting: AdaptiveLoad

1. Probleemstelling

Het trainen van grote video-diffusie-Transformers (zoals DiT en MMDiT) voor wereldmodellen staat voor twee kritieke knelpunten bij het verwerken van gemengde datasets met sterk heterogene sequentielengtes (variërend van statische afbeeldingen tot video's van lange duur):

• Rekenkundig onevenwicht door kwadratische complexiteit: Bestaande gedistribueerde trainingspijplijnen vertrouwen op strategieën voor "gelijk-token"-bucketing (waarbij $B \times S = \text{constant}$ wordt beperkt). Deze aanpak houdt geen rekening met de $O(S^2)$ rekenkundige complexiteit van self-attention-mechanismen. Bijgevolg worden buckets met lange sequenties "rekenkundige achterblijvers", wat leidt tot ernstige onevenwicht in de belasting. De globale synchronisatietijd wordt bepaald door de langzaamste GPU, wat aanzienlijke inactiviteit (synchronisatiebellen) voor andere knooppunten veroorzaakt. Experimentele data wijst op een zwakke correlatie ($R \approx 0,35$) tussen staplatentie en het totale aantal tokens, maar een sterke correlatie ($R \approx 0,92$) met het kwadraat van de sequentielengte.
• Geheugengrens en operator-overhead: In Multimodal DiT (MMDiT)-architecturen wordt Adaptive Layer Normalization (AdaLN) honderden keren per iteratie aangeroepen. Standaardimplementaties gebruiken discrete CUDA-kernels, wat frequente, overbodige heen-en-weer-bewegingen van tussenliggende tensors tussen Streaming Multiprocessors (SM's) en High Bandwidth Memory (HBM) veroorzaakt. Bovendien lijdt tijdens backpropagation de gradiëntreductie voor modulatieparameters onder inefficiënte gestreepte geheugentoegangspatronen (niet-continue toegang langs de sequentiedimensie $N$), wat hardware-level gecoalesceerde toegang verhindert en de geheugenbandbreedte ernstig onderbenut.

2. Methodologie: Het AdaptiveLoad-kader

Het artikel stelt AdaptiveLoad voor, een full-stack optimalisatiekader dat hoge-niveau-planning integreert met low-level kernel-fusie.

A. Adaptieve lastverdeling met dubbele beperking

Om het rekenkundige onevenwicht aan te pakken, vervangen de auteurs de lineaire "gelijk-token"-aanname door een planningslogica met dubbele beperking die gelijktijdig handhaaft:

1. Lineaire geheugenbeperking: $B \times S \leq M_{mem}$ (zorgend dat de batchgrootte in het GPU-geheugen past).
2. Polynomiale rekenkundige beperking: $B \times S^p \leq M_{comp}$ (rekening houdend met $O(S^2)$ complexiteit).

De uiteindelijke batchgrootte voor een bucket wordt bepaald door het snijpunt van deze beperkingen:
$$B_{shape} = \max(1, \min(\lfloor M_{mem}/S \rfloor, \lfloor M_{comp}/S^p \rfloor))$$

• Systeem voor kostenfitting: Om de empirische exponent $p$ en de rekenkundige bovengrens $M_{comp}$ te bepalen, maakt het systeem gebruik van een datagedreven Shape Benchmarking-workflow. Het voert synthetische pixel-scans uit in een live gedistribueerde omgeving (bijv. FSDP-paden) om de staplatentie te meten tegenover batchgrootte en sequentielengte. Met behulp van niet-lineaire regressie past het een kostenmodel aan ($\text{staptijd} \approx a + b \times B \times S^p$) om bucketconfiguraties dynamisch te kalibreren, waardoor de behoefte aan handmatige empirische afstelling wordt geëlimineerd.

B. Gefuseerde LayerNorm-Modulate CUDA-kernels

Om de "geheugengrens" aan te pakken, ontwerpen de auteurs een gefuseerde AdaLN CUDA-kernel die normalisatie- en modulatiewerkingen combineert in één kernel-start.

• Forward Pass: De kernel maakt gebruik van registers en Shared Memory om tussenliggende statistieken (gemiddelde en variantie) op te slaan, en voert twee-trapsreducties uit (Warp-level Shuffle en Block-level Reduction). Lineaire transformaties worden direct toegepast tijdens een derde doorloop van de feature-dimensie $D$, waardoor overbodige HBM-lezen/schrijven wordt geëlimineerd.
• Backward Pass (D-tile gecoalesceerde reductie): Om het knelpunt van gestreepte toegang bij gradiëntreductie op te lossen, introduceren de auteurs een D-tile gecoalesceerde reductie-strategie. In plaats van te reduceren langs de sequentiedimensie $N$ (wat gestreepte toegang veroorzaakt), wordt het Grid opgedeeld in een $(d_{tile}, n_{tile})-lay-out$. Threads zijn vastgekoppeld aan een feature-index $d$ en voeren verticale accumulatie uit langs sequentie-tiles. Dit zorgt ervoor dat geheugentoegangen binnen een Warp contigu zijn, wat gecoalesceerde HBM-transacties mogelijk maakt en het bandbreedtegebruik maximaliseert.
• Vereenvoudiging van het graf: Door de operator-keten te fuseren, wordt het rekenkundige graf samengevouwen van 5–8 discrete knooppunten tot één atomair knooppunt. Tussenliggende statistieken blijven in registers, waardoor het globale geheugenbeheer van het kader wordt omzeild en het geheugenverbruik voor activaties aanzienlijk wordt verminderd.

3. Belangrijkste bijdragen

1. Kwantificering en mitigeren van lastmismatch: Het artikel kwantificeert rigoureus de lastmismatch in video-diffusietraining en stelt een adaptief bucket-planningsalgoritme met dubbele beperking voor. Deze aanpak verlaagt de onevenwichtsgraad in de belasting met maximaal 50% door actief beheer van de $O(S^2)$ complexiteit van lange sequenties.
2. Datagedreven kostenmodellering: De auteurs ontwerpen een end-to-end-systeem dat de configuratie van gedistribueerde training verschuift van heuristisch-gedreven naar datagedreven. Door middel van geautomatiseerde kostenmodel-fitting via Shape Benchmarking, optimaliseert het systeem de clusterbrede resourcebenutting zonder handmatige tussenkomst.
3. Hoogpresterende gefuseerde kernels: Er is een gespecialiseerde gefuseerde AdaLN CUDA-kernel ontwikkeld voor MMDiT-ruggengraten. Door D-tile gecoalesceerde reductie in te voeren, wordt het knelpunt van gestreepte toegang in backpropagation opgelost, wat de trainingsdoorvoer met meer dan 25% verhoogt terwijl het tegelijkertijd het geheugenverbruik voor activaties verlaagt.

4. Experimentele resultaten

Experimenten werden uitgevoerd op het Wan 2.1-wereldmodel met een gemengd corpus van 10 miljoen samples (WebDataset en Koala-36m) op clusters van 8 GPU's en 16 GPU's.

• Verbetering van doorvoer:
  • Cluster van 8 GPU's: De gemiddelde trainingsdoorvoer steeg van 14.383 naar 18.069 tokens/sec (+25,6%).
  • Cluster van 16 GPU's: De gemiddelde doorvoer steeg van 30.170 naar 38.372 tokens/sec (+27,18%).
  • De AdaptiveLoad-curve handhaafde een hogere doorvoerondergrens vergeleken met de baseline, zelfs bij worst-case sequentieverdelingen.
• Efficiëntie van lastverdeling:
  • De onevenwichtsgraad in de rekenbelasting (gemeten door de Variatiecoëfficiënt van de rekenbelasting) werd verlaagd van 39% naar 18,9%.
  • De synchronisatie-overhead (CV van staptijd) daalde aanzienlijk, waarbij de configuratie met 16 GPU's een reductie zag van een baseline-gemiddelde van 18,7% naar 10,4% onder AdaptiveLoad.
• Geheugen- en kernelprestaties:
  • Pieken VRAM-gebruik: Verbeterd met 22,7%.
  • Snelheidswinst van kernels: De gefuseerde forward pass behaalde een snelheidswinst van 3,21× tot 3,39×. De backward pass (D-tile reductie) behaalde tot 1,42× snelheidswinst bij een sequentielengte van 64k.
  • Geheugenbesparing: De gefuseerde kernel verlaagde het geheugenverbruik voor activaties met ongeveer 61,9% (bijvoorbeeld ~4,8 GB besparen bij $N=64.000$).
  • Lastcapaciteit: Het geoptimaliseerde systeem ondersteunde een maximale sequentielengte van 52.800 tokens (een toename van 10% ten opzichte van de limiet van 48.000 tokens van de baseline) zonder het piekgeheugenverbruik te verhogen.
• Trainingsstabiliteit: De loss-curves voor AdaptiveLoad bleven zeer congruent met de baseline, wat aantoont dat de optimalisaties numerieke trouw en convergentiekwaliteit behielden.

5. Betekenis

Het artikel beweert dat AdaptiveLoad een robuste basis biedt voor het schalen van modellen voor video-generatie met ultra-lange sequenties. Door geheugenveiligheid te verzoenen met rekenkundige pariteit via planning met dubbele beperking en de "geheugengrens" aan te pakken via hardware-bewuste operator-fusie, elimineert het kader effectief de synchronisatieknelpunten die inherent zijn aan heterogene videodatasets. De resultaten tonen aan dat het overstappen van lineaire token-telling-heuristieken naar planning die rekening houdt met kwadratische complexiteit essentieel is voor het maximaliseren van de efficiëntie van grote GPU-clusters in taken voor video-generatie.