4-SFT

4-SFT#

在预训练阶段后，我们应该能够获得一个下一词预测模型，此时的模型已经掌握了大量的知识。不过，仅仅具备下一词预测能力是不够的，我们希望大模型能够获得问答能力，这一能力便是在有监督微调（Supervised Fine Tuning，SFT）阶段获得的.

在这个笔记本中，我们仅对 SFT 的训练流程进行展示和学习，因此只给出必要的代码片段，如 wandb 和 ddp 不会在此笔记本中涉及.

此笔记本的完整实现见主仓库 ../raw/data/Minimind/train_full_sft.py

# -------------------- 导入标准和第三方库 --------------------
# 导入 os 库，用于与操作系统交互
import os
# 导入 platform 库，用于获取系统信息
import platform
# 导入 argparse 库，用于解析命令行参数
import argparse
# 导入 time 库，用于计时
import time
# 导入 math 库，用于数学运算
import math
# 导入 warnings 库，用于控制警告信息
import warnings

# 导入 pandas 库，用于数据分析
import pandas as pd
# 导入 PyTorch 核心库
import torch
# 导入 PyTorch 的 functional 模块，包含常用函数
import torch.nn.functional as F
# 导入 PyTorch 的分布式训练库
import torch.distributed as dist
# 导入一个上下文管理器，用于在不同条件下执行代码块
from contextlib import nullcontext

# 导入 PyTorch 的优化器和神经网络模块
from torch import optim, nn
# 导入 PyTorch 的分布式数据并行工具
from torch.nn.parallel import DistributedDataParallel
# 导入 PyTorch 的数据加载器和分布式采样器
from torch.utils.data import DataLoader, DistributedSampler
# -------------------- 导入 Hugging Face 和自定义模块 --------------------
# 从 transformers 库导入 AutoTokenizer (自动加载分词器) 和 AutoModelForCausalLM (自动加载因果语言模型)
from transformers import AutoTokenizer, AutoModelForCausalLM
# 从我们自己的 model 文件夹中导入之前编写好的模块
from model.model import MiniMindLM # 我们的语言模型
from model.LMConfig import LMConfig # 模型配置类
from model.dataset import SFTDataset # 我们为 SFT 准备的数据集类

# 这行代码设置警告过滤器，让程序忽略所有的警告信息
# 这有助于保持输出的整洁，让我们能专注于关键的训练日志
warnings.filterwarnings('ignore')

可选参数设置#

首先，查看训练的可选参数，这些参数在实际使用时通过命令行导入，为了保持笔记本的易用性，选择用 class 进行包装.

# 使用一个 class 来模拟命令行参数，集中管理所有超参数
class args:
    # 训练过程超参数
    epochs: int = 1                  # 整个数据集要被训练多少遍
    batch_size: int = 2              # 每一步训练多少个样本。SFT 玩具数据集也只有2个样本
    learning_rate: float = 5e-4      # 学习率。在微调阶段，通常会使用比预训练更小的学习率，但为演示方便这里保持一致
    
    # 硬件和精度设置
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu' # 自动检测 GPU
    dtype: str = 'bfloat16'          # 使用 bfloat16 半精度进行训练
    
    # 日志和数据加载
    wandb_project: str = 'MiniMind-Notebook' 
    num_workers: int = 1             # DataLoader 使用的子进程数
    
    # 高级训练技巧
    accumulation_steps: int = 1      # 梯度累积步数
    grad_clip: float = 1.0           # 梯度裁剪阈值
    warmup_iters: int = 0            # 学习率预热步数
    log_interval: int = 1            # 每隔多少步打印一次日志
    
    # 分布式训练相关
    local_rank: int = 1              # 当前 GPU 编号

    # 模型结构超参数 (与预训练阶段的模型结构保持一致)
    dim: int = 512                   # 模型内部维度
    n_layers: int = 2                # Transformer Block 的层数
    max_seq_len: int = 512           # 最大序列长度
    use_moe: bool = False            # 是否使用混合专家(MoE)

    # 数据路径 (关键变化！)
    # 这里我们指向 SFT 的数据文件，而不是预训练数据
    data_path: str = '../raw/data/Minimind/toydata/sft_data.jsonl'

# 打印出 `args` 中设置的 `device`，确认程序将在哪个设备上运行（CPU 或 GPU）
print(f'查看工作设备 {args.device}')

查看工作设备 cuda

初始化训练#

接下来，我们对一些重要模块进行初始化，我们已经了解过，分词器，模型和数据集是大模型的基本组件，我们对其进行初始化.

注意与预训练阶段不同的是在 sft 阶段我们实际上是在上一阶段训练获得的模型的基础上修改数据集进行接续训练因此需要载入上一阶段的模型权重出于展示的目的载入权重的代码在此笔记本中只作展示并不执行

# 定义一个函数来封装 SFT 阶段的模型和分词器初始化过程
def init_model(lm_config):
    # 1. 加载分词器，这和预训练阶段完全一样
    tokenizer = AutoTokenizer.from_pretrained('../raw/data/Minimind/model/minimind_tokenizer')
    
    # 2. 根据相同的配置，重新创建一个与预训练阶段结构完全相同的模型
    model = MiniMindLM(lm_config).to(args.device)
    
    # --- 关键区别：加载预训练权重 ---
    # 下面的代码被注释掉了，但在真实的 SFT 流程中至关重要
    moe_path = '_moe' if lm_config.use_moe else ''
    # a. 定义预训练模型权重文件的路径
    # ckp = f'./out/pretrain_{lm_config.dim}{moe_path}.pth' 
    # b. 使用 torch.load 加载权重文件。`map_location` 确保权重被加载到正确的设备上
    # state_dict = torch.load(ckp, map_location=args.device)
    # c. 使用 `load_state_dict` 方法将加载的权重应用到我们新创建的模型上
    #    `strict=False` 表示如果模型结构和权重文件有轻微不匹配（比如某些层的名字不同），程序不会报错
    # model.load_state_dict(state_dict, strict=False)
    # ---------------------------------
    
    # 3. 计算并打印模型的总参数量
    print(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
    
    # 4. 确保模型在正确的设备上（虽然前面已经 .to() 过，但再次确认无害）
    model = model.to(args.device)
    
    # 5. 返回初始化好的模型和分词器
    return model, tokenizer

# --- 1. 创建模型配置 ---
# 使用 `args` 中的参数创建 LMConfig 对象，确保模型结构与预训练阶段一致
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)

# --- 2. 初始化模型和分词器 ---
# 调用 `init_model` 函数。注意：因为我们注释了加载权重的代码，所以这里的 `model` 实际上是一个随机初始化的新模型
model, tokenizer = init_model(lm_config)

# --- 3. 初始化数据集 (关键变化！) ---
# 这里我们使用 `SFTDataset` 类来加载 SFT 数据
train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)

# --- 4. 初始化数据加载器 (DataLoader) ---
# 这部分和预训练阶段的设置完全相同
train_loader = DataLoader(
    train_ds,
    batch_size=args.batch_size,
    pin_memory=True,
    drop_last=False,
    shuffle=False,
    num_workers=args.num_workers,
)

# --- 5. 打印确认信息 ---
print(f'模型位于设备：{model.device}, 词表长度：{tokenizer.vocab_size}, DataLoader：{train_loader}')

LLM总参数量：8.915 百万
模型位于设备：cuda:0, 词表长度：6400, DataLoader：<torch.utils.data.dataloader.DataLoader object at 0x000002260D603CA0>

# 这段代码用来检查 SFT 的 DataLoader

# 1. 将 DataLoader 转换成一个迭代器
loader = iter(train_loader)

# 2. 取出第一个 batch 的数据并打印
#    返回的仍然是 (X, Y, loss_mask) 三个张量，但它们的内容是由 SFTDataset 生成的
print(f'打印一个 iter 的数据:\n{next(loader)}\n')

# 3. 打印数据集和 DataLoader 的大小
#    我们的 SFT 玩具数据集也只有2个样本，batch_size=2，所以 DataLoader 的大小也是1
print(f'数据集大小：{len(train_ds)}, DataLoader 大小：{len(train_loader)}')

打印一个 iter 的数据:
[tensor([[  1,  85, 736,  ...,   0,   0,   0],
        [  1,  85, 736,  ...,   0,   0,   0]]), tensor([[ 85, 736, 201,  ...,   0,   0,   0],
        [ 85, 736, 201,  ...,   0,   0,   0]]), tensor([[0, 0, 0,  ..., 0, 0, 0],
        [0, 0, 0,  ..., 0, 0, 0]])]

数据集大小：2, DataLoader 大小：1

我们发现，train loader 的每一个 iter 都包含一个长度为 3 的张量列表，这是因为 train_dataset 每一次取数据都会返回三个张量，分别为:

样本 X: 包含 <bos> 在内的输入 conversation
标签 Y: 包含 <eos> 在内的输出 conversation
掩码 loss_mask: 指示需要计算损失的 token 位置

由于我们的数据集只有两条数据，而 batch size 设置为 2，因此我们的 dataloader 只有一个 iter.

启动训练#

训练一个深度学习模型，还涉及到了优化器，损失函数和学习率调度. 接下来，我们查看 MiniMind 训练部分的代码，并进行一轮简单的训练.

不难发现 pretrain 阶段和 sft 阶段的训练主体差不多因为这两个阶段的差异体现在数据集格式而数据集在经过 chat template 格式化后差异小了很多

# 这部分代码与预训练阶段完全相同，我们重新设置了训练所需的辅助工具

# --- 1. 定义学习率调度函数 ---
# 同样使用余弦退火学习率
def get_lr(current_step, total_steps, lr):
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

# --- 2. 设置混合精度训练的梯度缩放器 (Scaler) ---
# GradScaler 用于防止半精度训练时的梯度下溢
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))

# --- 3. 初始化优化器 ---
# 同样使用 AdamW 优化器
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

# --- 4. 设置自动混合精度上下文 (Autocast) ---
# 根据设备类型决定是否启用混合精度
device_type = "cuda" if "cuda" in args.device else "cpu"
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()

接下来，我们来看看 MiniMind 的训练函数

# 这个训练函数 `train_epoch` 的结构和代码与预训练阶段几乎完全一样
# 这展示了 PyTorch 训练流程的模块化特性：
# 只要你的 Dataset 输出格式是 (输入, 标签, 掩码)，那么训练循环的核心逻辑就可以保持不变。
# 唯一的区别在于，这次传入的 `loss_mask` 是 SFT 特有的（只在 assistant 回答处为1），
# 而不是预训练时的（在非 padding 处都为1）。

def train_epoch(epoch):
    loss_fct = nn.CrossEntropyLoss(reduction='none') # 交叉熵损失
    start_time = time.time()
    
    # 遍历 SFT DataLoader
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        # 将数据移动到 GPU
        X = X.to(args.device)
        Y = Y.to(args.device)
        loss_mask = loss_mask.to(args.device)

        # 更新学习率
        current_total_step = epoch * iter_per_epoch + step
        total_steps = args.epochs * iter_per_epoch
        lr = get_lr(current_total_step, total_steps, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # 在混合精度上下文中进行前向传播和损失计算
        with ctx:
            res = model(X)
            loss = loss_fct(res.logits.view(-1, res.logits.size(-1)), Y.view(-1)).view(Y.size())
            # 关键步骤：应用 SFT 的 loss_mask，只计算 assistant 回答部分的损失
            # `loss_mask.sum()` 可能会是0，如果一个 batch 恰好没有 assistant 的回答
            # 在实际应用中需要处理这种情况，例如加一个极小的数到分母
            if loss_mask.sum() > 0:
                loss = (loss * loss_mask).sum() / loss_mask.sum()
            else: # 如果这个 batch 没有需要计算损失的部分，则损失为0
                loss = torch.tensor(0.0).to(args.device)
            loss += res.aux_loss
            loss = loss / args.accumulation_steps

        # 反向传播和权重更新（与预训练完全相同）
        scaler.scale(loss).backward()
        if (step + 1) % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)

        # 打印日志（与预训练完全相同）
        if step % args.log_interval == 0:
            spend_time = time.time() - start_time
            print(
                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} ...'.format(
                    epoch + 1, args.epochs,
                    step, iter_per_epoch,
                    loss.item() * args.accumulation_steps,
                    optimizer.param_groups[-1]['lr'],
                    # ...
                )
            )

准备完毕，我们尝试一轮长度 1 个 iter 的训练.

# 计算每个 epoch 的迭代次数
iter_per_epoch = len(train_loader) # 结果是 1

# 开始主训练循环
# 由于 `args.epochs` 设置为 1，这个循环只会执行一次
for epoch in range(args.epochs):
    # 调用 `train_epoch` 函数，开始 SFT 训练
    train_epoch(epoch)

Epoch:[1/1](0/1) loss:9.002 lr:0.000550000000 epoch_Time:0.0min:

# SFT 训练演示结束后，删除模型对象以释放 GPU 显存
del model

4-SFT

Contents

4-SFT#

可选参数设置#

初始化训练#

启动训练#