【大模型】- 指令微调（SFT）

from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

class InstructionDataset(Dataset):
    def __init__(self, data, tokenizer, max_length=512):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        item = self.data[idx]
        
        # 格式化提示
        if item["input"]:
            prompt = f"### 指令:\n{item['instruction']}\n\n### 输入:\n{item['input']}\n\n### 回答:\n"
        else:
            prompt = f"### 指令:\n{item['instruction']}\n\n### 回答:\n"
        
        full_text = prompt + item["output"] + self.tokenizer.eos_token
        
        # 分词
        encodings = self.tokenizer(
            full_text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        # 创建标签（只在回答部分计算损失）
        input_ids = encodings["input_ids"].squeeze()
        labels = input_ids.clone()
        
        # 将指令部分的标签设为 -100
        prompt_length = len(self.tokenizer(prompt)["input_ids"])
        labels[:prompt_length] = -100
        
        return {
            "input_ids": input_ids,
            "attention_mask": encodings["attention_mask"].squeeze(),
            "labels": labels
        }

from peft import LoraConfig, get_peft_model, TaskType

def apply_lora(model):
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=8,           # LoRA 秩
        lora_alpha=32, # 缩放因子
        lora_dropout=0.1,
        target_modules=["q_proj", "v_proj"]  # 要适配的层
    )
    
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    # trainable params: 4,718,592 || all params: 124,439,808 || 3.79%
    
    return model

from transformers import get_linear_schedule_with_warmup

def train_sft(model, train_dataloader, val_dataloader, config):
    model = apply_lora(model)
    model = model.to(config["device"])
    
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config["lr"],
        weight_decay=config["weight_decay"]
    )
    
    # 学习率调度
    num_training_steps = len(train_dataloader) * config["epochs"]
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=config["warmup_steps"],
        num_training_steps=num_training_steps
    )
    
    for epoch in range(config["epochs"]):
        model.train()
        total_loss = 0
        
        for batch in train_dataloader:
            # 移到设备
            input_ids = batch["input_ids"].to(config["device"])
            attention_mask = batch["attention_mask"].to(config["device"])
            labels = batch["labels"].to(config["device"])
            
            # 前向传播
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )
            
            loss = outputs.loss
            
            # 反向传播
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            
            total_loss += loss.item()
        
        # 验证
        val_loss = evaluate(model, val_dataloader, config)
        
        print(f"Epoch {epoch + 1}:")
        print(f"  Train Loss: {total_loss / len(train_dataloader):.4f}")
        print(f"  Val Loss: {val_loss:.4f}")
    
    return model

def evaluate(model, dataloader, config):
    model.eval()
    total_loss = 0
    
    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(config["device"])
            attention_mask = batch["attention_mask"].to(config["device"])
            labels = batch["labels"].to(config["device"])
            
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels
            )
            
            total_loss += outputs.loss.item()
    
    return total_loss / len(dataloader)

def generate_response(model, tokenizer, instruction, input_text=""):
    model.eval()
    
    if input_text:
        prompt = f"### 指令:\n{instruction}\n\n### 输入:\n{input_text}\n\n### 回答:\n"
    else:
        prompt = f"### 指令:\n{instruction}\n\n### 回答:\n"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True,
            top_p=0.9
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### 回答:\n")[-1].strip()

# 保存 LoRA 权重
model.save_pretrained("./sft_model")

# 加载
from peft import PeftModel
base_model = AutoModelForCausalLM.from_pretrained("base_model")
model = PeftModel.from_pretrained(base_model, "./sft_model")

# 可选：合并权重
model = model.merge_and_unload()