El entrenamiento y ajuste de modelos grandes tiene altos requisitos de memoria de video, y el estado del optimizador es uno de los principales gastos de la memoria de video. Recientemente, el equipo de Zhu Jun y Chen Jianfei de la Universidad de Tsinghua propuso un optimizador de 4 bits para el entrenamiento de redes neuronales, que ahorra la sobrecarga de memoria del entrenamiento del modelo y puede lograr una precisión comparable a la de un optimizador de precisión total.
El optimizador de 4 bits se experimentó en numerosas tareas de preentrenamiento y ajuste y puede reducir la sobrecarga de memoria del ajuste fino del LLaMA-7B hasta en un 57 % manteniendo la precisión.
Documento: https://arxiv.org/abs/2309.01507
Código: https://github.com/thu-ml/low-bit-optimizers
Cuello de botella de memoria en el entrenamiento de modelos.
Desde GPT-3, Gopher hasta LLaMA, se ha convertido en un consenso de la industria que los modelos grandes tienen mejor rendimiento. Sin embargo, por el contrario, el tamaño de la memoria de video de una sola GPU ha crecido lentamente, lo que hace que la memoria de video sea el principal cuello de botella para el entrenamiento de modelos grandes. Cómo entrenar modelos grandes con memoria de GPU limitada se ha convertido en un problema importante.
Para hacer esto, primero necesitamos identificar las fuentes de consumo de memoria de video. De hecho, existen tres tipos de fuentes, a saber:
1. La "memoria de datos" incluye los datos de entrada y la salida del valor de activación de cada capa de la red neuronal. Su tamaño se ve directamente afectado por el tamaño del lote y la resolución de la imagen/longitud del contexto;
2. "Memoria del modelo", incluidos los parámetros del modelo, gradientes y estados del optimizador, su tamaño es proporcional al número de parámetros del modelo;
3. La "memoria de video temporal" incluye la memoria temporal y otros cachés utilizados en los cálculos del kernel de la GPU. A medida que aumenta el tamaño del modelo, la proporción de la memoria de vídeo del modelo aumenta gradualmente, convirtiéndose en un importante cuello de botella.
El tamaño del estado del optimizador está determinado por el optimizador que se utiliza. Actualmente, el optimizador AdamW se usa a menudo para entrenar Transformers, que necesitan almacenar y actualizar dos estados del optimizador durante el proceso de entrenamiento, es decir, el primer y segundo momento. Si el número de parámetros del modelo es N, entonces el número de estados del optimizador en AdamW es 2N, lo que obviamente supone una enorme sobrecarga de memoria gráfica.
Tomando LLaMA-7B como ejemplo, el número de parámetros de este modelo es de aproximadamente 7 B. Si se utiliza el optimizador AdamW de precisión total (32 bits) para ajustarlo, el tamaño de la memoria de video ocupada por el estado del optimizador es de aproximadamente 52,2 GB. Además, aunque el ingenuo optimizador SGD no requiere estados adicionales y guarda la memoria ocupada por el estado del optimizador, el rendimiento del modelo es difícil de garantizar. Por lo tanto, este artículo se centra en cómo reducir el estado del optimizador en la memoria del modelo y al mismo tiempo garantizar que el rendimiento del optimizador no se vea comprometido.
Formas de ahorrar memoria del optimizador
Actualmente, en términos de algoritmos de entrenamiento, existen tres métodos principales para ahorrar la sobrecarga de memoria del optimizador:
1. Realizar una aproximación de bajo rango del estado del optimizador mediante la idea de descomposición de bajo rango (factorización);
2. Evite guardar la mayoría de los estados del optimizador entrenando solo un pequeño conjunto de parámetros, como LoRA;
3. Un método basado en compresión que utiliza un formato numérico de baja precisión para representar el estado del optimizador.
En particular, Dettmers et al. (ICLR 2022) propusieron un optimizador de 8 bits correspondiente para SGD con impulso y AdamW, mediante el uso de técnicas de cuantización por bloques y formato numérico exponencial dinámico, logrando resultados que coinciden con el optimizador de precisión total original en tareas como modelado de lenguaje, clasificación de imágenes, aprendizaje autosupervisado y traducción automática.
En base a esto, este artículo reduce aún más la precisión numérica del estado del optimizador a 4 bits, propone un método de cuantificación para diferentes estados del optimizador y finalmente propone un optimizador AdamW de 4 bits. Al mismo tiempo, este artículo explora la posibilidad de combinar métodos de compresión y descomposición de bajo rango y propone un optimizador de factor de 4 bits, que disfruta de un buen rendimiento y una mejor eficiencia de la memoria. Este artículo evalúa el optimizador de 4 bits en una serie de tareas clásicas, incluida la comprensión del lenguaje natural, la clasificación de imágenes, la traducción automática y el ajuste fino de instrucciones de modelos grandes.
En todas las tareas, el optimizador de 4 bits logra un rendimiento comparable al optimizador de precisión total y ocupa menos memoria.
Configuración de preguntas
Un marco para optimizadores eficientes en memoria basados en compresión
Primero, debemos comprender cómo introducir operaciones de compresión en optimizadores de uso común, lo cual viene dado por el Algoritmo 1. donde A es un optimizador basado en gradiente (como SGD o AdamW). El optimizador ingresa los parámetros existentes w, gradiente g y estado del optimizador s, y genera nuevos parámetros y estado del optimizador. En el algoritmo 1, el s_t de precisión total existe temporalmente, mientras que el s_t de baja precisión (s_t ) ̅ persistirá en la memoria de la GPU. La razón importante por la que este método puede ahorrar memoria de video es que los parámetros de las redes neuronales a menudo se unen a partir de los vectores de parámetros de cada capa. Por lo tanto, la actualización del optimizador también se realiza capa por capa/tensor. Según el algoritmo 1, el estado del optimizador de como máximo un parámetro se deja en la memoria en forma de precisión total, y los estados del optimizador correspondientes a otras capas están comprimidos. estado. . 
Método de compresión principal: cuantificación.
La cuantificación es una tecnología que utiliza valores numéricos de baja precisión para representar datos de alta precisión. Este artículo desacopla la operación de cuantificación en dos partes: normalización y mapeo, lo que permite un diseño más liviano y la experimentación de un nuevo método de cuantificación. Las dos operaciones de normalización y mapeo se aplican secuencialmente a los datos de precisión total por elementos. La normalización es responsable de proyectar cada elemento del tensor al intervalo unitario, donde la normalización del tensor (normalización por tensor) y la normalización por bloques (normalización por bloques) se definen de la siguiente manera: Por lo tanto, cómo diseñar una T adecuada 
Juega un papel importante en la reducción de errores de cuantificación. Este artículo considera principalmente el mapeo lineal (lineal) y el mapeo exponencial dinámico (exponente dinámico). Finalmente, el proceso de descuantización consiste en aplicar los operadores inversos de mapeo y normalización en secuencia.
Método de compresión de momento de primer orden.
A continuación se proponen principalmente diferentes métodos de cuantificación para el estado del optimizador (momento de primer orden y momento de segundo orden) de AdamW. Para momentos de primer orden, el método de cuantificación en este artículo se basa principalmente en el método de Dettmers et al. (ICLR 2022), utilizando normalización de bloques (el tamaño de bloque es 2048) y mapeo exponencial dinámico.
En experimentos preliminares, redujimos directamente el ancho de bits de 8 bits a 4 bits y descubrimos que el momento de primer orden es muy robusto a la cuantificación y ha logrado efectos de coincidencia en muchas tareas, pero también sufrió pérdidas de rendimiento en algunas tareas. Para mejorar aún más el rendimiento, estudiamos cuidadosamente el patrón de los primeros momentos y descubrimos que hay muchos valores atípicos en un solo tensor.
En trabajos anteriores se han realizado algunas investigaciones sobre los patrones atípicos de parámetros y valores de activación: la distribución de parámetros es relativamente suave, mientras que los valores de activación tienen las características de distribución de canales. Este artículo encontró que la distribución de valores atípicos en el estado del optimizador es compleja: algunos tensores tienen valores atípicos distribuidos en filas fijas y otros tensores tienen valores atípicos distribuidos en columnas fijas.
Para tensores con distribución de valores atípicos por columnas, la normalización del bloque de primera fila puede encontrar dificultades. Por lo tanto, este artículo propone utilizar bloques más pequeños, con un tamaño de bloque de 128, lo que puede reducir el error de cuantificación y al mismo tiempo mantener la sobrecarga de memoria adicional dentro de un rango controlable. La siguiente figura muestra el error de cuantificación para diferentes tamaños de bloque. 
Método de compresión de momento de segundo orden.
En comparación con el momento de primer orden, la cuantificación del momento de segundo orden es más difícil y traerá inestabilidad al entrenamiento. Este artículo determina que el problema del punto cero es el principal cuello de botella en la cuantificación de momentos de segundo orden.Además, se propone un método de normalización mejorado para distribuciones atípicas mal condicionadas: la normalización de rango 1. Este artículo también intenta el método de descomposición (factorización) del momento de segundo orden.
problema de punto cero
En la cuantificación de parámetros, valores de activación y gradientes, el punto cero suele ser indispensable y también es el punto con la frecuencia más alta después de la cuantificación. Sin embargo, en la fórmula iterativa de Adam, el tamaño de la actualización es proporcional a la potencia -1/2 del segundo momento, por lo que cambiarlo en el rango cercano a cero afectará en gran medida el tamaño de la actualización, provocando así inestabilidad. 
La siguiente figura muestra la distribución del segundo momento de Adam -1/2 potencia antes y después de la cuantificación en forma de histograma, es decir, h (v) = 1/(√v+10^(-6)). Si se incluyen los ceros (Figura b), la mayoría de los valores se elevan a 10^6, lo que genera grandes errores de aproximación. Una solución sencilla es eliminar los ceros en el mapa exponencial dinámico. Después de hacer esto (figura c), la aproximación al segundo momento se vuelve más precisa. En situaciones reales, para utilizar eficazmente la capacidad expresiva de los valores numéricos de baja precisión, propusimos utilizar un mapeo lineal que elimina los puntos cero y logramos buenos resultados en los experimentos.
Normalización de rango 1
Basado en la compleja distribución de valores atípicos de momentos de primer orden y momentos de segundo orden, e inspirado en el optimizador SM3, este artículo propone un nuevo método de normalización denominado normalización de rango 1. Para un tensor matricial no negativo x∈R^(n×m), su estadística unidimensional se define como: la 
normalización de rango 1 utiliza la información unidimensional del tensor de una manera más detallada y puede ser más Maneje de manera inteligente y eficiente los valores atípicos en filas o columnas. Además, la normalización de rango 1 se puede generalizar fácilmente a tensores de alta dimensión y, a medida que aumenta el tamaño del tensor, la sobrecarga de memoria adicional que genera es menor que la de la normalización de bloques.
Además, este artículo encontró que el método de descomposición de bajo rango para momentos de segundo orden en el optimizador Adafactor puede evitar efectivamente el problema del punto cero, por lo que también se exploró la combinación de métodos de cuantificación y descomposición de bajo rango. La siguiente figura muestra una serie de experimentos de ablación para momentos de segundo orden, lo que confirma que el problema del punto cero es el cuello de botella para cuantificar los momentos de segundo orden. También verifica la efectividad de los métodos de normalización de rango 1 y descomposición de rango bajo. 
Resultados experimentales
Con base en los fenómenos observados y los métodos de uso, el estudio finalmente propuso dos optimizadores de baja precisión: AdamW de 4 bits y Factor de 4 bits, y los comparó con otros optimizadores, incluidos AdamW, Adafactor y SM3 de 8 bits. Se seleccionaron estudios para su evaluación en una amplia gama de tareas, incluida la comprensión del lenguaje natural, la clasificación de imágenes, la traducción automática y el ajuste fino de las instrucciones de modelos grandes. La siguiente tabla muestra el rendimiento de cada optimizador en diferentes tareas. 
Se puede ver que en todas las tareas de ajuste fino, incluidas NLU, QA y NLG, el optimizador de 4 bits puede igualar o incluso superar el AdamW de 32 bits. Al mismo tiempo, en todas las tareas previas al entrenamiento, CLS, MT y el optimizador de 4 bits alcanzan el mismo nivel que el AdamW completo de 32 bits. Nivel de precisión comparable. Se puede ver en la tarea de ajuste de instrucciones que AdamW de 4 bits no destruye la capacidad de los modelos previamente entrenados y, al mismo tiempo, puede permitirles obtener la capacidad de cumplir con las instrucciones. whatosoft aiot http://143ai.com
Después de eso, probamos la memoria y la eficiencia computacional del optimizador de 4 bits y los resultados se muestran en la siguiente tabla. En comparación con el optimizador de 8 bits, el optimizador de 4 bits propuesto en este artículo puede ahorrar más memoria, con un ahorro máximo del 57,7% en el experimento de ajuste fino LLaMA-7B. Además, proporcionamos una versión de operador de fusión de 4 bits de AdamW, que puede ahorrar memoria sin afectar la eficiencia computacional. Para la tarea de ajuste de instrucciones de LLaMA-7B, AdamW de 4 bits también aporta efectos de aceleración al entrenamiento debido a la reducción de la presión de la memoria caché. La configuración experimental detallada y los resultados se pueden encontrar en el enlace del artículo. 
Reemplace una línea de código para usar en PyTorch
import lpmm
optimizer = lpmm.optim.AdamW (model.parameters (), lr=1e-3, betas=(0.9, 0.999))
Proporcionamos un optimizador de 4 bits listo para usar. Solo necesita reemplazar el optimizador original con un optimizador de 4 bits. Actualmente, se admiten versiones de baja precisión de Adam y SGD. Al mismo tiempo, también proporcionamos una interfaz para modificar los parámetros de cuantificación para admitir escenarios de uso personalizados.