Minimind 系列笔记 03：训练全流程

源码：jingyaogong/minimind（trainer/ 目录，11 个训练脚本约 3000 行）。

MiniMind 的训练管线是一条完整的大模型后训练流水线缩影：

Pretrain（学语言与知识）
  → Full SFT（学会对话格式，融入 mid-training 数据）
      → LoRA（低成本垂域适配）
      → 蒸馏（白盒 KL 对齐教师分布）
      → DPO（离线偏好对齐，off-policy）
      → PPO（在线 RL，Actor+Critic 双网络）
      → GRPO / CISPO（在线 RL，组内相对优势，单网络）
      → Agentic RL（多轮工具调用 + 延迟奖励）

本文先讲所有训练脚本共享的工程骨架（值得单拎出来学），再逐个拆解训练阶段的算法。

一、所有训练脚本共享的骨架

把 train_pretrain.py 抄一遍，几乎就是其他所有脚本的模板：

1.1 学习率调度

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

余弦退火（cosine decay）的手写版：训练开始 lr × 1.0（0.1 + 0.45×2），结束时 lr × 0.1（0.1 + 0.45×0）。从不衰减到 0——末端保留 10% 的学习率继续微调，这是与"归零版"余弦的实用差异（末端小步长继续收敛，但不至于完全停止学习）。每个 step 手动写回 optimizer.param_groups，不用 torch.optim.lr_scheduler（PPO/GRPO 例外，用了 CosineAnnealingLR）。

1.2 混合精度训练（AMP）

dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
...
with autocast_ctx:
    res = model(input_ids, labels=labels)
    loss = res.loss + res.aux_loss
    loss = loss / args.accumulation_steps
scaler.scale(loss).backward()
if step % 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)

默认 bfloat16（不是 fp16），这条选型值得展开：

• fp16：动态范围小（指数 5 位），前向易上溢/下溢，必须配 GradScaler——把 loss 先放大 N 倍，反传梯度同倍放大，optimizer.step 前再缩回来检查 inf/nan。所以 fp16 路径上 scaler.scale→unscale_→step→update 一套都不能少。
• bf16：指数位与 fp32 相同（8 位），动态范围大，不需要 GradScaler（enabled=(dtype=='float16') 关掉了它），损失部分精度但换来完全免调参。现代 GPU（A100/3090 Ampere 之后）都支持，已是训练默认。

哪些计算不该进低精度（01 篇已见到，这里集中归纳）：norm 的统计量（.float() 进出）、softmax（.float() 进出）、loss 累加。模式统一为"低精度算矩阵乘，fp32 算 reduction"。

梯度裁剪必须在 unscale_ 之后——scale 状态下的梯度全部被放大了 100 倍量级，直接裁剪等于没裁（永远达不到 1.0 阈值之外的尺度）。clip_grad_norm_(1.0)。

zero_grad(set_to_none=True)：把梯度置 None 而不是 0，省一次 memset 且跳过后续的梯度更新判断，小优化。

1.3 梯度累积

loss = loss / args.accumulation_steps    # 先除
scaler.scale(loss).backward()             # 每步都反传（梯度自然累加）
if step % args.accumulation_steps == 0:   # 每 N 步才 step 一次
    scaler.unscale_(optimizer); clip; step; update; zero_grad()

显存不够大 batch 时的标准解法：micro-batch 跑 N 次，梯度在 .grad 里自动累加（backward 不会清零），loss 预先除以 N 保证等效于大 batch 的平均损失，每 N 步才做一次 optimizer.step。默认 accumulation_steps=8，等效 batch = 32×8 = 256。

train_epoch 结尾的尾巴逻辑（last_step % accumulation_steps != 0 时补一次 step）处理"最后一个不满 N 步的窗口"——否则这批梯度会白白丢掉。

1.4 DDP 与随机种子

local_rank = init_distributed_mode()          # 检测 RANK 环境变量，有则 init_process_group
if dist.is_initialized(): args.device = f"cuda:{local_rank}"
setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))
...
if dist.is_initialized():
    model = DistributedDataParallel(model, device_ids=[local_rank])
for epoch in range(start_epoch, args.epochs):
    train_sampler and train_sampler.set_epoch(epoch)

DDP（DistributedDataParallel）是数据并行：每张卡一份完整模型，各吃各的数据，反向时梯度 all-reduce 求平均。三个实践点：

1. 种子按 rank 偏移（seed + rank）：否则所有卡初始化和数据顺序完全相同，数据并行退化成单卡×N 浪费。每 epoch 再偏移（seed + epoch）保证轮间 shuffle 不同。
2. set_epoch 必须调：DistributedSampler 靠它改变各 epoch 的分片方式，不调的话每个 epoch 所有卡吃到相同的数据排列。
3. 保存权重只存 rank 0（is_main_process()），且从 model.module（剥掉 DDP 壳）再剥 _orig_mod（剥掉 torch.compile 壳）拿原始模型。

03 篇 GRPO/PPO 里还有两个 DDP 的深坑注释（"反向传播必须经过 DDP 包装后的模块，直接调 .module 会跳过 prepare_for_backward，梯度不会 all-reduce，各卡静默发散"）——DDP 的 hook 挂在包装层的 forward 上，绕过包装层调用参数，all-reduce 时序就错了，而且不报错、只是默默训废，属于最阴险的一类 bug。

1.5 断点续训

# lm_checkpoint 的双文件设计：
ckp_path   = f'{save_dir}/{weight}_{hidden_size}.pth'         # 纯权重（推理用）
resume_path = f'..._resume.pth'                                # 完整训练状态
resume_data = {
    'model': ..., 'optimizer': optimizer.state_dict(),
    'epoch': epoch, 'step': step,
    'world_size': ..., 'wandb_id': ...,
    # kwargs 里还可以塞 scaler / scheduler / critic 等任意组件
}

恢复 optimizer 状态是续训的关键：AdamW 的动量 ${\beta }_1$、二阶矩 ${\beta }_2$、当前 t 都在 state_dict 里，只恢复 model 不恢复 optimizer，等于重启了自适应学习率的记忆，loss 会跳变。两处细节：

• world_size 转换：续训时 GPU 数变了（8 卡→1 卡），step 坐标要按比例换算（step * saved_ws // current_ws）才能对上数据进度；
• 原子写入（.tmp 写完 os.replace）：训练中途 kill 不会留下半个损坏的 checkpoint。

权重保存统一转 fp16 存盘（v.half().cpu()）——训练精度 bf16/fp32，发布权重 fp16，一份权重两套精度需求分开满足。

完整源码：train_pretrain.py 的 train_epoch（所有训练脚本的骨架模板）

def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
    start_time = time.time()
    last_step = start_step
    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
        input_ids = input_ids.to(args.device)
        labels = labels.to(args.device)
        last_step = step
        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with autocast_ctx:
            res = model(input_ids, labels=labels)
            loss = res.loss + res.aux_loss
            loss = loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if step % 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 or step == iters:
            spend_time = time.time() - start_time
            current_loss = loss.item() * args.accumulation_steps
            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
            current_logits_loss = current_loss - current_aux_loss
            current_lr = optimizer.param_groups[-1]['lr']
            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})

        if (step % args.save_interval == 0 or step == iters) and is_main_process():
            model.eval()
            moe_suffix = '_moe' if lm_config.use_moe else ''
            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
            raw_model = getattr(raw_model, '_orig_mod', raw_model)
            state_dict = raw_model.state_dict()
            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
            model.train()
            del state_dict

        del input_ids, labels, res, loss

    if last_step > start_step and last_step % args.accumulation_steps != 0:
        # 最后一个不满 accumulation_steps 的窗口：补一次 step，避免这批梯度白白丢掉
        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)

以及主流程的启动段（DDP / 种子 / 混合精度 / 断点恢复的组织顺序）：

if __name__ == "__main__":
    # ========== 1. 初始化环境和随机种子 ==========
    local_rank = init_distributed_mode()
    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
    setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0))

    # ========== 2. 配置目录、模型参数、检查ckp ==========
    os.makedirs(args.save_dir, exist_ok=True)
    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None

    # ========== 3. 设置混合精度 ==========
    device_type = "cuda" if "cuda" in args.device else "cpu"
    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)

    # ========== 4. 配wandb ==========
    wandb = None
    if args.use_wandb and is_main_process():
        import swanlab as wandb
        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
        resume = 'must' if wandb_id else None
        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)

    # ========== 5. 定义模型、数据、优化器 ==========
    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
    train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

    # ========== 6. 从ckp恢复状态 ==========
    start_epoch, start_step = 0, 0
    if ckp_data:
        model.load_state_dict(ckp_data['model'])
        optimizer.load_state_dict(ckp_data['optimizer'])
        scaler.load_state_dict(ckp_data['scaler'])
        start_epoch = ckp_data['epoch']
        start_step = ckp_data.get('step', 0)

    # ========== 7. 编译和分布式包装 ==========
    if args.use_compile == 1:
        model = torch.compile(model)
    if dist.is_initialized():
        model = DistributedDataParallel(model, device_ids=[local_rank])

    # ========== 8. 开始训练 ==========
    for epoch in range(start_epoch, args.epochs):
        train_sampler and train_sampler.set_epoch(epoch)
        setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist()
        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
        ...

    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized():
        dist.barrier()
        dist.destroy_process_group()

完整源码：trainer_utils.py 的 lm_checkpoint（断点续训）与 SkipBatchSampler

def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
    os.makedirs(save_dir, exist_ok=True)
    moe_path = '_moe' if lm_config.use_moe else ''
    ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth'
    resume_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}_resume.pth'

    if model is not None:
        raw_model = model.module if isinstance(model, DistributedDataParallel) else model
        raw_model = getattr(raw_model, '_orig_mod', raw_model)
        state_dict = raw_model.state_dict()
        state_dict = {k: v.half().cpu() for k, v in state_dict.items()}
        ckp_tmp = ckp_path + '.tmp'
        torch.save(state_dict, ckp_tmp)
        os.replace(ckp_tmp, ckp_path)          # 原子替换：kill 在写一半也不会损坏 ckp
        wandb_id = None
        if wandb:
            if hasattr(wandb, 'get_run'):
                run = wandb.get_run()
                wandb_id = getattr(run, 'id', None) if run else None
            else:
                wandb_id = getattr(wandb, 'id', None)

        resume_data = {
            'model': state_dict,
            'optimizer': optimizer.state_dict(),
            'epoch': epoch,
            'step': step,
            'world_size': dist.get_world_size() if dist.is_initialized() else 1,
            'wandb_id': wandb_id
        }
        for key, value in kwargs.items():
            if value is not None:
                if hasattr(value, 'state_dict'):
                    raw_value = value.module if isinstance(value, DistributedDataParallel) else value
                    raw_value = getattr(raw_value, '_orig_mod', raw_value)
                    resume_data[key] = raw_value.state_dict()
                else:
                    resume_data[key] = value

        resume_tmp = resume_path + '.tmp'
        torch.save(resume_data, resume_tmp)
        os.replace(resume_tmp, resume_path)
        del state_dict, resume_data
        torch.cuda.empty_cache()
    else:  # 加载模式
        if os.path.exists(resume_path):
            ckp_data = torch.load(resume_path, map_location='cpu')
            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:
                # GPU 数量变化时，step 坐标按比例换算
                ckp_data['step'] = ckp_data['step'] * saved_ws // current_ws
                Logger(f'GPU数量变化({saved_ws}→{current_ws})，step已自动转换为{ckp_data["step"]}')
            return ckp_data
        return None


class SkipBatchSampler(Sampler):
    """续训时跳过已训练过的 batch，让 DataLoader 从断点继续"""
    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
                yield batch
                batch = []
        if len(batch) > 0 and skipped >= self.skip_batches:
            yield batch

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

二、阶段一：Pretrain

数据与目标见 02 篇。训练侧没有任何算法花活，就是骨架 + 最朴素的超参：

超参	值
epochs / batch_size / accum	2 / 32 / 8（等效 256）
learning_rate	5e-4
max_seq_len	340
loss	res.loss + res.aux_loss（MoE 时）

预训练 lr 5e-4 是所有阶段里最大的——模型从随机初始化开始，参数需要大幅移动；后面 SFT 降到 1e-5（50 倍），保护预训练学到的能力。

三、阶段二：Full SFT

脚本与 pretrain 几乎逐行相同，差异全在数据侧（SFTDataset + loss mask）。超参对照：

	Pretrain	Full SFT
lr	5e-4	1e-5
batch × accum	32×8	16×1
max_seq_len	340	768
from_weight	none	pretrain

lr 降 50 倍的原因：预训练权重已经"很好"了，SFT 只是在流形上做局部微调，大 lr 会把预训练的能力灾难性遗忘掉。from_weight 链条（pretrain → full_sft → dpo/grpo/...）就是大模型训练的 stage 化管理——每个阶段都站在前一阶段的权重上。

四、阶段三：LoRA

4.1 实现：monkey-patch 旁路

class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        self.A = nn.Linear(in_features, rank, bias=False)     # 降维
        self.B = nn.Linear(rank, out_features, bias=False)    # 升维
        self.A.weight.data.normal_(0, 0.02)   # A 高斯初始化
        self.B.weight.data.zero_()            # B 零初始化

def apply_lora(model, rank=16):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
            lora = LoRA(module.in_features, module.out_features, rank=rank)
            setattr(module, "lora", lora)
            original_forward = module.forward
            def forward_with_lora(x, layer1=original_forward, layer2=lora):
                return layer1(x) + layer2(x)
            module.forward = forward_with_lora     # 替换 forward

LoRA 假设：微调时的权重更新 $\mathrm{\Delta }W$ 是低秩的，可以分解为 $\mathrm{\Delta }W=BA$（$B:d\times r$，$A:r\times d$，$r\ll d$）。前向变成 $Wx+BAx$，只训 A、B，冻结原权重。

三处实现细节：

1. B 零初始化：训练开始时 $BA=0$，模型行为与原模型完全一致——从"已训练好的状态"出发做增量，而不是从一个随机扰动的状态出发。
2. 默认路径闭包陷阱的规避：def forward_with_lora(x, layer1=..., layer2=...) 用默认参数把变量当场定值。若写成闭包引用外层变量，循环里所有层会共享最后一次的值（Python 闭包晚绑定的经典坑）。
3. 只 patch 方阵 Linear（in_features == out_features）：q/k/v/o/gate/up/down 都是 768→768 或 2432→2432 的方阵，lm_head（768→6400）不 patch——输出投影动了会直接改写词表分布，容易破坏通用能力。

训练时只把 lora 参数给 optimizer：optim.AdamW(lora_params, lr=1e-4)。lr 是 1e-4，比 full SFT 大 10 倍——可训练参数少了 99% 以上，需要更大的步长才能学动。LoRA 参数占比实测：约 1.2M / 64M ≈ 1.9%。

4.2 保存与合并

def save_lora(model, path):        # 只存 lora 权重（几 MB）
def merge_lora(model, lora_path, save_path):
    ...
    state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()

merge_lora 把 $BA$ 加回原权重，推理零开销（没有额外分支）。注意 merge 在 fp16 下做矩阵乘有精度损失（BA 的乘积先 half 再加），教学代码可接受，生产实现一般 fp32 merge。

完整源码：model_lora.py（LoRA 全部 65 行）

import torch
from torch import optim, nn


# 定义Lora网络结构
class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        super().__init__()
        self.rank = rank  # LoRA的秩（rank），控制低秩矩阵的大小
        self.A = nn.Linear(in_features, rank, bias=False)  # 低秩矩阵A
        self.B = nn.Linear(rank, out_features, bias=False)  # 低秩矩阵B
        # 矩阵A高斯初始化
        self.A.weight.data.normal_(mean=0.0, std=0.02)
        # 矩阵B全0初始化
        self.B.weight.data.zero_()

    def forward(self, x):
        return self.B(self.A(x))


def apply_lora(model, rank=16):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
            lora = LoRA(module.in_features, module.out_features, rank=rank).to(model.device)
            setattr(module, "lora", lora)
            original_forward = module.forward

            # 显式绑定（默认参数当场定值，规避闭包晚绑定）
            def forward_with_lora(x, layer1=original_forward, layer2=lora):
                return layer1(x) + layer2(x)

            module.forward = forward_with_lora


def load_lora(model, path):
    state_dict = torch.load(path, map_location=model.device)
    state_dict = {(k[7:] if k.startswith('module.') else k): v for k, v in state_dict.items()}

    for name, module in model.named_modules():
        if hasattr(module, 'lora'):
            lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
            module.lora.load_state_dict(lora_state)


def save_lora(model, path):
    raw_model = getattr(model, '_orig_mod', model)
    state_dict = {}
    for name, module in raw_model.named_modules():
        if hasattr(module, 'lora'):
            clean_name = name[7:] if name.startswith("module.") else name
            lora_state = {f'{clean_name}.lora.{k}': v.cpu().half() for k, v in module.lora.state_dict().items()}
            state_dict.update(lora_state)
    torch.save(state_dict, path)


def merge_lora(model, lora_path, save_path):
    load_lora(model, lora_path)
    raw_model = getattr(model, '_orig_mod', model)
    state_dict = {k: v.cpu().half() for k, v in raw_model.state_dict().items() if '.lora.' not in k}
    for name, module in raw_model.named_modules():
        if isinstance(module, nn.Linear) and '.lora.' not in name:
            state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
            if hasattr(module, 'lora'):
                state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
    torch.save(state_dict, save_path)

五、阶段四：知识蒸馏

def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()
    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    kl = F.kl_div(student_log_probs, teacher_probs, reduction=reduction)
    return (temperature ** 2) * kl

蒸馏原理（Hinton 2015）：教师 softmax 的输出是"暗知识"（dark knowledge）——"猫 0.9、狗 0.08、汽车 0.02"里狗和猫的相似度信息比 one-hot 标签丰富得多。温度 T 软化分布：T 越大，小概率的次优 token 被放大，暗知识越显形；T² 是梯度尺度补偿因子（logits/T 的梯度量级会缩 1/T²，乘回来保持 loss 量级与 CE 可比）。

总损失是混合的：loss = alpha * ce_loss + (1-alpha) * distill_loss（默认 alpha=0.5，T=1.5）：

$$\mathcal{L}=\alpha \cdot {\mathcal{L}}_{CE}+(1-\alpha )\cdot T^2\mathrm{KL}(p_{teacher}^T\parallel p_{student}^T)$$

CE 锚住 ground truth（防纯蒸馏漂移），KL 传递教师的分布偏好。两个实现要点：

• 只在 loss mask 的 token 上算 KL：student_logits.view(-1, v)[loss_mask_flat == 1] 先过滤——prompt 位置的教师分布没有学习价值；
• 词表裁剪：teacher_logits[..., :vocab_size_student]——MoE 教师（198M）蒸馏 dense 学生（64M）时若词表不同，取前 N 对齐。README 的演示配置正是 moe→dense。

与 SFT 的关系：SFT 教"正确答案是什么"（one-hot 蒸馏，黑盒），白盒蒸馏教"每个候选 token 的相对倾向"。README 区分了黑盒蒸馏（拿强模型输出做 SFT——R1 蒸馏就是这个路线）和白盒蒸馏，这个脚本是后者的参考实现。

完整源码：train_distillation.py 的 train_epoch 核心（CE + KL 混合损失）

def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):
    with torch.no_grad():
        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()

    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

    kl = F.kl_div(
        student_log_probs,
        teacher_probs,
        reduction=reduction
    )
    return (temperature ** 2) * kl


def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_step=0, wandb=None, alpha=0.0, temperature=1.0):
    start_time = time.time()
    last_step = start_step

    if teacher_model is not None:
        teacher_model.eval()
        teacher_model.requires_grad_(False)

    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
        last_step = step
        input_ids = input_ids.to(args.device)
        labels = labels.to(args.device)
        loss_mask = (labels[..., 1:] != -100).float()
        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # 前向传播（学生模型）
        with autocast_ctx:
            res = model(input_ids)
            student_logits = res.logits[..., :-1, :].contiguous()

        # 教师模型前向传播（只在eval & no_grad）
        if teacher_model is not None:
            with torch.no_grad():
                teacher_logits = teacher_model(input_ids).logits[..., :-1, :].contiguous()
                vocab_size_student = student_logits.size(-1)
                teacher_logits = teacher_logits[..., :vocab_size_student]   # 词表对齐（moe→dense）

        # ========== 计算损失 ==========
        # 1) Ground-Truth CE Loss
        shift_labels = labels[..., 1:].contiguous()
        loss_mask_flat = loss_mask.view(-1)
        ce_loss = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
            reduction='none'
        )
        ce_loss_raw = torch.sum(ce_loss * loss_mask_flat) / (loss_mask_flat.sum() + 1e-8)
        if lm_config_student.use_moe: ce_loss = ce_loss_raw + res.aux_loss
        else: ce_loss = ce_loss_raw

        # 2) Distillation Loss（只在 loss mask 的 token 上算）
        if teacher_model is not None:
            distill_loss = distillation_loss(
                student_logits.view(-1, student_logits.size(-1))[loss_mask_flat == 1],
                teacher_logits.view(-1, teacher_logits.size(-1))[loss_mask_flat == 1],
                temperature=temperature
            )
        else:
            distill_loss = torch.tensor(0.0, device=args.device)

        # 3) 总损失 = alpha * CE + (1-alpha) * Distill
        loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps

        scaler.scale(loss).backward()

        if step % 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)
        # ...（日志与保存与 pretrain 骨架一致，略）

六、阶段五：DPO

6.1 原理

RLHF 的原始管线（PPO 式）需要：训奖励模型 → 用奖励模型在线指导策略。DPO（Direct Preference Optimization）的推导证明：KL 约束的奖励最大化目标存在闭式解 $r(x,y)=\beta \log \frac{\pi (y|x)}{{\pi }_{ref}(y|x)}$，代回偏好模型后可以跳过显式的奖励模型，直接从偏好对优化策略：

$${\mathcal{L}}_{DPO}=-\mathbb{E}[\log \sigma \left(\beta [\log \frac{{\pi }_{\theta }(y_w|x)}{{\pi }_{ref}(y_w|x)}-\log \frac{{\pi }_{\theta }(y_l|x)}{{\pi }_{ref}(y_l|x)}]\right)]$$

直觉读法：拉大"策略模型认为 chosen 好于 rejected 的程度"与"参考模型认为的差距"。参考模型 ${\pi }_{ref}$ 充当锚——防止策略无约束地把 rejected 概率压到零的同时把 chosen 也搞崩（那不是变好，是分布漂移）。

6.2 实现

def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
    ref_log_probs = (ref_log_probs * mask).sum(dim=1)      # 整段回答的 logprob 求和
    policy_log_probs = (policy_log_probs * mask).sum(dim=1)
    # batch 前半是 chosen，后半是 rejected
    chosen_ref = ref_log_probs[:batch_size // 2]
    reject_ref = ref_log_probs[batch_size // 2:]
    ...
    pi_logratios = chosen_policy - reject_policy          # log π(y_w)/π(y_l)
    ref_logratios = chosen_ref - reject_ref               # log πref(y_w)/πref(y_l)
    logits = pi_logratios - ref_logratios
    loss = -F.logsigmoid(beta * logits)
    return loss.mean()

实现与公式逐项对应：

• logits_to_log_probs 用 gather 抠出每个真实 token 的 logprob，mask 后 sum 得到整段回答的对数概率 $\log \pi (y|x)$（序列级似然，不是逐 token 平均）；
• chosen / rejected 在 batch 维拼成一个大 batch 过模型（2 次前向 → 1 次），减少 kernel launch；
• $\beta$（默认 0.15）控制"贴多紧参考模型"：越大越保守。lr 是全场最小（4e-8，脚本注释：建议 ≤5e-8 避免遗忘）——DPO 的梯度集中在压低 rejected 上，步子大了通用能力崩得极快。

模型侧：ref_model 与 policy 完全同构，requires_grad_(False) + eval()，等于存了两份 SFT 权重。

完整源码：train_dpo.py 的 loss 函数与 train_epoch 核心

def logits_to_log_probs(logits, labels):
    # logits shape: (batch_size, seq_len, vocab_size)
    # labels shape: (batch_size, seq_len)
    # log_probs shape: (batch_size, seq_len)
    log_probs = F.log_softmax(logits, dim=2)
    log_probs_per_token = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)
    return log_probs_per_token


def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
    # ref_log_probs 和 policy_log_probs 都是 shape: (batch_size, seq_len)
    ref_log_probs = (ref_log_probs * mask).sum(dim=1)
    policy_log_probs = (policy_log_probs * mask).sum(dim=1)

    # 将 chosen 和 rejected 数据分开
    batch_size = ref_log_probs.shape[0]
    chosen_ref_log_probs = ref_log_probs[:batch_size // 2]
    reject_ref_log_probs = ref_log_probs[batch_size // 2:]
    chosen_policy_log_probs = policy_log_probs[:batch_size // 2]
    reject_policy_log_probs = policy_log_probs[batch_size // 2:]

    pi_logratios = chosen_policy_log_probs - reject_policy_log_probs
    ref_logratios = chosen_ref_log_probs - reject_ref_log_probs
    logits = pi_logratios - ref_logratios
    loss = -F.logsigmoid(beta * logits)
    return loss.mean()


def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=None, beta=0.1):
    ...
    for step, batch in enumerate(loader, start=start_step + 1):
        last_step = step
        x_chosen = batch['x_chosen'].to(args.device)
        x_rejected = batch['x_rejected'].to(args.device)
        y_chosen = batch['y_chosen'].to(args.device)
        y_rejected = batch['y_rejected'].to(args.device)
        mask_chosen = batch['mask_chosen'].to(args.device)
        mask_rejected = batch['mask_rejected'].to(args.device)
        # chosen 和 rejected 拼成一个大 batch，2 次前向 → 1 次
        x = torch.cat([x_chosen, x_rejected], dim=0)
        y = torch.cat([y_chosen, y_rejected], dim=0)
        mask = torch.cat([mask_chosen, mask_rejected], dim=0)

        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with autocast_ctx:
            with torch.no_grad():
                ref_outputs = ref_model(x)
                ref_logits = ref_outputs.logits
            ref_log_probs = logits_to_log_probs(ref_logits, y)

            outputs = model(x)
            logits = outputs.logits
            policy_log_probs = logits_to_log_probs(logits, y)

            dpo_loss_val = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta)
            loss = dpo_loss_val + outputs.aux_loss
            loss = loss / args.accumulation_steps

        scaler.scale(loss).backward()
        # ...（step / 日志 / 保存与 pretrain 骨架一致，略）

6.3 定位

README 的总结很到位：DPO 是 off-policy（静态偏好数据集、可多轮 epoch）、免奖励模型、实现简单；但它不做在线探索，更多用于"偏好/安全"对齐，对"能不能做对题"的能力提升有限。

七、阶段六：PPO / GRPO / CISPO

7.1 统一视角

README 给了一个所有 Policy Optimization 算法的统一目标：

$${\mathcal{J}}_{PO}={\mathbb{E}}_{q\sim P(Q),o\sim {\pi }_{\theta }}[\underset{\text{策略项}}{\underbrace{\mathrm{\Phi }(r_t,A_t)}}-\underset{\text{正则项}}{\underbrace{h({\text{KL}}_t)}}]$$

三组件 = 策略项（概率比 $r_t$ 和优势 $A_t$ 怎么组合）× 优势项（$A_t$ 怎么估）× 正则项（怎么约束别跑偏）。DPO / PPO / GRPO / CISPO 全部是这个框架的实例化：

算法	策略项	优势项	正则	模型数
DPO	$\log r_w-\log r_l$	无（偏好对隐式）	隐含在 β	1（前向 2）
PPO	$min(rA,\text{clip}(r,1\pm \epsilon )A)$	Critic + GAE	$\beta \cdot \text{KL}$	2（+Critic）
GRPO	同 PPO	组内归一化 $\frac{R-\mu }{\sigma }$	$\beta \cdot {\text{KL}}_t$	1
CISPO	$min(r,{\epsilon }_{high})\cdot A\cdot \log \pi$	同 GRPO	$\beta \cdot {\text{KL}}_t$	1

7.2 共享的基建：Rollout 引擎

PPO/GRPO/AgentRL 共享一个抽象（rollout_engine.py），这是理解现代 RL 训练系统的关键：

class RolloutEngine(ABC):
    @abstractmethod
    def rollout(self, prompt_ids, attention_mask, num_generations, max_new_tokens, temperature=0.8) -> RolloutResult: ...
    @abstractmethod
    def update_policy(self, model): ...

class TorchRolloutEngine(RolloutEngine):     # 直接调模型 generate
class SGLangRolloutEngine(RolloutEngine):    # HTTP 调远程推理服务

RolloutResult 携带五件套：output_ids（完整序列）、completion_ids（生成部分）、per_token_logps（每个生成 token 的 logprob——RL 需要）、completions（文本）、completion_mask。训推分离：训练侧只管优化 policy，rollout 侧负责高吞吐采样；每批更新完 update_policy 把新权重同步过去（SGLang 版走 /update_weights_from_disk HTTP 接口）。当前实现是同步模式（采完一批更一批），还没有异步 rollout buffer——但 openrlhf/verl 的雏形已经在了。

奖励函数（三份脚本共用逻辑）：

rewards[i] += 0.5 if 20 <= len(response) <= 800 else -0.5        # 长度合规
if '</think>' in response:
    rewards[i] += 1.0 if 20 <= len(thinking) <= 300 else -0.5    # 思考长度
    rewards[i] += 0.25 if response.count('</think>') == 1 else -0.25  # 格式闭合
rewards[i] -= rep_penalty(answer)                                 # 重复惩罚
score = reward_model.get_score(messages, answer)                  # InternLM2-1.8B-Reward 打分

规则奖励管格式（长度、思考标签闭合、重复度），奖励模型管内容质量。为什么用模型打分而不是规则二元奖励？README 讲透了小模型 RL 的根本困难——奖励稀疏：0.1B 模型做数学题几乎全错，二元 0/1 奖励全零 → 组内优势 $A=r-b\approx 0$ → 梯度信号消失。连续分数奖励（-2.5~+3.0）保证"再差也有区分度"，让小模型在能力边界内仍有梯度可学。

完整源码：rollout_engine.py（训推分离的抽象与两个实现）

# ===== 计算每个 token 的 logprob =====
def compute_per_token_logps(model, input_ids: Tensor, n_keep: int, attention_mask: Optional[Tensor] = None) -> Tensor:
    if n_keep <= 0:
        return input_ids.new_empty((input_ids.size(0), 0), dtype=torch.float32)
    unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
    input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
    logits = unwrapped(input_ids, attention_mask=attention_mask, logits_to_keep=n_keep + 1).logits[:, :-1, :]
    per_token_logps = []
    for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
        ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
        per_token_logps.append(
            torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)
        )
    return torch.stack(per_token_logps)


# ===== Rollout 结果 =====
@dataclass
class RolloutResult:
    output_ids: Tensor          # 完整序列 (B*num_gen, P+R)
    completion_ids: Tensor      # 生成部分
    per_token_logps: Tensor     # 每个生成 token 的 logprob（RL 需要）
    completions: List[str]      # 文本
    prompt_lens: Tensor
    completion_mask: Tensor


# ===== Rollout 引擎抽象基类 =====
class RolloutEngine(ABC):
    tokenizer = None

    @abstractmethod
    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
        pass

    @abstractmethod
    def update_policy(self, model: torch.nn.Module):
        pass


# ===== PyTorch 原生推理引擎 =====
class TorchRolloutEngine(RolloutEngine):
    def __init__(self, policy_model: torch.nn.Module, tokenizer, device: str = "cuda", autocast_ctx=None):
        self.policy_model = policy_model
        self.tokenizer = tokenizer
        self.device = device
        self.autocast_ctx = autocast_ctx

    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
        model = self.policy_model.module if isinstance(self.policy_model, DistributedDataParallel) else self.policy_model
        ctx = self.autocast_ctx if self.autocast_ctx else nullcontext()
        with torch.no_grad(), ctx:
            output_ids = model.generate(
                input_ids=prompt_ids.repeat_interleave(num_generations, dim=0),
                attention_mask=attention_mask.repeat_interleave(num_generations, dim=0),
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=temperature,
                num_return_sequences=1,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            ).clone()  # [B*num_gen, P+R]
            prompt_len = prompt_ids.size(1)
            completion_ids = output_ids[:, prompt_len:]  # [B*num_gen, R]
            full_mask = torch.cat([attention_mask.repeat_interleave(num_generations, dim=0), attention_mask.new_ones(output_ids.size(0), completion_ids.size(1))], dim=1)
            per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask)
        completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
        return RolloutResult(output_ids, completion_ids, per_token_logps, completions,
                             prompt_ids.new_full((output_ids.size(0),), prompt_len),
                             attention_mask.new_ones(output_ids.size(0), completion_ids.size(1)))

    def update_policy(self, model: torch.nn.Module):
        self.policy_model = model


# ===== SGLang HTTP API 推理引擎（节选核心 rollout / update_policy）=====
class SGLangRolloutEngine(RolloutEngine):
    def __init__(self, base_url: str, model_path: str, shared_ckpt_path: str = "./sglang_ckpt", timeout: int = 120):
        self.base_url = base_url.rstrip('/')
        self.shared_ckpt_path = shared_ckpt_path
        self.timeout = timeout
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.http = requests

    def rollout(self, prompt_ids, attention_mask, num_generations, max_new_tokens, temperature=0.8):
        # 去除左侧 padding tokens，只保留有效 token
        input_ids_list = []
        for ids, mask in zip(prompt_ids, attention_mask):
            valid_ids = ids[mask.bool()].tolist()
            input_ids_list.append(valid_ids)
        all_input_ids = [ids for ids in input_ids_list for _ in range(num_generations)]

        payload = {
            "input_ids": all_input_ids,
            "sampling_params": {
                "temperature": temperature,
                "max_new_tokens": max_new_tokens,
                "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [],
            },
            "return_logprob": True,
        }
        resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout)
        resp.raise_for_status()
        results = resp.json()
        # ...（解析 output_ids / output_token_logprobs 并 pad 成 tensor，略）

    def update_policy(self, model: torch.nn.Module):
        # rank 0 存权重到磁盘 → HTTP 通知推理服务热更新
        ok = True
        if not dist.is_initialized() or dist.get_rank() == 0:
            unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
            unwrapped = getattr(unwrapped, '_orig_mod', unwrapped)
            state_dict = {k: v.detach().half().cpu() for k, v in unwrapped.state_dict().items()}
            unwrapped.save_pretrained(os.path.abspath(self.shared_ckpt_path), state_dict=state_dict, safe_serialization=False)
            self.tokenizer.save_pretrained(os.path.abspath(self.shared_ckpt_path))
            resp = self.http.post(f"{self.base_url}/update_weights_from_disk", json={"model_path": os.path.abspath(self.shared_ckpt_path)}, timeout=self.timeout)
            ok = resp.status_code == 200
        if dist.is_initialized():
            ok_t = torch.tensor(int(ok), device=next(model.parameters()).device)
            dist.broadcast(ok_t, src=0); dist.barrier(); ok = bool(ok_t.item())
        if not ok: raise RuntimeError("SGLang update_policy failed")
        return ok


# ===== 工厂函数 =====
def create_rollout_engine(engine_type: str = "torch", policy_model=None, tokenizer=None, device="cuda",
                          autocast_ctx=None, sglang_base_url=None, sglang_model_path=None, sglang_shared_path=None):
    if engine_type == "torch":
        return TorchRolloutEngine(policy_model, tokenizer, device, autocast_ctx)
    elif engine_type == "sglang":
        return SGLangRolloutEngine(sglang_base_url, sglang_model_path, sglang_shared_path)
    else:
        raise ValueError(f"不支持的引擎类型: {engine_type}")

7.3 PPO：双网络 + GAE

PPO 涉及四个模型，是小模型上能跑的最复杂实现：

• Actor：策略模型，要训练；
• Critic：价值估计网络（CriticModel = MiniMind 主干 + value_head: 768→1），要训练；
• Ref：冻结的 SFT 权重，算 KL 正则；
• Reward Model：冻结的 InternLM2-1.8B，给整段回答打分。

核心循环里最值得学的是 GAE（Generalized Advantage Estimation）：

token_rewards = torch.zeros_like(old_resp_logp)
token_rewards[arange(B), last_idx] += rewards   # 序列级奖励放到最后一个 token

lastgaelam = 0; advs_rev = []
for t in reversed(range(gen_len)):              # 时间反序递推
    nv = old_resp_values[:, t+1] if t < gen_len-1 else 0.0
    delta = token_rewards[:, t] + gamma * nv - old_resp_values[:, t]   # TD 误差
    lastgaelam = delta + gamma * lam * lastgaelam
    advs_rev.append(lastgaelam)
advantages = torch.stack(advs_rev[::-1], dim=1)
returns = advantages + old_resp_values

奖励是稀疏的（整段回答一个分，放在最后 token 上），GAE 通过 TD 误差 ${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$ 的指数加权递推，把这个终局奖励沿序列反传分布到每个 token，$\lambda$ 控制偏差（纯 TD，低方差高偏差）与方差（纯蒙特卡洛，高方差低偏差）的权衡。

PPO 的 clip 目标：

log_ratio = mb_resp_logp - old_resp_logp
ratio = torch.exp(log_ratio)
policy_loss = max(-adv*ratio, -adv*clamp(ratio, 1-ε, 1+ε))    # 保守更新
value_loss = 0.5 * max((V-returns)², (clamp(V, old±0.2)-returns)²)  # value 也 clip

old_resp_logp 是 rollout 时的 logprob（重要性采样的"旧策略"）；ratio 偏离 1 太多（新策略相对采样策略漂移过大）的部分被 clip 掉，防止单步更新过激。整批 rollout 还可以复用 ppo_update_iters=2 轮（PPO 唯一允许的数据复用——importance sampling 修正了分布偏移）。

还有三个工程防御，都是实战血泪：

1. 早停 KL：approx_kl > 0.25 时本轮停止更新（loss 乘 0 保住 forward-backward 闭环，DDP 不挂）；
2. all_reduce 同步 approx_kl：各卡独立判断早停会导致某卡 break 其它卡继续 → DDP 死锁；
3. debug_log_ratio：验证"首轮第一个 minibatch 的 ratio≈1"——不≈1 说明 rollout 与训练的实现有数值错位（dropout 未关、token 对不上等），这是 RL 实现最常见的静默 bug。

README 的实测结论：PPO reward 提升缓慢——Critic 要先收敛才能给 Actor 提供靠谱优势估计，两者相互依赖；显存也是单网络方法的 1.5-2 倍。

完整源码：train_ppo.py 的 Critic 模型、奖励与 GAE 计算

# 自定义的Critic模型，继承自MiniMindLM
class CriticModel(MiniMindForCausalLM):
    def __init__(self, params):
        super().__init__(params)
        # 替换lm_head为输出单一价值的线性层
        self.value_head = nn.Linear(params.hidden_size, 1)

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        # 使用基础模型获取隐藏状态
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
        hidden_states = self.model.norm(outputs[0])
        # 使用value_head获取价值估计
        values = self.value_head(hidden_states).squeeze(-1)
        return values


def ppo_train_epoch(...):
    ...
    for step, batch in enumerate(loader, start=start_step + 1):
        ...
        rollout_result = rollout_engine.rollout(
            prompt_ids=enc.input_ids,
            attention_mask=enc.attention_mask,
            num_generations=1,
            max_new_tokens=args.max_gen_len,
            temperature=0.8,
        )
        gen_out = rollout_result.output_ids
        completion_ids = rollout_result.completion_ids
        prompt_lens = rollout_result.prompt_lens.to(args.device)
        responses_text = rollout_result.completions
        old_resp_logp = rollout_result.per_token_logps.to(args.device)
        rewards = calculate_rewards(prompts, responses_text, reward_model)  # [B]
        ...
        with torch.no_grad():  # Rollout阶段只需推理获取old_logp和old_values，切断梯度省显存
            critic_for_rollout = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
            values_seq = critic_for_rollout(input_ids=gen_out, attention_mask=full_mask)
            old_resp_values = values_seq.gather(1, logp_pos) * resp_value_mask

            ref_resp_logp = F.log_softmax(ref_model(input_ids=gen_out, attention_mask=full_mask).logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
            token_rewards = torch.zeros_like(old_resp_logp)
            last_idx = resp_lengths - 1  # [B]
            token_rewards[torch.arange(B, device=args.device)[valid_resp], last_idx[valid_resp]] += rewards[valid_resp]  # 末尾加外部奖励

            # ===== GAE：时间反序递推 =====
            gen_len = old_resp_values.size(1); lastgaelam = torch.zeros(B, device=args.device); advs_rev = []
            for t in reversed(range(gen_len)):
                nv = old_resp_values[:, t + 1] if t < gen_len - 1 else 0.0
                delta = token_rewards[:, t] + args.gamma * nv - old_resp_values[:, t]      # TD 误差
                lastgaelam = delta + args.gamma * args.lam * lastgaelam
                advs_rev.append(lastgaelam)
            advantages = torch.stack(advs_rev[::-1], dim=1)  # [B, R]
            returns = advantages + old_resp_values  # [B, R]

            # 优势标准化（batch 级）
            adv_mean = (advantages * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
            adv_var = ((advantages - adv_mean) ** 2 * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
            advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8) * resp_policy_mask

        # ===== 多 epoch 复用 + clip 目标 =====
        for ppo_epoch in range(args.ppo_update_iters):
            ...
            ratio = torch.exp(log_ratio)
            policy_loss = ((torch.max(-advantages[inds] * ratio,
                                      -advantages[inds] * torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon))
                           * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
                           + args.kl_coef * kl_ref_penalty)
            value_loss = 0.5 * (torch.max((mb_resp_values - returns[inds]) ** 2,
                                          (torch.clamp(mb_resp_values, old_resp_values[inds] - args.cliprange_value,
                                                       old_resp_values[inds] + args.cliprange_value) - returns[inds]) ** 2)
                                * resp_value_mask[inds]).sum() / resp_value_mask[inds].sum().clamp(min=1)

            # 早停时必须保证 forward-backward 闭环，故只截断 loss 不中断 DDP 通信
            if stop_ppo:
                loss = (policy_loss + args.vf_coef * value_loss + aux_loss) * 0.0
            else:
                loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps

            loss.backward()
        ...

7.4 GRPO：去掉 Critic

GRPO 的核心一步：同一个问题采 N 个回答，用组内平均分当 baseline：

rollout_result = rollout_engine.rollout(prompt_ids, ..., num_generations=6, ...)  # 每题 6 采样
rewards = calculate_rewards(...)                     # (B * 6)

grouped_rewards = rewards.view(-1, num_generations)  # (B, 6) 分组
mean_r = grouped_rewards.mean(dim=1).repeat_interleave(num_generations)
std_r  = grouped_rewards.std(dim=1).repeat_interleave(num_generations)
advantages = (rewards - mean_r) / (std_r + 1e-4)     # 组内标准化 = 优势

$$A_{i,j}=\frac{R_{i,j}-{\mu }_i}{{\sigma }_i+ϵ}$$

同一题的 6 个回答互为参照：高于组均值的被鼓励、低于的被抑制。Critic 网络整个不需要了——组均值就是 value 的蒙特卡洛估计。loss 是 PPO 式的 clip 目标 + token 级 KL（对 ref 逐 token 算 $\exp (\mathrm{\Delta })-\mathrm{\Delta }-1$，这是 k3 估计量，恒正、低方差，比直接 $\mathrm{\Delta }$ 平滑）：

kl_div = ref_per_token_logps - per_token_logps
per_token_kl = torch.exp(kl_div) - kl_div - 1
ratio = torch.exp(per_token_logps - old_per_token_logps)
clipped = torch.clamp(ratio, 1-ε, 1+ε)
per_token_loss = -(min(ratio*adv, clipped*adv) - β*per_token_kl)
policy_loss = ((per_token_loss * completion_mask).sum(1) / completion_mask.sum(1)).mean()

completion_mask 是 EOS 之前的有效区域（EOS 之后的 padding token 不该参与 loss）：

is_eos = (completion_ids == eos_token_id) & completion_pad_mask
eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]   # 第一个 EOS
completion_mask = (arange(R) <= eos_idx.unsqueeze(1)) & completion_pad_mask

实测曲线：GRPO 的 reward 稳定上升且收敛上限高于 PPO（单网络、无 Critic 依赖）。GRPO 的软肋是退化组——如果某题 N 个回答得分几乎一样，$\sigma \to 0$（分母有 1e-4 兜底但分子同样趋零），整组学习信号≈0。小模型能力弱，退化组出现频率更高，所以训练必须限定在能力边界内的任务。

完整源码：train_grpo.py 的 train_epoch 核心（组内优势 + token 级 KL + CISPO 分支）

def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model, start_step=0, wandb=None, use_sglang=False):
    for step, batch in enumerate(loader, start=start_step + 1):
        prompts = batch['prompt']  # list[str], length B
        prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
                                  padding_side="left", add_special_tokens=False).to(args.device)
        if args.max_seq_len:
            prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
            prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]

        rollout_result = rollout_engine.rollout(
            prompt_ids=prompt_inputs["input_ids"],
            attention_mask=prompt_inputs["attention_mask"],
            num_generations=args.num_generations,
            max_new_tokens=args.max_gen_len,
            temperature=0.8,
        )
        outputs = rollout_result.output_ids
        completion_ids = rollout_result.completion_ids
        completions = rollout_result.completions
        old_per_token_logps = rollout_result.per_token_logps.to(args.device).detach()
        prompt_lens = rollout_result.prompt_lens.to(args.device)
        full_mask = (outputs != tokenizer.pad_token_id).long()
        logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
        full_mask.scatter_(1, logp_pos + 1, rollout_result.completion_mask.to(args.device, dtype=full_mask.dtype))

        rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]

        with autocast_ctx:
            # 反向传播必须经过 DDP 包装后的模块：直接调用 .module 会跳过
            # DDP 的 prepare_for_backward，梯度不会 all-reduce，各卡静默发散。
            res = model(outputs, attention_mask=full_mask)
            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
            per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)

        with torch.no_grad():
            ref_per_token_logps = F.log_softmax(ref_model(outputs, attention_mask=full_mask).logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
        ...
        # ===== 组内相对优势 =====
        grouped_rewards = rewards.view(-1, args.num_generations)  # [B, num_gen]
        mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
        std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)  # [B*num_gen]
        advantages = (rewards - mean_r) / (std_r + 1e-4)  # [B*num_gen]

        # ===== 只保留第一个 EOS 之前的有效区域 =====
        completion_pad_mask = rollout_result.completion_mask.to(args.device).bool()
        is_eos = (completion_ids == tokenizer.eos_token_id) & completion_pad_mask  # [B*num_gen, R]
        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1) - 1, dtype=torch.long, device=args.device)
        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
        completion_mask = ((torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)) & completion_pad_mask).int()

        # ===== token 级 k3 KL + GRPO/CISPO 两种 loss =====
        kl_div = ref_per_token_logps - per_token_logps
        per_token_kl = torch.exp(kl_div) - kl_div - 1  # [B*num_gen, R] k3 估计量
        ratio = torch.exp(per_token_logps - old_per_token_logps)  # [B*num_gen, R]
        if args.loss_type == "cispo":
            clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
            per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
        else:
            clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
            per_token_loss1 = ratio * advantages.unsqueeze(1)
            per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
            per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
        policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()
        loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
        loss.backward()

        if step % args.accumulation_steps == 0:
            if args.grad_clip > 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        # ...（日志 / 保存 / rollout_engine.update_policy(model)，略）

7.5 CISPO：改一行 loss

CISPO（Clipped Importance Sampling PO）观察到一个别扭的问题：PPO/GRPO 的 clip 把 ratio 截断成常数后，梯度流也被硬切断了——被 clip 的 token 在这次更新里完全失聪。CISPO 把策略项改写成"截断权重 × log 概率"：

$${\mathcal{L}}_{CISPO}=-\mathbb{E}[min(r_t,{\epsilon }_{high})\cdot A_t\cdot \log {\pi }_{\theta }(a_t|s)-\beta \cdot {\text{KL}}_t]$$

ratio 即使被截断，也只作为系数，梯度走 $\log {\pi }_{\theta }$ 路径继续回传。实现上就是 GRPO 脚本里的一个 if：

if args.loss_type == "cispo":
    clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()   # 上界截断，detach 掉 clip 分支
    per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
else:
    ...  # GRPO 原版

注意 .detach() 加在 clamp 后的 ratio 上——权重是常数，梯度只流过 per_token_logps。epsilon_high=5.0 只设上界（ratio 很大时截断防爆炸，ratio 很小时保留梯度让它继续学）。MiniMind 没有为 CISPO 单独写脚本，直接作为 GRPO 的 loss 变体——这也印证了"XXPO 都是同一框架的实例化"。

7.6 Agentic RL：多轮工具调用

train_agent.py 把 RL 从"单轮问答"扩展到"多轮交互"。数据 (x, T, gt)：问题 + 工具集 + ground truth。优化对象不再是单个回答，而是整条轨迹：

$$\tau =(a_1,o_1,a_2,o_2,\dots ,a_T),a_t\sim {\pi }_{\theta }(\cdot \mid s_t,\mathcal{T})$$

rollout 过程（rollout_single）：

模型生成 → 含 <tool_call>？→ 是：解析 JSON → 执行工具（mock 环境）→
结果包进 <tool_response> 拼回上下文 → 继续 rollout（最多 max_turns 轮）
→ 无 tool_call 且给出最终回答 → 结束

奖励对整条轨迹联合打分：

$$R(\tau )=R_{\text{answer}}+R_{\text{tool}}+R_{\text{format}}+R_{\text{rm}}-R_{\text{unfinished}}$$

答案命中 gt（Agentic RL 终于用上了规则可验证奖励——工具调用场景下数学题有唯一答案，不存在奖励稀疏问题）、工具调用合法性、格式闭合、未完成惩罚、RM 分数。梯度回传在整个轨迹结束后才发生（延迟奖励），并且要把多轮拼接的每段 logprob 都收集起来（response_ids.extend(new_ids) 逐轮累积）。

README 对这套实现的定位："模板组织、工具执行、多轮 rollout、延迟奖励、训推分离"的最小串联——工业级 Agent 训练框架（verl/slime）的骨架已具备，只是每处都取了最简形态。

完整源码：train_agent.py 的多轮 rollout_single（含 observation 边界对齐）

def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
    all_outputs = []
    prompt_ids = None
    response_ids = []
    response_mask = []
    response_old_logps = []
    final_context = ""
    unfinished = False
    open_thinking = random.random() < thinking_ratio
    for turn in range(max_turns):
        context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
        if prompt_ids is None:
            prompt_ids = tokenizer(context, add_special_tokens=False)["input_ids"]
        input_ids = torch.tensor([prompt_ids + response_ids], device=device)
        rollout_result = rollout_engine.rollout(
            prompt_ids=input_ids,
            attention_mask=torch.ones_like(input_ids),
            num_generations=1,
            max_new_tokens=max_new_tokens,
            temperature=0.8,
        )
        valid_len = int(rollout_result.completion_mask[0].sum().item())
        new_ids = rollout_result.completion_ids[0, :valid_len].tolist()
        new_logps = rollout_result.per_token_logps[0, :valid_len].tolist()
        if len(new_ids) != len(new_logps):
            raise RuntimeError(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
        new_text = rollout_result.completions[0]
        all_outputs.append(new_text)
        response_ids.extend(new_ids)
        response_mask.extend([int(t != tokenizer.eos_token_id) for t in new_ids])
        response_old_logps.extend(new_logps)
        final_context = context + new_text
        calls = parse_tool_calls(new_text)
        if not calls:                       # 没有 tool_call：轨迹结束
            break
        unfinished = turn == max_turns - 1
        assistant_message = {"role": "assistant", "content": new_text}
        messages.append(assistant_message)
        for call in calls:                  # 执行每个工具，结果包进 tool 消息
            name, raw = call.get("name", ""), call.get("arguments", {})
            if isinstance(raw, str):
                try: raw = json.loads(raw)
                except: raw = {}
            result = execute_tool(name, raw)
            result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048]  # 防止天文数字撑爆tokenizer
            messages.append({"role": "tool", "content": result_str})

        # ===== observation 边界对齐 =====
        # 用唯一 marker 标记 assistant 内容末尾，重新渲染模板后定位"模板自动拼出的
        # <tool_response> 段"的精确 token 范围——这段是环境给的（非模型生成），
        # 计入 response_ids 但 mask=0、logp=0，不参与策略梯度。
        marker = f"<|agent_observation_{id(messages)}_{len(response_ids)}|>"
        assistant_message["content"] += marker
        marked_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
        assistant_message["content"] = new_text
        _, found, observation = marked_context.partition(marker)
        if not found:
            raise RuntimeError("chat template did not preserve the assistant content boundary")
        obs_delta = tokenizer(observation, add_special_tokens=False)["input_ids"]
        if new_ids and new_ids[-1] == tokenizer.eos_token_id and obs_delta[:1] == [tokenizer.eos_token_id]:
            obs_delta = obs_delta[1:]
        observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
        response_ids.extend(obs_delta)
        response_mask.extend([0] * len(obs_delta))
        response_old_logps.extend([0.0] * len(obs_delta))
        final_context = observe_context

    final_output = all_outputs[-1] if all_outputs else ""
    prompt_ids = prompt_ids or []
    return final_output, final_context, prompt_ids, response_ids, response_mask, response_old_logps, list(all_outputs), unfinished

完整源码：train_agent.py 的轨迹奖励 calculate_rewards（gt 校验 + 格式分 + RM 分）

def validate_gt_in_text(text, gt_list):
    text, text_num = str(text), str(text).replace(',', '')
    nums = [float(x) for x in re.findall(r'(?<![\w.])[-+]?\d+(?:\.\d+)?(?![\w.])', text_num)]
    return {g for g in gt_list if ((s := str(g).strip()) and s.lower() in text.lower())
            or (re.fullmatch(r'[-+]?\d+(?:\.\d+)?', str(g).strip().replace(',', ''))
                and any(abs(float(str(g).strip().replace(',', '')) - n) < 1e-6 for n in nums))}

def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, reward_model=None, device="cuda", turn_outputs_batch=None, unfinished_batch=None):
    rewards = torch.zeros(len(completions), device=device)
    for idx, response in enumerate(completions):
        reward, answer = 0.0, response
        sample_idx = idx // num_gen
        tools = tools_batch[sample_idx]
        turn_outputs = turn_outputs_batch[idx] if turn_outputs_batch is not None else [response]
        unfinished = unfinished_batch[idx] if unfinished_batch is not None else False
        turn_answers = [turn.split('</think>', 1)[-1].strip() if '</think>' in turn else turn.strip() for turn in turn_outputs]
        answer = turn_answers[-1] if turn_answers else response.strip()
        valid_names = {t['function']['name'] for t in tools} if tools else set()
        tool_calls = []
        for turn_answer in turn_answers: tool_calls.extend(parse_tool_calls(turn_answer))  # 解析tool调用
        reward -= 0.5 * sum(abs(turn.count('<tool_call>') - turn.count('</tool_call>')) for turn in turn_answers)  # 标签闭合扣分
        # -------- 无工具调用：格式+reward奖励 --------
        if not tool_calls:
            reward += 0.5 if 5 <= len(response.strip()) <= 800 else -0.5  # 长度分
            if '</think>' in response:
                think, answer = response.split('</think>', 1)
                reward += 1.0 if 20 <= len(think.strip()) <= 300 else -0.5  # 思考长度分
                reward += 0.25 if response.count('</think>') == 1 else -0.25  # 思考闭合分
                answer = answer.strip()
            if reward_model is not None:
                prompt = prompts[sample_idx]
                pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
                matches = re.findall(pattern, prompt, re.DOTALL)
                messages = [{"role": role, "content": content.strip()} for role, content in matches]
                score = reward_model.get_score(messages, answer)
                reward += score  # RM分
            reward -= rep_penalty(answer)
            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
        # -------- 有工具调用：执行结果奖励 --------
        else:
            gt = gt_batch[sample_idx]
            valid_call_count = 0
            for tool_call in tool_calls:
                name, raw = tool_call.get("name", ""), tool_call.get("arguments", {})
                if isinstance(raw, str):
                    try: raw = json.loads(raw)
                    except: raw = {}
                check = CHECK_ARGS.get(name)
                valid_call_count += int(bool(name in valid_names and check and check(raw)))
            tool_gap = abs(valid_call_count - len(gt)) + max(0, len(tool_calls) - valid_call_count)  # tool数差值
            reward += 0.5 if tool_gap == 0 else -0.5 * tool_gap  # tool对齐分

            final_text = "" if unfinished else (answer.split('</tool_call>')[-1] if '</tool_call>' in answer else answer)
            verified = validate_gt_in_text(final_text, gt) if gt else set()
            if gt: reward += 2.5 * len(verified) / len(gt)  # GT命中分
            if unfinished: reward -= 0.5  # 未完成扣分
            reward -= rep_penalty(final_text if final_text else answer)
            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
    return rewards

八、训练成本与观察

README 实测（单卡 3090，约 1.3￥/h）：

阶段	minimind-3 (64M)	minimind-3-moe (198M-A64M)
Pretrain	1.21h / 1.57￥	1.69h / 2.20￥
SFT	1.10h / 1.43￥	1.54h / 2.00￥
Tool Call	0.9h	1.26h
RLAIF	1.1h	1.54h

从 0 到能对话（pretrain + sft）约 2.3 小时 3 块钱——"3 块钱训个 LLM"的由来。

两条值得记录的实验观察（都来自 README 实测曲线）：

1. PPO vs GRPO：PPO reward 缓慢爬升（双网络互锁），GRPO 稳定上升且上限更高（单网络）；
2. RL 的代价：RL 权重的 reward score 更高，但通用能力和知识面会退化（reward hacking / capability trade-off）——RL 优化的是奖励函数的代理，不是真实能力。发布的 transformers 版模型一般从 full_sft 而不是 RL 权重转换，就是这个原因。

下一篇 04-多模态VLM与Omni：同一个 64M 主干，怎么长出眼睛和嗓子。