在LLM的训练/微调过程中,开发者通常会遇到一些瓶颈,包括GPU显存不够,经常遇到OOM,GPU使用率100%而且非常慢等。
Liger Kernel是专为 LLM 训练设计的 Triton 内核集合。它可以有效地将多 GPU训练吞吐量提高 20%,并将内存使用量降低 60%。
Liger Kernel为什么选择Triton:
-Triton编程更简单:比 CUDA 更快地完成内核开发。
-用 Numpy 思考:用vector(块)而不是Element(线程)思考。
-与 AI 研究人员的协作更简单:AI 研究人员可以轻松掌握和扩展。
-Python 原生:一个内核不需要五种不同的文件类型。
-干净的依赖关系:Triton 在大多数情况下都可以正常工作。
项目地址:
https://github.com/linkedin/Liger-Kernel
使用 Liger Kernel 增强模型
只需如上的一行代码,Liger Kernel 就能将吞吐量提高 20% 以上,并将内存使用量降低 60%,从而实现更长的上下文长度、更大的批次大小和海量词汇。
Liger Kernel 主要特点
-易于使用:只需用一行代码使用Liger Kernel SDK加速内核来增强模型。
-时间和内存效率高:与 Flash-Attn 秉承同样的精神,但适用于RMNSNorm、RoPE、SwiGLU和CrossEntropy等层!通过内核融合、In-place 替换和分块技术,可将多 GPU 训练吞吐量提高 20%,并将内存使用量降低多达 60% 。
-精确:计算精确——无近似值!前向和后向传递均通过严格的单元测试实现,并在没有 Liger Kernel 的情况下针对训练运行进行收敛测试,以确保准确性。
-轻量级:Liger Kernel 的依赖性极小,只需要 Torch 和 Triton — 无需额外的库!告别依赖性烦恼!
-支持多 GPU:兼容多 GPU 设置(PyTorch FSDP、DeepSpeed、DDP 等)。
-训练器框架集成:Axolotl、LLaMa-Factory、SFTTrainer、Hugging Face Trainer、MS-Swift。
最佳实践
基于Huggingface Trainer做Qwen2全参数微调(双卡A100)
训练脚本:
from dataclasses import dataclass import datasets import torch import transformers from callback import EfficiencyCallback from trl import DataCollatorForCompletionOnlyLM, SFTTrainer from modelscope.msdatasets import MsDataset from liger_kernel.transformers import AutoLigerKernelForCausalLM @dataclass class CustomArguments: model_name: str = "meta-llama/Meta-Llama-3-8B" dataset: str = "OmniData/alpaca" max_seq_length: int = 512 use_liger: bool = False def formatting_prompts_func(example): return example["text"] def train(): parser = transformers.HfArgumentParser( (transformers.TrainingArguments, CustomArguments) ) training_args, custom_args = parser.parse_args_into_dataclasses() tokenizer = transformers.AutoTokenizer.from_pretrained( custom_args.model_name, padding_side="left", truncation_side="left", ) tokenizer.pad_token = tokenizer.eos_token dataset = MsDataset.load(custom_args.dataset, subset_name='default', split='train').train_test_split( test_size=0.1 ) train_dataset = dataset["train"] eval_dataset = dataset["test"] response_prompt = tokenizer.encode("### Response:\n", add_special_tokens=False) collator = DataCollatorForCompletionOnlyLM( tokenizer=tokenizer, response_template=response_prompt, pad_to_multiple_of=16, ) if custom_args.use_liger: model = AutoLigerKernelForCausalLM.from_pretrained( custom_args.model_name, trust_remote_code=True, use_cache=False, torch_dtype=torch.bfloat16, # These args will get passed to the appropriate apply_liger_kernel_to_* function # to override the default settings # cross_entropy=True, # fused_linear_cross_entropy=False, ) else: model = transformers.AutoModelForCausalLM.from_pretrained( custom_args.model_name, trust_remote_code=True, use_cache=False, torch_dtype=torch.bfloat16, ) trainer = SFTTrainer( model=model, args=training_args, data_collator=collator, max_seq_length=custom_args.max_seq_length, train_dataset=train_dataset, eval_dataset=eval_dataset, formatting_func=formatting_prompts_func, callbacks=[EfficiencyCallback()], ) trainer.train() if __name__ == "__main__": train()
训练参数:设置use_liger为 True
torchrun --nnodes=1 --nproc-per-node=2 training.py \ --model_name "/mnt/workspace/cherry/Qwen2-1.5B" \ --bf16 \ --num_train_epochs 1 \ --per_device_train_batch_size 8 \ --per_device_eval_batch_size 8 \ --eval_strategy "no" \ --save_strategy "no" \ --learning_rate 6e-6 \ --weight_decay 0.05 \ --warmup_ratio 0.1 \ --lr_scheduler_type "cosine" \ --logging_steps 1 \ --include_num_input_tokens_seen \ --report_to none \ --fsdp "full_shard auto_wrap" \ --fsdp_config config/fsdp_config.json \ --seed 42 \ --use_liger True \ --output_dir alpaca_finetuning
性能对比,显存节约49%:
- 使用liger的total_peak_memory_reserved_MB: 19322.0
- 不使用liger的total_peak_memory_reserved_MB: 37314.0,
llama-factory支持-仅需9G显存(单卡A10)
在魔搭社区免费算力使用llama-factory+liger微调Qwen2
安装依赖
!git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git %cd LLaMA-Factory !pip install -e . !pip install transformers -U !pip install liger-kernel
加载数据集
import json %cd /mnt/workspace/LLaMA-Factory/ NAME = "Qwen-2" AUTHOR = "LLaMA Factory" with open("data/identity.json", "r", encoding="utf-8") as f: dataset = json.load(f) for sample in dataset: sample["output"] = sample["output"].replace("{{"+ "name" + "}}", NAME).replace("{{"+ "author" + "}}", AUTHOR) with open("data/identity.json", "w", encoding="utf-8") as f: json.dump(dataset, f, indent=2, ensure_ascii=False)
微调模型
import json from modelscope import snapshot_download model_dir = snapshot_download("LLM-Research/Qwen2-7B-Instruct-bnb-4bit") args = dict( stage="sft", # do supervised fine-tuning do_train=True, model_name_or_path=model_dir, # use bnb-4bit-quantized Qwen2-7B-4bit model dataset="identity,alpaca_en_demo", # use alpaca and identity datasets template="qwen", # use llama3 prompt template finetuning_type="lora", # use LoRA adapters to save memory lora_target="all", # attach LoRA adapters to all linear layers output_dir="qwen2_lora", # the path to save LoRA adapters per_device_train_batch_size=2, # the batch size gradient_accumulation_steps=4, # the gradient accumulation steps lr_scheduler_type="cosine", # use cosine learning rate scheduler logging_steps=10, # log every 10 steps warmup_ratio=0.1, # use warmup scheduler save_steps=1000, # save checkpoint every 1000 steps learning_rate=5e-5, # the learning rate num_train_epochs=3.0, # the epochs of training max_samples=500, # use 500 examples in each dataset max_grad_norm=1.0, # clip gradient norm to 1.0 loraplus_lr_ratio=16.0, # use LoRA+ algorithm with lambda=16.0 fp16=True, # use float16 mixed precision training enable_liger_kernel=True, # use liger kernel for efficient training ) json.dump(args, open("train_qwen2.json", "w", encoding="utf-8"), indent=2) %cd /mnt/workspace/LLaMA-Factory !llamafactory-cli train train_qwen2.json
性能对比,显存节约24%
使用liger
不使用liger
对比图
魔搭的SWIFT框架也已经集成了Liger Kernel,首先确保swift和liger-kernel都安装了:
pip install ms-swift -U pip install liger-kernel
通过如下swift命令行,就可以启用Liger-Kernel,使用aplaca-en数据集来微调qwen2-7b-instruct。在这个例子里,我们采用的是非量化版本的qwen2-7b模型。同时为了保证训练过程中显存的稳定性,命令行里指定了packing参数,确保每个训练sequence都是等长的。
# 为了保证显存的稳定性,我们使用了--packing参数,保证每个sequence都是等长的 # --use_liger true/false 即可切换使用Liger Kernel swift sft \ --model_type qwen2-7b-instruct \ --dataset alpaca-en#5000 \ --use_liger true \ --sft_type lora \ --target_modules ALL \ --learning_rate 1e-5 \ --packing true \ --max_length 2048
使用Liger Kernel的显存占用情况:
不使用Liger Kernle的显存占用情况:
对比图
可以看到对于非量化版本的qwen2-7b训练,Liger Kernel同样可以提供稳定的显存节省。