Unsloth afirma entrenamiento de LLM 5 veces más rápido con nuevos kernels Triton y empaquetado automático

Unsloth, la biblioteca de ajuste fino de código abierto popular entre los profesionales de LLM, ha lanzado una actualización importante que afirma poder ofrecer velocidades de entrenamiento hasta 5 veces más rápidas mientras reduce el uso de VRAM entre un 30% y un 90%. Las mejoras provienen de un conjunto de kernels Triton personalizados y un nuevo sistema de empaquetado automático inteligente que el equipo dice que funciona con cualquier backend de atención, FlashAttention 3, xFormers o SDPA, y cualquier familia de modelos.

La adición principal es un kernel fusionado QK RoPE. Anteriormente, cada consulta y clave necesitaba dos kernels Triton separados; la nueva implementación los fusiona en uno solo, permitiendo RoPE de longitud variable. En microbenchmarks, el kernel fusionado es 2.3 veces más rápido en longitudes de contexto más largas y 1.9 veces más rápido en las más cortas. El kernel es completamente in situ, eliminando operaciones de clonación y transposición contigua, reduciendo aún más la presión sobre la memoria de la GPU.

Unsloth también abordó una limitación de larga data para el entrenamiento de contexto largo: errores de desbordamiento de CUDA durante ejecuciones de 500K tokens causados por la indexación predeterminada int32 en Triton y CUDA. La solución introduce una constante de tiempo de compilación LONG_INDEXING que permite al compilador Triton especializar kernels para contextos más cortos y más largos sin sobrecarga.

“Cambiando la proporción de 50% secuencias cortas, y suponiendo que tenemos más secuencias cortas, por ejemplo, 20% secuencias largas y 80% secuencias cortas, obtenemos un entrenamiento 5 veces más rápido”, escribió el equipo de Unsloth en el comunicado. “Esto significa que la aceleración del empaquetado depende de cuántas filas cortas tenga tu conjunto de datos; cuantas más cortas, más rápido”.

El sistema de empaquetado automático (habilitado por defecto) agrupa automáticamente múltiples secuencias cortas en un solo tensor unidimensional, manteniendo metadatos de longitud de secuencia para evitar fugas de atención entre muestras. El resultado: una mejora de rendimiento reportada de 1.7 a 3 veces en una ejecución de ajuste fino de Qwen3-32B en el conjunto de datos yahma/alpaca-cleaned con tamaños de lote variables, con algunas configuraciones viendo 5 veces o más. Unsloth enfatiza que las pérdidas de entrenamiento coinciden exactamente con las ejecuciones sin empaquetado.

Los resultados de referencia publicados por el equipo muestran que, sin empaquetado, los tokens de relleno pueden representar casi el 50% de un lote con un tamaño de lote de 8. Con empaquetado, la utilización de tokens válidos se mantiene cerca del 100% independientemente del tamaño del lote. El uso de memoria para Qwen3-8B y Qwen3-32B disminuyó en un 60% mientras que el rendimiento se duplicó.

La actualización funciona para todos los métodos de entrenamiento, ajuste fino completo, preentrenamiento y RL, y es compatible con versiones anteriores de GPUs como Tesla T4 y RTX 2080, así como con las más nuevas como H100 y B200. Los usuarios solo necesitan actualizar Unsloth para tener procesamiento por lotes sin relleno por defecto; el empaquetado explícito se puede habilitar con packing = True en el SFTConfig.

El lanzamiento se basa en el trabajo anterior de Unsloth sobre kernels RoPE de diciembre de 2023 y un PR de la comunidad en marzo de 2024 que hizo que el entrenamiento de extremo a extremo fuera un 1-2% más rápido al optimizar el kernel para lanzar bloques para grupos de cabezas.

Unsloth, que se ha posicionado como una alternativa ligera al Hugging Face Trainer completo para ajuste fino, ahora afirma ser compatible con modelos que incluyen la familia Qwen3, Llama 3.1 y otros. Los cuadernos de la biblioteca en Google Colab se han actualizado automáticamente.