Databricks ukázal, ako obmedziť straty pri zlyhaní veľkých tréningov v PyTorch

Databricks zverejnil technický postup pre odolnejší tréning modelov v PyTorch, ktorý spája distribuované asynchrónne checkpointy so zachovaním presného stavu dátového toku. Nejde o nový model ani samostatný benchmark, ale o dôležitú infraštruktúrnu tému: pri dlhých úlohách na desiatkach či stovkách GPU už zlyhanie nie je výnimočná nehoda. Rozhodujúce je, koľko už vykonanej práce sa po poruche stratí, ako rýchlo sa výpočet obnoví a či pokračuje s rovnakými dátami. Databricks preto navrhuje sledovať „goodput“, teda podiel času, počas ktorého akcelerátory vykonávajú užitočný výpočet namiesto čakania, zápisu checkpointu alebo obnovy po chybe.

Zlyhanie sa pri veľkom klastri mení na bežnú prevádzku

Pravdepodobnosť prerušenia rastie spolu s počtom zariadení a dĺžkou tréningu. Databricks používa ilustračný model s približne jednopercentnou ročnou mierou zlyhania jedného GPU. Podľa jeho výpočtu má úloha na 256 GPU bežiaca 30 dní asi 19-percentnú pravdepodobnosť, že zaznamená poruchu; pri 1 024 GPU má pravdepodobnosť vzrásť na 57 percent. Tieto čísla nie sú univerzálnou predpoveďou pre každý klaster, pretože výsledok ovplyvňuje hardvér, sieť, ovládače aj spôsob plánovania. Dobre však ukazujú, prečo sa pri veľkom tréningu neoplatí navrhovať systém s predpokladom, že celý beh dokončí bez prerušenia.

Praktický problém sa nekončí výmenou chybného uzla. Ak sa stav ukladá iba raz denne, náhodná porucha znamená v priemere zopakovanie približne dvanástich hodín výpočtu. Pri drahých akcelerátoroch môže byť cena stratenej práce vyššia než náklady na samotné odstránenie chyby. Častejšie checkpointy tento interval skracujú, klasický zápis však môže zastaviť všetky pracovné procesy a počas prenosu veľkého súboru nechať GPU nečinné. Frekvencia ukladania preto nemôže byť oddelená od formátu checkpointu a od spôsobu, akým sa dáta zapisujú do lokálneho alebo vzdialeného úložiska.

Namiesto jedného súboru zapisuje každý proces vlastný diel

Jednoduchý prístup s torch.save často zhromaždí stav na procese s poradím nula a následne ho zapíše ako jeden monolitický súbor. Pri distribuovanom tréningu sa tak jedno miesto môže stať úzkym hrdlom pre pamäť, sieť aj úložisko. Rozhranie PyTorch Distributed Checkpoint tento postup obracia: jednotlivé procesy zapisujú svoje časti stavu paralelne a malý metadátový súbor opisuje ich spoločné usporiadanie. Checkpoint sa následne môže načítať aj pri inom počte procesov, pretože PyTorch naplánuje nové rozdelenie uložených bajtov. To je užitočné napríklad vtedy, keď sa úloha po výpadku musí obnoviť na menšom klastri.

Databricks upozorňuje, že distribuovaný checkpoint má význam aj pri klasickom Distributed Data Parallel, kde každý proces drží rovnakú repliku parametrov. Mechanizmus môže stav pri zápise rozdeliť medzi pracovníkov a súčasne vytvára cestu k neskoršiemu prechodu na FSDP alebo tensorový paralelizmus. Tím tak nemusí prerábať celú vrstvu obnovy práve vo chvíli, keď model prerastie možnosti jednoduchšej konfigurácie. Dokumentácia PyTorch zároveň označuje Distributed Checkpoint za distribuované rozhranie: pri ukladaní predpokladá dostupnosť viacerých procesov a vývojári musia správne koordinovať kolektívne operácie i životný cyklus procesnej skupiny.

Druhou časťou návrhu je async_save. Tréning najprv skopíruje stav do pomocného priestoru a samotný zápis pokračuje na pozadí, zatiaľ čo GPU môžu spracúvať ďalšie dávky. Databricks uvádza dve vlastné merania. Pri 2,8-miliardovom modeli trénovanom cez DDP na 32 GPU H100 nameral zrýchlenie ukladania z 66 na 36 sekúnd, teda približne 1,8-násobok. Pri 20-miliardovom modeli s FSDP uvádza 522 sekúnd oproti deviatim sekundám, čo zodpovedá 58-násobnému rozdielu. Ide o výsledky konkrétnych konfigurácií a nemožno ich automaticky preniesť na iné siete, úložiská či veľkosti modelov.

Obnoviť treba aj dátový tok, nielen váhy modelu

Najdôležitejšia pripomienka sa týka stavu dát. Uloženie parametrov modelu, optimalizátora a plánovača učenia samo osebe nezaručuje korektné pokračovanie. Loader musí vedieť, ktoré vzorky už spracoval, aké náhodné premiešanie použil a na akom mieste sa nachádzal každý pracovník. Ak sa po obnove začne epocha od začiatku, časť záznamov sa zopakuje. Ak sa preskočí nesprávny počet dávok, iné záznamy môžu chýbať. Tréning pritom môže naďalej bežať bez viditeľnej chyby, takže problém sa prejaví až zmenou kvality modelu alebo slabšou reprodukovateľnosťou experimentu.

Databricks preto odporúča ukladať spolu s modelom aj pozíciu dátového procesu a používať deterministické premiešanie. Pri veľkých dátových súboroch má loader podporovať efektívny návrat na uloženú pozíciu bez toho, aby musel od začiatku čítať všetky predchádzajúce vzorky. Rovnako dôležité je zásobovanie GPU: pomalá serializácia, vzdialené úložisko alebo nedostatočný počet pracovníkov môže akcelerátory vyhladovať aj bez poruchy. Odolnosť a priepustnosť teda nie sú dve oddelené disciplíny. Rovnaký návrh dátovej vrstvy rozhoduje o bežnom využití hardvéru aj o tom, či bude obnovený beh pokračovať korektne.

Pre slovenské tímy je táto téma relevantná aj bez vlastného tisícového klastra. Viacdňový tréning na ôsmich či desiatkach prenajatých GPU môže prerušiť chyba uzla, spotová kapacita alebo krátky výpadok úložiska. Pri cloudovom účtovaní sa každá hodina opakovaného výpočtu priamo premieta do rozpočtu. Rozumným základom je preto automatický test obnovy: úlohu zámerne prerušiť, načítať checkpoint, porovnať krok optimalizátora a poradie nasledujúcich dávok a overiť, že neúplný zápis systém nepovažuje za platný. Samotná existencia súboru ešte neznamená, že je checkpoint konzistentný a bezpečne čitateľný.

Čo ešte treba overiť v konkrétnom prostredí

Publikované výsledky zatiaľ neodpovedajú na všetky prevádzkové otázky. Nie je z nich možné odvodiť rovnakú úsporu pre pomalšie objektové úložisko, inú topológiu siete, lacnejšie GPU alebo veľmi časté checkpointy, pri ktorých sa zápisy začnú prekrývať. Asynchrónne ukladanie navyše potrebuje pomocnú pamäť a mechanizmus, ktorý pred ukončením úlohy počká na dokončenie zápisu. Prevádzkovateľ musí riešiť aj retenciu starších verzií, kontrolu integrity, šifrovanie a situáciu, keď zlyhanie nastane počas vytvárania checkpointu. Užitočným ďalším údajom by bolo porovnanie celkového goodputu vrátane reálnych porúch počas dlhých behov, nielen času jedného zápisu.

Hlavný prínos postupu preto nie je v jednom pôsobivom násobku, ale v zmene architektonického predpokladu. Tréningová úloha má byť od začiatku pripravená na prerušenie, checkpoint má byť lacný natoľko, aby sa vytváral často, a obnovený dátový tok musí byť rovnako presný ako obnovené váhy. PyTorch už na to poskytuje všeobecné distribuované rozhranie, zatiaľ čo Databricks opisuje jeho integráciu s vlastným AI Runtime a úložiskom Unity Catalog. Pre iné platformy zostáva potrebné zmerať správanie v ich konkrétnej infraštruktúre, princíp však platí široko: pri drahom distribuovanom tréningu je schopnosť spoľahlivo pokračovať súčasťou výkonu, nie iba poistkou pre prípad havárie.