大语言模型监督式微调(SFT)实战指南

1. 项目概述与背景

在大语言模型（LLM）领域，监督式微调（Supervised Fine-Tuning, SFT）是模型适应特定任务的关键步骤。这个项目基于Human-Like-DPO数据集，对SmolLM2-135M-Instruct模型进行微调，为后续的PPO强化学习训练做准备。整个过程涉及数据处理、模型训练和生成测试三个核心环节。

提示：虽然项目使用的是小型模型（135M参数），但同样的方法可以应用于更大的模型，只需调整batch size和梯度累积步数以适应显存限制。

2. 环境准备与数据加载

2.1 环境配置

首先需要设置Hugging Face镜像源以加速数据集下载。国内用户经常会遇到数据集下载缓慢的问题，通过修改环境变量可以显著提升加载速度：

python复制import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

2.2 数据集加载

项目使用了Human-Like-DPO数据集，这是一个专门为强化学习设计的数据集，包含prompt以及对应的优选(chosen)和非优选(rejected)回答：

python复制from datasets import load_dataset

all_data = load_dataset("HumanLLMs/Human-Like-DPO-Dataset")
train_data = all_data["train"]

数据集结构如下：

code复制Dataset({
    features: ['prompt', 'chosen', 'rejected'],
    num_rows: 10884
})

这个数据集的设计非常适合后续的强化学习训练，因为它已经标注了哪些回答更符合人类偏好。

3. 数据预处理

3.1 格式转换

为了与ChatML格式兼容，需要对数据进行格式化处理：

python复制def format_to_chatml(example):
    return {
        "chosen": f'<|im_start|>user\n{example["prompt"]}<|im_end|>\n<|im_start|>assistant\n{example["chosen"]}<|im_end|>',
        "rejected": f'<|im_start|>user\n{example["prompt"]}<|im_end|>\n<|im_start|>assistant\n{example["rejected"]}<|im_end|>',
    }

original_columns = train_data.column_names
train_data = train_data.map(format_to_chatml, remove_columns=original_columns)

这种格式添加了对话标记，明确区分用户输入和助手回复，有助于模型理解对话结构。

3.2 数据批处理

创建DataLoader时需要注意几个关键点：

1. 批处理函数：需要正确处理padding和标签生成
2. 注意力掩码：标识哪些位置是真实token，哪些是padding
3. 标签生成：将输入右移一位作为预测目标

python复制from torch.utils.data import DataLoader

batch_size = 2
gradient_accumulation_steps = 8

def collate_batch(batch):
    # 详细实现见下文
    pass

dataloader = DataLoader(
    train_data,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=collate_batch
)

4. 模型准备与训练

4.1 模型加载

使用HuggingFace提供的预训练模型SmolLM2-135M-Instruct：

python复制import torch
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer

device = torch.device("cuda")
torch.set_default_dtype(torch.bfloat16)  # 与模型训练精度一致

base_model_name = "HuggingFaceTB/SmolLM2-135M-Instruct"

config = AutoConfig.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    config=config,
).to(device)

tokenizer = AutoTokenizer.from_pretrained(base_model_name)

4.2 训练配置

训练采用以下关键配置：

• 优化器：AdamW
• 学习率：9e-6
• 学习率调度：余弦退火
• 梯度累积：8步（有效batch size=16）

python复制num_epochs = 1
num_steps = math.ceil(len(dataloader) / gradient_accumulation_steps)

optimizer = torch.optim.AdamW(
    params=base_model.parameters(),
    lr=9.0e-6,
    betas=(0.9, 0.999),
    eps=1e-08,
)

# 学习率调度器实现
def _get_cosine_schedule(current_step, num_training_steps, num_warmup_steps=0, linear_warmup=False, min_value=0.0):
    # 实现细节...
    pass

scheduler = LambdaLR(optimizer, lr_lambda=functools.partial(
    _get_cosine_schedule,
    num_training_steps=num_epochs*num_steps,
    min_value=0.3,
))

4.3 训练循环

训练过程中需要注意的几个关键点：

1. 梯度累积的实现
2. 损失计算（交叉熵）
3. 模型状态的保存

python复制for epoch in range(num_epochs):
    base_model.train()
    optimizer.zero_grad()
    record_loss = []

    for i, (inputs, labels, masks) in enumerate(dataloader):
        with torch.set_grad_enabled(True):
            outputs = base_model(input_ids=inputs, attention_mask=masks)
            loss = F.cross_entropy(outputs.logits.transpose(1,2), labels)
            record_loss.append(loss.item())

            loss.backward()
            if ((i + 1) % gradient_accumulation_steps == 0) or (i + 1 == len(dataloader)):
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

    # 保存检查点
    base_model.save_pretrained("./llm_sft")

5. 模型测试与生成

5.1 文本生成实现

使用动态缓存(DynamicCache)来加速生成过程：

python复制from transformers import DynamicCache
from torch.nn import functional as F

def generate_token_by_policy(chat_data, model, tokenizer, max_seq_len):
    # 初始化缓存和输入
    batch_size = chat_data["input_ids"].shape[0]
    cur_iids = chat_data["input_ids"]
    cur_mask = chat_data["attention_mask"]
    proceed_flag = torch.ones(batch_size, dtype=bool).to(device)
    cache_position = None
    past_key_values = DynamicCache()

    while(torch.any(proceed_flag)):
        # 生成逻辑...
        pass
    
    return cur_iids, cur_mask

5.2 生成测试

测试模型的生成能力：

python复制max_seq_len = 768

messages = [
    "What do you most want to do right now?",
    "What is the best gift to give a friend who loves the outdoors?",
    "How do you relax after something bad happens?",
]

# 准备输入
inputs = [f"<|im_start|>user\n{m}<|im_end|>\n<|im_start|>assistant\n" for m in messages]
input_batch = tokenizer(
    inputs,
    padding=True,
    padding_side="left",
    return_tensors="pt").to(device)

# 加载微调后的模型
base_model = AutoModelForCausalLM.from_pretrained("./llm_sft").to(device)
base_model.eval()

# 生成回复
with torch.no_grad():
    iids, mask = generate_token_by_policy(
        input_batch,
        base_model,
        tokenizer,
        max_seq_len,
    )

# 解码并打印结果
iids = iids[:,input_seq_len:]
outputs = tokenizer.batch_decode(iids, skip_special_tokens=True)

6. 常见问题与解决方案

6.1 中文支持问题

从测试结果看，模型对中文的支持不够理想。这可能是因为：

1. 预训练数据中中文占比低
2. 微调数据集中没有中文样本
3. tokenizer对中文的分词效率低

解决方案：

• 在微调数据中加入中文样本
• 使用支持多语言的tokenizer
• 考虑专门针对中文进行继续预训练

6.2 训练不稳定

小模型训练时容易出现不稳定的情况，可以尝试：

1. 降低学习率
2. 增加梯度累积步数
3. 使用学习率warmup
4. 尝试不同的优化器参数

6.3 显存不足

对于资源有限的开发者：

1. 减小batch size
2. 使用梯度检查点
3. 尝试混合精度训练
4. 考虑使用LoRA等参数高效微调方法

7. 性能优化技巧

7.1 缓存利用

生成过程中使用DynamicCache可以显著减少计算量：

python复制past_key_values = DynamicCache()
# 在模型调用时传递cache参数
outputs = model(
    input_ids=inputs,
    attention_mask=attention_mask,
    past_key_values=past_key_values,
    use_cache=True,
)

7.2 批处理生成

通过保持批次维度进行生成，可以充分利用GPU并行计算能力：

python复制# 同时处理多个输入
messages = ["Q1", "Q2", "Q3"]
input_batch = tokenizer(messages, padding=True, return_tensors="pt").to(device)
outputs = model.generate(**input_batch)

7.3 量化推理

训练完成后，可以考虑将模型量化为8位或4位以减少推理时的资源消耗：

python复制from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    "./llm_sft",
    quantization_config=quantization_config,
)

8. 后续改进方向

1. 强化学习微调：在SFT基础上进行PPO训练，进一步优化模型输出
2. 多轮对话支持：扩展模型处理多轮对话的能力
3. 领域适配：针对特定领域（如客服、教育）进行专项优化
4. 评估体系：建立系统的评估方法，量化模型改进效果

这个项目完整展示了从数据准备到模型微调的全流程，为后续更高级的强化学习训练打下了坚实基础。在实际应用中，可以根据具体需求调整模型规模、训练数据和微调策略。