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训练领域的重要资源。