Сократите потери в моделях языка с большим словарным запасом.

Cut Your Losses in Large-Vocabulary Language Models

November 13, 2024

Авторы: Erik Wijmans, Brody Huval, Alexander Hertzberg, Vladlen Koltun, Philipp Krähenbühl

cs.AI

Аннотация

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

English

As language models grow ever larger, so do their vocabularies. This has shifted the memory footprint of LLMs during training disproportionately to one single layer: the cross-entropy in the loss computation. Cross-entropy builds up a logit matrix with entries for each pair of input tokens and vocabulary items and, for small models, consumes an order of magnitude more memory than the rest of the LLM combined. We propose Cut Cross-Entropy (CCE), a method that computes the cross-entropy loss without materializing the logits for all tokens into global memory. Rather, CCE only computes the logit for the correct token and evaluates the log-sum-exp over all logits on the fly. We implement a custom kernel that performs the matrix multiplications and the log-sum-exp reduction over the vocabulary in flash memory, making global memory consumption for the cross-entropy computation negligible. This has a dramatic effect. Taking the Gemma 2 (2B) model as an example, CCE reduces the memory footprint of the loss computation from 24 GB to 1 MB, and the total training-time memory consumption of the classifier head from 28 GB to 1 GB. To improve the throughput of CCE, we leverage the inherent sparsity of softmax and propose to skip elements of the gradient computation that have a negligible (i.e., below numerical precision) contribution to the gradient. Experiments demonstrate that the dramatic reduction in memory consumption is accomplished without sacrificing training speed or convergence.