用于训练的 Kernels

自定义内核针对矩阵乘法、注意力和归一化等特定算子进行优化，以提升其执行速度。将多个算子融合到单个内核中，可以减少对 GPU 内存的读写次数，从而降低内存带宽占用，并削减每个算子的启动开销。

Hub 内核

Hub 托管了社区提供的内核，您可以使用 KernelConfig 进行加载。将该配置作为 kernel_config 参数传入 from_pretrained()。内核加载完成后，即可在训练过程中生效。有关所有可用选项，请参阅加载内核指南。

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 将 RMSNorm、RoPE、SwiGLU、CrossEntropy 和 FusedLinearCrossEntropy 等层融合为单个 Triton 内核。它与 FlashAttention、FSDP 和 DeepSpeed 兼容，在降低内存占用的同时提升多 GPU 训练吞吐量，从而使得更大的词表、批大小和上下文长度更具可行性。

pip install liger-kernel

在 TrainingArguments 中设置 use_liger_kernel=True，即可使用 Liger 的内核修补相应的模型层。

有关支持模型的完整列表，请参阅 patching 页面。

from transformers import TrainingArguments

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

要控制哪些层被修补，请将 liger_kernel_config 作为字典传入。可用选项因模型而异，包括：rope、swiglu、cross_entropy、fused_linear_cross_entropy、rms_norm 等。

from transformers import TrainingArguments

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

后续步骤

有关 FlashAttention 等可降低内存占用的内核的详细信息，请参阅 Attention backends 指南。
请参阅 torch.compile 指南，了解如何编译整个训练步骤的前向和反向传播。

Update on GitHub

Transformers

用于训练的 Kernels

Hub 内核

Liger

后续步骤