实战：ColossalChat 的 SFT → RM → PPO 流水线

ColossalChat 是 ColossalAI 团队 2023 年开源、首批中文友好的 InstructGPT 复刻项目。 整套 RLHF 三段式 pipeline （SFT → RM → PPO）的代码都集中在 applications/ColossalChat/， 本章一段一段拆给你看。

9.1项目结构总览

applications/ColossalChat/
├── examples/
│   ├── training_scripts/
│   │   ├── train_sft.py
│   │   ├── train_rm.py
│   │   ├── train_ppo.py             ★ PPO 主脚本
│   │   ├── train_dpo.py
│   │   ├── train_orpo.py
│   │   ├── train_kto.py
│   │   └── train_grpo.py
│   ├── data_preparation_scripts/
│   │   ├── prepare_sft_dataset.py
│   │   ├── prepare_preference_dataset.py
│   │   └── prepare_prompt_dataset.py
│   └── inference/
├── coati/
│   ├── trainer/
│   │   ├── sft.py
│   │   ├── rm.py
│   │   ├── ppo.py                    ★ PPOTrainer
│   │   ├── dpo.py
│   │   └── kto.py
│   ├── models/
│   │   ├── reward_model.py
│   │   ├── critic.py
│   │   └── lora.py
│   ├── experience_maker/             ★ rollout 引擎
│   │   └── naive.py
│   ├── experience_buffer/
│   ├── dataset/
│   └── utils/
└── tests/

三段式数据流：

9.2段 1：SFT

SFT 跟第 8 章基本一样，重点是用 chat template。看 train_sft.py：

colossalai run --nproc_per_node 8 train_sft.py \
    --pretrain meta-llama/Meta-Llama-3-8B \
    --tokenizer_dir meta-llama/Meta-Llama-3-8B \
    --dataset ./sft_data/ \
    --plugin zero2 \
    --batch_size 4 --lr 5e-6 \
    --max_len 2048 \
    --save_path ./out_sft \
    --conversation_template_config ./conf/conversation_template/llama3.json

核心代码（简化）：

from coati.trainer import SFTTrainer

trainer = SFTTrainer(
    model=model,
    booster=booster,
    optim=optimizer,
    lr_scheduler=scheduler,
    max_epochs=1,
    accumulation_steps=1,
    use_flash_attn=True,
)
trainer.fit(train_dataloader, eval_dataloader)

9.3段 2：Reward Model

RM 把语言模型的 lm_head 换成一个标量 head：

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base = base_model              # transformer body
        d = base_model.config.hidden_size
        self.value_head = nn.Linear(d, 1)   # 输出单一标量

    def forward(self, input_ids, attention_mask):
        h = self.base(input_ids, attention_mask).last_hidden_state
        # 取最后一个非 pad token 的 hidden
        last_idx = attention_mask.sum(dim=1) - 1
        last_hidden = h[torch.arange(h.size(0)), last_idx]
        return self.value_head(last_hidden).squeeze(-1)

训练 loss 是 Bradley-Terry pairwise：

$$\mathcal{L}_{\text{RM}} = -\mathbb{E}_{(x,y_w,y_l) \sim D}\Big[\log \sigma\big(R(x,y_w) - R(x,y_l)\big)\Big]$$

其中 $y_w$ 是被偏好的回答、$y_l$ 是被拒绝的回答。

colossalai run --nproc_per_node 8 train_rm.py \
    --pretrain ./out_sft \
    --dataset ./preference_data/ \
    --plugin zero2 \
    --batch_size 4 --lr 5e-6 \
    --max_len 2048 \
    --save_path ./out_rm

9.4段 3：PPO 主流程

PPO 同时维护四个模型：

显存压力是 SFT 的 4 倍，所以 PPO 通常需要 TP / Gemini / LoRA。

PPO objective

$$\mathcal{L}_\text{PPO} = -\mathbb{E}_t\Big[\min\!\big(r_t \hat A_t,\,\text{clip}(r_t,\,1\!-\!\epsilon,\,1\!+\!\epsilon)\hat A_t\big)\Big] + \beta\,\text{KL}\big(\pi_\theta\|\pi_\text{ref}\big)$$

$r_t = \pi_\theta(a_t|s_t)/\pi_{\theta_\text{old}}(a_t|s_t)$ 是策略比、$\hat A_t$ 是 GAE advantage、$\beta$ 是 KL 惩罚强度。

每个 step 干什么

启动 PPO

colossalai run --nproc_per_node 8 train_ppo.py \
    --pretrain ./out_sft \
    --rm_pretrain ./out_rm \
    --prompt_dataset ./prompt_data/ \
    --pretrain_dataset ./sft_data/ \    # 防遗忘的 PT-Mixin
    --plugin hybrid_parallel \
    --tp 2 --zero_stage 1 \
    --max_length 2048 \
    --num_episodes 1 \
    --num_collect_steps 4 \
    --num_update_steps 4 \
    --train_batch_size 4 \
    --experience_batch_size 8 \
    --lr 1e-6 \
    --kl_coef 0.1 \
    --save_path ./out_ppo \
    --use_flash_attn

9.5experience_maker：rollout 引擎

看 coati/experience_maker/naive.py：

class NaiveExperienceMaker:
    def __init__(self, actor, critic, reward_model, ref_model, tokenizer, kl_coef, ...):
        self.actor = actor
        self.critic = critic
        self.reward_model = reward_model
        self.ref_model = ref_model

    @torch.no_grad()
    def make_experience(self, prompts):
        # 1) Actor 生成
        sequences = self.actor.generate(prompts, max_length=...)
        # 2) 计算 log-probs（actor 和 ref）
        action_log_probs = compute_log_probs(self.actor, sequences)
        ref_log_probs    = compute_log_probs(self.ref_model, sequences)
        # 3) Reward
        rewards = self.reward_model(sequences)
        # 4) KL 惩罚（per-token）
        kl = action_log_probs - ref_log_probs
        rewards_with_kl = rewards - self.kl_coef * kl.sum(-1)
        # 5) Value
        values = self.critic(sequences)
        # 6) GAE advantage
        advantages = compute_gae(rewards_with_kl, values, gamma, lam)
        return Experience(sequences, action_log_probs, advantages, ...)

9.6显存 / 时间预算

LLaMA-3-8B PPO、8×A100-80GB、TP=2、ZeRO-1、seqlen=2048：

显存不够时常用招：

• 把 reference 和 reward 合并部署到一组卡（共用 base 权重）；
• actor 上 LoRA，只更新 LoRA 部分；
• 开 gradient checkpointing。

9.7DPO / GRPO 的位置

ColossalChat 同时支持 DPO（去 RM 直接训）和 GRPO（去 critic）：

显存紧时优先用 DPO（一份模型 + reference 就够），效果接近 PPO 但工程复杂度低一个数量级。

9.8常见踩坑

9.9这章你需要带走的

• InstructGPT RLHF 三段式：SFT → RM → PPO，ColossalChat 全部提供官方实现；
• PPO 维护 4 个模型（Actor / Critic / RM / Reference），显存压力是 SFT 的 4 倍；
• experience_maker 是核心 rollout 引擎，串起来 generate / reward / KL / GAE；
• 显存紧选 DPO；rule-based reward 选 GRPO；
• 典型问题 = reward hacking，靠 kl_coef + reward normalization 控制；
• 训练时间最大头是 rollout（~50% 时间）。