Prime Intellect prezentuje Prime Mega MoE: 3,11x szybszy forward pass dla treningu MoE

Prime Intellect wydała Prime Mega MoE, GPU kernel zaprojektowany do efektywnego trenowania modeli mixture-of-experts. Rozszerzenie to dodaje do DeepSeek'owego kernela Mega MoE szybki, nakładający się backward pass, zamieniając narzędzie dostosowane wyłącznie do inferencji w pełnoprawne rozwiązanie treningowe.

Kernel został przetestowany na modelach Qwen3-30B-A3B i Qwen3-235B-A22B, oba z 128 ekspertami i routingiem top-8. Działa na GPU NVIDIA Hopper i Blackwell w precyzji BF16. Wyniki pokazują imponujące przyspieszenia - pełne kroki treningowe są 32-34 procent szybsze niż baseline Prime-RL i 48-55 procent szybsze niż TorchTitan. Na pojedynczej warstwie MoE Qwen3-30B-A3B forward pass przyspieszył się aż 3,11-krotnie, z 12 ms do 3,9 ms.

Problema, którą Prime Mega MoE rozwiązuje, leży w naturze architektur MoE. Każdy token trafia do małego podzbioru ekspertów rozmieszczonych na różnych GPU, co wymaga komunikacji między urządzeniami zarówno w forward passie, jak i backward passie. Prime Intellect osiągnął przyspieszenie poprzez fuzję kernela i nakładanie operacji komunikacyjnych, czym zmniejszył wąskie gardło komunikacyjne. Dla zespołów trenujących duże modele ekspertów, to bezpośrednia redukcja kosztu GPU hours.