第2章 微调你的专属模型：从原理到实战的完全指南#

[BOS] 用户: 头痛怎么办? [SEP] 助手: 建议多休息，必要时就医。[EOS]

import torch

def compute_sft_loss(logits, labels, attention_mask):
    """
    Args:
        logits: [batch, seq_len, vocab_size]
        labels: [batch, seq_len]  # -100 表示忽略该位置
        attention_mask: [batch, seq_len]
    """
    # 只计算labels != -100的位置
    loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100)
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1)
    )
    return loss

# 实际使用：构造labels
def create_sft_labels(input_ids, answer_start_idx):
    """只在答案部分计算loss"""
    labels = input_ids.clone()
    labels[:, :answer_start_idx] = -100  # 掩码指令部分
    return labels

指令部分（掩码）         答案部分（计算Loss）
┌──────────────┐        ┌──────────────┐
│ 用户: 头痛... │ -100   │ 建议多休息... │ Loss
└──────────────┘        └──────────────┘
      ↓                        ↓
  不反向传播                 更新梯度

# 伪代码：AdamW状态
for param in model.parameters():
    optimizer_state = {
        'm': torch.zeros_like(param),      # 一阶矩（均值） - 4 bytes (FP32)
        'v': torch.zeros_like(param),      # 二阶矩（方差） - 4 bytes (FP32)
        'grad': torch.zeros_like(param)    # 梯度缓存 - 4 bytes (FP32)
    }
    # 总计：4 + 4 + 4 = 12 bytes/参数

# LoRA可训练参数量
def lora_params(model_dim, rank, num_layers):
    params_per_layer = 2 * model_dim * rank  # W_A + W_B
    return params_per_layer * num_layers

# Llama-3-8B: 4096维度, 32层, rank=16
lora_params(4096, 16, 32) = 8,388,608 ≈ 8M (仅Attention层)

import torch
import torch.nn as nn
import math

class LoRALinear(nn.Module):
    """LoRA线性层：完整PyTorch原生实现"""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        alpha: float = 16,
        dropout: float = 0.0,
        bias: bool = True
    ):
        super().__init__()

        # 冻结的预训练权重
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        self.linear.weight.requires_grad = False
        if bias:
            self.linear.bias.requires_grad = False

        # LoRA参数
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank

        # 低秩矩阵
        self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))

        # Dropout（可选）
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

        # 初始化
        self.reset_parameters()

    def reset_parameters(self):
        """Kaiming初始化A，零初始化B"""
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播: h = W0*x + (α/r)*B*A*x

        Args:
            x: [batch, seq_len, in_features]
        Returns:
            h: [batch, seq_len, out_features]
        """
        # 基础输出（冻结）
        result = self.linear(x)

        # LoRA增量
        lora_out = self.dropout(x) @ self.lora_A.T  # [*, rank]
        lora_out = lora_out @ self.lora_B.T         # [*, out_features]
        lora_out = lora_out * self.scaling

        return result + lora_out

    def merge_weights(self):
        """合并LoRA权重到主权重（推理优化）"""
        if self.rank > 0:
            delta_w = self.lora_B @ self.lora_A * self.scaling
            self.linear.weight.data += delta_w
            # 清空LoRA参数以节省显存
            self.lora_A.data.zero_()
            self.lora_B.data.zero_()

    def get_lora_state_dict(self):
        """只保存LoRA参数"""
        return {
            'lora_A': self.lora_A,
            'lora_B': self.lora_B
        }

# 使用示例
layer = LoRALinear(in_features=4096, out_features=4096, rank=16, alpha=32)

# 统计参数量
frozen_params = sum(p.numel() for p in layer.linear.parameters())
trainable_params = sum(p.numel() for p in [layer.lora_A, layer.lora_B])

print(f"冻结参数: {frozen_params:,} ({frozen_params * 2 / 1e9:.2f} GB)")
print(f"可训练参数: {trainable_params:,} ({trainable_params * 2 / 1e9:.4f} GB)")
print(f"参数比例: {trainable_params / frozen_params * 100:.2f}%")

# 输出：
# 冻结参数: 16,777,216 (0.03 GB)
# 可训练参数: 131,072 (0.0003 GB)
# 参数比例: 0.78%

# NF4量化点（从标准正态分布的分位数推导）
NF4_QUANTILES = torch.tensor([
    -1.0000, -0.6962, -0.5251, -0.3949,
    -0.2844, -0.1848, -0.0911,  0.0000,
     0.0796,  0.1609,  0.2461,  0.3379,
     0.4407,  0.5626,  0.7230,  1.0000
])

def quantize_nf4(tensor):
    """NF4量化"""
    # 标准化到[-1, 1]
    absmax = tensor.abs().max()
    normalized = tensor / (absmax + 1e-8)

    # 查找最近的量化点
    distances = (normalized.unsqueeze(-1) - NF4_QUANTILES).abs()
    indices = distances.argmin(dim=-1)

    # 量化值（4bit索引）
    quantized = NF4_QUANTILES[indices]

    # 反量化
    dequantized = quantized * absmax
    return dequantized, indices, absmax

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

class DoRALinear(nn.Module):
    """DoRA: Direction + Magnitude分解"""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        rank: int = 8,
        alpha: float = 16,
        dropout: float = 0.0
    ):
        super().__init__()

        # 冻结的预训练权重
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.linear.weight.requires_grad = False

        # LoRA参数
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank

        self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))

        # DoRA特有：幅度向量
        self.magnitude = nn.Parameter(torch.ones(out_features, 1))

        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

        # 初始化
        self.reset_parameters()

    def reset_parameters(self):
        """初始化策略"""
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

        # 幅度初始化为W0的Row Norm
        with torch.no_grad():
            weight_norm = self.linear.weight.norm(p=2, dim=1, keepdim=True)
            self.magnitude.copy_(weight_norm)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播: h = m * (V + ΔV) / ||V + ΔV||

        Args:
            x: [batch, seq_len, in_features]
        Returns:
            h: [batch, seq_len, out_features]
        """
        # 计算方向更新
        lora_out = self.dropout(x) @ self.lora_A.T
        lora_out = lora_out @ self.lora_B.T
        lora_out = lora_out * self.scaling

        # 组合权重
        combined_weight = self.linear.weight + (self.lora_B @ self.lora_A)

        # 标准化方向
        weight_norm = combined_weight.norm(p=2, dim=1, keepdim=True)
        normalized_weight = combined_weight / (weight_norm + 1e-8)

        # 应用幅度
        final_weight = self.magnitude * normalized_weight

        # 计算输出
        return F.linear(x, final_weight)

    def extra_repr(self) -> str:
        return f'in_features={self.linear.in_features}, out_features={self.linear.out_features}, rank={self.rank}'

# 性能对比测试
def benchmark_dora_vs_lora():
    torch.manual_seed(42)
    x = torch.randn(32, 128, 4096)  # [batch, seq, dim]

    lora = LoRALinear(4096, 4096, rank=16)
    dora = DoRALinear(4096, 4096, rank=16)

    # 前向传播时间
    import time

    start = time.time()
    _ = lora(x)
    lora_time = time.time() - start

    start = time.time()
    _ = dora(x)
    dora_time = time.time() - start

    print(f"LoRA时间: {lora_time:.4f}s")
    print(f"DoRA时间: {dora_time:.4f}s")
    print(f"DoRA额外开销: {(dora_time - lora_time) / lora_time * 100:.1f}%")

import torch
from torch.optim.optimizer import Optimizer

class GaLoreAdamW(Optimizer):
    """GaLore优化器：梯度低秩投影"""

    def __init__(
        self,
        params,
        lr: float = 1e-3,
        rank: int = 128,
        update_proj_gap: int = 200,  # 每200步更新投影矩阵
        scale: float = 1.0,
        proj_type: str = "std"  # "std", "reverse", "right", "left"
    ):
        defaults = dict(lr=lr, rank=rank, update_proj_gap=update_proj_gap,
                       scale=scale, proj_type=proj_type)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                state = self.state[p]

                # 初始化状态
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(grad)
                    state['exp_avg_sq'] = torch.zeros_like(grad)
                    state['projector'] = None

                state['step'] += 1

                # 更新投影矩阵
                if state['step'] % group['update_proj_gap'] == 1:
                    state['projector'] = self._get_projector(
                        grad, group['rank'], group['proj_type']
                    )

                # 投影梯度
                if state['projector'] is not None:
                    grad_proj = self._project(grad, state['projector'])
                else:
                    grad_proj = grad

                # AdamW更新
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = 0.9, 0.999

                exp_avg.mul_(beta1).add_(grad_proj, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad_proj, grad_proj, value=1 - beta2)

                denom = exp_avg_sq.sqrt().add_(1e-8)
                step_size = group['lr'] * group['scale']

                p.add_(exp_avg / denom, alpha=-step_size)

    def _get_projector(self, grad, rank, proj_type):
        """计算SVD投影矩阵"""
        matrix = grad if grad.dim() == 2 else grad.flatten(1)

        # SVD分解（只取前r个奇异向量）
        U, S, Vt = torch.linalg.svd(matrix, full_matrices=False)
        U_r = U[:, :rank]
        Vt_r = Vt[:rank, :]

        if proj_type == "left":
            return U_r
        elif proj_type == "right":
            return Vt_r
        else:  # "std"
            return (U_r, Vt_r)

    def _project(self, grad, projector):
        """应用投影"""
        if isinstance(projector, tuple):
            U_r, Vt_r = projector
            # 双边投影: U_r @ U_r^T @ G @ Vt_r^T @ Vt_r
            return U_r @ (U_r.T @ grad @ Vt_r.T) @ Vt_r
        else:
            # 单边投影
            return projector @ (projector.T @ grad)

# 使用示例
model = MyTransformer()
optimizer = GaLoreAdamW(model.parameters(), lr=1e-3, rank=128)

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# LoRA+参数组设置
optimizer = torch.optim.AdamW([
    {'params': [p for n, p in model.named_parameters() if 'lora_A' in n],
     'lr': 1e-4},
    {'params': [p for n, p in model.named_parameters() if 'lora_B' in n],
     'lr': 1e-3},  # 10x更大
])

def ewc_loss(model, fisher_matrix, old_params, lambda_ewc=0.4):
    loss = 0
    for name, param in model.named_parameters():
        if name in fisher_matrix:
            # Fisher值越大，惩罚越重
            _loss = fisher_matrix[name] * (param - old_params[name]) ** 2
            loss += _loss.sum()
    return lambda_ewc * loss

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import torch

# ========== 1. 加载模型和Tokenizer ==========
model_name = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# ========== 2. 配置LoRA ==========
lora_config = LoraConfig(
    r=16,                            # 秩
    lora_alpha=32,                   # 缩放因子（通常是r的2倍）
    target_modules=[                 # 目标层
        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
        "gate_proj", "up_proj", "down_proj"      # FFN
    ],
    lora_dropout=0.05,
    bias="none",                     # 不训练bias
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出：trainable params: 20,971,520 || all params: 3,213,978,624 || trainable%: 0.65%

# ========== 3. 数据处理 ==========
dataset = load_dataset("tatsu-lab/alpaca", split="train[:1000]")

# 关键：设置padding方向
tokenizer.padding_side = "right"  # SFT必须用right padding！
tokenizer.pad_token = tokenizer.eos_token

# 应用聊天模板
def format_chat(example):
    messages = [
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]}
    ]
    return {"text": tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )}

dataset = dataset.map(format_chat)

# ========== 4. 数据Collator（关键：只计算答案Loss）==========
# 自动识别response_template，只在助手回复部分计算loss
response_template = "<|start_header_id|>assistant<|end_header_id|>"
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template,
    tokenizer=tokenizer,
    mlm=False
)

# ========== 5. 训练配置 ==========
training_args = SFTConfig(
    output_dir="./llama3-lora-sft",

    # 基础设置
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,      # 等效batch size = 16

    # 优化器
    optim="adamw_torch",
    learning_rate=2e-4,
    weight_decay=0.01,
    warmup_ratio=0.03,

    # 学习率调度
    lr_scheduler_type="cosine",

    # 显存优化
    gradient_checkpointing=True,
    bf16=True,                          # 使用BF16混合精度

    # 日志
    logging_steps=10,
    save_strategy="epoch",

    # NEFTune（噪声注入技巧）
    neftune_noise_alpha=5,              # 在嵌入层加噪声

    # 其他
    max_seq_length=512,
    dataset_text_field="text",          # 数据集中的文本字段
)

# ========== 6. 开始训练 ==========
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    data_collator=collator,
)

trainer.train()

# ========== 7. 保存模型 ==========
trainer.save_model("./final_lora_model")

# Llama-3的聊天格式
"""
<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>

用户问题<|eot_id|><|start_header_id|>assistant<|end_header_id|>

助手回答<|eot_id|>
"""

# 自动应用模板
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "什么是量子纠缠？"},
    {"role": "assistant", "content": "量子纠缠是..."}
]

formatted = tokenizer.apply_chat_template(
    messages,
    tokenize=False,           # 返回字符串而非token ids
    add_generation_prompt=False  # 训练时不加生成提示符
)

# ❌ 错误：left padding会破坏因果注意力
tokenizer.padding_side = "left"

# 示例数据
texts = ["短句", "这是一个很长很长的句子"]
tokens = tokenizer(texts, padding=True)

# 结果：
# [PAD][PAD][PAD] 短 句
# 这 是 一 个 很 长 ...

# 问题：模型会在PAD位置计算loss！

# ✅ 正确：right padding保持因果结构
tokenizer.padding_side = "right"

# 结果：
# 短 句 [PAD][PAD][PAD]
# 这 是 一 个 很 长 ...

# 只在真实token上计算loss

class NEFTuneEmbedding(nn.Module):
    """NEFTune嵌入层包装器"""

    def __init__(self, embedding_layer, alpha=5):
        super().__init__()
        self.embedding = embedding_layer
        self.alpha = alpha

    def forward(self, input_ids):
        embeddings = self.embedding(input_ids)

        if self.training:
            # 计算噪声标准差
            seq_len, emb_dim = embeddings.shape[1], embeddings.shape[2]
            noise_std = self.alpha / (seq_len * emb_dim) ** 0.5

            # 添加噪声
            noise = torch.randn_like(embeddings) * noise_std
            embeddings = embeddings + noise

        return embeddings

# TRL自动支持（在SFTConfig中设置）
training_args = SFTConfig(
    neftune_noise_alpha=5,  # 自动应用NEFTune
    ...
)

from peft import PeftModel

# ========== 方法1：保持LoRA分离（灵活） ==========
base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B")
model = PeftModel.from_pretrained(base_model, "./final_lora_model")

# 推理
inputs = tokenizer("用户问题", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100)

# ========== 方法2：合并权重（快速） ==========
# 合并LoRA到基础模型
merged_model = model.merge_and_unload()

# 保存合并后的模型
merged_model.save_pretrained("./merged_model")
tokenizer.save_pretrained("./merged_model")

# 之后可以当作普通模型加载
model = AutoModelForCausalLM.from_pretrained("./merged_model")

import torch
import numpy as np

def slerp(t, v0, v1, dot_threshold=0.9995):
    """
    球面线性插值
    Args:
        t: 插值系数 [0, 1]
        v0, v1: 两个形状相同的张量
    """
    # 归一化
    v0_norm = v0 / torch.norm(v0)
    v1_norm = v1 / torch.norm(v1)

    # 计算夹角余弦
    dot = torch.sum(v0_norm * v1_norm)

    # 如果两个向量几乎平行，直接线性插值
    if torch.abs(dot) > dot_threshold:
        return (1 - t) * v0 + t * v1

    # 计算夹角
    theta_0 = torch.acos(dot)
    sin_theta_0 = torch.sin(theta_0)

    theta_t = theta_0 * t
    sin_theta_t = torch.sin(theta_t)

    s0 = torch.sin(theta_0 - theta_t) / sin_theta_0
    s1 = sin_theta_t / sin_theta_0

    return s0 * v0 + s1 * v1