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 訓練領域的重要資源。