Technische Samenvatting: Begrip en Verbetering van Communicatieprestaties bij Multi-node LLM-inferentie

Probleemstelling

Naarmate Large Language Models (LLM's) groter worden, overschrijdt hun geheugenvoetafdruk vaak de capaciteit van één enkele GPU, wat gedistribueerde inferentie over meerdere nodes noodzakelijk maakt. Hoewel strategieën voor modelparallelisme, zoals Tensor Parallelism (TP) en Pipeline Parallelism (PP), goed bestudeerd zijn voor inferentie op één node, blijft hun prestatie in multi-node omgevingen onderbelicht. Multi-node inferentie introduceert aanzienlijke uitdagingen door hogere latenties tussen nodes in vergelijking met verbindingen binnen een node (bijvoorbeeld NVLink).

De auteurs stellen vast dat TP vaak superieur is voor werklasten die geheugengebonden en decode-zwaar zijn, maar dat het lijdt aan aanzienlijke communicatie- overheads door frequente all-reduce-bewerkingen. Omgekeerd vermindert Hybrid Parallelism (HP), dat TP en PP combineert, de communicatie- overhead, maar faalt het om de rekentijd effectief te verminderen in decode-zware regimes vanwege de aard van kleine matrixvermenigvuldigingen. Een kritieke bottleneck die wordt geïdentificeerd, is de suboptimale prestatie van standaard NCCL all-reduce-implementaties voor de kleine berichtgroottes (128 KB tot 2 MB) die typerend zijn voor LLM-decodefasen, met name bij schaling over nodes.

Methodologie

De auteurs voerden een systematische prestatiestudie uit met een combinatie van bestaande state-of-the-art inferentie-engines en een aangepast onderzoeksprototype:

1. Experimentele Opstelling:

   • Hardware: Experimenten werden uitgevoerd op het Perlmutter-systeem (NVIDIA A100 GPU's, HPE Slingshot-11 interconnect) en het Vista-systeem (NVIDIA GH200 GPU's, InfiniBand interconnect).
   • Modellen: Llama 3.1 70B en 405B (Instruct) in bf16-precisie.
   • Werklasten: Twee primaire regimes werden getest: "Prefill-heavy" (grote promptlengtes) en "Decode-heavy" (lange generatiesequenties).
   • Engines: De studie gebruikte vLLM en SGLang naast YALIS, een aangepaste open-source inferentie-engine ontwikkeld door de auteurs om gecontroleerde experimenten in multi-node HPC-omgevingen mogelijk te maken. YALIS beschikt over een uniforme modeldefinitieslaag, uitvoering op basis van Torch Compile en tensorparallelisme op basis van AxoNN.

2. Analysebenadering:

   • De auteurs voerden strong-scaling-experimenten uit, waarbij ze de end-to-end batch-latentie maten naarmate het aantal GPU's toenam van single-node tot 32-node configuraties.
   • Ze splitsten de uitvoeringstijd op in berekening (Matmul, Overige Berekening), communicatie en wachttijd om bottlenecks te identificeren.
   • Synthetische GEMM-benchmarks werden gebruikt om de effecten van micro-batching (PP) versus tensor splitsing (TP) op matrixvermenigvuldigingskernen te isoleren.

3. Voorgestelde Oplossing (NVRAR):

   • Om de communicatiebottleneck aan te pakken, ontwikkelden de auteurs NVRAR, een hiërarchisch all-reduce-algoritme gebaseerd op recursief verdubbelen met behulp van NVSHMEM.
   • Ontwerp: NVRAR werkt in drie fasen:
     1. Intra-node Reduce-Scatter: Gebruikt NCCL om data binnen een node te reduceren.
     2. Inter-node Recursief Verdubbelen: Gebruikt niet-blokkerende RMA-primitieven van NVSHMEM (put_nbi) om all-reduce over nodes uit te voeren met een topologie van recursief verdubbelen.
     3. Intra-node All-Gather: Gebruikt NCCL om de gereduceerde data terug te verzamelen naar alle GPU's binnen een node.
   • Optimalisaties:
     • Geblokkeerde Niet-blokkerende Communicatie: Overlapt berekening en communicatie door data in chunks te verwerken.
     • Gefuseerde Payloads: Combineert data en synchronisatievlaggen in één payload van 8 bytes om dure expliciete signaleringsprimitieven op Slingshot/InfiniBand te vermijden.
     • Synchronisatie op Basis van Sequencenummers: Gebruikt unieke sequencenummers voor globale synchronisatie om de latentie- overhead van globale barrières (fence/quiet) te vermijden, terwijl de correctheid wordt gewaarborgd.

Belangrijkste Bijdragen

1. Systematische Prestatiestudie: Het artikel biedt een gedetailleerde uiteenzetting van de schaalgedragingen van TP versus HP in multi-node omgevingen, waarbij wordt geconstateerd dat TP over het algemeen superieur is voor decode-zware werklasten, maar beperkt wordt door communicatie- overheads, terwijl HP moeite heeft om de rekentijd in deze regimes te verminderen.
2. YALIS-engine: De ontwikkeling van YALIS, een op onderzoek gerichte inferentie-engine ontworpen voor instrumenteerbaarheid en gecontroleerde experimenten in multi-node HPC-omgevingen.
3. NVRAR-algoritme: De voorstelling en implementatie van NVRAR, een aangepast all-reduce-algoritme dat is geoptimaliseerd voor het regime van kleine berichten dat typerend is voor LLM-decodefasen. Het maakt gebruik van NVSHMEM om een lagere latentie te bereiken dan standaard NCCL-implementaties op Slingshot- en InfiniBand-netwerken.
4. Integratie en Evaluatie: De integratie van NVRAR in YALIS en vLLM, waarbij end-to-end prestatieverbeteringen worden aangetoond voor grote modellen (Llama 3.1 405B) en Mixture-of-Experts (MoE)-modellen (Qwen3-235B).

Resultaten

• Schaalgedrag: Zowel TP als HP vertonen een slechte strong scaling in multi-node omgevingen. HP presteert beter dan TP in rekengebonden (prefill-heavy) regimes, terwijl TP beter presteert dan HP in geheugengebonden (decode-heavy) regimes.
• Communicatiebottlenecks: NCCL all-reduce schaalt slecht over nodes voor kleine berichten (128 KB–2 MB), en presteert in deze specifieke regimes vaak slechter dan MPI-gebaseerde oplossingen.
• NVRAR-prestaties:
  • Microbenchmarks: NVRAR bereikt 1,9×–3,6× lagere latentie dan NCCL voor berichtgroottes tussen 128 KB en 2 MB op HPE Slingshot- en InfiniBand-interconnects.
  • End-to-End Inferentie: Geïntegreerd in YALIS en vLLM vermindert NVRAR de end-to-end batch-latentie met maximaal 1,72× voor het Llama 3.1 405B-model in decode-zware werklasten met gebruik van tensorparallelisme.
  • MoE-modellen: NVRAR verbetert ook de doorvoer voor MoE-modellen (bijvoorbeeld Qwen3-235B) door het TP-component van de parallelisme-strategie te versnellen, met een doorvoer die maximaal 1,14× hoger is dan NCCL-gebaseerde configuraties.

Betekenis

Het artikel stelt dat naarmate LLM's groter worden en inferentie verschuift naar multi-node clusters, communicatie-efficiëntie de primaire determinant van prestaties wordt. De auteurs tonen aan dat standaard communicatiebibliotheken (NCCL) niet zijn geoptimaliseerd voor de specifieke berichtgroottes en latentie-eisen van LLM-decodefasen. Door NVRAR in te voeren, bieden ze een gerichte oplossing die de communicatie- overhead aanzienlijk vermindert, waardoor een efficiëntere schaling van tensorparallelisme voor grote modellen mogelijk wordt. Dit werk benadrukt de noodzaak van aangepaste communicatieprimitieven voor gedistribueerde inferentiesystemen van de volgende generatie en valideert het nut van hiërarchische, op NVSHMEM gebaseerde benaderingen boven standaard collectieve communicatiebibliotheken in specifieke HPC-interconnect-omgevingen.