自己训练大模型？MiniMind 全流程解析 (二) 监督微调SFT

MiniMind 监督微调（SFT）全流程解析

MiniMind 是一个高效、灵活的大语言模型框架，旨在提供完整的模型训练、微调和推理解决方案。本教程详细解析 MiniMind 的监督微调（SFT）流程，涵盖从数据准备到模型保存的完整技术实现。

一、整体流程概述

二、SFT核心技术详解

1. 对话模板处理

SFT阶段的核心是将原始对话数据转换为模型可理解的格式：

源码位置：dataset/lm_dataset.py 第78-87行

def _create_chat_prompt(self, conversations):
    """构建符合ChatML格式的对话"""
    messages = []
    for i, turn in enumerate(conversations):
        role = 'user' if i % 2 == 0 else 'assistant'
        messages.append({"role": role, "content": turn['content']})
    return self.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=False
    )

训练数据如下

{"conversations": [{"role": "user", "content": "请告诉我在中国古代的“四大发明”是什么？"}, {"role": "assistant", "content": "中国古代的“四大发明”是指造纸术、印刷术、火药和指南针。这四项发明对世界文明的发展产生了深远的影响：\n\n1. **造纸术**：据史书记载，东汉时期的蔡伦改进了造纸工艺，使得纸张的生产更加便捷、成本更低，质量也更加优良。这一发明极大地促进了文化的传播和保存。\n\n2. **印刷术**：中国古代的印刷术最早可以追溯到唐代的雕版印刷，到了宋代发展出了活字印刷技术。印刷术的发明极大地促进了知识的传播，降低了书籍的成本，对教育和文化的发展起到了重要的推动作用。\n\n3. **火药**：火药最初是在唐代被发现的，最初可能用于医疗或炼金术。到了宋代，人们开始将火药用于军事目的，发明了各种火器。火药的发明改变了战争的面貌，对世界军事技术的发展产生了重大影响。\n\n4. **指南针**：指南针最初在中国被用于风水测量，后来逐渐发展成为航海导航的重要工具。这一发明极大地促进了海上航行技术的发展，对于新航路的开辟和世界地理大发现起到了关键作用。\n\n这四项发明不仅在中国历史上占有重要地位，而且对全世界的科技进步和文明发展都产生了深远的影响。"}]}

优势：

• 标准化对话格式，支持多轮对话
• 与模型输入格式完全兼容
• 自动处理角色标识和分隔符

2. 损失掩码机制

SFT的关键创新是只对助手回复部分计算损失：

源码位置：dataset/lm_dataset.py 第89-108行

def _generate_loss_mask(self, input_ids):
    loss_mask = [0] * len(input_ids)
    i = 0
    while i < len(input_ids):
        if input_ids[i:i + len(self.bos_id)] == self.bos_id:
            start = i + len(self.bos_id)
            end = start
            while end < len(input_ids):
                if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                    break
                end += 1
            for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
                loss_mask[j] = 1
            i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
        else:
            i += 1
    return loss_mask

核心原理：

• 仅对助手回复部分（BOS到EOS之间）计算损失
• 忽略用户输入部分的梯度更新
• 显著提高训练效率和对话质量

3. 预训练模型加载

SFT与预训练的关键区别在于模型初始化方式：

源码位置：train_full_sft.py 第116-128行

def init_model(lm_config):
    tokenizer = AutoTokenizer.from_pretrained('../model')
    model = MiniMindForCausalLM(lm_config)
    
    # 加载预训练权重
    moe_path = '_moe' if lm_config.use_moe else ''
    ckp = f'{args.save_dir}/pretrain_{lm_config.hidden_size}{moe_path}.pth'
    state_dict = torch.load(ckp, map_location=args.device)
    model.load_state_dict(state_dict, strict=False)
    
    Logger(f'LLM可训练总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
    return model.to(args.device), tokenizer

特点：

• 加载预训练权重作为初始化
• 支持部分参数加载（strict=False）
• 自动计算可训练参数数量
• 支持MoE和标准模型切换

代码所涉及的系统具备能够在两种不同模型架构（混合专家模型（Mixture of Experts，MoE ）和标准模型）之间进行切换的能力，具体可从以下几个方面理解：
模型架构的区别
标准模型：是一种常见的神经网络架构，模型中各个神经元或模块以固定的方式进行连接和计算，在处理任务时，所有的输入数据都经过相同的网络结构进行处理。例如常见的 Transformer 架构模型，在自然语言处理任务中，无论是处理一句话还是一段文本，数据都按照固定的 Transformer 结构的前向传播路径进行计算。
混合专家模型（MoE）：是一种更复杂且灵活的模型架构，它将模型的功能拆分成多个 “专家” 模块，每个 “专家” 模块负责处理输入数据的一部分。在处理输入时，根据输入数据的特点，动态地选择一个或多个 “专家” 模块来处理数据，这样可以更高效地处理不同类型的数据，在大规模数据和复杂任务上表现出更好的性能。比如在一个处理多种语言文本的任务中，不同语言相关的文本可能由不同的 “专家” 模块来处理。
代码中实现切换的方式
在代码中，通过 lm_config.use_moe 这个配置项来控制模型架构的选择：
当 lm_config.use_moe 为 True 时，会构建基于 MoE 架构的模型。从代码中的 moe_path = ‘_moe’ if lm_config.use_moe else ‘’ 这一行可以看出，通过设置 moe_path 变量来标识这是一个 MoE 模型，后续加载预训练权重时，可能会根据这个标识去加载特定的 MoE 架构的预训练权重文件。
当 lm_config.use_moe 为 False 时，构建的就是标准模型。此时不会按照 MoE 架构相关的逻辑去处理，而是按照标准模型的方式加载预训练权重并进行后续的模型初始化等操作。
切换的意义
灵活性：在不同的任务场景下，不同的模型架构可能会有不同的表现。通过支持 MoE 和标准模型切换，用户或开发者可以根据实际任务需求（如数据规模、任务复杂度、计算资源等），灵活地选择合适的模型架构，以达到更好的性能表现。
实验和优化：对于研究人员和算法开发者来说，这种切换功能提供了便利，可以方便地对比 MoE 模型和标准模型在同一任务上的效果差异，从而进行实验和算法优化，探索更优的模型方案。

4. SFT专用损失计算

SFT阶段的损失计算与预训练有显著差异：

源码位置：train_full_sft.py 第60-75行

def train_epoch(epoch, wandb):
    loss_fct = nn.CrossEntropyLoss(reduction='none')
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        X, Y, loss_mask = X.to(args.device), Y.to(args.device), loss_mask.to(args.device)
        
        with ctx:
            res = model(X)
            loss = loss_fct(
                res.logits.view(-1, res.logits.size(-1)),
                Y.view(-1)
            ).view(Y.size())
            # 关键：应用损失掩码
            loss = (loss * loss_mask).sum() / loss_mask.sum()
            loss += res.aux_loss  # 添加MoE辅助损失
            loss = loss / args.accumulation_steps

与预训练的区别：

• 使用损失掩码只计算助手回复部分
• 保持MoE辅助损失的计算
• 损失归一化方式不同

三、SFT数据格式与处理

1. SFT数据格式

所有SFT数据文件采用统一的JSONL格式：

{
    "conversations": [
        {"role": "user", "content": "你好"},
        {"role": "assistant", "content": "你好！很高兴为您服务。"},
        {"role": "user", "content": "请介绍一下你自己"},
        {"role": "assistant", "content": "我是MiniMind，一个基于Transformer架构的大语言模型..."}
    ]
}

2. 数据集选择建议

根据训练需求和GPU资源选择合适的数据集：

• 快速验证：sft_mini_512.jsonl (~1.2GB) ✨推荐
• 标准训练：sft_512.jsonl (~7.5GB)
• 高质量对话：sft_1024.jsonl (~5.5GB)
• 长文本处理：sft_2048.jsonl (~9GB)
• 领域定制：lora_identity.jsonl、lora_medical.jsonl

3. SFT数据预处理流程

源码位置：dataset/lm_dataset.py

class SFTDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = self._load_data(data_path)
        self.bos_id = tokenizer.encode('<|im_start|>')
        self.eos_id = tokenizer.encode('<|im_end|>')
        
    def __getitem__(self, index):
        conversations = self.data[index]['conversations']
        prompt = self._create_chat_prompt(conversations)
        input_ids = self.tokenizer.encode(prompt)
        loss_mask = self._generate_loss_mask(input_ids)
        
        # 截断和填充
        if len(input_ids) > self.max_length:
            input_ids = input_ids[:self.max_length]
            loss_mask = loss_mask[:self.max_length]
            
        return torch.tensor(input_ids), torch.tensor(input_ids), torch.tensor(loss_mask)

关键特点：

• 自动识别对话角色和边界
• 生成精确的损失掩码
• 支持多轮对话处理

四、SFT训练流程详解

1. 初始化与配置

SFT的参数配置与预训练有所不同：

源码位置：train_full_sft.py 第130-150行

def main():
    parser = argparse.ArgumentParser(description="MiniMind Full SFT")
    parser.add_argument("--out_dir", type=str, default="../out")
    parser.add_argument("--epochs", type=int, default=2)  # SFT通常需要更少轮次
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--learning_rate", type=float, default=5e-7)  # 更小的学习率
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--dtype", type=str, default="bfloat16")
    parser.add_argument("--use_wandb", action="store_true")
    parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
    # ... 其他参数 ...
    args = parser.parse_args()

2. SFT数据加载器配置

源码位置：train_full_sft.py 第151-165行

model, tokenizer = init_model(lm_config)
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if ddp else None
train_loader = DataLoader(
    train_ds,
    batch_size=args.batch_size,
    pin_memory=True,
    drop_last=False,
    shuffle=False,
    num_workers=args.num_workers,
    sampler=train_sampler
)

与预训练数据加载的区别：

• 使用SFTDataset而非PretrainDataset
• 返回三元组：(input_ids, labels, loss_mask)
• 专门处理对话格式数据

3. SFT模型保存策略

SFT模型保存与预训练略有不同：

源码位置：train_full_sft.py 第103-114行

if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
    model.eval()
    moe_path = '_moe' if lm_config.use_moe else ''
    # 注意：保存为full_sft而非pretrain
    ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
    
    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
        state_dict = model.module.state_dict()
    else:
        state_dict = model.state_dict()
    
    state_dict = {k: v.half() for k, v in state_dict.items()}  # 半精度保存
    torch.save(state_dict, ckp)
    model.train()

保存文件命名规则：

• 预训练：pretrain_{hidden_size}.pth
• SFT：full_sft_{hidden_size}.pth
• LoRA：lora_{hidden_size}.pth

五、SFT训练启动与调优

1. 基础训练命令

# 单GPU训练
python train_full_sft.py \
    --data_path ./dataset/sft_mini_512.jsonl \
    --max_seq_len 512 \
    --batch_size 16 \
    --learning_rate 5e-7 \
    --epochs 2 \
    --save_interval 100

# 多GPU分布式训练
torchrun --nproc_per_node 2 train_full_sft.py \
    --data_path ./dataset/sft_1024.jsonl \
    --max_seq_len 1024 \
    --batch_size 8 \
    --learning_rate 3e-7 \
    --use_wandb

# 启用wandb监控
python train_full_sft.py \
    --use_wandb \
    --wandb_project "MiniMind-SFT" \
    --wandb_run_name "sft-512-experiment"

2. SFT专用参数调优

参数	预训练推荐值	SFT推荐值	说明
learning_rate	5e-4	5e-7 ~ 1e-6	SFT需要更小学习率
epochs	1-6	2-5	SFT通常需要更少轮次
batch_size	32	8-32	根据对话长度调整
max_seq_len	512	512-2048	对话通常更长

3. SFT显存优化策略

SFT特有的显存挑战：

• 对话数据通常比预训练数据更长
• 损失掩码增加额外内存开销
• 多轮对话增加序列复杂度

优化方案：

# 针对长对话的优化
python train_full_sft.py \
    --max_seq_len 1024 \
    --batch_size 4 \
    --accumulation_steps 8 \
    --dtype bfloat16

# 针对短对话的快速训练
python train_full_sft.py \
    --max_seq_len 512 \
    --batch_size 32 \
    --accumulation_steps 1

六、SFT效果评估与调试

1. SFT训练监控指标

除了通用的训练指标外，SFT还需要关注：

# 在wandb中记录SFT专用指标
if wandb and step % args.log_interval == 0:
    wandb.log({
        "train_loss": loss.item() * args.accumulation_steps,
        "learning_rate": optimizer.param_groups[-1]['lr'],
        "effective_tokens": loss_mask.sum().item(),  # 有效训练token数
        "mask_ratio": loss_mask.sum().item() / loss_mask.numel(),  # 掩码比例
        "epoch": epoch,
        "step": step
    })

关键指标说明：

• effective_tokens：实际参与损失计算的token数量
• mask_ratio：损失掩码的覆盖比例，反映助手回复占比

2. SFT常见问题与解决方案

问题1：模型只会重复或胡言乱语

可能原因：

• 预训练模型加载失败
• 学习率过大导致灾难性遗忘
• 数据质量问题

解决方案：

# 检查预训练模型是否正确加载
python -c "
import torch
state_dict = torch.load('./out/pretrain_512.pth')
print('预训练模型参数数量:', len(state_dict))
print('模型键值示例:', list(state_dict.keys())[:5])
"

# 降低学习率
python train_full_sft.py --learning_rate 1e-7

# 检查数据格式
head -n 1 ./dataset/sft_mini_512.jsonl | python -m json.tool

问题2：损失不下降或下降缓慢

可能原因：

• 损失掩码设置错误
• 有效训练数据太少
• 学习率过小

调试方法：

# 检查损失掩码是否正确
from dataset.lm_dataset import SFTDataset
dataset = SFTDataset('./dataset/sft_mini_512.jsonl', tokenizer)
input_ids, labels, loss_mask = dataset[0]
print(f"序列长度: {len(input_ids)}")
print(f"掩码覆盖率: {loss_mask.sum().item() / len(loss_mask):.2%}")
print(f"有效token数: {loss_mask.sum().item()}")

问题3：对话格式不正确

症状：模型输出格式混乱，不遵循对话模板

解决方案：

• 检查tokenizer的chat_template设置
• 验证BOS/EOS token的正确性
• 确保数据预处理正确

# 验证对话模板
python -c "
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('./model')
messages = [{'role': 'user', 'content': '你好'}]
result = tokenizer.apply_chat_template(messages, tokenize=False)
print('对话模板结果:', result)
"

七、SFT与其他微调方法对比

1. Full SFT vs LoRA

特性	Full SFT	LoRA
参数更新	全部参数	低秩矩阵
显存需求	高	低
训练速度	慢	快
效果质量	最佳	良好
适用场景	充足资源	资源受限

2. 选择建议

• 使用Full SFT的情况：

  • GPU显存充足（>8GB）
  • 追求最佳对话效果
  • 有充足的高质量数据

• 使用LoRA的情况：

  • GPU显存受限（<8GB）
  • 快速原型验证
  • 领域特定微调

八、SFT最佳实践总结

1. 数据准备最佳实践

• 数据质量：优先选择高质量、多样化的对话数据
• 数据长度：根据GPU显存选择合适的序列长度
• 数据平衡：确保用户和助手回复的平衡性

2. 训练策略最佳实践

• 学习率：从预训练学习率的1/100开始尝试
• 训练轮次：通常2-5个epoch即可，避免过拟合
• 保存策略：定期保存，便于回滚到最佳状态

3. 监控与调试最佳实践

• 实时监控：使用wandb跟踪损失和学习率变化
• 定期验证：人工检查模型输出质量
• 渐进式训练：从小数据集开始，逐步扩大

通过以上SFT专用的技术解析和实践指导，你可以高效地完成MiniMind模型的监督微调，获得优秀的对话能力。SFT是从预训练模型到实用对话模型的关键步骤，掌握这些技术要点将帮助你构建高质量的对话AI系统。