Эффективное использование памяти при обучении больших языковых моделей

По мере того как языковые модели становятся всё больше, растут и их словари. Это привело к непропорциональному увеличению объема памяти, используемого LLMs во время обучения, особенно в одном слое: при вычислении кросс-энтропии в функции потерь. Кросс-энтропия создает матрицу логитов с элементами для каждой пары входных токенов и элементов словаря и, для небольших моделей, потребляет на порядок больше памяти, чем остальная часть LLM вместе взятая. Мы предлагаем метод Cut Cross-Entropy (CCE), который вычисляет потери кросс-энтропии без материализации логитов для всех токенов в глобальной памяти. Вместо этого, CCE вычисляет логит только для правильного токена и оценивает log-sum-exp по всем логитам на лету. Мы реализовали специальное ядро, которое выполняет матричные умножения и редукцию log-sum-exp по словарю во флэш-памяти, делая потребление глобальной памяти для вычислений кросс-энтропии незначительным. Это имеет драматический эффект. Например, в случае модели Gemma 2 (2B), CCE снижает объем памяти, необходимый для вычисления потерь, с 24 ГБ до 1 МБ, а общее потребление памяти при обучении головы классификатора - с 28 ГБ до 1 ГБ. Для улучшения пропускной способности CCE мы используем врожденную разреженность softmax и предлагаем пропускать элементы вычисления градиента, которые имеют пренебрежимо малый вклад в градиент (то есть ниже численной точности). Эксперименты показывают, что драматическое сокращение потребления памяти достигается без ущерба для скорости обучения или сходимости.