Métodos de entrenamiento de baja precisión

La aparición de nuevos tipos de hardware dio lugar a nuevos paradigmas de entrenamiento que los aprovechan mejor. Actualmente, esto se manifiesta en el entrenamiento con precisión de 8 bits mediante paquetes como TransformersEngine (TE) o torchao (FP8 nativo de PyTorch).

Para una introducción a los temas tratados hoy, recomendamos revisar la guía de uso de baja precisión, ya que esta documentación hará referencia a ella con regularidad.

Un gráfico rápido

A continuación se muestra un gráfico rápido que muestra las diferentes precisiones de bits para cada solución durante el entrenamiento:

TransformersEngine

TransformersEngine es la primera solución que intenta entrenar en punto flotante de 8 bits. Funciona mediante el uso de capas de reemplazo directo para ciertas capas de un modelo, aprovechando su motor FP8 para reducir el número de bits (por ejemplo, de 32 a 8) sin degradar la precisión final del modelo.

Específicamente, Accelerate encontrará y reemplazará las siguientes capas con versiones de TransformersEngine:

• nn.LayerNorm para te.LayerNorm
• nn.Linear para te.Linear

Como resultado, obtenemos un modelo que tiene la mayoría de sus capas en BF16, mientras que algunas capas están en FP8, lo que reduce parte de la memoria.

De forma anecdótica, hemos observado que las mejoras de rendimiento no se hacen evidentes al usar TransformerEngine hasta que una gran mayoría de las capas del modelo están compuestas por esas dos capas a reemplazar. Como resultado, solo los modelos más grandes han mostrado mejoras de rendimiento cuando el número de parámetros es de varios miles de millones o más.

El TransformerEngine puede recibir muchos argumentos diferentes que personalizan cómo realiza los cálculos en FP8 y qué hacen. A continuación se muestra una lista completa de los argumentos:

• margin: El margen a utilizar para la escalado del gradiente.
• interval: El intervalo a utilizar para la frecuencia con la que se recalcula el factor de escalado.
• fp8_format``: The format to use for the FP8 recipe. Must be one ofHYBRIDorE4M3. (GenerallyHYBRIDfor training,E4M3` for evaluation)
• amax_history_len: La longitud del historial que se utilizará para el cálculo del factor de escalado
• amax_compute_algo: El algoritmo que se utilizará para el cálculo del factor de escalado. Debe ser uno de max o most_recent.
• override_linear_precision: Si se ejecutan o no los GEMM de fprop, dgrad y wgrad en mayor precisión.

Puedes personalizar cada uno de estos elementos como parte de utils.TERecipeKwargs para ayudar a optimizar el rendimiento de tus modelos.

Si observamos el gráfico mencionado anteriormente, TE simplemente convierte las capas de cómputo a FP8, mientras que todo lo demás permanece en FP32. Como resultado, esto termina utilizando la mayor cantidad de memoria, pero lo hace con la ventaja de garantizar la menor pérdida de precisión final durante el entrenamiento.