Naukowcy z arXiv przedstawili block parallelism (BP) - nowy wymiar paralelizmu rozproszonego dla efektywnego treningu modeli dyfuzji na długie konteksty. Problem polegał na tym, że dotychczasowe podejście context parallelism (CP) wymagało intensywnej komunikacji między GPU przy przetwarzaniu zarówno czystych jak i uszkodzonych sekwencji, co stanowiło wąskie gardło przy skalowaniu.
Klucz do rozwiązania tkwi w obserwacji, że funkcja celu w block diffusion language models rozszczepia się po poszczególnych blokach. Zespół zaproponował przypisanie każdego obliczenia dla uszkodzonego bloku jednemu rankowi GPU, a następnie rozszerzył to podejście do context-sharded block parallelism (CSBP), które dodatkowo rozdziela czystą sekwencję wśród tych rankingów. Rezultat jest elegancki - uszkodzonym K/V i ich gradienty pozostają lokalnie, eliminując replikowane prefeksy, jednocześnie zachowując semantykę treningu.
Wyniki są imponujące. Na 16 GPU H200 z kontekstem 256K, CSBP uzyskuje przyspieszenie 1.18-1.45x dla supervised fine-tuningu i 1.27-1.33x dla konwersji modeli autoregresywnych, przy czym wykorzystanie pamięci HBM jest równe lub niższe. Przy kontekście 512K całkowite przyspieszenie modelu sięga 1.61x. Na ośmiu GPU H100 technika przyśpiesza trening dekodera spekulacyjnego DFlash2 aż 7.59x przy kontekście 1M. W praktycznym teście na benchmarkach SWE-bench Verified i Terminal-Bench Lite model DiffusionGemma 26B trenowany z CSBP osiąga wyższe wskaźniki sukcesu na każdym checkpocie.