Reducción del tiempo de entrenamiento de Qwen en 6x en AI Runtime

El entrenamiento o ajuste fino (finetuning) de un modelo base de varios miles de millones de parámetros requiere un cómputo significativo, que a menudo abarca de cientos a miles de horas de entrenamiento en GPU. Sin embargo, optimizar el rendimiento del entrenamiento para un LLM es todo un desafío, y las métricas de alto nivel, como la utilización de la GPU, podrían no reflejar la utilización real. En esta publicación de blog, mostramos cómo atravesar las capas de abstracción y acelerar 6 veces el ajuste fino del modelo Qwen 3.8 27B más reciente en AI Runtime. Mediante una optimización sistemática del rendimiento, incluso cuando la utilización de la GPU ya estaba en el 98 %, pasamos de 7 horas y 44 minutos a 1 hora y 17 minutos al realizar el ajuste fino de un conjunto de datos fijo.

Profundizaremos en las técnicas detrás de estas mejoras, desde cómo medir lo que importa hasta los detalles de implementación y los compromisos de diseño que implica cada optimización.

Si quieres empezar de inmediato, ve directamente al código de ejemplo en AI Runtime y pruébalo.

Configuración del modelo

Comenzamos con un bucle de entrenamiento estándar para Qwen 3.8 27B utilizando la librería transformers y el conjunto de datos de entrenamiento SFT HuggingFaceH4/ultrachat_200k.

• Aplicamos FSDP2, la técnica más reciente de fragmentación de modelos (sharding) de torch. Esto es necesario para hacer caber los pesos del modelo y el estado del optimizador en las GPU. Un enfoque puramente paralelo de datos distribuidos no lograría ajustar los pesos del modelo, con este tamaño, en una sola GPU.
• También habilitamos el punto de control de activación (activation checkpointing), una técnica que intercambia espacio de memoria por cómputo en la pasada hacia atrás (backwards pass) para permitir un tamaño de lote más grande por GPU.
• El dataloader está empaquetando lotes 2D con un máximo de 8192 tókenes por lote, aplicando relleno (padding) hasta el tamaño máximo del ejemplo de los lotes.

Puntos de referencia de entrenamiento iniciales

Nuestro entrenamiento de referencia del modelo sobre el conjunto de datos completo tarda 7 horas y 44 minutos, y durante el entrenamiento la GPU muestra un 98 % de utilización promedio. El resultado es un entrenamiento de 9,396 tok/s.

Esta utilización de la GPU parece alta, pero como se mencionó anteriormente, ese número es engañoso. La métrica estándar de utilización de la GPU solo cuenta el porcentaje de tiempo que la GPU estuvo haciendo algo, por lo que no nos dice qué tan bien estaba usando el modelo la GPU mientras estaba activa. Para entender esto, necesitamos mirar más a fondo.

Una forma de hacer esto es el mecanismo de traza de torch; podemos ajustar nuestro bucle de entrenamiento de la siguiente manera (al tiempo que limitamos el número de pasos para obtener una respuesta rápida):

Para un paso de muestra, esto nos da los siguientes datos:

Tabla 1: CPU: 14,4 s con 1 714 382 llamadas. La tabla contiene las operaciones principales por tiempo total.

Tabla 2: GPU: 10,3 s con 269 448 operaciones. La tabla contiene las operaciones principales por tiempo total

Áreas de mejora

A partir de estas tablas, notamos algunos problemas de alto nivel con el modelo actual:

• La CPU tarda más en configurar cada paso de entrenamiento de lo que la GPU tarda en ejecutarlo. Esto es inesperado ya que, por lo general, el lado de la CPU de la ejecución en la GPU es asíncrono, es decir, la CPU debería ejecutarse bastante por delante de la GPU. El efecto principal de esto es que la GPU se detiene, por lo que tenemos cierta oportunidad de mejorar esto.
• ¡Hay demasiadas operaciones! En las tablas anteriores, tenemos cientos de miles de operaciones pequeñas; en este caso, el problema es que la CPU no puede encolar operaciones de la GPU lo suficientemente rápido para mantener ocupada la GPU. La idea clave es que debemos reducir la cantidad de operaciones que la CPU encola hacia la GPU. Existen algunas técnicas generales para esto, una se adapta a nuestra situación:
  • Los gráficos CUDA de torch registran los kernels CUDA observados y los reproducen con una sola llamada a la CPU. Aunque es potente, tiene muchas limitaciones y entra en conflicto con el dinamismo de FSDP2. Por lo tanto, por ahora, esto está descartado.
  • Torch.compile rastrea el código y usa Triton para fusionar automáticamente kernels cuando sea posible. En este ejemplo, usaremos kernels fusionados manualmente y, en el caso de los LLM, esto no deja margen de mejora para torch.compile.
  • Los kernels fusionados manualmente utilizan técnicas de fusión especializadas para operaciones particulares que el compilador de torch no puede aplicar por sí solo. Afortunadamente, ya existen muchos kernels preconstruidos establecidos en la comunidad. Los ejemplos de esto incluyen ‘Flash Attention’, que proporciona una pequeña mejora en el rendimiento en este caso.

Además, tenemos un problema con el exceso de relleno (padding) en el lote: el tensor de entrada debe ser rectangular y cada fila se rellena hasta el número máximo de tókenes para ese lote. Mostraremos un truco para los LLM que se puede usar para minimizarlo.

En las siguientes secciones, profundizaremos en cada tema técnico, cubriendo qué es, las mejores prácticas recomendadas y las principales ventajas y desventajas a considerar al aplicarlo.

Kernels fusionados

¡Es posible que hayas notado que algunos kernels se ejecutan muchísimo! En particular, para Qwen, el kernel de multiplicación de matrices en lotes se ejecuta más de 30 000 veces en un solo paso.

En este caso, primero exploraremos los kernels fusionados. Un kernel fusionado toma varios kernels de GPU más pequeños y los combina en una sola unidad. Esto puede mejorar el rendimiento al evitar transferencias costosas entre la memoria de alto ancho de banda de una GPU y su memoria en chip de ancho de banda aún mayor. Asimismo, también puede permitir una programación más eficiente de los núcleos de la GPU.

Con tantos kernels adicionales, aunque torch.compile podría funcionar, llevará mucho tiempo ejecutarse. Por lo tanto, reducir la cantidad directa de ops con kernels fusionados manualmente es un buen primer paso. Además, fusiones bien conocidas como Flash Attention no son algo que torch.compile encontraría por sí mismo.

La biblioteca transformers utiliza de forma predeterminada el SDPA de torch para las capas de atención estándar. Sin embargo, resulta que muchas de estas pequeñas operaciones provienen de las capas no estándar de gated delta network. Tenemos suerte de que para estas capas existen dos kernels, causal-conv1d y flash-linear-attention, que se escribieron para fusionar muchas de estas operaciones por capa. ¿Cómo descubrimos esto? La implementación de la biblioteca transformers de Qwen ya nos muestra los kernels que admite:

• flash-linear-attention
• causal-conv1d

Hay dos métodos de alto nivel para obtener estos kernels:

• Instalar los kernels por tu cuenta. Esto puede ir desde:
  • Un simple pip install (por ejemplo, para flash-linear-attention) desde PyPi
  • Un git clone y la compilación del repositorio correspondiente
  • Buscar otra fuente para un paquete precompilado en caso de que la compilación tome demasiado tiempo
• Usar Kernel Hub de Hugging Face para obtener kernels precompilados.

Por lo general, recomendamos la primera opción, especialmente combinada con nuestro soporte de contenedores de Docker: crea tu contenedor con todos tus kernels instalados en el momento de la compilación para reducir el tiempo que tarda tu trabajo en iniciarse.

Aunque kernelhub es conveniente, dado que los kernels a menudo dependen explícitamente de la versión instalada de torch y cuda, es posible que no haya un kernel disponible para tu configuración específica y el modelo de transformers recurrirá silenciosamente a un kernel no optimizado con bajo rendimiento.

Después de instalar los paquetes de Python flash-linear-attention y causal-conv1d, mejoramos nuestro tiempo de extremo a extremo en un factor de 3.7x, bajando a 2 horas y 5 minutos. En las siguientes tablas, puedes ver los nuevos kernels (resaltados), que ocupan el lugar de las funciones elementwise de las tablas anteriores. Además, el número total de operaciones ejecutadas se ha reducido enormemente, en un factor de 15x:

Tabla 3: CPU: 3.9 s con 118,316 llamadas. La tabla contiene las principales operaciones por tiempo total.

Tabla 4: GPU: 3.9 s con 18,060 llamadas. La tabla contiene las principales operaciones por tiempo total.

Advertencia: durante el entrenamiento necesitarás un kernel tanto para el pase forward como para el backward, y algunos kernels son solo forward, ya que solo se han implementado para acelerar la inferencia.

Conclusiones clave

• Considera usar kernels bien establecidos y altamente optimizados que estén fácilmente disponibles para modelos comunes. Aprovechar estos kernels puede ofrecer mejoras significativas de rendimiento con relativamente poco esfuerzo. En este caso específico vimos una mejora de 3.7x en la velocidad de entrenamiento, reduciendo nuestro tiempo de entrenamiento de 7 horas y 44 minutos a 2 horas y 5 minutos.

Batching, tokenización y padding

Nuestro modelo Qwen ahora es bastante rápido, pero podemos hacerlo aún más rápido. Por defecto, nuestra configuración de batching utiliza tensores rectangulares en 2D de 8,192 tokens, ya que es el tamaño más grande que podemos ajustar sin quedarnos sin memoria. La desventaja de este enfoque es que tendemos a desperdiciar una cantidad considerable de espacio en el tensor con padding.

Si Qwen solo utilizara mecanismos de atención estándar, el NestedTensor de torch ahora es muy bien compatible y ayudaría a reducir este costo al eliminar por completo el padding. Pero flash-linear-attention solo admite un enfoque de tensor empaquetado en 1D más tradicional:

Pero no podemos hacer esto de manera arbitraria; pasar simplemente el tensor en 1D al modelo puede ser aceptable para el preentrenamiento, pero para el ajuste fino (finetuning), debemos asegurarnos de que, dentro de un ejemplo, los tokens solo atiendan a tokens del mismo ejemplo. Para solucionar esto, debemos pasar los position_ids de ejemplo correctos al modelo. A su vez, el modelo los pasará a los kernels de atención fusionados y este aplicará la máscara de atención correcta. Si observas el código de ejemplo aquí, puedes ver cómo calculamos las entradas necesarias para que esto funcione.

Con este cambio de apariencia sencilla, nuestro tiempo de entrenamiento pasa de 2 horas y 5 minutos a 1 hora y 17 minutos.

Esto se debe a dos razones:

1. La razón principal es que este empaquetado aumenta la cantidad de secuencias por lote y reduce la cantidad total de lotes en aproximadamente un 36%. Ten en cuenta que esto depende de la distribución y la longitud de los ejemplos en tu conjunto de datos.
2. Una razón secundaria es que los kernels de flash linear attention deben compilarse justo a tiempo (just-in-time) para cada forma de entrada y la forma antes era dinámica, lo que provocaba unos minutos adicionales de compilación.

Compromiso (trade-off): Este tipo de cambio hace que la canalización de entrada sea más compleja y existe una alta probabilidad de introducir errores que afecten la calidad del modelo. Pero en este caso, el ahorro del 36% vale la pena. Existen otras técnicas para reducir el padding, como mantener el mecanismo de batching original pero ordenar todas las entradas por longitud para minimizar el padding, pero este orden puede introducir correlaciones inesperadas en muestras adyacentes, lo que sesga el modelo.

Conclusiones clave: Al entrenar LLM en un conjunto de datos con entradas más cortas que el máximo de tokens entrenables (por ejemplo, en este caso 8192), empaquetar los datos de manera eficiente en lotes puede ahorrar una cantidad significativa de tiempo. Para este conjunto de datos, vimos una mejora de 1.6x en la velocidad de entrenamiento, reduciendo nuestro tiempo de entrenamiento de 2 horas y 5 minutos a 1 hora y 17 minutos.

Conclusiones

En este blog, te mostramos cómo multiplicar por 6 la eficiencia de tu entrenamiento. Se redujo a (1) primero medir lo que importa, no solo la utilización de la GPU, sino también en qué parte el modelo pasa tiempo entrenando, (2) optimizar y fusionar kernels, utilizando bibliotecas conocidas como Flash Linear Attention, y (3) adaptar cómo se alimentan los lotes de datos al modelo.

Para hacer esto a escala, necesitas una plataforma que haga de la ruta de alto rendimiento la opción predeterminada. En AIR, muchas de estas mejoras de rendimiento vienen integradas. Incluimos las últimas versiones de CUDA y flash attention en nuestros entornos base, rastreamos automáticamente métricas clave como el uso de la GPU, la temperatura y su estado por ti, y nos aseguramos de que las versiones de torch (y los perfiladores) estén actualizadas. Nuestro agente interno, Genie Code, también está ajustado para buscar oportunidades como estas y optimizar aún más tu entrenamiento. Además, la lectura de datos desde Databricks está optimizada a través de nuestro UCVolumeDataset.

¡Prueba AIR para tu próxima carga de trabajo de entrenamiento! ¿Listo para probarlo? Consulta el recetario de AI Runtime, con detalles sobre cómo ejecutar este modelo tú mismo.

(Esta entrada del blog ha sido traducida utilizando herramientas basadas en inteligencia artificial) Publicación original