KutatásHír
Prime Mega MoE: tanításra is jó a DeepSeek MoE kernele
🇬🇧 Read this article in English →
A Prime Intellect backward passt írt a DeepSeek fused Mega MoE kerneléhez, így az eddig csak inferenciára használható kernel tanításra is jó lett. Az eredmény akár 3,11-szer gyorsabb MoE kernel és akár 34 százalékkal nagyobb end-to-end tanítási throughput.
Írta: Mikola Máté Patrik · · 2 perc olvasás
Mikola Máté Patrik
Mi a probléma?
Egy mixture-of-experts modellben minden token csak néhány experthez jut el, ezek az expertek viszont különböző GPU-kon vannak. Tanításkor ezért minden rétegben kétszer kell adatot mozgatni: a tokeneket el kell küldeni az experthez tartozó GPU-ra, majd az eredményt vissza. Ha ezek a lépések külön műveletek, a GPU egy része mindig vár, ez pedig elvesztegetett számítási idő (Prime Intellect).
A DeepSeek nyílt forrású DeepGEMM gyűjteményében lévő Mega MoE kernel ezt egyetlen fused kernellel oldja meg: egy hívásban routolja a tokeneket, lefuttatja az expert MLP-t és visszaküldi az eredményt, közben átfedi a kommunikációt a számítással (GitHub). A gond az volt, hogy csak forward implementáció létezett, tehát tanításra nem volt használható.
Mit csinált a Prime Intellect?
A csapat gyors, szintén átfedéses backward kernelt írt hozzá, majd a standard MoE rétegeket ezekre a kernelekre cserélte a saját Prime-RL tanítókódjában. A támogatott adattípusok: bfloat16, MXFP8 (csak Blackwell), illetve blockwise FP8 (csak Hopper).
- Akár 3,11-szer gyorsabb MoE kernel a PyTorch alapú alapimplementációhoz képest. A mérés egy node-on, 8 H200-on, egyetlen MoE rétegen futott, 128 experttel, top-8 routinggal, FP8-ban.
- Akár 1,34-szeres end-to-end tanítási sebesség, vagyis 34 százalékkal nagyobb throughput a Prime-RL alapverzióhoz képest.
- Az end-to-end teszt Qwen3-30B-A3B modellel 4 node-on, Qwen3-235B-A22B modellel 16 node-on futott, node-onként 8 H200-zal, 16 384 tokenes szekvenciával, EP=8 és FSDP mellett, FP8-ban.
Mit jelent ez neked?
Ha nagy MoE modellt tanítasz vagy RL-lel utótanítasz, a GPU-k közötti kommunikáció könnyen szűk keresztmetszet lesz. Egy 34 százalékos throughput-javulás ugyanannyi GPU-órából ennyivel több tanítást jelent, ami nagy futásoknál komoly pénz. A projekt azt is mutatja, hogy a DeepSeek nyíltan közzétett infrastruktúra-kódjára mások érdemben tudnak építeni. Érdemes figyelni, hogy a Prime Intellect mikor és milyen formában teszi elérhetővé a kernelt a Prime-RL-ben.
MI segítségével készült, szerkesztő ellenőrizte. Hogyan készülnek a cikkeink
Szerkesztőségünk ellenőrizte és szerkesztette. Rólunk
