Chunkovaná KL strata umožnila destilovať LLM s 32-tisícovým kontextom na jednom H200

Výskumníci z Multiverse Computing podrobnejšie opísali postup, ktorým znižujú pamäťové nároky pri destilácii veľkých jazykových modelov. Nejde o nový model ani o úplne nový princíp učenia, ale o kombináciu dvoch systémových úprav: výstupy učiteľského modelu sa vypočítajú vopred a samotná Kullbackova–Leiblerova strata sa spracúva po blokoch bez vytvorenia celej matice logitov. V experimente s približne 3,2-miliardovým študentom táto kombinácia umožnila trénovať s kontextom 32 768 tokenov na jednom GPU NVIDIA H200. Toto tvrdenie treba čítať presne: výsledok sa týka konkrétnej zostavy, nie ľubovoľného LLM a ľubovoľného 32-tisícového tréningu.

Dva odlišné zdroje úspory

Pri online destilácii zostáva počas tréningu v pamäti učiteľ aj študent. Učiteľ pri každom kroku opakovane vytvára pravdepodobnostné rozdelenie nad celým slovníkom, hoci sa jeho parametre nemenia. Autori namiesto toho spustili učiteľa vopred a pre každú pozíciu uložili iba 100 najpravdepodobnejších tokenov a ich hodnoty. Študent sa potom učí z tejto cache. V porovnaní pri kontexte 8 192 tokenov klesla uvádzaná špičková spotreba z približne 103 na 78 GB, čas iterácie z 25,9 na 18,5 sekundy a priepustnosť stúpla z 237 na 331 TFLOP/s. Ide o merania autorov na jednom H200, nie o nezávislý benchmark.

Druhá úprava rieši iný problém. Aj po odstránení učiteľa môže výstupná vrstva študenta vytvárať obrovský tenzor s rozmermi dávka krát dĺžka sekvencie krát veľkosť slovníka. Pri dlhom kontexte a slovníku so státisícmi položiek môže tento medzivýsledok spotrebovať desiatky až stovky gigabajtov. Fused chunked KL loss preto spája výstupnú projekciu s výpočtom straty. Logity vytvorí iba pre malý blok pozícií, započíta ich do výsledku a zahodí ich. Pri spätnom priechode príslušný blok znova vypočíta. Platí sa teda dodatočným výpočtom, no v pamäti netreba držať logity celej sekvencie.

Tieto dve úspory sa nemajú zamieňať. Offline cache odstraňuje učiteľa z tréningovej slučky, zatiaľ čo chunkovanie odstraňuje najväčší dočasný tenzor študenta. Pri 8-tisícovom kontexte bola dopredne chunkovaná verzia dokonca o niečo rýchlejšia než plne spojená verzia, pretože nemusela logity pri spätnom priechode prepočítavať. Plne chunkovaný variant sa presadzuje najmä tam, kde je rozhodujúca pamäť alebo veľmi dlhá sekvencia. Podľa práce znížil špičku pri 8-tisícovom kontexte na približne 58 GB a pri 32 768 tokenoch zostal okolo 128 GB, teda pod oficiálnou kapacitou 141 GB modelu H200.

Čo dopĺňajú podrobnejšie experimenty

Dôležitým spresnením oproti stručnému súhrnu je rozsah testov. Hlavné porovnanie používalo Llama 3.1 8B Instruct ako učiteľa a kompaktný model s približne 3,2 miliardy parametrov ako študenta. Offline variant uchovával top-100 výstupov učiteľa bez opätovnej normalizácie odrezanej distribúcie. Autori tvrdia, že online a offline postup mali takmer zhodné krivky tréningovej straty. To podporuje technickú správnosť implementácie v tomto nastavení, ale ešte nedokazuje rovnakú výslednú kvalitu pri inom učiteľovi, dátach, jazyku alebo agresívnejšej kompresii.

Samostatný syntetický test izoloval iba výstupnú projekciu a KL stratu, bez transformerových blokov, pozornosti, optimalizátora či načítavania dát. Pri 32-tisícovej sekvencii v ňom klesla pamäť z 85,2 GiB pri hustej KL strate na 5,45 GiB pri plne chunkovanej verzii. Hustý variant od 64-tisíc tokenov zlyhával pre nedostatok pamäte; pri 256-tisíc tokenoch používala plne chunkovaná verzia 11,6 GiB. Tieto čísla dobre ukazujú správanie jadra, nemožno ich však vydávať za spotrebu celého LLM. Samotná práca na toto obmedzenie výslovne upozorňuje.

Autori uvádzajú aj experiment s GPT-OSS-20B pri kontexte 32 768 tokenov. Uvoľnená pamäť podľa nich umožnila znížiť nasadenie zo štyroch výpočtových uzlov na jeden uzol s viacerými H200, nie na jediný GPU. Odstránenie komunikácie medzi uzlami skrátilo krok z 57 na 12,23 sekundy a zvýšilo priepustnosť na GPU zo 74,2 na 345,7 TFLOP/s. Rozlíšenie medzi jedným GPU a jedným viac-GPU uzlom je podstatné pre odhad ceny aj reprodukovateľnosti; titulkové tvrdenie o jedinom H200 patrí ku kompaktnému 3,2-miliardovému študentovi.

Praktický význam pre menšie tímy

Optimalizácia je zaujímavá pre firmy, ktoré po prerezaní alebo kvantizácii modelu obnovujú jeho kvalitu destiláciou. Cache učiteľa možno opakovane použiť pri porovnávaní stratových funkcií, dátových zmesí alebo hyperparametrov, takže najdrahší dopredný priechod učiteľa netreba zakaždým opakovať. Chunkovaná strata zasa môže rozhodnúť, či experiment potrebuje niekoľko GPU alebo sa zmestí na jeden akcelerátor. Úspora preto nemusí znamenať iba nižší účet za hardvér; menší počet zariadení obmedzuje aj komunikačnú réžiu a zjednodušuje prevádzku.

Pre slovenské a európske tímy je dôležitý najmä on-premise rozmer. Organizácie pracujúce s citlivými podnikovými, zdravotníckymi alebo verejnými dátami často nemôžu posielať tréningové vstupy do ľubovoľnej externej služby. Menší lokálny model je realistickejší, jeho prispôsobenie však stále naráža na cenu akcelerátorov. Postup, ktorý znižuje potrebný počet H200, môže rozšíriť okruh realizovateľných experimentov v univerzitných centrách, spoločných európskych výpočtových kapacitách aj menších firmách. Stále však ide o optimalizáciu tréningu; sama osebe nerieši licencie učiteľských modelov, dostupnosť dát ani náklady následnej inferencie.

Kde zostáva dôvod na opatrnosť

Technický smer nadväzuje na nezávislý výskum pamäťovo úsporných stratových funkcií. Práca Cut Cross-Entropy, prijatá na ICLR 2025, už ukázala, že pri klasickej krížovej entropii možno vynechať materializáciu plnej matice logitov. Nový preprint prenáša podobnú myšlienku na destiláciu, kde cieľom nie je jediný správny token, ale riedka distribúcia učiteľa. Zverejnený repozitár obsahuje tri implementácie, testy ekvivalencie a syntetický benchmark, čo uľahčuje kontrolu kódu. Nejde však o nezávislú reprodukciu výsledkov, pretože repozitár aj blog pochádzajú od rovnakého tímu ako preprint.

Otvorených otázok zostáva viacero. Hlavné end-to-end porovnanie pokrýva jedinú dvojicu učiteľ–študent a jedinú hodnotu K rovnu 100. Cache je viazaná na konkrétne dáta, tokenizer a učiteľa, zaberá diskový priestor a pri zmene dát ju treba vytvoriť znova. Odrezanie zvyšku distribúcie môže pri iných modeloch alebo úlohách poškodiť jemné informácie, hoci v uvedenom experimente bola tréningová strata podobná. Preprint zatiaľ neprešiel štandardným nezávislým recenzným procesom a uvádza čakajúcu patentovú prihlášku, takže pred produkčným použitím treba preveriť aj licenčné podmienky. Najpresnejší záver preto je, že tím ukázal sľubnú a kontrolovateľnú implementáciu, ktorej všeobecnosť ešte musia potvrdiť cudzie laboratóriá.