LinkedInが開発した効率的なTritonカーネルの集合体で、大規模言語モデルのトレーニングに特化して最適化されており、トレーニング速度を20%向上させ、メモリ使用量を60%削減します。
Liger-Kernel プロジェクト詳細
プロジェクト概要
Liger-Kernelは、LinkedInが開発したTritonカーネルの集合であり、大規模言語モデル(LLM)のトレーニング専用に設計されています。このプロジェクトは、マルチGPUトレーニングのスループットを20%向上させ、メモリ使用量を60%削減できます。プロジェクト名「Liger」は「LinkedIn GPU Efficient Runtime」の略で、その効率的なGPUランタイムというコアコンセプトを体現しています。
コア特性
性能優位性
- トレーニング速度の向上: カーネル融合、インプレース置換、およびチャンク化技術により、マルチGPUトレーニングのスループットが20%向上
- メモリ効率: メモリ使用量が60%削減され、より長いコンテキスト長、より大きなバッチサイズ、および大規模な語彙をサポート
- ポストトレーニング最適化: アライメントおよび蒸留タスク向けに、ポストトレーニングカーネルは最大80%のメモリを節約可能
技術的実現
- 正確な計算: 近似計算なしで、フォワードおよびバックワード伝播は厳密なユニットテスト済み
- 軽量な依存関係: TorchとTritonのみが必要で、追加のライブラリ依存関係なし
- 強力な互換性: すぐに使用可能で、Flash Attention、PyTorch FSDP、およびMicrosoft DeepSpeedと互換性あり
サポートされるモデルと操作
サポートされるモデルアーキテクチャ
プロジェクトは、以下を含むさまざまな主流の大規模言語モデルアーキテクチャをサポートしています。
- LLaMAシリーズ: LLaMA 2、LLaMA 3、LLaMA 3.2-Vision
- Mistralシリーズ: Mistral、Mixtral
- Gemmaシリーズ: Gemma1、Gemma2、Gemma3
- Qwenシリーズ: Qwen2、Qwen2.5、Qwen2-VL、Qwen3など
- その他のモデル: Phi3、Granite、OLMo2、GLM-4など
コアカーネル操作
プロジェクトは、さまざまな最適化されたカーネル操作を実装しています。
基礎カーネル
LigerRMSNorm
: RMS正規化LigerLayerNorm
: 層正規化liger_rotary_pos_emb
: 回転位置エンコーディング(RoPE)LigerSwiGLUMLP
: SwiGLU活性化関数LigerGEGLUMLP
: GeGLU活性化関数LigerCrossEntropyLoss
: 交差エントロピー損失LigerFusedLinearCrossEntropyLoss
: 融合線形交差エントロピー損失
ポストトレーニングカーネル
さまざまなアライメントおよび選好最適化損失関数をサポートします。
LigerFusedLinearDPOLoss
: DPO損失LigerFusedLinearORPOLoss
: ORPO損失LigerFusedLinearCPOLoss
: CPO損失LigerFusedLinearSimPOLoss
: SimPO損失LigerFusedLinearKTOLoss
: KTO損失
使用方法
1. 自動統合方式
from liger_kernel.transformers import AutoLigerKernelForCausalLM
model = AutoLigerKernelForCausalLM.from_pretrained("path/to/some/model")
2. 手動パッチ方式
import transformers
from liger_kernel.transformers import apply_liger_kernel_to_llama
apply_liger_kernel_to_llama()
apply_liger_kernel_to_llama(
rope=True,
swiglu=True,
cross_entropy=True,
fused_linear_cross_entropy=False,
rms_norm=False
)
model = transformers.AutoModelForCausalLM("path/to/llama/model")
3. 低レベルAPI方式
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
import torch.nn as nn
import torch
model = nn.Linear(128, 256).cuda()
loss_fn = LigerFusedLinearCrossEntropyLoss()
input = torch.randn(4, 128, requires_grad=True, device="cuda")
target = torch.randint(256, (4, ), device="cuda")
loss = loss_fn(model.weight, input, target)
loss.backward()
4. ポストトレーニング損失の使用例
from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss
orpo_loss = LigerFusedLinearORPOLoss()
y = orpo_loss(lm_head.weight, x, target)
インストール方法
安定版のインストール
pip install liger-kernel
開発版のインストール
pip install liger-kernel-nightly
ソースコードからのインストール
git clone https://github.com/linkedin/Liger-Kernel.git
cd Liger-Kernel
pip install -e .
開発環境のインストール
pip install -e ".[dev]"
システム要件
NVIDIA GPU環境
torch >= 2.1.2
triton >= 2.3.0
AMD GPU環境
torch >= 2.5.0
triton >= 3.0.0
その他の依存関係
transformers >= 4.x
: transformersモデルパッチAPIを使用する場合
性能ベンチマークテスト
ベンチマークテスト条件:
- モデル:LLaMA 3-8B
- バッチサイズ:8
- データ型:bf16
- オプティマイザ:AdamW
- 勾配チェックポイント:有効
- 分散戦略:FSDP1、8つのA100 GPU
テスト結果:
- Hugging Faceモデルは4Kコンテキスト長でメモリ不足になり始めましたが、Hugging Face + Liger Kernelは16Kまで拡張可能
- トレーニングスループットが20%以上向上
- メモリ使用量が60%削減
フレームワーク統合
Liger-Kernelは、複数の主流トレーニングフレームワークに統合されています。
- Axolotl
- LLaMa-Factory
- SFTTrainer
- Hugging Face Trainer
- SWIFT
- oumi
技術原理
カーネル融合技術
複数の操作を単一のカーネルに融合することにより、GPUメモリへのアクセス回数を減らし、計算効率を向上させます。
チャンク化計算
メモリ集約的な操作の場合、チャンク処理技術を採用し、大規模な計算を小さなチャンクに分解して、ピーク時のメモリ使用量を削減します。
インプレース操作
可能な限りインプレース操作を使用し、追加のメモリ割り当てを回避して、メモリ効率をさらに最適化します。
まとめ
Liger-Kernelは、大規模言語モデルのトレーニング最適化における重要な進歩を表しています。慎重に設計されたTritonカーネル、メモリ最適化技術、および幅広いモデルサポートを通じて、研究者とエンジニアに強力で使いやすいツールを提供し、トレーニング効率を大幅に向上させ、計算コストを削減できます。プロジェクトのオープンソースの性質と活発なコミュニティサポートにより、LLMトレーニング分野における重要なリソースとなっています。