SkewAdam redukuje zapotrzebowanie na pamięć podczas trenowania modeli Mixture-of-Experts o ponad 60% i zapewnia zauważalny wzrost dokładności, umożliwiając programistom zmieszczenie modelu MoE o 6,78 miliarda parametrów na pojedynczym procesorze GPU o pojemności 40 GB.

Dlaczego trenowanie MoE napotyka barierę pamięci

Architektury Mixture-of-Experts przekierowują każdy wejściowy sygnał przez mały podzbiór podsieci „eksperckich”, zachowując jednocześnie gęstą strukturę bazową (dense backbone). Korzyścią jest model, który może skalować się do miliardów parametrów bez proporcjonalnego wzrostu zapotrzebowania na moc obliczeniową. W praktyce to optymalizator — a konkretnie wariant AdamW, z którego korzysta większość zespołów — pochłania większość pamięci RAM procesora GPU. W przypadku modelu o 6,78 mld parametrów same tensory pędu (momentum) i wariancji optymalizatora wymagają około 50 GB. Po dodaniu aktywacji i wag modelu, szczytowe zużycie pamięci przekracza 81,4 GB, co wymusza stosowanie konfiguracji wieloprocesorowych (multi-GPU) lub wolniejszych harmonogramów trenowania.

Co SkewAdam robi inaczej

SkewAdam nie wymyśla nowej reguły uczenia. Zmienia on sposób rozmieszczenia momentów (moments) algorytmu Adam:

  • Gęsta struktura bazowa (dense backbone) zachowuje pęd o pełnej precyzji, co zapewnia płynne aktualizacje niezbędne dla warstw gęstych.
  • Bank ekspertów przechowuje faktoryzowaną aproksymację wariancji, co redukuje ilość danych przechowywanych dla każdego eksperta.
  • Router, który decyduje o aktywacji poszczególnych ekspertów, zachowuje dokładną estymację drugiego momentu.

Traktując te trzy komponenty jako oddzielne „poziomy” (tiers), optymalizator rezygnuje ze zbędnej precyzji tam, gdzie jest ona mniej istotna, i zachowuje ją tam, gdzie ma największe znaczenie.

Mierzalny wpływ

Uruchomienie tego samego modelu MoE o 6,78 mld parametrów z użyciem SkewAdam daje wyniki:

  • Szczytowe zużycie pamięci GPU: 31,3 GB (spadek z 81,4 GB)
  • Ślad pamięci stanu optymalizatora: 1,29 GB (spadek z 50 GB)

Zredukowany stan bez problemu mieści się w pojedynczym akceleratorze o pojemności 40 GB.

Zysk w dokładności, a nie tylko oszczędności

Redukcja zużycia pamięci często pogarsza jakość modelu, ale SkewAdam poprawia perplexity — standardową metrykę modeli językowych — z poziomu 126,8 (AdamW) do 108,4. Przewyższa on również Muon (120,2) oraz Lion (393,7) w tym samym zadaniu. Autorzy twierdzą, że wzrost wynika z zachowania pędu na wszystkich trzech poziomach; metody, które całkowicie rezygnują z pędu, takie jak Adafactor, wypadają gorzej.

Otwarte pytania

Wyniki SkewAdam zostały zaprezentowane na modelu o 6,78 mld parametrów. Pozostaje niejasne, czy te same proporcje stanu pamięci zachowają się w przypadku modeli o rząd wielkości większych lub w zadaniach wykraczających poza modelowanie języka.

Na co zwrócić uwagę

  • Benchmarki na większych konfiguracjach MoE (dziesiątki miliardów parametrów).
  • Adaptacja przez frameworki open-source i uwzględnienie w popularnych skryptach trenujących.
  • Opinie społeczności na temat ukrytych kompromisów, takich jak szybkość zbieżności lub stabilność przy różnych harmonogramach tempa uczenia (learning-rate schedules).

Źródło: post dewelopera szczegółowo opisujący projekt optymalizatora i wyniki empiryczne.