TRL GRPOTrainer：5090 单卡实跑

为什么单卡框架层选 TRL

这条路径参考的两个工业框架都跑不了单卡（slime 最小 8×H100，Miles 的硬件列表是 A100/H100/B200/MI300 系列）。所以框架层实跑用 TRL：

选 TRL 的关键理由是它的抽象层次刚好：GRPOTrainer 帮你管住训练循环和 vLLM 协作，但 reward 函数还是你自己写的普通 Python。你 L2 写的 rewards.py 可以原样接进来。

最小可跑脚本

# scripts/trl_grpo/train.py
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer

from scripts.grpo_scratch.rewards import correctness_reward, format_reward

MODEL = "./models/Qwen3-0.6B"


# ---------- reward：签名是固定的 ----------
def reward_correct(completions, answer, **kwargs):
    """
    completions: list[str]，本批所有生成
    answer:      list[str]，数据集里同名列会自动按批传进来
    返回:         list[float]，长度与 completions 一致
    """
    return [correctness_reward(c, a) for c, a in zip(completions, answer)]


def reward_format(completions, **kwargs):
    return [format_reward(c) for c in completions]


def main():
    ds = load_dataset("openai/gsm8k", "main", split="train")

    def prep(row):
        # prompt 列会被自动套 chat template
        return {
            "prompt": [{"role": "user", "content": row["question"]}],
            "answer": row["answer"].split("####")[-1].strip(),
        }

    ds = ds.map(prep, remove_columns=ds.column_names)

    cfg = GRPOConfig(
        output_dir="out/trl-grpo-01",

        # ---- 对应你手写脚本里的变量 ----
        num_generations=8,              # ← GROUP_SIZE
        per_device_train_batch_size=8,  # ← PROMPTS_PER_STEP（见下方说明）
        gradient_accumulation_steps=4,
        max_completion_length=512,      # ← max_tokens
        max_prompt_length=256,
        temperature=1.0,
        top_p=1.0,
        beta=0.0,                       # ← KL 系数
        epsilon=0.2,                    # ← clip 下界
        epsilon_high=0.28,              # ← clip 上界（非对称）
        learning_rate=1e-6,
        adam_beta2=0.98,
        max_grad_norm=1.0,
        max_steps=200,

        # ---- 单卡 colocate ----
        use_vllm=True,
        vllm_mode="colocate",
        vllm_gpu_memory_utilization=0.30,

        # ---- 显存 ----
        bf16=True,
        gradient_checkpointing=True,

        # ---- 日志 ----
        logging_steps=1,
        save_steps=20,
        report_to="none",               # 想接 wandb 就改这里
    )

    trainer = GRPOTrainer(
        model=MODEL,
        args=cfg,
        train_dataset=ds,
        reward_funcs=[reward_correct, reward_format],
        reward_weights=[1.0, 1.0],      # 权重原则同 L2：便宜的分给得少
    )
    trainer.train()


if __name__ == "__main__":
    main()

python -c "import trl; print(trl.__version__)"
python -c "from trl import GRPOConfig; import inspect; print(inspect.signature(GRPOConfig.__init__))"

参数对照表

这张表是这一节的核心。左边是你 L2 亲手写的，右边是 TRL 的名字：

接自己的 reward

TRL 的 reward 函数约定：

def my_reward(completions, **kwargs) -> list[float]:
    ...

三条规则：

1. completions 是本批所有生成的文本列表（不是单条）
2. 返回列表长度必须与 completions 一致
3. 数据集里的其它列会按名字通过 kwargs 传进来 —— 上面 answer=... 就是这么拿到的

多个 reward 函数用 reward_weights 加权求和，等价于你 L2 里的 total_reward。

def reward_correct(completions, answer, **kwargs):
    out = []
    for c, a in zip(completions, answer):
        try:
            out.append(correctness_reward(c, a))
        except Exception as e:
            print(f"[reward] 跳过一条：{e}")
            out.append(0.0)
    return out

跑起来，和 L2 的曲线放一起比

python -m scripts.trl_grpo.train 2>&1 | tee logs/trl-01.log

TRL 实跑演练。重点是把它的指标和你 L2 的曲线对上。
goals 看目标，hint 要提示。

分数对不上时先看哪三个地方

如果 TRL 的结果和你手写版差很多，按这个顺序查：

一、batch 语义

per_device_train_batch_size 是序列条数不是 prompt 数。算错会导致实际 batch 差好几倍，而 batch 大小直接影响梯度噪声。

查法：看日志里 epoch 的推进速度，或直接打印每步的实际样本数。

二、loss_type（序列平权 vs token 平权）

这是上面演示里出现的差异来源。两者不等价，长度曲线会明显不同。

查法：print(cfg.loss_type)，和你手写实现里的归一化方式对齐。

三、importance sampling 修正

TRL 默认开 vllm_importance_sampling_correction，它在修正 vLLM 与 transformers 的概率差异。你手写版（如果用的是「多跑一次前向拿 logp_old」的方案 A）不需要这个修正，两边的 ratio 分布会有细微不同。

查法：把它关掉再跑一次对比。

这节课的落点

• 单卡框架层选 TRL：抽象层次刚好，reward 还是你自己的 Python
• 你 L2 写的 rewards.py 能原样接进 reward_funcs
• reward 函数三条约定：收批量、返回等长列表、数据集其它列走 kwargs
• per_device_train_batch_size 是序列条数不是 prompt 数，最容易算错 batch
• loss_type 默认是 token 平权，与手写的序列平权不等价，长度曲线会分岔
• 用 inspect.signature 问你装的那个版本，不要照抄博客
• 分数对不上先查三处：batch 语义、loss_type、importance sampling 修正