Technische Samenvatting: SwiftFusion

Probleemstelling

Diffusie-Transformers (DiT's) worden steeds vaker gebruikt voor het genereren van beelden en video's van hoge kwaliteit. Het genereren van beelden met hoge resolutie of video's met lange duur resulteert echter in lange sequentielengtes, waardoor activatietensors lineair groeien met de sequentielengte. Dit leidt tot geheugenbeperkingen en hoge latentie op enkele GPU's, wat gedistribueerde inferentie via Sequentieparallelisme (SP) noodzakelijk maakt.

Bestaande SP-technieken, zoals Ring Attention, Ulysses Attention en hun combinatie Unified Sequence Parallelism (USP), staan voor drie primaire beperkingen bij schaling over meerdere GPU-machines:

1. Suboptimale Topologiebenutting: Huidige frameworks houden vaak geen rekening met de aanzienlijke bandbreedteverschillen tussen netwerken binnen een machine (bijv. NVSwitch) en netwerken tussen machines (bijv. EFA/RDMA). Zo gebruikt USP bijvoorbeeld Ring Attention voor communicatie tussen machines, wat leidt tot een hoog communicatievolume dat niet afneemt met meer machines, waardoor een knelpunt ontstaat.
2. Niet-overlappende Communicatie: Technieken die vertrouwen op Ulysses Attention maken gebruik van all-to-all-operatoren die worden behandeld als atomaire bewerkingen. Deze bewerkingen blokkeren de berekening, waardoor latentiebellen ontstaan; dit is vooral problematisch voor communicatie tussen machines waar de bandbreedte lager is.
3. Synchronisatie- en Berekeningsoverhead: Bestaande implementaties vertrouwen op tweezijdige communicatiebibliotheken (bijv. NCCL), die strikte zender-ontvangersynchronisatie vereisen. Dit introduceert impliciete barrières en verbruikt GPU Streaming Multiprocessors (SM's) voor communicatie op basis van kernels, wat leidt tot concurrentie met modelberekeningen.

Methodologie

De auteurs stellen SwiftFusion voor, een topologiebewuste, efficiënte DiT-servengine die is ontworpen om deze uitdagingen aan te pakken door middel van drie sleutelinnovaties:

1. Topologiebewuste Communicatieplanning

SwiftFusion heroverweegt de toewijzing van SP-technieken op basis van netwerktopologie. In tegenstelling tot USP, dat Ring Attention gebruikt voor communicatie tussen machines, keert SwiftFusion deze strategie om:

• Tussen machines: Het maakt gebruik van Ulysses Attention. Hoewel Ulysses in sommige contexten een hoog communicatievolume per machine heeft, neemt het all-to-all-volume af naarmate het aantal machines toeneemt, waardoor het geschikter is voor langzamere netwerken tussen machines bij schaling.
• Binnen een machine: Het maakt gebruik van Ring Attention. Omdat netwerken binnen een machine (bijv. NVSwitch) hoge bandbreedte en lage latentie bieden, vormt het hogere communicatievolume van Ring Attention hier minder een knelpunt.
  Deze planning minimaliseert het totale communicatievolume tussen machines en benut de hoogwaardige bandbreedte van links binnen een machine voor de meer communicatie-intensieve Ring-bewerkingen.

2. Torus Attention (Overlappende Communicatie en Berekening)

Om de latentie van de all-to-all-operatoren die nodig zijn voor Ulysses Attention (die in standaardimplementaties niet-overlappend zijn) te mitigeren, introduceert SwiftFusion Torus Attention.

• Chunking: De all-to-all-operatoren worden opgesplitst in meerdere chunks.
• Pipelining: Het systeem identificeert "stationaire" elementen in de tensors die niet hoeven te worden overgedragen tijdens specifieke fasen van de all-to-all-bewerking. Deze elementen worden direct gebruikt voor attentieberekening.
• Overlap: Terwijl de stationaire elementen worden berekend, start het systeem op de achtergrond de overdracht van de resterende chunks. Dit creëert een pijplijn waarbij de communicatie van één chunk overlapt met de berekening van de eerder ontvangen chunk, waardoor de latentie van communicatie tussen machines effectief wordt verborgen.

3. Implementatie met Éénzijdige Communicatie

SwiftFusion vervangt traditionele tweezijdige bibliotheken (zoals NCCL) door bibliotheken voor éénzijdige communicatie (bijv. NVSHMEM).

• Mechanisme: In plaats van send-receive-primitieven die wederzijdse synchronisatie vereisen, maakt het gebruik van push- en pull-primitieven.
• Optimalisatie: De auteurs ontwerpen een unificerend communicatiealgoritme dat Torus, Ulysses en Ring Attention verenigt. Deze aanpak minimaliseert synchronisatiebarrières en vereist alleen synchronisatie binnen een machine en twee specifieke synchronisaties tussen machines (aan het begin en einde van een laag).
• Efficiëntie: Door synchronisatie per overdracht te verwijderen en communicatie op basis van CUDA-kernels te vermijden die concurreren om SM's, vermindert het systeem zowel synchronisatie- als berekeningsoverhead.

Belangrijkste Bijdragen

Het paper vat de bijdragen als volgt samen:

1. Topologiebewuste Planning: Een nieuwe strategie die moderne GPU-machine-topologieën beter benut door Ulysses Attention toe te wijzen aan communicatie tussen machines en Ring Attention aan communicatie binnen een machine.
2. Torus Attention: Een nieuw gedistribueerd attentiealgoritme dat all-to-all-operatoren opbreekt om communicatie te pipelinen met berekening, waardoor latentie tussen machines wordt verborgen.
3. Implementatie met Éénzijdige Communicatie: Een gespecialiseerde implementatie die Torus, Ulysses en Ring Attention verenigt met behulp van bibliotheken voor éénzijdige communicatie om synchronisatie- en berekeningsoverhead te minimaliseren.
4. Empirische Evaluatie: Omvattende evaluatie op werklasten voor beeldgeneratie (Flux) en videogeneratie (CogVideoX) die state-of-the-art prestaties aantoont.

Resultaten

De auteurs hebben SwiftFusion geëvalueerd op AWS p4de.24xlarge-instanties (8x NVIDIA A100 GPU's per machine, onderling verbonden via NVSwitch en EFA) in vergelijking met het state-of-the-art USP-framework.

• Algemene Prestaties: SwiftFusion behaalde een gemiddelde snelheidswinst van 1,35× ten opzichte van USP, met een maximale snelheidswinst van 1,77× over verschillende gedistribueerde configuraties en werklasten (Flux 3072×3072/4096×4096, CogVideoX 20s/40s).
• Componentanalyse:
  • De topologiebewuste planning alleen (TAS) leverde een snelheidswinst van 1,27× op ten opzichte van USP bij gebruik van meer dan twee GPU-machines, voornamelijk door het verminderen van het communicatievolume tussen machines.
  • De toevoeging van Torus Attention en éénzijdige communicatie (SwiftFusion) leverde een extra snelheidswinst van 1,35× op ten opzichte van TAS, toegeschreven aan het overlappen van communicatie en berekening en verminderde synchronisatieoverhead.
• Schaalbaarheid: SwiftFusion presteerde consistent beter dan USP over verschillende aantallen GPU-machines (tot 4 machines) en verschillende gedistribueerde configuraties (variërende graden van Ulysses- en Ring-parallelisme).
• Geheugen: Het framework introduceerde geen hoger geheugengebruik in vergelijking met USP, met gebruik van maximaal één kopiebuffer voor tensors.
• Generalisatie Werklast: Prestatiewinsten waren consistent over verschillende sequentielengtes, batchgroottes en hoofd-dimensies, hoewel de winsten iets duidelijker waren bij grotere hoofd-dimensies waar berekeningssaturatie gemakkelijker te bereiken is.

Betekenis

Het paper stelt dat SwiftFusion een nieuw paradigma biedt voor efficiënte en schaalbare gedistribueerde DiT-inferentie. Door expliciet in te spelen op het mismatch tussen netwerktopologie en communicatiepatronen, en door mechanismen in te voeren om communicatielatenie te verbergen en synchronisatieoverhead te verminderen, overwint SwiftFusion de door communicatie beperkte beperkingen van huidige SP-technieken. De auteurs stellen dat deze inzichten cruciaal zijn voor het schalen van DiT-inferentie naar de eisen van hoge resolutie en lange duur van toekomstige generatieve AI-toepassingen.