Liger Kernel(ライガーカーネル)
Hugging Faceが2024年に公開したLLM学習向けCUDA融合カーネル集。CrossEntropy・RMSNorm・RoPE・SwiGLUなどをメモリ効率最適化実装し、学習スループットを向上させる。
Liger Kernelとは
Hugging Faceが2024年9月にOSSとして公開した、LLM学習の主要演算を最適化するCUDA融合カーネル(fused kernel)集。Tritonで実装されており、PyTorchの標準実装と完全互換でモデルコードを変更することなくapplyできる。名前はLLMとTigerを組み合わせた造語。
背景: LLM学習の演算ボトルネック
大規模LLMの学習では以下の演算が繰り返し実行される。
- RMSNorm(各Transformer layerの正規化)
- RoPE(Rotary Position Embedding、位置情報付与)
- SwiGLU/GeGLU(GLU系活性化関数)
- CrossEntropy Loss(語彙サイズが大きいほど計算コスト増大)
PyTorchの標準実装ではこれらを個別のCUDAカーネルとして実行するため、中間結果のHBM書き込み・読み込みが多発する。Liger KernelはこれらをTritonで融合(fuse)することでHBMアクセスを大幅削減する。
実装された主要カーネル
Chunked CrossEntropy(最大の効果)
大語彙LLMでは語彙サイズVが128k〜150kに達する。CrossEntropyの前段logits行列はshape=(batch×seq_len, V)で、Llama-3 8Bの場合4096×128256のfloat32行列が生成される。これだけで数十GBのVRAMを消費する。
Liger KernelのChunked CrossEntropyはlogitsをチャンク単位で計算し、中間行列をHBMに保存しない。VRAMピーク消費を最大60%削減する。
Fused RMSNorm
RMSNorm(Root Mean Square Normalization)は各Transformer layerで実行される。Liger KernelはForward + Backward + weight gradientを1カーネルに融合し、HBMアクセスを削減。
Fused RoPE
Rotary Position EmbeddingのForward/Backwardを融合カーネルで実装。特に長いシーケンス長で効果が大きい。
Fused SwiGLU
SwiGLU(Swish × GLU)はLlama/Mistral系で標準的な活性化関数。Gate・Up・Downの3線形変換を含む処理を融合実装。
インストールと適用
pip install liger-kernel
from liger_kernel.transformers import apply_liger_kernel_to_llama
# モデルロード前に1行追加するだけ
apply_liger_kernel_to_llama() # monkey patchでLlama実装を置換
from transformers import LlamaForCausalLM
model = LlamaForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
Mistral・Gemma・Phi・Qwen・Mixtral など主要アーキテクチャに対応したapply関数が用意されている。
効果(公式ベンチマーク)
Llama-3 8B、A100 80GB、seq_len=4096、batch=4での測定:
| 指標 | PyTorch標準 | Liger Kernel | 改善率 |
|---|---|---|---|
| GPU Memory | 33.1 GB | 17.6 GB | -47% |
| Training throughput | 3,220 tok/s | 4,040 tok/s | +25% |
| Convergence | 基準 | 同等 | 精度劣化なし |
Axolotl・TRLとの統合
Axolotlではuse_liger_kernel: trueの1行で有効化可能。TRL SFTTrainerでもuse_liger_kernel=Trueで対応。主要ファインチューニングフレームワークでのサポートが進んでいる。
対応アーキテクチャ(2025年時点)
Llama 2/3/3.1/3.2・Mistral・Mixtral・Gemma/Gemma 2・Phi 3/3.5・Qwen 2/2.5・ChatGLM・LlaVA(マルチモーダル)等。HuggingFace公式のtransformersリポジトリとの密接な連携で継続的に拡張されている。
まとめ
Liger KernelはTritonで実装されたLLM学習向けCUDA融合カーネル集であり、VRAMを最大47%削減しスループットを25%向上させる。モデルコードの変更不要でmonkey patchにより透過的に適用でき、Axolotl・TRLなど主要フレームワークと統合されている。特にChunked CrossEntropyによる大語彙でのVRAMピーク削減効果が大きく、8B〜70BクラスのLLMファインチューニングに実用価値が高い。