LoRA / QLoRA / Unsloth 微调笔记

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

class LoRALinear(nn.Module):    """
    带 LoRA 的线性层。
    原始权重 W_0 冻结，只训练低秩矩阵 A 和 B。
    """
    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        alpha: float = 16.0,
        dropout_rate: float = 0.0,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.scaling = alpha / rank  # α/r 缩放因子

        # 原始权重：冻结
        self.weight = nn.Parameter(
            torch.randn(out_features, in_features) * 0.02,
            requires_grad=False,  # 冻结！
        )

        # LoRA 矩阵：可训练
        # A 用 Kaiming-uniform 初始化（与 nn.Linear 默认初始化一致，PEFT 库标准做法）
        # 不用高斯 * 0.02，因为 Kaiming-uniform 能更好地保持各层激活值的方差
        self.lora_A = nn.Parameter(torch.empty(rank, in_features))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))  # B 初始化为 0，保证 ΔW=0

        # 可选的 dropout（防止 LoRA 过拟合）
        self.lora_dropout = nn.Dropout(dropout_rate) if dropout_rate > 0 else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 原始路径（冻结权重）
        base_output = F.linear(x, self.weight)

        # LoRA 路径：x → A → B，乘以缩放因子
        lora_output = F.linear(self.lora_dropout(x), self.lora_A)  # [*, rank]
        lora_output = F.linear(lora_output, self.lora_B)            # [*, out_features]

        return base_output + self.scaling * lora_output

    def merge_weights(self) -> nn.Linear:
        """
        将 LoRA 权重合并到原始权重，返回标准 Linear 层（推理用）。
        合并后推理无额外开销。
        """
        merged_weight = self.weight + self.scaling * (self.lora_B @ self.lora_A)
        merged_layer = nn.Linear(self.in_features, self.out_features, bias=False)
        merged_layer.weight = nn.Parameter(merged_weight)
        return merged_layer


def apply_lora_to_attention(model: nn.Module, rank: int = 8, alpha: float = 16.0) -> nn.Module:
    """
    将模型中所有名为 q_proj、k_proj、v_proj、o_proj 的线性层替换为 LoRA 版本。
    其余参数全部冻结。
    """
    # 先冻结所有参数
    for param in model.parameters():
        param.requires_grad = False

    # 替换目标层为 LoRA 版本
    target_layer_names = {"q_proj", "k_proj", "v_proj", "o_proj"}
    for module_name, module in model.named_modules():
        for layer_name in target_layer_names:
            if hasattr(module, layer_name):
                original_layer = getattr(module, layer_name)
                if isinstance(original_layer, nn.Linear):
                    lora_layer = LoRALinear(
                        in_features=original_layer.in_features,
                        out_features=original_layer.out_features,
                        rank=rank,
                        alpha=alpha,
                    )
                    # 复制原始权重（冻结）
                    lora_layer.weight.data = original_layer.weight.data.clone()
                    setattr(module, layer_name, lora_layer)

    return model


def count_trainable_params(model: nn.Module) -> tuple[int, int]:
    """返回 (可训练参数量, 总参数量)"""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total


# ── 验证 ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    torch.manual_seed(42)

    # 验证 LoRALinear 的前向传播和权重合并
    lora_layer = LoRALinear(in_features=512, out_features=512, rank=8, alpha=16.0)
    x = torch.randn(2, 10, 512)

    output = lora_layer(x)
    print(f"LoRA 输出形状: {output.shape}")  # [2, 10, 512]

    # 验证合并前后输出一致
    merged_layer = lora_layer.merge_weights()
    output_merged = merged_layer(x)
    max_diff = (output - output_merged).abs().max().item()
    print(f"合并前后最大误差: {max_diff:.2e}")  # 应接近 0

    # 统计参数量
    trainable_params = sum(p.numel() for p in lora_layer.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in lora_layer.parameters())
    frozen_params = total_params - trainable_params
    print(f"可训练参数: {trainable_params:,}  ({trainable_params/total_params*100:.1f}%)")
    print(f"冻结参数:   {frozen_params:,}  ({frozen_params/total_params*100:.1f}%)")

NF4_levels = [
    -1.0,
    -0.6961928009986877,
    -0.5250730514526367,
    -0.39491748809814453,
    -0.28444138169288635,
    -0.18477343022823334,
    -0.09105003625154495,
     0.0,
     0.07958029955625534,
     0.16093020141124725,
     0.24611230194568634,
     0.33791524171829224,
     0.44070982933044434,
     0.5626170039176941,
     0.7229568362236023,
     1.0,
]

from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset
import torch

# 1. 加载模型（自动应用 QLoRA + Unsloth 优化）
# 模型名称末尾带 "unsloth-bnb-4bit" 的是 Unsloth 动态 4bit 量化版，精度优于标准 bnb-4bit
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/llama-3.1-8b-unsloth-bnb-4bit",
    max_seq_length=2048,
    dtype=None,          # None = 自动选择（新 GPU 用 BF16，旧 GPU 用 FP16）
    load_in_4bit=True,   # QLoRA：4bit 量化基础模型
)

# 2. 添加 LoRA 适配器
model = FastLanguageModel.get_peft_model(
    model,
    r=16,                # LoRA rank，从 16 开始试
    target_modules=[     # 应用 LoRA 的层（包含 FFN 效果更好）
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha=16,       # α = r，scaling = 1.0，简化调参
    lora_dropout=0.0,    # Unsloth 的 Triton kernel 假设无 dropout，设为 0
    bias="none",         # 不训练 bias
    use_gradient_checkpointing="unsloth",  # Unsloth 优化版梯度检查点，比标准版省更多显存
    random_state=42,
)

# 3. 准备数据集
# Unsloth 推荐用 formatting_func 动态格式化，避免提前 map 整个数据集
alpaca_prompt = """Below is an instruction that describes a task.
Write a response that appropriately completes the request.

### Instruction:
{}

### Response:
{}"""

def formatting_func(examples):
    """将数据集的每条样本格式化为训练文本，末尾加 EOS token。"""
    texts = []
    for instruction, output in zip(examples["instruction"], examples["output"]):
        text = alpaca_prompt.format(instruction, output) + tokenizer.eos_token
        texts.append(text)
    return texts

dataset = load_dataset("yahma/alpaca-cleaned", split="train")

# 4. 训练
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=formatting_func,  # 动态格式化，比 dataset_text_field 更灵活
    max_seq_length=2048,
    args=TrainingArguments(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,  # 等效 batch_size = 8
        warmup_steps=5,
        num_train_epochs=1,
        learning_rate=2e-4,
        fp16=not torch.cuda.is_bf16_supported(),
        bf16=torch.cuda.is_bf16_supported(),
        logging_steps=1,
        optim="adamw_8bit",             # 8bit Adam，进一步节省显存
        output_dir="outputs",
    ),
)

trainer.train()

# 5. 保存 LoRA adapter（只有几十 MB，不含基础模型）
model.save_pretrained("lora_adapter")
tokenizer.save_pretrained("lora_adapter")

# 6. 推理：加载 adapter 并切换到推理模式
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="lora_adapter",
    max_seq_length=2048,
    load_in_4bit=True,
)
FastLanguageModel.for_inference(model)  # 必须调用，启用 Unsloth 的推理加速

inputs = tokenizer(
    [alpaca_prompt.format("Explain LoRA in one sentence.", "")],
    return_tensors="pt",
).to("cuda")

outputs = model.generate(**inputs, max_new_tokens=128)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 保存：只有几十 MB
model.save_pretrained("lora_adapter")

# 加载：需要同时加载基础模型 + adapter
from transformers import AutoModelForCausalLM
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
model = PeftModel.from_pretrained(base_model, "lora_adapter")

# 合并 LoRA 权重到基础模型（推理时无额外开销）
merged_model = model.merge_and_unload()
merged_model.save_pretrained("merged_model")  # 完整模型，几 GB

# Unsloth 支持直接导出 GGUF，供 llama.cpp / Ollama 使用
model.save_pretrained_gguf(
    "model_gguf",
    tokenizer,
    quantization_method="q4_k_m",  # 4bit 量化，平衡精度和速度
)