Kernels para treinamento

Kernels personalizados visam operações específicas, como multiplicações de matrizes, atenção e normalização, para executá-las mais rapidamente. A fusão de múltiplas operações em um único kernel reduz o uso de largura de banda de memória, pois lê e grava na memória da GPU menos vezes, e diminui o overhead de lançamento por operação.

Kernels do Hub

O Hub hospeda kernels da comunidade que você pode carregar com KernelConfig. Passe a configuração para kernel_config em from_pretrained(). Uma vez que o kernel é carregado, ele fica ativo para o treinamento. Leia o guia Carregando kernels para todas as opções disponíveis.

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 funde camadas como RMSNorm, RoPE, SwiGLU, CrossEntropy e FusedLinearCrossEntropy em kernels Triton únicos. É compatível com FlashAttention, FSDP e DeepSpeed, e melhora o throughput de treinamento em múltiplas GPUs, ao mesmo tempo em que reduz o uso de memória, tornando vocabulários maiores, tamanhos de lote maiores e comprimentos de contexto maiores mais viáveis.

pip install liger-kernel

Defina use_liger_kernel=True em TrainingArguments para aplicar os kernels do Liger às camadas correspondentes do modelo.

Consulte a página de patching para obter a lista completa de modelos suportados.

from transformers import TrainingArguments

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

Para controlar quais camadas serão substituídas, passe liger_kernel_config como um dicionário. As opções disponíveis variam conforme o modelo e incluem: 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óximas etapas

• Consulte o guia Backends de atenção para obter detalhes sobre kernels como FlashAttention que reduzem o uso de memória.
• Consulte o guia de torch.compile para aprender a compilar o passo de forward e o passo de backward para todo o seu passo de treinamento.