Scalable Training of Mixture-of-Experts Models with Megatron Core

Dit paper introduceert Megatron Core, een schaalbaar en productieklaar open-source framework dat geïntegreerde optimalisaties voor geheugen, communicatie en berekening combineert om de training van Mixture-of-Experts-modellen tot triljoenen parameters op duizenden GPU's efficiënt te maken.

Zijie Yan (NVIDIA), Hongxiao Bai (NVIDIA), Xin Yao (NVIDIA), Dennis Liu (NVIDIA), Tong Liu (NVIDIA), Hongbin Liu (NVIDIA), Pingtian Li (NVIDIA), Evan Wu (NVIDIA), Shiqing Fan (NVIDIA), Li Tao (NVIDIA), Robin Zhang (NVIDIA), Yuzhong Wang (NVIDIA), Shifang Xu (NVIDIA), Jack Chang (NVIDIA), Xuwen Chen (NVIDIA), Kunlun Li (NVIDIA), Yan Bai (NVIDIA), Gao Deng (NVIDIA), Nan Zheng (NVIDIA), Vijay Anand Korthikanti (NVIDIA), Abhinav Khattar (NVIDIA), Ethan He (NVIDIA), Soham Govande (NVIDIA), Sangkug Lym (NVIDIA), Zhongbo Zhu (NVIDIA), Qi Zhang (NVIDIA), Haochen Yuan (NVIDIA), Xiaowei Ren (NVIDIA), Deyu Fu (NVIDIA), Tailai Ma (NVIDIA), Shunkang Zhang (NVIDIA), Jiang Shao (NVIDIA), Ray Wang (NVIDIA), Santosh Bhavani (NVIDIA), Xipeng Li (NVIDIA), Chandler Zhou (NVIDIA), David Wu (NVIDIA), Yingcan Wei (NVIDIA), Ashwath Aithal (NVIDIA), Michael Andersch (NVIDIA), Mohammad Shoeybi (NVIDIA), Jiajie Yao (NVIDIA), June Yang (NVIDIA)

Gepubliceerd Tue, 10 Ma
📖 5 min leestijd🧠 Diepgaand

Each language version is independently generated for its own context, not a direct translation.

De "Megatron-Core MoE": Hoe we enorme AI-hersens bouwen zonder dat ze ontploffen

Stel je voor dat je een superintelligente AI wilt bouwen die zo groot is dat hij alle kennis van de mensheid in zich draagt. In het verleden bouwden we deze "hersens" als één gigantisch, zwaar blok (een dense model). Maar dat is als proberen een heel orkest te laten spelen in één kleine kamer: het wordt rommelig, traag en er is niet genoeg ruimte.

NVIDIA heeft een nieuwe manier bedacht, genaamd Mixture of Experts (MoE). In plaats van één grote hersenstam, bouwen we een gebouw met duizenden kleine, gespecialiseerde kamers (de "experts"). Als een vraag binnenkomt, kijkt een slimme conciërge (de router) alleen naar de kamers die relevant zijn voor die vraag. De rest van het gebouw blijft rustig. Dit maakt de AI veel slimmer en sneller.

Maar er is een probleem: als je duizenden experts bouwt, wordt het beheer van dat gebouw een nachtmerrie. De paper beschrijft hoe NVIDIA met Megatron-Core deze nachtmerrie heeft opgelost. Ze hebben drie grote muren moeten doorbreken: de Geheugenmuur, de Communicatiemuur en de Rekenkrachtmuur.

Hier is hoe ze dat deden, vertaald naar alledaagse analogieën:

1. De Geheugenmuur: "De koelkast die te vol zit"

Stel je voor dat je een koelkast (het geheugen van je computer) hebt die vol zit met ingrediënten voor een gigantisch diner. Bij een MoE-model heb je niet alleen de ingrediënten die je nu gebruikt, maar ook de recepten voor alle mogelijke gerechten die je ooit zou kunnen maken. Dat past niet in je koelkast.

  • Het probleem: De AI moet alle "experts" (recepten) in het geheugen houden, zelfs als ze op dat moment niet gebruikt worden.
  • De oplossing:
    • Slimme opslag: In plaats van alles in zware glazen potten (FP16) te bewaren, gebruiken ze lichte plastic bakjes (FP8 of FP4). Het smaakt bijna hetzelfde, maar bespaart enorm veel ruimte.
    • Herrekenen: In plaats van alle tussenresultaten op te slaan, laten ze de kok (de GPU) die resultaten even opnieuw berekenen als ze nodig zijn. Het kost een seconde extra tijd, maar bespaart veel koelkastruimte.
    • Uitlenen: Als de koelkast echt vol zit, schuiven ze tijdelijke ingrediënten naar de schuur (het CPU-geheugen) en halen ze ze er pas weer uit als ze echt nodig zijn.

2. De Communicatiemuur: "De verkeersfile in de tunnel"

In een MoE-model moeten de experts met elkaar praten. Als de conciërge een vraag naar expert A stuurt, moet expert A het antwoord terugsturen. Bij duizenden experts is dit als een file in een tunnel: iedereen probeert tegelijkertijd te praten, en de tunnel (de verbinding tussen de chips) raakt verstopt.

  • Het probleem: De tijd die de AI kwijt is aan het wachten op berichten, is vaak langer dan de tijd die hij besteedt aan het denken.
  • De oplossing:
    • Snellere postbodes: Ze hebben nieuwe, supersnelle postbodes (DeepEP en HybridEP) bedacht die berichten niet één voor één, maar in bundels en via slimme routes bezorgen.
    • Tijdverschuiving: Ze laten de postbodes werken terwijl de kok al aan het koken is. De "verkeersfile" wordt verborgen achter het werk dat al gedaan wordt. Zo lijkt het alsof er geen file is.
    • Parallelle banen: Ze zorgen ervoor dat de experts binnen één gebouw (een server) met elkaar praten via een snelle interne lift (NVLink), in plaats van dat ze naar een ander gebouw moeten rijden over een trage weg.

3. De Rekenkrachtmuur: "De trage chef-kok"

Soms is de keuken zo groot dat de chef-kok (de GPU) niet snel genoeg nieuwe orders kan aannemen. De computer moet wachten op de chef, terwijl de machines (de GPU's) al klaar staan om te werken. Dit heet "host overhead".

  • Het probleem: De computer doet te veel kleine dingen tegelijk, waardoor de grote machines vaak stil staan.
  • De oplossing:
    • Groepsbestellingen: In plaats van 100 kleine orders één voor één te geven, bundelen ze ze in één grote order (Grouped GEMM). De kok kan dan in één keer 100 borden klaarmaken in plaats van 100 keer te beginnen en te stoppen.
    • De "Vooraf-geplande" menukaart: Ze gebruiken een truc genaamd CUDA Graphs. In plaats van dat de chef elke seconde een nieuwe vraag krijgt ("Klaar? Nee, wacht even... Klaar?"), krijgen ze één keer een volledig menu voor de hele avond. De chef kan dan ononderbroken werken zonder te hoeven wachten op instructies.
    • Slimme verdeling: Als één expert te veel werk heeft en een andere niets doet, huren ze tijdelijk een tweede kok in (ECHO) om het werk te verdelen, zodat niemand stilzit.

De Grote Innovatie: "Parallel Folding" (De Slimme Vouwtechniek)

Stel je voor dat je een gebouw hebt met twee soorten kamers:

  1. De Lezingzaal: Hier moeten mensen in grote groepen zitten (Dense layers).
  2. De Werkkamers: Hier werken mensen in kleine, gespecialiseerde groepjes (MoE experts).

Vroeger moesten beide kamers dezelfde indeling hebben. Dat was dom: de lezingzaal had grote tafels nodig, de werkkamers kleine bureau's. Als je de indeling aanpaste voor de lezingzaal, werden de werkkamers onbruikbaar, en andersom.

Megatron-Core introduceert Parallel Folding. Dit is alsof je het gebouw kunt vouwen. Je kunt de lezingzaal en de werkkamers volledig onafhankelijk van elkaar indelen. Je gebruikt de beste indeling voor de lezingzaal én de beste indeling voor de werkkamers, zonder dat ze elkaar in de weg zitten. Hierdoor kunnen ze duizenden experts tegelijk laten werken zonder dat het systeem instort.

Wat levert dit op?

Dankzij deze slimme trucs kunnen ze nu modellen trainen die biljoenen parameters groot zijn (zoals DeepSeek-V3 en Qwen3).

  • Op de nieuwste NVIDIA-chips (GB300/GB200) werken ze zo snel dat ze in één seconde meer berekeningen doen dan er sterren in het melkwegstelsel zijn.
  • Het is alsof ze van een fiets (oude methoden) zijn overgestapt op een raket (Megatron-Core).

Kortom: NVIDIA heeft een systeem gebouwd dat enorme AI-modellen niet alleen mogelijk maakt, maar ook efficiënt, snel en betaalbaar houdt, door slim om te gaan met geheugen, communicatie en rekenkracht. Het is de "besturingssysteem" voor de toekomst van superintelligente AI.