Prime Intellect dopisał backward pass do kernela Mega MoE od DeepSeeka. Trening modeli MoE szybszy o jedną trzecią

Trening modeli Mixture-of-Experts w dużej mierze ogranicza komunikacja między GPU, a nie same obliczenia. Każdy token trafia do kilku ekspertów rozmieszczonych na różnych kartach, więc w każdej warstwie dane są wysyłane, liczone i odsyłane z powrotem. Dzieje się to w forward passie i jeszcze raz w backwardzie. Przy standardowym podejściu to osobne operacje, a GPU między nimi czeka bezczynnie.

DeepSeek skleił ten łańcuch w jedno wywołanie. Kernel Mega MoE, dodany do biblioteki DeepGEMM w kwietniu 2026 roku, łączy expert-parallel dispatch, dwie warstwy liniowe ekspertów, SwiGLU i combine. Gdy jedna grupa tokenów jest w transporcie, kernel liczy na tej, która już dotarła. Był jednak ograniczony do forward passu, więc nadawał się do inferencji, ale nie do treningu.

Tę lukę zamknął Prime Intellect. Firma napisała brakujący backward kernel, również nakładający obliczenia na komunikację, i nazwała całość Prime Mega MoE. Obsługiwane są BF16 oraz FP8: blockwise na kartach Hopper i MXFP8 na Blackwell. W konfiguracji prime-rl jest też tryb fp8_bwd, czyli forward w BF16 i backward w FP8.

Co pokazują pomiary

Wszystkie liczby pochodzą od samego Prime Intellect. Testy objęły Qwen3-30B-A3B i Qwen3-235B-A22B (128 ekspertów, routing top-8) na węzłach 8×H200, w FP8, z sekwencjami po 16 384 tokenów i pełnym activation checkpointingiem. Wynik to średnia z kroków 6–20.

• Qwen3-30B-A3B na 32 H200: pełny krok trwa 2,08 s wobec 2,78 s w bazowym prime-rl (czas krótszy o około 25%) i 3,07 s w TorchTitan (czas krótszy o około 32%). Przepustowość względem TorchTitan rośnie z 5 337 do 7 886 tokenów/s na GPU.
• Qwen3-235B-A22B na 128 H200: 9,65 s wobec 12,72 s w prime-rl (czas krótszy o około 24%) i 14,98 s w TorchTitan (czas krótszy o około 36%). Przepustowość rośnie z 1 094 do 1 697 tokenów/s na GPU.

Bazowy prime-rl używa torchowego all-to-all i ekspertów FP8 z DeepGEMM, więc to porównanie jest bliższe praktyce niż zestawienie z gołym frameworkiem.

Nagłówkowe „do 3,11×” dotyczy wyłącznie forward passu pojedynczej warstwy MoE w modelu 30B (12,0 ms → 3,9 ms). Backward daje mniej: 1,87× w tej samej warstwie (15,4 → 8,2 ms), a w warstwie modelu 235B tylko 1,24× (30,5 → 24,6 ms). Forward w modelu 235B przyspiesza 2,17×. Pełna warstwa (dwa forwardy plus backward, bo włączony jest checkpointing) skraca się 2,47× w modelu 30B i 1,64× w 235B. Pomiary warstw wykonano na jednym węźle 8×H200 z losowym routingiem, więc nie przekładają się wprost na krok treningowy na wielu węzłach.

Porównania i dostępność

Jak zauważa MarkTechPost, w publikowanych zestawieniach pojawia się też wynik TorchTitan z DeepEP: do 41% szybszy trening DeepSeek-V3 671B. Pochodzi jednak z innego układu (B200, model 671B, bazowy BF16), więc nie da się go zestawić z liczbami Prime Intellect. Podobną drogę – fuzję dispatchu z GEMM-ami – opisuje też dokumentacja AMD Primus dla kart MI355X, ale bez opublikowanych przyspieszeń.

Prime Mega MoE działa jako backend dispatchu w prime-rl. Wystarczy ustawić type = "mega_moe" w sekcji [trainer.model.moe.dispatch] i wybrać precyzję: bf16, mxfp8, fp8 lub fp8_bwd. Kernele treningowe są w gotowej paczce DeepGEMM z repozytorium prime-kernels. Standardowy build DeepGEMM ich nie zawiera. Prime-rl ma licencję Apache 2.0, a DeepGEMM – MIT.

Fused kernel Mega MoE trafił też do ekosystemu Hugging Face: dokumentacja Transformers opisuje backend deepgemm_megamoe dla Blackwella. Wcześniej oba podejścia, fused i overlapping, były w praktyce dostępne głównie na ścieżce inferencyjnej. Wersja Prime Intellect przenosi je do pętli treningowej, ale na razie sprawdzono ją tylko na dwóch modelach Qwen3 i na kartach H200.