Destylacja wiedzy AI taniej o 50% dzięki nowej metodzie cache'owania

Naukowcy z Multiverse Computing opublikowali nowe badanie pokazujące, jak drastycznie obniżyć koszty destylacji wiedzy w dużych modelach językowych. Ich metoda pozwala na trening na pojedynczej karcie graficznej zamiast setek GPU. Szczegóły w ich publikacji na Hugging Face.

Kluczowe wnioski

• Tradycyjna destylacja wiedzy wymaga około 250GB pamięci VRAM, co przekracza możliwości nawet najnowszych kart H200 czy B200.
• Nowa metoda obniża zużycie pamięci do 128GB poprzez cache'owanie top-100 tokenów i chunked KL loss.
• Model Kimi-K3 z 2,8 biliona parametrów potrzebuje około 3TB VRAM tylko do załadowania, co pokazuje skalę problemu.
• Offline distillation eliminuje potrzebę trzymania modelu-nauczyciela w pamięci podczas treningu.

Problem kosztów standardowej destylacji

Destylacja wiedzy to proces uczenia mniejszego modelu-ucznia, by dorównywał wydajności większego modelu-nauczyciela. W standardowym podejściu oba modele muszą być jednocześnie załadowane w pamięci. Przy każdym kroku treningowym nauczyciel wykonuje pełny przebieg, generując rozkład prawdopodobieństwa dla całego słownika, a uczeń uczy się go naśladować.

Przykład pokazuje skalę problemu: model gpt-oss-120b ma słownik 201,088 tokenów. Przy długości sekwencji 32K i batch size 4, tensor prawdopodobieństw nauczyciela ma wymiary 4 × 201,088 × 32,768. W formacie bfloat16 to już około 50GB VRAM dla jednego tensora. Dodając gradienty, aktywacje, wagi modelu i stany optymalizatora, pojedyncza iteracja treningowa może osiągnąć szczyt około 250GB VRAM.

Dwie kluczowe zmiany systemowe

Offline distillation zastępuje ponowne obliczanie nauczyciela w każdym kroku. Zamiast tego oblicza jego wyniki raz, cache'uje top-100 najbardziej prawdopodobnych tokenów na pozycję i trenuje ucznia przeciwko tej pamięci podręcznej. Nauczyciel nie musi już siedzieć w pamięci podczas treningu.

Fused chunked KL loss to nowa implementacja funkcji straty. Standardowe podejście buduje ogromną siatkę: jeden wiersz na każdy wpis w słowniku i jedną kolumnę na każdą pozycję w sekwencji. Dla słownika 100K+ słów i długiej sekwencji ta siatka jest enormous.

Autorzy porównują trzy metody obliczania tej samej straty:

• Dense KL (podejście podręcznikowe) — odbudowuje pełną, gęstą siatkę prawdopodobieństw
• Forward-chunked KL — przetwarza stratę kawałek po kawałku
• Fused chunked KL — łączy projekcję wyjściową modelu bezpośrednio z obliczaniem straty, przetwarzając jeden chunk sekwencji na raz

Nowa metoda nigdy nie produkuje pełnej siatki logitów ucznia, tylko przetwarza jeden fragment sekwencji end-to-end, rzutuje ukryte stany na logity dla tego fragmentu, włącza wynik do bieżącej straty i odrzuca fragment przed przejściem do następnego.

Te optymalizacje sprawiają, że destylacja wiedzy staje się na tyle tania, by umożliwić eksperymenty na dużą skalę i trenowanie długich kontekstów na pojedynczej karcie graficznej.