先问一个问题

训练一个 LLM 需要：多 GPU 并行、断点续训、学习率调度、混合精度……

6 个训练器（pretrain / sft / lora / dpo / grpo / ppo）各自实现一遍吗？

本章路线图

阶段 11 覆盖支撑训练的基础设施代码。先建立全局视野：

基础设施总览

trainer/trainer_utils.py 提供了训练流程中的公共基础设施：

9 步训练模式

定义：所有训练器都遵循完全相同的 9 步初始化→训练→清理流程。

白话：6 个训练器的"骨架"一模一样，就像同一个工厂流水线，只是装的"零件"不同。

以下是训练脚本的9个关键步骤，分为初始化（步骤1-5）和训练循环（步骤6-9）两个阶段。

各训练器差异

理解了 9 步框架，每个训练器你只需要关注它的差异部分：

混合精度训练

定义：混合精度用 FP16 做前向/反向计算（快），同时维护一份 FP32 的 master weights（准）。

白话：像用草稿纸算数学题（FP16，快但不精确），但把最终答案工整抄到正式本上（FP32，精确）。

参数统计

定义：get_model_params() 区分 MoE 模型的"总参数"和"活跃参数"。

白话：MoE 模型像个大公司——员工很多（总参数），但每次只派一小撮人干活（活跃参数）。

代码——来自 trainer_utils.py：

def get_model_params(model, config):
    total = sum(p.numel() for p in model.parameters()) / 1e6
    n_routed = config.n_routed_experts    # 总专家数
    n_active = config.num_experts_per_tok  # 每次激活几个
    expert = sum(p.numel() for n, p in model.named_parameters()
                 if 'mlp.experts.0.' in n) / 1e6
    shared_expert = sum(p.numel() for n, p in model.named_parameters()
                        if 'mlp.shared_experts.0.' in n) / 1e6
    base = total - (expert * n_routed) - (shared_expert * n_shared)
    active = base + (expert * n_active) + (shared_expert * n_shared)
    if active < total:
        Logger(f'Model Params: {total:.2f}M-A{active:.2f}M')
    else:
        Logger(f'Model Params: {total:.2f}M')

Model Params: 198.42M-A63.91M
# 总参数 198.42M，但每次推理只激活 63.91M

MoE 参数分布

MoE 模型里，不是所有参数每次都被用到——理解"总参数"和"活跃参数"的差别：

余弦学习率

定义：get_lr() 实现余弦退火调度，让学习率在训练过程中平滑变化。

白话：起步时学习率较低（0.1倍），保证训练稳定；中段学习率最高（约0.55倍），是学习最积极的阶段；末段学习率再次降低，精细收敛不冲过头。

代码：

def get_lr(current_step, total_steps, lr):
    return lr * (0.1 + 0.45 * (1 + math.cos(
        math.pi * current_step / total_steps)))

lr 调度详解

随机种子

定义：setup_seed() 统一设置所有随机源，保证每次运行结果可复现。

白话：深度学习的每一步都有随机性（数据打乱、Dropout、采样）。固定种子就是"抽签前先洗好牌"，保证下次抽到一样的。

代码：

def setup_seed(seed: int):
    random.seed(seed)               # Python 内置随机
    np.random.seed(seed)            # NumPy 随机
    torch.manual_seed(seed)         # PyTorch CPU 随机
    torch.cuda.manual_seed(seed)    # 当前 GPU 随机
    torch.cuda.manual_seed_all(seed)  # 所有 GPU 随机
    torch.backends.cudnn.deterministic = True   # cuDNN 确定性
    torch.backends.cudnn.benchmark = False       # 关闭自动调优

主进程日志

定义：Logger() 确保在多 GPU 训练时，日志只在主进程打印一次。

白话：4 张 GPU 跑同样的代码，不加判断的话，一条日志打印 4 遍，终端就炸了。

代码：

def is_main_process():
    return not dist.is_initialized() or dist.get_rank() == 0

def Logger(content):
    if is_main_process():   # 只在 rank=0 的主 GPU 上打印
        print(content)

检查点生命周期

定义：lm_checkpoint() 是训练工程中最复杂的组件，负责保存/恢复训练状态。

白话：训练一个模型可能跑几天。如果中间停电崩溃了，没有检查点就只能从头来。检查点就是"游戏存档"。

检查点内容

一个完整的检查点文件包含训练恢复所需的全部状态：

代码——保存模式的核心逻辑：

resume_data = {
    'model': state_dict,
    'optimizer': optimizer.state_dict(),
    'epoch': epoch,
    'step': step,
    'world_size': dist.get_world_size(),
    'wandb_id': wandb_id
}
resume_tmp = resume_path + '.tmp'
torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path)  # 原子操作！

原子写入

定义：原子写入指文件操作"要么全部完成，要么完全不变"，不存在中间状态。

白话：先写到临时文件，写完了才把名字改成正式的。写一半断了？临时文件坏了就坏了，正式文件好好的。

GPU 数量适配

定义：断点续训时如果换了 GPU 数量，需要按比例换算 step。

白话：8 张卡一起训练时，每一步吃 8 个 batch。换到 2 张卡，每步只吃 2 个。要看等量的数据，step 得翻倍。

代码：

saved_ws = ckp_data.get('world_size', 1)
current_ws = dist.get_world_size() if dist.is_initialized() else 1
if saved_ws != current_ws:
    ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws
    Logger(f'GPU数量变化({saved_ws}→{current_ws})，'
           f'step已自动转换为{ckp_data["step"]}')

断点续训跳过

定义：SkipBatchSampler 在断点续训时跳过已经处理过的 batch，避免数据重复。

白话：假设你在第 500 步崩了。重启后，前 500 个 batch 已经训练过了，再训一遍就是浪费。这个工具帮你跳过它们。

代码——完整实现：

class SkipBatchSampler(Sampler):
    def __init__(self, sampler, batch_size, skip_batches=0):
        self.sampler = sampler
        self.batch_size = batch_size
        self.skip_batches = skip_batches

    def __iter__(self):
        batch, skipped = [], 0
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                if skipped < self.skip_batches:
                    skipped += 1
                    batch = []
                    continue      # 跳过已处理的 batch！
                yield batch
                batch = []

    def __len__(self):
        total = (len(self.sampler) + self.batch_size - 1) // self.batch_size
        return max(0, total - self.skip_batches)

断点续训：完整故事

把检查点保存、恢复、GPU适配、数据跳过串起来：

奖励模型封装

定义：LMForRewardModel 封装了外部奖励模型（如 internlm2-1_8b-reward），在 RL 训练中给模型回答打分。

白话：RL 训练需要一个"裁判"来打分。这个类就是裁判的封装——给它对话和回答，它返回一个分数。

代码——完整实现：

class LMForRewardModel:
    def __init__(self, model_path, device="cuda", dtype=torch.float16):
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(
            model_path, torch_dtype=dtype, trust_remote_code=True)
        self.model = self.model.to(device).eval()

    @torch.no_grad()
    def get_score(self, messages, response):
        # 把对话历史和新问题拼成评估输入
        history_text = "\n".join(
            [f"{m['role']}: {m['content']}" for m in messages[:-1]])
        last_query = messages[-1]['content'] if messages else ""
        eval_messages = [
            {"role": "user", "content": f"{history_text}\n{last_query}"},
            {"role": "assistant", "content": response}
        ]
        score = self.model.get_score(self.tokenizer, eval_messages)
        return max(min(score, 3.0), -3.0)  # 裁剪到 [-3, 3]

奖励模型在 RL 中的角色

打分只是第一步——分数如何反馈给模型？

模型初始化入口

定义：init_model() 是所有训练器创建模型的统一入口。

白话：不管哪个训练器，"造模型"这一步都调它。参数决定了是随机初始化还是加载已有权重。

代码——完整实现：

def init_model(lm_config, from_weight='pretrain',
               tokenizer_path='../model', save_dir='../out', device='cuda'):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    model = MiniMindForCausalLM(lm_config)

    if from_weight != 'none':
        # 加载已有权重（微调 / 继续训练）
        moe_suffix = '_moe' if lm_config.use_moe else ''
        weight_path = f'{save_dir}/{from_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
        weights = torch.load(weight_path, map_location=device)
        model.load_state_dict(weights, strict=False)
    # from_weight == 'none' 时不加载，随机初始化（从头训练）

    get_model_params(model, lm_config)
    Logger(f'Trainable Params: '
           f'{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f}M')
    return model.to(device), tokenizer

DDP 初始化流程

定义：分布式数据并行（DDP）让多个 GPU 同时训练，每个 GPU 处理不同数据，然后同步梯度。

白话：一个人搬砖慢，找 4 个人同时搬。每个人搬自己的那份，搬完对一下账（同步梯度），保证每个人拿到的结果一样。

代码：

def init_distributed_mode():
    if int(os.environ.get("RANK", -1)) == -1:
        return 0  # 单卡模式，不初始化 DDP

    dist.init_process_group(backend="nccl")  # 用 NVIDIA 的 NCCL 通信
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)  # 绑定当前进程的 GPU
    return local_rank

DDP 数据流

DDP 关键概念

常见问题诊断

课后练习

学完本阶段后

参考答案