4. 实现前缀微调（Prefix Tuning）

任务说明

本案例为模型实现前缀微调（Prefix Tuning）：冻结主干模型参数，仅训练插入到每一层注意力前的可学习前缀向量，使模型在特定提示风格下生成结果发生可控变化。完成后，同一提示词在微调前后的生成文本应呈现明显差异，且训练过程仅更新前缀参数。

环境依赖

pip install torch transformers datasets

完整实现

import torch
from torch import nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer

class PrefixTuningGPT2(nn.Module):
    def __init__(self, model_name="", prefix_length=20, num_layers=12, hidden_size=768):
        super().__init__()
        self.model = GPT2LMHeadModel.from_pretrained(model_name)
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        self.tokenizer.pad_ = self.tokenizer.eos_
        self.prefix_length = prefix_length
        self.num_layers = num_layers
        self.hidden_size = hidden_size

        # 冻结主干模型
        for param in self.model.parameters():
            param.requires_grad = False

        # 可学习前缀嵌入（每层独立前缀）
        self.prefix_embeds = nn.Parameter(
            torch.randn(num_layers, 2, prefix_length, hidden_size) * 0.02
        )

    def forward(self, input_ids, attention_mask=None, labels=None):
        batch_size = input_ids.shape[0]
        past_key_values = []

        for layer_idx in range(self.num_layers):
            # 每层前缀拆分为 key/value 状态
            prefix_k = self.prefix_embeds[layer_idx, 0].unsqueeze(0).expand(batch_size, -1, -1)
            prefix_v = self.prefix_embeds[layer_idx, 1].unsqueeze(0).expand(batch_size, -1, -1)
            # 构造与  匹配的 past_key_values 格式
            past_key_values.append((prefix_k, prefix_v))

        # 将前缀作为 past_key_values 传入，实现前缀参与注意力计算
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            labels=labels,
            use_cache=False
        )
        return outputs

    def generate(self, prompt, max_length=50):
        self.eval()
        inputs = self.tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"].to(next(self.parameters()).device)
        attention_mask = inputs["attention_mask"].to(next(self.parameters()).device)

        with torch.no_grad():
            # 生成时同样注入前缀
            batch_size = input_ids.shape[0]
            past_key_values = []
            for layer_idx in range(self.num_layers):
                prefix_k = self.prefix_embeds[layer_idx, 0].unsqueeze(0).expand(batch_size, -1, -1)
                prefix_v = self.prefix_embeds[layer_idx, 1].unsqueeze(0).expand(batch_size, -1, -1)
                past_key_values.append((prefix_k, prefix_v))

            outputs = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                max_length=max_length,
                pad__id=self.tokenizer.eos__id,
                do_sample=True,
                temperature=0.8
            )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

def train_prefix(model, train_texts, epochs=3, lr=1e-3):
    """在给定文本上训练前缀向量"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    optimizer = torch.optim.AdamW([model.prefix_embeds], lr=lr)
    model.train()

    for epoch in range(epochs):
        total_loss = 0
        for text in train_texts:
            inputs = model.tokenizer(text, return_tensors="pt", truncation=True, max_length=64)
            input_ids = inputs["input_ids"].to(device)
            attention_mask = inputs["attention_mask"].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
            loss = outputs.loss
            total_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_texts):.4f}")

if __name__ == "__main__":
    # 1. 初始化模型
    model = PrefixTuningGPT2(model_name="", prefix_length=20)

    # 2. 微调前生成（观察基线效果）
    print("=== 微调前生成 ===")
    print(model.generate("The weather today is"))

    # 3. 构造训练数据：希望模型输出的风格
    train_data = [
        "The weather today is sunny and warm, perfect for a walk in the park.",
        "The weather today is rainy and cold, remember to bring an umbrella.",
        "The weather today is cloudy with a chance of snow in the evening.",
        "The weather today is windy and dry, ideal for flying kites."
    ]

    # 4. 训练前缀
    print("\n=== 开始训练前缀 ===")
    train_prefix(model, train_data, epochs=5, lr=1e-3)

    # 5. 微调后生成（对比效果）
    print("\n=== 微调后生成 ===")
    print(model.generate("The weather today is"))

运行与验证

运行方式：直接执行 python prefix_tuning.py。首次运行会自动下载模型（约500MB），需保持网络畅通。

验证方法：

1. 观察训练输出：Loss应随epoch逐步下降，说明前缀向量正在学习训练文本的分布特征。
2. 对比微调前后的生成结果：微调前模型输出为通用续写；微调后输出应明显偏向训练数据中的天气描述风格（如包含sunny/rainy/cloudy等词汇）。
3. 参数冻结检查：训练过程中仅 model.prefix_embeds 的 requires_grad 为True，主干模型参数保持冻结。可通过 print(model.model.lm_head.weight.requires_grad) 验证输出为False。

通过条件：微调后生成文本中，天气相关词汇出现频率显著高于微调前，且训练Loss较初始值下降超过30%。

常见问题

• 显存不足：可减小 prefix_length 至10或使用 -small 替代默认模型。
• 生成结果无变化：增大训练epoch至10或调高学习率至5e-3，确保前缀向量充分更新。
• past_key_values格式错误：的past_key_values要求每层包含(key, value)二元组，维度为(batch, num_heads, seq_len, head_dim)，本实现已按768维hidden_size自动匹配。