Reducir sus pérdidas en modelos de lenguaje de gran vocabulario.

Resumen

A medida que los modelos de lenguaje crecen cada vez más, también lo hacen sus vocabularios. Esto ha desplazado de manera desproporcionada la huella de memoria de los LLM durante el entrenamiento a una sola capa: la entropía cruzada en el cálculo de la pérdida. La entropía cruzada construye una matriz de logit con entradas para cada par de tokens de entrada e ítems de vocabulario y, para modelos pequeños, consume un orden de magnitud más de memoria que el resto del LLM combinado. Proponemos Cut Cross-Entropy (CCE), un método que calcula la pérdida de entropía cruzada sin materializar los logit para todos los tokens en la memoria global. En su lugar, CCE solo calcula el logit para el token correcto y evalúa la suma logarítmica exponencial sobre todos los logit sobre la marcha. Implementamos un kernel personalizado que realiza las multiplicaciones de matrices y la reducción de la suma logarítmica exponencial sobre el vocabulario en la memoria flash, haciendo que el consumo de memoria global para el cálculo de la entropía cruzada sea insignificante. Esto tiene un efecto dramático. Tomando el modelo Gemma 2 (2B) como ejemplo, CCE reduce la huella de memoria del cálculo de pérdida de 24 GB a 1 MB, y el consumo total de memoria en tiempo de entrenamiento de la cabeza clasificadora de 28 GB a 1 GB. Para mejorar el rendimiento de CCE, aprovechamos la esparsidad inherente de softmax y proponemos omitir elementos del cálculo del gradiente que tienen una contribución despreciable (es decir, por debajo de la precisión numérica) al gradiente. Los experimentos demuestran que la reducción dramática en el consumo de memoria se logra sin sacrificar la velocidad de entrenamiento o la convergencia.

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.

Reducir sus pérdidas en modelos de lenguaje de gran vocabulario.

Cut Your Losses in Large-Vocabulary Language Models

Resumen

Support