写 rollout：让模型自己生成一批答案

这一阶段要写的东西

四节课拼出一个文件：

scripts/grpo_scratch/
├── rollout.py   ← 本节：让模型生成一批答案
├── rewards.py   ← 下一节：打分
├── grpo.py      ← 第三节：优势、ratio、clip、KL
└── train.py     ← 第四节：串起来跑

总共约 200 行。跑完你会有一个完全属于自己的 RL 训练脚本。

rollout 的输入输出，先定死

写代码前先把接口定清楚，后面三节都按这个接口对接：

输入：一批 prompt（比如 8 个数学题）
     ↓
输出：一个列表，每个元素是
     {
       prompt_ids     : list[int]   ← prompt 的 token
       completion_ids : list[int]   ← 生成部分的 token
       text           : str         ← 生成的文本（打分和肉眼检查用）
       group_id       : int         ← 属于哪个 prompt，算优势时要按组分
     }

prompts × G 条记录。8 个 prompt、G=8 就是 64 条。

ids = [15496, 11, 1917]
text = tok.decode(ids)          # "Hello, world"
ids2 = tok.encode(text)         # 可能变成 [15496, 11, 995] —— 长度或内容都可能变

写 rollout

# scripts/grpo_scratch/rollout.py
from dataclasses import dataclass
from vllm import LLM, SamplingParams


@dataclass
class Sample:
    prompt_ids: list[int]
    completion_ids: list[int]
    text: str
    group_id: int


def build_engine(model_path: str, util: float = 0.3, max_len: int = 1024) -> LLM:
    return LLM(
        model=model_path,
        gpu_memory_utilization=util,   # colocate：只给 vLLM 这么多
        max_model_len=max_len,
        dtype="bfloat16",
        enable_sleep_mode=True,        # 训练时能把显存还回来，第四节用
    )


def rollout(
    engine: LLM,
    tokenizer,
    questions: list[str],
    group_size: int = 8,
    max_tokens: int = 512,
    temperature: float = 1.0,
) -> list[Sample]:
    # 用 chat template 拼 prompt。注意 add_generation_prompt=True，
    # 否则模型不知道该轮到它说话了
    prompts = [
        tokenizer.apply_chat_template(
            [{"role": "user", "content": q}],
            tokenize=False,
            add_generation_prompt=True,
        )
        for q in questions
    ]

    params = SamplingParams(
        n=group_size,
        temperature=temperature,
        top_p=1.0,
        max_tokens=max_tokens,
    )
    outputs = engine.generate(prompts, params)

    samples: list[Sample] = []
    for gid, out in enumerate(outputs):
        for cand in out.outputs:
            samples.append(
                Sample(
                    prompt_ids=list(out.prompt_token_ids),
                    completion_ids=list(cand.token_ids),   # ← 关键：拿 id
                    text=cand.text,
                    group_id=gid,
                )
            )
    return samples

不到 40 行。关键点都在注释里。

采样参数怎么设

temperature：不能是 0

temperature = 0    # ❌ 绝对不行

GRPO 的基线是组内均值。temperature=0 时 8 条回答完全一样，reward 完全一样，std = 0，整组优势归零，一个梯度都没有。

推荐值：

top_p：设 1.0

推理时常用 top_p=0.9 来提升质量。但训练时不要裁剪分布：

被 top_p 砍掉的 token 概率在采样时是 0，但训练侧算 logprob 时用的是完整分布。两边不一致，会引入偏差。

保持 top_p = 1.0，让采样分布和训练分布严格对齐。

max_tokens

直接决定 KV cache 和采样时间。数学题任务 512 起步就够。

跑一次，肉眼看看

在写任何 reward 之前，先把 8 条回答打印出来看。这一步不能跳。

from transformers import AutoTokenizer

MODEL = "./models/Qwen3-0.6B"
tok = AutoTokenizer.from_pretrained(MODEL)
engine = build_engine(MODEL, util=0.3)

q = "小明有 5 个苹果，他给了小红 2 个，又买了 7 个。他现在有几个苹果？"
samples = rollout(engine, tok, [q], group_size=8, max_tokens=512)

for i, s in enumerate(samples):
    print(f"\n===== #{i}  ({len(s.completion_ids)} token) =====")
    print(s.text.strip())

你会看到 8 条不一样的回答，质量参差不齐。这个参差就是 GRPO 的信息来源。

rollout 检查演练。跑完采样后有三件事必须确认，按顺序来。
goals 看目标，hint 要提示。

拼 batch：pad 的方向

采样出来的序列长度不一，喂进训练前要 pad 成矩阵。这里有个容易错的地方。

右 pad（训练用这个）           左 pad（生成时用）
[prompt][completion][pad][pad]   [pad][pad][prompt][completion]

• 生成时用左 pad：这样所有序列的「最后一个 token」在同一列，方便取下一步的 logits。vLLM 内部自己处理，你不用管。
• 训练时用右 pad：位置索引从 0 开始连续，attention mask 和 loss mask 都好算。

同时要造两个 mask：

attention_mask   # 哪些位置是真实 token（pad 之外全是 1）
completion_mask  # 哪些位置要算 loss（只有 completion 部分是 1）

这节课的落点

• 接口定死：(prompt_ids, completion_ids, text, group_id)，共 prompts × G 条
• token id 一路带到训练，中间绝不重新 tokenize（TITO）
• temperature 必须 > 0，默认 1.0；top_p 保持 1.0 让采样与训练分布一致
• 观察「顶到 max_tokens 的比例」，超过 15% 要处理
• 生成用左 pad（vLLM 自理），训练用右 pad
• loss mask 只覆盖 completion，且注意与 logits 错一位对齐
• 写 reward 前先把 8 条回答打印出来肉眼看一遍