Эффективное использование памяти при обучении больших языковых моделей
В последние годы большие языковые модели (LLM) значительно увеличились в размерах, что привело к росту их словарных запасов. Этот рост словарного запаса, в свою очередь, привел к тому, что память, используемая для обучения этих моделей, стала распределяться неравномерно, особенно в одном слое — слое кросс-энтропии в вычислении потерь. В этой статье мы рассмотрим метод, известный как Cut Cross-Entropy (CCE), который предлагает новый подход к вычислению кросс-энтропии, значительно сокращая использование памяти без ущерба для скорости обучения или сходимости.
Введение в проблему
LLM, такие как Gemma 2 (2B), используют кросс-энтропию для оценки вероятности следующего токена в последовательности. Кросс-энтропийная потеря требует вычисления логит-матрицы, где каждая запись соответствует паре входного токена и элемента словаря. Для моделей с большими словарями это может потребовать огромного количества памяти, что делает обучение неэффективным и дорогим.
Кросс-энтропия и её проблемы
Кросс-энтропийная потеря вычисляется следующим образом:
[ \ell_i(x) = \log \text{softmax}{x_i} \left( C^T E_i \right) = C^T{x_i} E_i - \log \sum_j \exp(C^T_j E_i) ]
где ( E_i ) — это вложение (embedding) для последовательности токенов, ( C ) — классификатор, а ( x_i ) — правильный токен. В традиционном подходе вся матрица логитов ( C^T E ) материализуется в памяти, что для больших словарей может потребовать до 90% всей памяти, используемой на этапе обучения.
Метод Cut Cross-Entropy (CCE)
CCE предлагает вычислять кросс-энтропию без материализации полной матрицы логитов в глобальной памяти. Вместо этого:
-
Индексированное умножение матриц: Вычисляется только логит для правильного токена, используя индексированное умножение матриц.
-
Логарифмическое суммирование экспонент: Операция log-sum-exp вычисляется на лету, что позволяет избежать хранения всех логитов в памяти.
Реализация CCE
Для реализации CCE используется набор специализированных ядер CUDA, которые выполняют:
- Индексированное умножение матриц: Используется для вычисления логита для правильного токена.
- Блочное вычисление log-sum-exp: Выполняется в SRAM, что уменьшает использование глобальной памяти.
Эти ядра работают с блоками данных, что позволяет эффективно использовать кэш GPU и сократить общее потребление памяти.
Преимущества CCE
- Сокращение использования памяти: CCE сокращает потребление памяти от 24 ГБ до 1 МБ для модели Gemma 2 (2B).
- Сохранение скорости обучения: Несмотря на уменьшение использования памяти, CCE не замедляет процесс обучения.
- Сходимость: Эксперименты показывают, что CCE не влияет на сходимость модели.
Анализ производительности
При анализе производительности CCE в сравнении с традиционными методами, такими как PyTorch, torch.compile, Torch Tune и Liger Kernels, CCE показывает выдающиеся результаты:
- Память: CCE использует значительно меньше памяти, чем другие методы, что позволяет увеличить размер батча в обучении.
- Время выполнения: CCE имеет конкурентоспособное время выполнения, особенно при вычислении потерь и градиентов.
Градиентный фильтр
CCE также использует градиентный фильтр, который пропускает вычисления для элементов градиента, вносящих незначительный вклад. Это позволяет еще больше сократить время вычислений, не влияя на точность.
Заключение
Метод Cut Cross-Entropy представляет собой значительный шаг вперед в оптимизации обучения больших языковых моделей. Он позволяет эффективно использовать память, что открывает возможности для обучения моделей с еще большими словарными запасами и увеличивает возможности для параллельных вычислений. CCE не только сокращает потребление памяти, но и сохраняет или даже улучшает скорость и качество обучения, делая его привлекательным решением для будущих разработок в области LLM.