Kernels para el entrenamiento

Los kernels personalizados se dirigen a operaciones específicas, como multiplicaciones de matrices, atención y normalización, para ejecutarlas más rápido. Fusionar varias operaciones en un solo kernel reduce el uso de ancho de banda de memoria al leer y escribir la memoria de la GPU menos veces, y elimina el sobrecoste de lanzamiento por operación.

Kernels de Hub

El Hub aloja kernels de la comunidad que puedes cargar con KernelConfig. Pasa la configuración a kernel_config en from_pretrained(). Una vez que el kernel se carga, queda activo para el entrenamiento. Consulta la guía Carga de kernels para ver todas las opciones disponibles.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": "kernels-community/liger-kernels:LigerRMSNorm",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
)

Liger

Liger Kernel fusiona capas como RMSNorm, RoPE, SwiGLU, CrossEntropy y FusedLinearCrossEntropy en kernels de Triton individuales. Es compatible con FlashAttention, FSDP y DeepSpeed, y mejora el rendimiento del entrenamiento multi-GPU al tiempo que reduce el uso de memoria, lo que hace más viables vocabularios, tamaños de lote y longitudes de contexto más grandes.

pip install liger-kernel

Establezca use_liger_kernel=True en TrainingArguments para parchear las capas del modelo correspondientes con los kernels de Liger.

Consulte la página de patching para obtener una lista completa de los modelos compatibles.

from transformers import TrainingArguments

training_args = TrainingArguments(
    ...,
    use_liger_kernel=True
)

Para controlar qué capas se parchean, pase liger_kernel_config como un diccionario. Las opciones disponibles varían según el modelo e incluyen: rope, swiglu, cross_entropy, fused_linear_cross_entropy, rms_norm, etc.

from transformers import TrainingArguments

training_args = TrainingArguments(
    ...,
    use_liger_kernel=True,
    liger_kernel_config={
        "rope": True,
        "cross_entropy": True,
        "rms_norm": False,
        "swiglu": True,
    }
)

Próximos pasos

• Consulte la guía de Backends de atención para obtener detalles sobre kernels como FlashAttention que reducen el uso de memoria.
• Consulte la guía de torch.compile para aprender a compilar el paso hacia adelante y el paso hacia atrás de su paso de entrenamiento completo.