Cut Your Losses in Large-Vocabulary Language Modelsавторы вспоминают вычисление кросс энтропии
loss = F.cross_entropy((e @ c.T).float(), targets)
где e - финальные эмбеддинги, с - матрица, которая переводит из эмбеддингов в vocab_size вектор. и эта самая матрица
e @ c.T, как оказывается, кушает оочень многоо памяти, а при увеличении вокаб сайза затрат по памяти становится еще больше
и чтобы снизить память, авторы сделали свой тритоновский кернел, который не засовывает сразу все логиты в глобальную память, а вычисляет логиты для верных токенов (во flash memory из-за концепции teacher forcing) и log-sum-exp на лету (по аналогии с онлайн софтмаксом из флеш аттн), т.е. не аллоцирует память на
e @ c.T
к тому же еще при бекворде они заметили, что меньше чем 0.02% от нормализованных логитов являются ненулевыми (по большей части из-за действительно большого вокаб сайза) → не вычисляют градиент для элементов, которые не проходят трешхолд в 2е-12 для бф16 точности, что тоже снижает по памяти и прибавляет по скорости
оверолл голова ллмок начала теперь кушать по памяти в среднем гигабайт вместо 28 при батч сайзе в 65к (при том экспы ставили на модельках от 1.3B do 70B за что респект). есть правда вопросы насколько стабильно будет этот метод работать для претрена с нуля, ибо авторы только “файнтюнили”
очень понятно описано решение их кернела, как и сам код, рекомендуем к прочтению
выглядит интересно и прикольно, на последней картинке только с осторожностью относился бы к абсолютным числам, которые у них получились на экспах по замеру времени
👀LINK 
