Unsloth 高效微调 — 比普通 LoRA 快 2-5 倍、省 80% 显存的微调神器¶

# 方式1：pip 安装（最简单）
pip install unsloth                          # 自动检测 CUDA 版本

# 方式2：uv 安装（官方推荐，最快）
pip install uv                               # 先装 uv 包管理器
uv venv unsloth_env --python 3.11           # 创建虚拟环境
source unsloth_env/bin/activate             # 激活环境
uv pip install unsloth --torch-backend=auto # 自动选最优 torch 版本

# 更新到最新版（重要，Unsloth 更新频率极高）
pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth
pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth_zoo

# 验证安装
python -c "import unsloth; print(unsloth.__version__)"

from unsloth import FastLanguageModel                # 导入 Unsloth 的快速模型加载器
import torch

# 第一步：加载模型（比 transformers 快很多）
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Llama-3.2-8B-Instruct",  # 模型名（支持 HuggingFace 上的 unsloth 优化版）
    max_seq_length = 2048,                           # 最大序列长度
    dtype = None,                                    # 自动检测（None = bf16 或 fp16）
    load_in_4bit = True,                             # 4bit 量化，大幅节省显存
)

# 第二步：给模型加 LoRA 适配器
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,                                         # LoRA 秩（8-64 之间选择）
    target_modules = ["q_proj", "k_proj", "v_proj", # 对 Attention 层加 LoRA
                      "o_proj", "gate_proj",
                      "up_proj", "down_proj"],
    lora_alpha = 16,                                 # LoRA 缩放系数
    lora_dropout = 0,                                # Unsloth 推荐设为 0（已内置优化）
    bias = "none",                                   # 不训练 bias
    use_gradient_checkpointing = "unsloth",         # Unsloth 的特殊梯度检查点（更省显存）
    random_state = 3407,                             # 随机种子（可复现）
)

from trl import SFTTrainer                          # 使用 TRL 的 SFT 训练器
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported           # 检查硬件是否支持 bf16

# 配置训练参数
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,                         # HuggingFace Dataset 对象
    dataset_text_field = "text",                     # 数据集中文本字段名
    max_seq_length = 2048,                           # 最大序列长度
    dataset_num_proc = 2,                            # 数据处理线程数
    args = TrainingArguments(
        per_device_train_batch_size = 2,             # 每张 GPU 的 batch 大小
        gradient_accumulation_steps = 4,             # 梯度累积步数
        warmup_steps = 5,                            # 学习率预热步数
        max_steps = 60,                              # 总训练步数（用 max_steps 或 num_train_epochs 之一）
        learning_rate = 2e-4,                        # 学习率
        fp16 = not is_bfloat16_supported(),         # 如果不支持 bf16 则用 fp16
        bf16 = is_bfloat16_supported(),             # 优先用 bf16（更稳定）
        logging_steps = 1,                           # 每多少步打印一次 loss
        optim = "adamw_8bit",                        # 8bit Adam 优化器（节省显存）
        weight_decay = 0.01,                         # 权重衰减（防止过拟合）
        lr_scheduler_type = "linear",               # 学习率调度策略
        seed = 3407,                                 # 随机种子
        output_dir = "outputs",                      # 模型保存路径
    ),
)

# 完整微调流程：加载数据 → 格式化 → 训练 → 保存

from datasets import load_dataset                   # HuggingFace 数据集加载器

# 加载并格式化数据集
dataset = load_dataset("yahma/alpaca-cleaned", split = "train")

# 定义 Alpaca 格式的提示模板
alpaca_prompt = """下面是一个指令，描述了一个任务。请写出完成该任务的回复。

### 指令:
{}

### 输入:
{}

### 回复:
{}"""

def formatting_prompts_func(examples):              # 把数据集转换为模型输入格式
    instructions = examples["instruction"]
    inputs       = examples["input"]
    outputs      = examples["output"]
    texts = []
    for instruction, input, output in zip(instructions, inputs, outputs):
        text = alpaca_prompt.format(instruction, input, output) + tokenizer.eos_token
        texts.append(text)
    return {"text": texts}                          # 返回格式化后的文本列表

dataset = dataset.map(formatting_prompts_func, batched = True)

# 开始训练
trainer_stats = trainer.train()
print(f"训练完成！共耗时 {trainer_stats.metrics['train_runtime']:.2f} 秒")

# 保存 LoRA 权重
model.save_pretrained("lora_model")                # 只保存 LoRA 适配器（小文件）
tokenizer.save_pretrained("lora_model")

# 合并并保存完整模型（用于部署）
model.save_pretrained_merged("merged_model",
    tokenizer, save_method = "merged_16bit")        # 合并成 float16 全量模型