Liger kernel训练加速，一行代码训练吞吐量提高 20%，显存使用量降低 60%

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()

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

!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

pip install ms-swift -U
pip install liger-kernel

# 为了保证显存的稳定性，我们使用了--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