Experimentální backend v PyTorchi 2.14 umí postavit skupinu procesů po výpadku znovu
PyTorch vydal 2. září verzi 2.14 a v ní backend nccl2, port komunikační vrstvy z projektu torchcomms. Proti dosavadnímu backendu umí přestavět skupinu procesů: zboří komunikátor a postaví nový s jiným složením. Zapíná se proměnnou prostředí a domluvu o tom, kdo v nové skupině zůstane, si musí zařídit volající sám.
PyTorch vydal 2. září verzi 2.14. Balíčky se objevily na PyPI ve 13:42 UTC, poznámky k vydání na GitHubu o čtyři hodiny později. Mezi novinkami je celá nová vrstva, přes kterou si procesy jednoho trénovacího běhu předávají mezivýsledky: backend nccl2.
Není to náhrada za NCCL, tedy za knihovnu NVIDIE, o jejíž verzi 2.31 jsme psali v srpnu. NCCL zůstává pod tím; mění se vrstva v PyTorchi, která ho volá. A mění se hodně: v 2.13.0 z 8. července po ní není ani stopa, soubor distributed_c10d.py mezitím narostl z 6 908 na 8 055 řádků (vlastní srovnání obou značek).
Zapíná se ručně, výchozí zatím není
Registrace v kódu vypadá tak, že přibyla tři jména: nccl2, nccl-lazy a nccl-legacy. Poslední z nich ukazuje na dosavadní implementaci a jméno nccl na ni pořád míří taky – pokud v prostředí nestojí TORCH_DIST_USE_NCCL2=1. Pak se pod známým jménem vyrobí nová. Komentář ve zdrojovém souboru nový backend popisuje jako postavený na jádru z projektu torchcomms, který Meta otevřela loni v říjnu pod licencí BSD-3-Clause.
Autor návrhu, který nccl-legacy zavedl, k tomu píše, že jde o dílčí krok k přepnutí výchozí volby. Staré jméno tedy nevzniklo pro pořádek, ale jako záchranná cesta zpátky.
Přestavba skupiny chce domluvu mimo pásmo
Hlavní věc, kterou nový backend přináší, je reconfigure: skupina procesů se dá za běhu složit znovu, s jinými členy. Volá se to na dvakrát. Nejdřív init_process_group(..., enable_reconfigure=True), po kterém se komunikátor vůbec nezaloží – dokumentace k tomu říká, že se inicializace odkládá až na první reconfigure() a kolektivní operace před ním skončí chybou. Pak si každý proces vyžádá svůj popisovač funkcí _get_reconfigure_handle(); je to řetězec ve tvaru nccl2:<rank>:<uuid>:<adresa store>, kde rank je pořadové číslo procesu ve skupině.
Tady je ta nepříjemná část. Popisovače se musí posbírat od všech přeživších procesů a rozeslat je zpátky spolu s čerstvým číslem uuid, a tohle PyTorch nedělá. Ukázka v dokumentaci to odbývá voláním smyšlené funkce a poznámkou, že postup vyžaduje domluvu mimo pásmo. Kdo tedy rozhodne, že uzel vypadl a kdo zůstává, musí být někdo jiný – plánovač úlohy nebo knihovna nad ním. Všechna tři volání jsou navíc soukromá, mají v názvu podtržítko a u každého stojí varování, že se API může změnit nebo zmizet.
Rychlá cesta se nepřevzala schválně
V torchcomms existuje zkratka: když se skupina jen zmenšuje nebo roste, komunikátor se dá upravit, místo aby se zahodil. Do PyTorche se nedostala, a soubor ReconfigureNCCL.cpp u toho rovnou vysvětluje proč: ta cesta přiřazuje ranky v pořadí, v jakém je vyskládá NCCL, což se rozchází s pravidlem c10d, že pořadí popisovačů určuje pořadí ranků.
Přestavba proto pokaždé projde celým kolem: zruší rozdělanou práci, zastaví hlídací vlákno, vyprázdní frontu, ukončí komunikátor a přes úložiště klíčů pojmenované podle uuid postaví nový. Rank 0 si to číslo napřed zabere, aby se dvě přestavby nepopraly o tytéž klíče; opakované uuid skončí chybou. Zmenšení komunikátoru bez plné přestavby jde udělat zvlášť, metodou shrink(), ta ale chce NCCL 2.27 nebo novější.
Jednostranná okna umí zapsat, ne přečíst
Druhá novinka backendu jsou jednostranná okna. Rank zapíše data přímo do paměti souseda, aniž ten musí cokoli zavolat. NVIDIA tahle volání ohlásila v NCCL 2.29 a popisuje u nich dvě věci, kvůli kterým to lidi zajímá: přenos nezaměstná výpočetní jednotky karty, takže zůstanou volné pro počítání, a uvnitř uzlu ho odbaví kopírovací jednotky přes NVLink nebo PCIe. PyTorch to zpřístupňuje jako _new_window() s operacemi put, signal a wait_signal, které sedí přímo na ncclPutSignal, ncclSignal a ncclWaitSignal.
Opačný směr chybí. Metoda map_remote_tensor() funguje jen pro vlastní rank a na cizí vrátí chybu s odůvodněním, že NCCL ukazatele na okna sousedů nezveřejňuje. Backend se podle toho i popisuje: dotaz na vlastnosti okna hlásí režim SEPARATE, aby volající sáhli po zápisu se signálem místo přímého čtení cizí paměti.
Netýká se to jen karet NVIDIE
Přestavba skupiny není v 2.14 vázaná na jeden backend. Rozhraní sedí v ProcessGroup, tedy ve společné části, a druhý návrh ho ve stejném vydání doplnil do Gloo, backendu pro procesory. Popisovač má tam stejný tvar a testy jsou psané tak, aby je mohly převzít i další implementace. Backendy, které přestavbu neumí, příznak enable_reconfigure mlčky ignorují.
Kolik toho v tom vydání přibylo, se dá odhadnout z velikosti: adresář s novým backendem má 25 souborů a 265 637 bajtů (vlastní součet z výpisu rozhraní GitHubu). Na experimentální kus kódu, který se sám o sobě nezapne, je to slušná porce – a jméno nccl-legacy napovídá, že se u dobrovolné volby nemá zůstat.