01 - 高效微调技术（PEFT）¶

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class LoRALayer(nn.Module):
    """
    LoRA层的基本实现
    """
    def __init__(self, in_features, out_features, rank=16, lora_alpha=1):
        super().__init__()  # super()调用父类方法

        self.rank = rank
        self.lora_alpha = lora_alpha
        self.scaling = lora_alpha / rank

        # 低秩矩阵
        # 注：论文中 ΔW = BA，A ∈ ℝ^(r×k)（降维），B ∈ ℝ^(d×r)（升维）
        # 代码中存储为转置形式以便直接 x @ lora_A @ lora_B 计算
        # lora_A 对应 A^T ∈ ℝ^(k×r)，lora_B 对应 B^T ∈ ℝ^(r×d)
        self.lora_A = nn.Parameter(torch.zeros(in_features, rank))   # [k, r]
        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))  # [r, d]

        # 初始化：A随机初始化，B零初始化 → 保证初始ΔW=0
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    def forward(self, x):
        """
        x: [batch, seq_len, in_features]
        返回: [batch, seq_len, out_features]
        """
        # 计算 ΔWx = BAx
        # 等价于 x @ A^T @ B^T（使用存储的转置形式）
        return (x @ self.lora_A @ self.lora_B) * self.scaling

class LinearWithLoRA(nn.Module):
    """
    带LoRA的线性层
    """
    def __init__(self, linear_layer, rank=16, lora_alpha=1):
        super().__init__()

        # 原始线性层（冻结）
        self.linear = linear_layer
        self.linear.weight.requires_grad = False
        if self.linear.bias is not None:
            self.linear.bias.requires_grad = False

        # LoRA层
        self.lora = LoRALayer(
            linear_layer.in_features,
            linear_layer.out_features,
            rank=rank,
            lora_alpha=lora_alpha
        )

    def forward(self, x):
        # 原始输出 + LoRA输出
        return self.linear(x) + self.lora(x)

def inject_lora_to_model(model, target_modules, rank=16, lora_alpha=1):
    """
    将LoRA注入到模型的指定模块

    Args:
        model: 预训练模型
        target_modules: 要替换的模块名列表，如["q_proj", "v_proj"]
        rank: LoRA秩
        lora_alpha: LoRA alpha参数
    """
    for name, module in model.named_modules():
        # 检查是否是目标模块
        if any(target in name for target in target_modules):  # any()任一为True则返回True
            if isinstance(module, nn.Linear):  # isinstance检查类型
                # 获取父模块和属性名
                parent_name = '.'.join(name.split('.')[:-1])
                child_name = name.split('.')[-1]

                if parent_name == '':
                    parent = model
                else:
                    parent = model.get_submodule(parent_name)

                # 替换为带LoRA的线性层
                lora_layer = LinearWithLoRA(module, rank=rank, lora_alpha=lora_alpha)
                setattr(parent, child_name, lora_layer)  # setattr动态设置对象属性
                print(f"注入LoRA到: {name}")

    return model

# 使用示例
from transformers import AutoModelForCausalLM

# 加载预训练模型
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 注入LoRA
model = inject_lora_to_model(
    model,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    rank=16,
    lora_alpha=32
)

# 打印可训练参数
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
print(f"可训练参数: {trainable_params:,} ({100*trainable_params/total_params:.2f}%)")

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 4-bit量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,  # 双量化
    bnb_4bit_quant_type="nf4",       # NF4量化
    bnb_4bit_compute_dtype=torch.bfloat16
)

# 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto"
)

# 准备模型用于训练
model = prepare_model_for_kbit_training(model)

# 配置LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# 应用LoRA
model = get_peft_model(model, lora_config)

from peft import LoraConfig, get_peft_model
import torch.nn as nn

class DoRALayer(nn.Module):
    """DoRA层的简化实现"""
    def __init__(self, weight, rank=8):
        super().__init__()
        self.weight = nn.Parameter(weight.clone())
        out_features, in_features = weight.shape

        # 幅度向量
        self.magnitude = nn.Parameter(weight.norm(dim=0, keepdim=True))

        # LoRA矩阵
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01)

    def forward(self, x):
        # 计算更新后的权重
        delta_W = self.lora_B @ self.lora_A
        W_updated = self.weight + delta_W

        # 归一化并应用幅度
        W_normalized = W_updated / W_updated.norm(dim=0, keepdim=True)
        W_dora = self.magnitude * W_normalized

        return x @ W_dora.T

from peft import AdaLoraConfig, get_peft_model

# 配置AdaLoRA
adalora_config = AdaLoraConfig(
    init_r=12,              # 初始秩
    target_r=4,             # 目标秩
    beta1=0.85,             # 奇异值保留比例
    beta2=0.85,             # 奇异值保留比例
    tinit=200,              # 开始剪枝的步数
    tfinal=1000,            # 停止剪枝的步数
    deltaT=10,              # 剪枝间隔
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    task_type="CAUSAL_LM"
)

# 应用AdaLoRA
model = get_peft_model(model, adalora_config)

def print_trainable_parameters(model):
    """打印可训练参数"""
    trainable_params = 0
    all_params = 0

    for name, param in model.named_parameters():
        all_params += param.numel()
        if param.requires_grad:
            trainable_params += param.numel()
            print(f"可训练: {name}: {param.numel()}")

    print(f"\n可训练参数: {trainable_params:,} || "
          f"总参数: {all_params:,} || "
          f"比例: {100*trainable_params/all_params:.2f}%")

# 使用
print_trainable_parameters(model)

# 训练前保存LoRA权重
lora_weights_before = {name: param.clone() for name, param in model.named_parameters() if param.requires_grad}

# 训练...

# 训练后比较
lora_weights_after = {name: param.clone() for name, param in model.named_parameters() if param.requires_grad}

for name in lora_weights_before:
    diff = (lora_weights_after[name] - lora_weights_before[name]).abs().mean()
    print(f"{name} 平均变化: {diff:.6f}")

"""
QLoRA微调LLaMA的完整可运行代码
需要: pip install transformers peft bitsandbytes trl datasets
硬件: RTX 3090/4090 (24GB VRAM) 即可运行 7B 模型
"""
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer
from datasets import load_dataset

# ---------- Step 1: 量化加载模型 ----------
model_name = "meta-llama/Llama-3.1-8B"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model = prepare_model_for_kbit_training(model)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# ---------- Step 2: 配置 LoRA ----------
lora_config = LoraConfig(
    r=16,                          # 秩
    lora_alpha=32,                 # alpha = 2 * r 是常见选择
    target_modules=[               # LLaMA 的注意力和 FFN 权重
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 83,886,080 || all params: 8,114,212,864 || trainable%: 1.03%

# ---------- Step 3: 准备数据集 ----------
# 使用 Alpaca 格式的数据集作为示例
dataset = load_dataset("tatsu-lab/alpaca", split="train[:5000]")

def format_instruction(example):
    """将 Alpaca 格式转为对话模板"""
    if example.get("input", ""):
        text = (
            f"### Instruction:\n{example['instruction']}\n\n"
            f"### Input:\n{example['input']}\n\n"
            f"### Response:\n{example['output']}"
        )
    else:
        text = (
            f"### Instruction:\n{example['instruction']}\n\n"
            f"### Response:\n{example['output']}"
        )
    return {"text": text}

dataset = dataset.map(format_instruction)

# ---------- Step 4: 训练配置 ----------
training_args = TrainingArguments(
    output_dir="./lora-llama-output",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,     # 有效batch_size = 4 * 4 = 16
    learning_rate=2e-4,                # LoRA 学习率通常比全量微调高10倍
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,                         # 使用 bfloat16 混合精度
    logging_steps=10,
    save_strategy="steps",
    save_steps=200,
    max_grad_norm=0.3,                 # 梯度裁剪
    optim="paged_adamw_8bit",          # 分页优化器，节省显存
)

# ---------- Step 5: 使用 SFTTrainer 训练 ----------
# 注意：TRL v0.8+ 废弃了 dataset_text_field 参数，
# 需要将数据集预处理为包含 "text" 列或使用 formatting_func
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=training_args,
    processing_class=tokenizer,        # TRL v0.12+: tokenizer 参数更名为 processing_class
    max_seq_length=512,
    packing=True,                      # 将短样本打包，提高GPU利用率
    # 如果数据集列名不是 "text"，使用 formatting_func：
    # formatting_func=lambda examples: [ex["text"] for ex in examples],
)

trainer.train()

# ---------- Step 6: 保存 LoRA 权重 ----------
model.save_pretrained("./lora-llama-adapter")
# 只保存 LoRA 权重（~100MB vs 完整模型 ~16GB）

# ---------- Step 7: 加载和推理 ----------
from peft import PeftModel

# 重新加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)

# 加载 LoRA 权重
model = PeftModel.from_pretrained(base_model, "./lora-llama-adapter")

# 推理
prompt = "### Instruction:\n用Python实现快速排序\n\n### Response:\n"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# ---------- Step 8: 合并权重（可选，部署用）----------
merged_model = model.merge_and_unload()  # 将 LoRA 合并回基础模型
merged_model.save_pretrained("./llama-merged")
tokenizer.save_pretrained("./llama-merged")
# 现在 merged 模型可以直接用于推理，无需 peft 库

from peft import PeftModel

# 加载基础模型（一次）
base_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# 加载代码助手 LoRA
code_model = PeftModel.from_pretrained(base_model, "./lora-code-assistant")

# 加载翻译器 LoRA（同一个基座）
translation_model = PeftModel.from_pretrained(base_model, "./lora-translator")

# 运行时切换适配器（如果使用同一个 PeftModel）
model = PeftModel.from_pretrained(base_model, "./lora-code-assistant")
model.load_adapter("./lora-translator", adapter_name="translator")

model.set_adapter("default")     # 使用代码助手
model.set_adapter("translator")  # 切换到翻译器

# 禁用所有 LoRA（回到原始模型）
with model.disable_adapter():
    outputs = model.generate(...)  # 使用原始基础模型

"""微调效果诊断脚本"""

def diagnose_lora_training(model, training_log):
    """训练后诊断清单"""

    # 1. 检查可训练参数比例
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    ratio = trainable / total * 100
    print(f"[1] 可训练参数: {trainable:,} / {total:,} ({ratio:.2f}%)")
    assert ratio < 5, "警告：可训练参数过多，检查是否正确冻结了基础模型"  # assert断言：条件False时抛出AssertionError

    # 2. 检查 loss 是否下降
    losses = [log["loss"] for log in training_log if "loss" in log]
    if len(losses) > 1:
        improvement = (losses[0] - losses[-1]) / losses[0] * 100  # [-1]负索引取最后一个元素
        print(f"[2] Loss: {losses[0]:.4f} → {losses[-1]:.4f} (下降 {improvement:.1f}%)")

    # 3. 检查 LoRA 权重是否有变化
    for name, param in model.named_parameters():
        if "lora_B" in name and param.requires_grad:
            if param.abs().max() < 1e-6:
                print(f"[3] 警告：{name} 仍为零，LoRA可能未生效")
                break
    else:
        print("[3] LoRA B矩阵已更新 ✓")

    # 4. 检查梯度范数
    grad_norms = []
    for name, param in model.named_parameters():
        if param.requires_grad and param.grad is not None:
            grad_norms.append(param.grad.norm().item())
    if grad_norms:
        print(f"[4] 梯度范数: min={min(grad_norms):.6f}, max={max(grad_norms):.6f}")