Integración y uso de Kernels Hub

La biblioteca kernels permite cargar kernels de cálculo optimizados directamente desde el Hub.
Puede encontrar kernels en organizaciones dedicadas o buscando la etiqueta kernel dentro del Hub.

Los kernels son fragmentos de código optimizados que ayudan en el desarrollo, el entrenamiento y la inferencia de modelos. Aquí nos centraremos en su integración con TRL, pero consulta los recursos anteriores para obtener más información sobre ellos.

Instalación

Para usar kernels con TRL, es necesario instalar la biblioteca en tu entorno de Python:

pip install kernels

Uso de Kernels del Hub en TRL

Los kernels pueden reemplazar directamente las implementaciones de atención, eliminando la necesidad de compilar manualmente backends de atención como Flash Attention y aumentando la velocidad de entrenamiento simplemente obteniendo el kernel de atención correspondiente desde el Hub.

Puede especificar un kernel al cargar un modelo:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2"  # other options: kernels-community/vllm-flash-attn3, kernels-community/paged-attention
)

O al ejecutar un script de entrenamiento de TRL:

python sft.py ... --attn_implementation kernels-community/flash-attn2

O usando la CLI de TRL:

trl sft ... --attn_implementation kernels-community/flash-attn2

Ahora puedes aprovechar backends de atención más rápidos con un kernel preoptimizado para tu configuración de hardware desde el Hub, acelerando tanto el desarrollo como el entrenamiento.

Elegir y fijar una versión del kernel

Los repositorios de kernels en el Hub se versionan como ramas (v1, v2, v3, …), y Transformers selecciona una versión predeterminada para cada repositorio. Esa versión predeterminada puede cambiar de una versión de Transformers a la siguiente, por lo que el mismo valor de attn_implementation no siempre se resuelve en la misma compilación.

Para controlar qué compilación se carga, agregue una revisión al identificador del repositorio, ya sea una rama de versión o un SHA de commit:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2@v2",  # or a commit SHA to pin an exact build
)

Fijar la versión es útil para hacer que una ejecución de entrenamiento sea reproducible y para mantenerse en una compilación conocida y estable cuando una versión más reciente presenta regresiones. Tenga en cuenta que cada versión solo incluye compilaciones para un rango de versiones de Torch y CUDA, por lo que una versión fijada puede no tener una variante para su entorno.

attn_implementation="flash_attention_2" recurre al kernel de Hub kernels-community/flash-attn2 cuando el paquete flash-attn no está instalado, por lo que es posible que esté utilizando un kernel de Hub sin haberlo solicitado explícitamente.

Las compilaciones de v3 para CUDA 12.8 actualmente fallan en la pasada inversa para modelos que utilizan atención de consulta agrupada. Si su entorno utiliza esa versión de CUDA, fije @v2 hasta que se cierre huggingface/kernels-community#1085.

Comparación de implementaciones de atención

Evaluamos varias implementaciones de atención disponibles en transformers, junto con diferentes backends de kernel, utilizando TRL y SFT.
Los experimentos se ejecutaron en una sola H100 GPU con CUDA 12.9, aprovechando Qwen3-8B con un batch size de 8, acumulación de gradientes de 1 y precisión bfloat16.
Tenga en cuenta que los resultados mostrados aquí son específicos de esta configuración y pueden variar con diferentes configuraciones de entrenamiento.

La siguiente figura ilustra tanto la latencia (tiempo por paso de entrenamiento) como la memoria asignada máxima para las diferentes implementaciones de atención y backends de kernel.
Las implementaciones basadas en kernels rinden a la par que la atención instalada de forma personalizada, y aumentar la max_length del modelo mejora aún más el rendimiento. El consumo de memoria es similar en todas las implementaciones, sin diferencias significativas. Obtenemos el mismo rendimiento pero con menos fricción, como se describe en la siguiente sección.

Flash Attention vs. Núcleos de Hub

Compilar Flash Attention desde el código fuente puede llevar mucho tiempo, desde varios minutos hasta horas, dependiendo de tu hardware, la configuración de CUDA/PyTorch y de si hay ruedas precompiladas disponibles.

En cambio, Hugging Face Kernels ofrecen un flujo de trabajo mucho más rápido y fiable. Los desarrolladores no tienen que preocuparse por configuraciones complejas: todo se gestiona automáticamente. En nuestros benchmarks, los kernels estaban listos para usarse en unos 2.5 segundos, sin necesidad de compilación. Esto permite comenzar a entrenar casi de inmediato, acelerando significativamente el desarrollo. Basta con especificar la versión deseada y kernels se encarga del resto.