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トレーニング分野における重要なリソースとなっています。