Unsloth实战指南：3倍加速大模型训练，单卡轻松玩转70B参数

# 创建虚拟环境
conda create -n unsloth_env python=3.10 -y
conda activate unsloth_env

# 安装核心库（自动匹配CUDA版本）
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

# 安装扩展包
pip install flash-attn==2.5.8 triton==2.2.0

from unsloth import is_unsloth_available
print(f"Unsloth可用状态: {is_unsloth_available()}")
# 输出显存优化报告
import unsloth; unsloth.report_memory_optimization()

from unsloth import FastLanguageModel

# 加载7B模型（自动应用优化）
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = 4096,  # 支持动态长度
    dtype = torch.float16,  # 自动选择最优精度
    load_in_4bit = True,    # 4bit量化加载
)

# 定制化加载配置
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "mistralai/Mistral-7B-v0.1",
    max_seq_length = 32768,  # 支持32K上下文
    rope_scaling = {  # 动态RoPE缩放
        "type": "dynamic", 
        "factor": 2.0
    },
    attn_implementation = "flash_attention_3",  # FA3加速
    token = "hf_your_token",  # HuggingFace令牌
)

model = FastLanguageModel.get_peft_model(
    model,
    r = 16,  # LoRA秩
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"],  # 目标模块
    lora_alpha = 16,  # 缩放因子
    lora_dropout = 0.05,  # 防止过拟合
    use_gradient_checkpointing = True,  # 梯度检查点
    random_state = 3407,  # 随机种子
    max_seq_length = 4096,
    use_rslora = True,  # 启用稳定LoRA
    loftq_config = {},   # 量化感知训练
)

from torch.optim import AdamW

optimizer = AdamW(
    model.parameters(),
    lr = 2e-5 * (batch_size / 32),  # 动态学习率
    weight_decay = 0.01,
    betas = (0.9, 0.999),
)

# Unsloth专属Adafactor优化器
from unsloth import Adafactor
optimizer = Adafactor(
    model.parameters(),
    scale_parameter = True,
    relative_step = True,
    warmup_init = True,
)

trainer = TrainingArguments(
    per_device_train_batch_size = 4,  # 根据显存自动扩展
    gradient_accumulation_steps = 8,   # 等效batch_size=32
    warmup_ratio = 0.1,  # 10%步数用于预热
    learning_rate = 5e-5,
    fp16 = not torch.cuda.is_bf16_supported(),  # 自动精度选择
    bf16 = torch.cuda.is_bf16_supported(),
    logging_steps = 10,
    optim = "adafactor",
    weight_decay = 0.01,
    lr_scheduler_type = "cosine_with_restarts",  # 带重启的余弦退火
    seed = 42,
    output_dir = "outputs",
    save_strategy = "steps",
    save_steps = 1000,
    report_to = "none",  # 禁用wandb等监控
    ddp_find_unused_parameters = False,
)

# 推荐使用Arrow格式
dataset = {
    "instruction": "解释量子纠缠现象",
    "input": "",
    "output": "量子纠缠是...",  # 实际回答
    "category": "physics"
}

# 转换为模型输入
text = f"""<s>[INST] <<SYS>>
你是有帮助的AI助手
<</SYS>>

{dataset['instruction']} [/INST]
{dataset['output']} </s>"""

from unsloth import FastLanguageModel

# 创建优化后的DataLoader
data_loader = FastLanguageModel.get_train_dataloader(
    dataset = dataset,  # HuggingFace Dataset对象
    tokenizer = tokenizer,
    max_seq_length = 4096,
    batch_size = 4,
    packing = True,  # 动态序列打包（效率提升3倍）
    shuffle = True,
    num_proc = 4,    # 并行预处理进程
)

from trl import SFTTrainer

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = 4096,
    args = trainer,
    packing = True,  # 关键加速选项
)

# 启动训练（自动应用Unsloth优化）
trainer.train()

# 实时监控命令（另开终端）
watch -n 1 nvidia-smi  # 监控显存
gpustat -i  # 查看GPU利用率

# 动态批处理（自动调整）
model.enable_dynamic_batching(
    min_batch_size=2,
    max_batch_size=16,
    utilization_threshold=0.85,  # 当GPU利用率>85%时增大batch
)

# 梯度检查点配置
model.gradient_checkpointing_enable(
    strategy="aggressive",  # 激进模式节省40%显存
    checkpoint_interval=5   # 每5层设检查点
)

# 混合精度策略
model.configure_mixed_precision(
    forward_dtype=torch.bfloat16,  # 前向bf16
    backward_dtype=torch.float8,   # 反向fp8
    optimizer_dtype=torch.float32, # 优化器fp32
)

# 保存完整模型（包含LoRA权重）
model.save_pretrained_merged(
    "finetuned_model", 
    tokenizer,
    save_method = "merged_4bit",  # 4bit量化导出
    push_to_hub = True,
)

# 仅保存适配器（轻量部署）
model.save_pretrained_peft_adapter(
    "adapter_model",
    push_to_hub = True,
)

# 加载优化后的推理模型
from unsloth import FastLanguageModelForInference

model = FastLanguageModelForInference.from_pretrained(
    "finetuned_model",
    max_seq_length = 4096,
    dtype = torch.float16,
)

# 创建优化管道
pipeline = model.create_serving_pipeline(
    temperature = 0.7,
    repetition_penalty = 1.2,
    max_new_tokens = 512,
    do_sample = True,
)

# 执行推理（比原生快3倍）
outputs = pipeline("如何解释相对论？")
print(outputs[0]["generated_text"])

# 步骤1：环境检查
assert torch.cuda.get_device_properties(0).total_memory >= 80e9, "需80GB显存"

# 步骤2：4bit加载模型
model, tokenizer = FastLanguageModel.from_pretrained(
    "meta-llama/Meta-Llama-3-70B",
    load_in_4bit = True,
    attn_implementation = "flash_attention_3",
)

# 步骤3：注入LoRA
model = FastLanguageModel.get_peft_model(model, r=64)

# 步骤4：配置动态批处理
trainer_args = TrainingArguments(
    per_device_train_batch_size = 1,  # 物理batch_size
    gradient_accumulation_steps = 32, # 等效batch_size=32
    fp8_grads = True,  # 关键！启用FP8梯度
)

# 步骤5：启动训练
trainer = SFTTrainer(...)
trainer.train()  # 平均显存占用：72GB/80GB

# 启用更激进的量化
model, tokenizer = FastLanguageModel.from_pretrained(
    ...,
    load_in_4bit = True,
    quant_method = "bnb",  # 使用Bitsandbytes
    # 或
    quant_method = "gptq", # GPTQ量化
)

# 启用梯度检查点
model.gradient_checkpointing_enable(strategy="ultra")

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

# 2. 启用kernel融合
export UNSLOTH_KERNEL_FUSION_MODE=aggressive

# 3. 调整DataLoader配置
data_loader = FastLanguageModel.get_train_dataloader(
    num_workers = 0,  # Colab中设为0
    pin_memory = True,
)

# 学习率热重启配置
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer, 
    T_0 = 1000,     # 首次周期步数
    T_mult = 2,      # 周期倍增系数
    eta_min = 1e-7,  # 最小学习率
)

# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 损失缩放（FP16训练）
scaler = torch.cuda.amp.GradScaler()

# 终极单行代码体验
from unsloth import train; train("your_dataset")