Fire GPU-noder ble til én: Multiverse kutter minnebruken i destillasjon

250 gigabyte toppminne for én eneste treningsiterasjon: det er regnestykket Multiverse Computing legger fram i en artikkel på Hugging Face for å forklare hvorfor destillasjon av store språkmodeller vanligvis krever hundrevis av GPU-er. Med vokabularet på 201 088 tokens i gpt-oss-120b blir lærermodellens sannsynlighetstensor alene rundt 50 GB i bfloat16 ved 32K kontekst og batch 4, før gradienter, vekter og optimizer-tilstand er regnet med.

De to endringene er systemiske, ikke matematiske. Først kjøres lærermodellen én gang og de hundre mest sannsynlige tokenene per posisjon lagres i en cache, slik at læreren aldri må ligge i minnet ved siden av studenten og kan gjenbrukes på tvers av forsøk. Deretter fusjoneres modellens utgangsprojeksjon inn i selve tapsfunksjonen, som behandler én bit av sekvensen om gangen og kaster den før neste. Bakoverpasset regner hver bit på nytt i stedet for å lagre den. Prisen er én ekstra projeksjon, gevinsten er at toppminnet vokser lineært med sekvenslengden i stedet for å eksplodere med vokabular ganger sekvens.

På en enkelt H200 med Llama 3.1 8B som lærer og en student på 3,2 milliarder parametere havner alle variantene på nesten identisk treningstap, ifølge tallene i artikkelen. Det tette tapet feiler helt fra 64 000 tokens og oppover. For en destillasjon av GPT-OSS 20B ved 32 768 tokens kontekst betydde det frigjorte minnet at oppsettet krympet fra fire GPU-noder til én, med steg-tid ned fra 57 til 12,23 sekunder.

Studenten på 3,2 milliarder parametere beholder mesteparten av lærerens treffsikkerhet på BoolQ og HellaSwag og ligger rundt ni poeng bak på MMLU, med under halvparten så mange parametere. Implementasjonen av det chunkede tapet er lagt ut som åpen kildekode under navnet Full-Chunked-KL-Loss, så terskelen for å prøve teknikken på egne modeller er én GPU i stedet for en klynge.