学習用カーネル

カスタムカーネルは、行列積、アテンション、正規化などの特定のオペレーションに特化して、それらを高速に実行することを目的としています。複数のオペレーションを単一のカーネルに融合（フュージョン）することで、GPU メモリへの読み書き回数を減らし、メモリ帯域幅の消費を低減するとともに、オペレーションごとの起動オーバーヘッドを削減できます。

Hubカーネル

Hubには、KernelConfigで読み込めるコミュニティ製カーネルがホストされています。設定をfrom_pretrained()のkernel_configに渡します。カーネルの読み込みが完了すると、トレーニングで有効になります。利用可能なすべてのオプションについては、カーネルの読み込みガイドを参照してください。

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

use_liger_kernel=True を TrainingArguments に設定すると、対応するモデルのレイヤーに Liger のカーネルが適用されます。

サポートされているモデルの完全なリストについては、パッチングページを参照してください。

from transformers import TrainingArguments

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

パッチを適用するレイヤーを制御するには、liger_kernel_config を dict として渡します。利用可能なオプションはモデルによって異なり、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のガイドを参照し、トレーニングステップ全体のフォワードパスとバックワードパスをコンパイルする方法を学習してください。