Prime Intellect opublikował Prime Mega MoE, rozszerzenie do szkolenia złączonego jądra GPU Mega MoE firmy DeepSeek. Podczas gdy jądro DeepSeek obsługiwało tylko przebieg w przód, Prime Intellect dodaje do niego szybki przebieg wsteczny, nakładający się na inne obliczenia, dzięki czemu jądro można wykorzystać do szkolenia, a nie tylko do wnioskowania. Według firmy, w jej narzędziu szkoleniowym prime-rl pełne kroki szkolenia na modelach Qwen3 MoE działają nawet o 34% szybciej, a w porównaniu z TorchTitan nawet o 55% szybciej.
W modelu MoE (mixture of experts, czyli „mieszanka ekspertów ») każdy token jest wysyłany tylko do małego podzbioru ekspertów, rozproszonych na różnych GPU. Każda warstwa musi więc wysłać tokeny, kazać ekspertom obliczyć, a następnie odebrać wyniki, zarówno podczas przebiegu w przód, jak i przebiegu wstecznego, fazy wstecznej propagacji używanej do uczenia. W standardowej konfiguracji operacje te są rozdzielone: GPU sortuje tokeny, czeka na transport, oblicza, a potem znów czeka. Te czasy oczekiwania to okresy bezczynności obliczeniowej.
Złączone jądro, uzupełnione do uczenia
Mega MoE firmy DeepSeek, dodany do biblioteki DeepGEMM w kwietniu 2026 roku, zastępuje ten łańcuch jednym wywołaniem jądra. Jądro (kernel) to zoptymalizowany program działający bezpośrednio na GPU. To jądro łączy równoległe rozsyłanie do ekspertów, dwie liniowe warstwy ekspertów, SwiGLU i etap kombinacji. Podczas gdy jedna grupa tokenów jest w transporcie, jądro oblicza na grupie już dotarłej i unika tensorów pośrednich dzięki buforom roboczym znajdującym się wyłącznie na GPU. To jądro obejmowało jednak tylko implementację przebiegu w przód.
Prime Intellect napisał brakujący przebieg wsteczny, również nakładający się. Obsługiwane typy danych to BF16, MXFP8 (zarezerwowany dla GPU Blackwell) oraz FP8 blokowe (zarezerwowane dla GPU Hopper). Konfiguracja prime-rl oferuje również tryb fp8_bwd: przebieg w przód w BF16, przebieg wsteczny w FP8.
Dane z własnych pomiarów
Wszystkie wyniki zostały opublikowane przez samego Prime Intellect. Kroki szkolenia wykorzystują węzły z 8 H200, w FP8, z sekwencjami 16 384 tokenów i pełnym checkpointingiem aktywacji. Dla Qwen3-30B-A3B, rozproszonego na 32 H200, pełny krok trwa 2,78 sekundy z bazowym prime-rl, 3,07 sekundy z TorchTitan (FP8) i 2,08 sekundy z Prime Mega MoE, co daje wzrosty odpowiednio o 34% i 48%. Dla Qwen3-235B-A22B, rozproszonego na 128 H200, czasy wynoszą 12,72 sekundy, 14,98 sekundy i 9,65 sekundy, dając wzrosty o 32% i 55%. Przepustowość wzrasta dla modelu 30B z 5 337 do 7 886 tokenów na sekundę na GPU w porównaniu z TorchTitan, a dla modelu 235B z 1 094 do 1 697.
Zmierzony warstwa po warstwie na pojedynczym węźle, zysk jest większy na przodzie niż na tyle. Przebieg w przód modelu Qwen3-30B-A3B skraca się z 12,0 ms do 3,9 ms, czyli współczynnik 3,11. Przebieg wsteczny modelu Qwen3-235B-A22B skraca się z 30,5 ms do 24,6 ms, czyli 1,24. Liczba „nawet 3,11x » dotyczy więc tylko przebiegu w przód. TorchTitan z kolei deklaruje zysk do 41%, ale w innej konfiguracji (B200, model 671 miliardów parametrów, baza BF16), którą źródło uważa za nieporównywalną bezpośrednio.
Prime Mega MoE jest dystrybuowany na licencji Apache 2.0. W prime-rl aktywuje się jako backend rozsyłania: wystarczy ustawić type = « mega_moe » w sekcji [trainer.model.moe.dispatch], a następnie wybrać precyzję spośród bf16, mxfp8, fp8 lub fp8_bwd. Jądra szkoleniowe są dostarczane w prekompilowanym pakiecie DeepGEMM, dostępnym w repozytorium prime-kernels, ponieważ standardowa wersja DeepGEMM ich nie zawiera.
