FlashAttention-3 w BF16 psuje trening po cichu, dopiero po 25 miliardach tokenów

Przez pierwsze 25 miliardów tokenów trenowanie wyglądało zdrowo. Dwunastoosobowy zespół badaczy, w którym jest Eric P. Xing, uczył od zera transformer o 450 milionach parametrów na 50 miliardach tokenów i liczył mechanizm uwagi przez FlashAttention-3 w formacie BF16, w konfiguracji uważanej dziś za standard. W drugiej połowie biegu norma gradientu, która mówi, jak duże poprawki model dostaje w każdym kroku, urosła tysiąckrotnie. Na koniec wyszedł z tego wyraźnie słabszy model: jego funkcja straty, mierząca, jak bardzo myli się przy przewidywaniu kolejnych tokenów, skończyła wyżej niż w identycznym przebiegu, w którym uwagę liczono w pełnej precyzji FP32. Ani razu nie pojawił się przy tym NaN, a to właśnie na taką "nie-liczbę" reaguje większość alarmów w skryptach treningowych. Wszystko to autorzy opisali w pracy opublikowanej na arXiv 28 września.

Najbardziej niepokoi mnie w tej historii to, że błąd przez cały czas milczał. Ktoś, kto pilnuje tylko NaN-ów i spadającej krzywej straty, uznałby, że wszystko idzie dobrze, i dostałby na koniec słabszy model, nie wiedząc dlaczego. Dotyczy to jednak tylko tych, którzy sami trenują albo długo dostrajają modele na FlashAttention-3 w BF16. Kto korzysta z gotowych modeli przez API, nie musi niczego zmieniać, bo problem tkwi w liczeniu gradientów w trakcie trenowania.

Gdzie schował się błąd