Как сократить потребление памяти при обучении больших языковых моделей
Как сократить потребление памяти при обучении больших языковых моделей Авторы новой статьи предложили метод Cut Cross-Entropy (CCE) для решения проблемы высокого потребления памяти при обучении больших языковых моделей (LLM) с крупными словарями. Проблема заключается в том, что материализация логитов может занимать до 90% доступной памяти, что ограничивает размер батча и масштаб обучения. Метод CCE позволяет вычислять лосс без сохранения всех логитов в глобальной памяти. Вместо этого используется вычисление log-sum-exp «на лету» в SRAM. Например, при использовании Gemma-2 с 2 миллиардами параметров потребление памяти сокращается с 24 ГБ до 1 МБ для вычисления лосса и с 28 ГБ до 1 ГБ для classifier-head. CCE обеспечивает ускорение примерно в 3,5 раза на бэкворде и увеличивает достижимый размер батча на 16 GPU в 1,5–10 раз в зависимости от модели. При этом кривые обучения при файнтюнинге совпадают с torch.compile, а для претрейнинга точность выравнивается вариантом CCE-Kahan-FullC.