vime 使用指导书

Versatile Infrastructure for Model Evolution — 基于 Megatron + vLLM 的大规模 LLM/VLM 强化学习训练框架

目录

1. 项目概述

vime 是由 GLM 团队开发的大规模强化学习（RL）训练框架，专为 LLM 和 VLM 的后训练（post-training）设计，是开源 GLM/ChatGLM 模型系列的 RL 训练后端。

1.1 适用场景

vime 的设计目标是覆盖 LLM/VLM 后训练阶段的各类强化学习需求。以下按场景详细说明 vime 的适用范围与典型用法。

数学与推理强化

场景描述：通过规则驱动的奖励信号（如答案正确性验证），提升模型在数学、逻辑推理等任务上的表现。

典型工作流：

  1. 准备含问题-标签对的数据集（如 dapo-math-17k、AIME）
  2. 使用 GRPO 算法，每条 Prompt 采样多个回答，组内相对比较计算优势
  3. 选择规则型奖励（--rm-type deepscaler 或 --rm-type math），自动判别答案正确性
  4. 可选启用动态采样（DAPO），过滤全对或全错的 Prompt 组，提升训练数据质量

适用特征：

  • 奖励可由规则函数精确判定，无需训练奖励模型
  • GRPO 无需 Critic，资源效率高，适合快速迭代
  • 对应示例：scripts/run-qwen3-4B.sh
代码生成与编程 Agent

场景描述：训练模型自主编写代码、调用工具（如代码解释器、沙箱）、根据执行结果迭代修正，完成编程任务。

典型工作流：

  1. 准备编程题数据集，在 metadata 中传入测试用例或执行环境参数
  2. 编写自定义生成函数（--custom-generate-function-path），实现"生成代码 → 沙箱执行 → 获取输出 → 继续生成"的多轮循环
  3. 编写基于测试用例的奖励函数（--custom-rm-path），通过测试通过率计算奖励
  4. 正确设置 loss_mask：模型生成的代码 token 参与 loss，沙箱输出 token 不参与

适用特征：

  • 需要多轮交互与外部工具调用
  • 奖励信号来自环境反馈（测试用例通过/失败）
  • 对应示例：examples/coding_agent_rl/
多轮对话与工具调用（Agent RL）

场景描述：训练具备工具使用能力（搜索、API 调用、浏览器操作等）的 Agent 模型，在多轮交互中自主决策。

典型工作流：

  1. 将对话历史、工具定义、会话状态等存入 metadata 字段
  2. 自定义生成函数实现交互循环：模型输出 → 解析动作 → 调用工具 → 拼接观察 → 继续生成
  3. 自定义奖励函数评估完整轨迹（如任务完成度、工具使用效率）
  4. 必要时使用 --rollout-function-path 替换完整 Rollout 编排，实现多 Agent 协作

适用特征：

  • 单次推理无法完成，需要与环境/工具反复交互
  • loss_mask 需要精细区分"自主决策"与"被动接收"的 token
  • 支持多 Agent 配合（如 planner + executor 模式）
  • 对应示例：examples/multi_agent/、examples/fully_async/
视觉语言模型（VLM）强化学习

场景描述：对图文多模态模型进行 RL 后训练，提升视觉理解与视觉推理能力。

典型工作流：

  1. 准备含图像输入的数据集，vime 支持多模态 Sample 的数据加载
  2. 使用与文本模型相同的 RL 算法（GRPO/PPO 等），vime 自动处理多模态输入的 tokenization
  3. 可选使用自定义奖励函数处理视觉相关的评分逻辑

适用特征：

  • 数据中包含图像等多模态输入
  • 训练流程与纯文本场景基本一致，vime 内部处理多模态差异
  • 对应示例：examples/geo3k_vlm/、examples/geo3k_vlm_multi_turn/
大规模 MoE 模型训练

场景描述：对 Mixture-of-Experts 架构的大模型（如 DeepSeek-V3、Qwen3-30B-A3B）进行 RL 训练，需要处理专家并行、路由稳定性等挑战。

典型工作流：

  1. 配置 Expert Parallel（--expert-model-parallel-size）和专家 TP（--expert-tensor-parallel-size）
  2. 启用 MoE Routing Replay（--use-routing-replay / --use-rollout-routing-replay）确保前向-反向路由一致性
  3. 多机训练时使用 PD 分离或推测解码优化推理吞吐
  4. 可配合 vLLM 多模型部署（--vllm-config），同时部署 Actor 和冻结参考模型

适用特征：

  • 模型参数量大（30B+），需要多机多卡分布式训练
  • MoE 的稀疏激活需要特殊的并行和路由策略
  • 长时间训练需要容错机制（--use-fault-tolerance）
  • 对应示例：docs/zh/examples/qwen3-30B-A3B.md
通用 RL 后训练研究

场景描述：研究新型 RL 算法、奖励函数或训练策略，需要灵活替换训练流水线中的各环节。

典型工作流：

  1. 选择基础算法（--advantage-estimator）作为起点
  2. 通过 --custom-loss-function-path 实现自定义损失函数
  3. 通过 --custom-reward-post-process-path 实现奖励塑形
  4. 通过 --custom-tis-function-path 实现离线策略修正
  5. 通过 --custom-pg-loss-reducer-function-path 实现自定义策略梯度归约（如 Dr.GRPO）
  6. 使用 --rollout-data-postprocess-path 和 --custom-convert-samples-to-train-data-path 定制数据转换
  7. 利用 CPU contract 测试（tests/plugin_contracts/）验证自定义实现的接口合规性

适用特征：

  • vime 的 18 个 --*-path 扩展接口可覆盖训练流水线的每个环节
  • 无需修改框架代码，即插即用
  • 支持多种重要性采样和离线策略方法（TIS 等）
场景选择速查
你的需求 推荐算法 核心接口 关键特性
数学/推理强化 GRPO --rm-type 规则奖励，动态采样
代码 Agent GRPO --custom-generate-function-path + --custom-rm-path 沙箱执行，测试用例奖励
多轮工具调用 GRPO/PPO --custom-generate-function-path 或 --rollout-function-path loss_mask，多 Agent
VLM 强化 GRPO 默认接口 多模态数据加载
MoE 大模型 GRPO --vllm-config + --use-routing-replay EP，路由重放，PD 分离
算法研究 自选 --custom-loss-function-path 等 18 个扩展接口

1.2 核心能力

能力 说明
异步 PPO/GRPO 训练 训练与推理异步重叠，最大化 GPU 利用率
多算法支持 GRPO、GSPO、Reinforce++、PPO 等多种 RL 算法
MoE 支持 完整的 Expert Parallel、MoE Routing Replay 等优化
Megatron 后端 高效的大规模模型训练，支持 TP/PP/CP/EP 多维并行
vLLM 推理 高吞吐量的 rollout 数据生成，支持 PD 分离
高度可定制 通过 18 个 --*-path 参数替换训练流水线的各环节
多 Agent RL 支持多轮交互、工具调用、沙箱执行等 Agent 场景
推测解码 支持 MTP/EAGLE 等投机采样加速推理
容错恢复 Rollout Engine 健康检查、自动重启、Debug Replay

1.3 支持的模型

vime 已验证支持以下模型系列（完整列表参见 scripts/models/ 目录）：

  • Qwen 系列（Qwen3-4B、Qwen3-30B-A3B 等）
  • DeepSeek V3 系列
  • GLM 系列（GLM-4、GLM-5 等）
  • Llama 3 系列

2. 系统架构

┌─────────────────────────────────────────────────────────────┐
│                    train_async.py / train.py                 │
│                     （主入口，训练循环编排）                    │
├──────────────┬──────────────────────┬───────────────────────┤
│              │                      │                       │
│  Placement   │    Rollout Manager   │    Training Actors    │
│   Group      │   （vLLM 推理引擎）    │    （Megatron 训练）   │
│              │                      │                       │
│  GPU 资源    │  ┌──────────────┐    │  ┌─────────────────┐  │
│  跨节点调度  │  │ vLLM Engine 0│    │  │ Actor Model     │  │
│              │  │ vLLM Engine 1│    │  │ (Policy Network)│  │
│              │  │     ...      │    │  │ Critic Model    │  │
│              │  └──────────────┘    │  │ (仅 PPO)        │  │
│              │                      │  └─────────────────┘  │
│              │  vllm-router 负载均衡 │                       │
├──────────────┴──────────────────────┴───────────────────────┤
│                      Data Buffer                            │
│         （数据缓存、动态采样、Partial Rollout）                 │
└─────────────────────────────────────────────────────────────┘

核心流程

┌──────────┐    ┌──────────┐    ┌──────────┐    ┌──────────┐
│  Prompt   │───▶│  Rollout  │───▶│  Reward   │───▶│ Training  │
│  数据加载  │    │  (vLLM)  │    │  计算     │    │ (Megatron)│
└──────────┘    └──────────┘    └──────────┘    └──────────┘
                                          │
                                          ▼
                                   ┌──────────┐
                                   │  Weight   │
                                   │  Sync     │
                                   │ (Megatron │
                                   │  → vLLM)  │
                                   └──────────┘

异步模式：Rollout 生成与 Training 训练并行进行。当一个 batch 在训练时，下一个 batch 已在 Rollout 阶段生成，从而最大化 GPU 利用率。

3. 环境搭建

3.1 硬件支持

硬件平台 状态 说明
H 系列 (H100/H200) 官方支持 完整 CI 测试保护，推荐生产环境使用
B 系列 (B200/B300) 支持 基本功能稳定，暂无 CI 保护
GB200/GB300 支持 与 H 系列运行步骤相同

3.2 使用 Docker 镜像（推荐）

vime 可能包含针对 vLLM/Megatron 的临时补丁，强烈建议使用官方提供的 Docker 镜像，已预置所有依赖：

# 拉取最新镜像
docker pull inferactinc/public:vime-latest

# 启动容器
docker run --rm --gpus all --ipc=host --shm-size=16g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -it inferactinc/public:vime-latest /bin/bash

3.3 安装 vime

Docker 镜像中已预装 vime。如需更新到最新版本：

cd /root/vime
git pull
pip install -e . --no-deps

3.4 依赖说明

主要依赖包括：

  • PyTorch — 深度学习框架
  • Megatron-LM — 大规模模型训练后端（需加入 PYTHONPATH）
  • vLLM — 高性能推理引擎
  • Ray — 分布式计算框架
  • Transformers — HuggingFace 模型库
  • wandb / TensorBoard — 训练监控（可选）

重要：运行时需确保 Megatron-LM 在 PYTHONPATH 中：

export PYTHONPATH=/root/Megatron-LM/

4. 数据准备

4.1 数据格式

vime 当前仅支持 .jsonl 格式，每行一个 JSON 对象。

基础格式（单轮对话）：

{
  "prompt": [
    {
      "content": "Solve the following math problem step by step...",
      "role": "user",
      "step_loss_mask": 1
    }
  ],
  "label": "34"
}

多轮对话格式：

{
  "prompt": [
    {"content": "System instruction", "role": "system", "step_loss_mask": 0},
    {"content": "User question", "role": "user", "step_loss_mask": 1}
  ],
  "label": "answer"
}

附带元数据（用于自定义生成/奖励函数）：

{
  "prompt": [...],
  "label": "answer",
  "metadata": "{\"session_id\": \"sess_123\", \"tool_code\": \"code_A\"}"
}

4.2 字段说明

字段 说明 是否必需
prompt 输入对话列表，每项包含 content、role、step_loss_mask 是
label 评估标签，用于奖励计算 视 --rm-type 而定
metadata 附加信息，供自定义函数使用 否
step_loss_mask SFT 阶段使用。1 = 参与 loss 计算，0 = 不参与 否（默认 1）

4.3 数据下载示例

# 下载训练数据集 (dapo-math-17k)
hf download --repo-type dataset zhuzilin/dapo-math-17k \
  --local-dir /root/dapo-math-17k

# 下载评估数据集 (aime-2024)
hf download --repo-type dataset zhuzilin/aime-2024 \
  --local-dir /root/aime-2024

4.4 对应参数

--prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
--input-key prompt           # 指定 prompt 对应的 JSON key
--label-key label            # 指定 label 对应的 JSON key
--metadata-key metadata      # 指定 metadata 对应的 JSON key（可选）
--apply-chat-template        # prompt 为 OpenAI message 格式时需要开启

5. 权重准备与转换

Megatron 无法直接加载 HuggingFace 格式的检查点，需先进行权重转换。

5.1 HuggingFace → Megatron torch_dist 格式

第一步：加载模型配置。scripts/models/ 目录提供了常用模型的参数配置：

cd /root/vime
source scripts/models/qwen3-4B.sh

注意：请务必检查模型配置文件中的参数（如 --rotary-base）是否与您使用的模型版本完全匹配。如需修改，可在 source 之后直接覆盖：

source scripts/models/qwen3-4B.sh
MODEL_ARGS+=(--rotary-base 10000)

第二步：运行转换脚本：

PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /root/Qwen3-4B \
    --save /root/Qwen3-4B_torch_dist

对于更大的模型，可以使用 torchrun 多卡转换：

torchrun --nproc_per_node=8 tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /root/large-model \
    --save /root/large-model_torch_dist

5.2 Megatron torch_dist → HuggingFace 格式

将训练保存的 Megatron 权重转回 HuggingFace 格式：

PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \
  --input-dir /path/to/torch_dist_ckpt/iter_xxx/ \
  --output-dir /root/Qwen3-4B-iter_xxx \
  --origin-hf-dir /root/Qwen3-4B

注意：Megatron 会对 embedding 做 padding，可能导致转换后的权重 embedding 形状不匹配。此时需要设置 --vocab-size 参数。

5.3 检查点存储结构

/root/Qwen3-4B_torch_dist/
├── latest_checkpointed_iteration.txt
├── iter_0000100/
│    ├── _0_0.distcp
│    ├── _0_1.distcp
│    └── ...
├── iter_0000200/
└── iter_0000300/

加载时传递目录根路径，使用 --ckpt-step 指定训练步数（不指定则读取 latest_checkpointed_iteration.txt）。

6. 快速启动：GRPO 训练实战

以下以 Qwen3-4B + dapo-math-17k 数据集 为例，展示完整的 GRPO 训练流程。

6.1 准备工作

# 1. 下载模型与数据
hf download zai-org/Qwen3-4B --local-dir /root/Qwen3-4B
hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k
hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024

# 2. 转换权重
cd /root/vime
source scripts/models/qwen3-4B.sh
PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
    ${MODEL_ARGS[@]} \
    --hf-checkpoint /root/Qwen3-4B \
    --save /root/Qwen3-4B_torch_dist

6.2 启动训练

cd /root/vime
bash scripts/run-qwen3-4B.sh

6.3 脚本参数拆解

以 scripts/run-qwen3-4B.sh 为例，理解核心参数分组：

模型配置参数

通过 source scripts/models/qwen3-4B.sh 加载，包含 Megatron 所需的模型结构参数（层数、隐藏维度、注意力头等）。

检查点参数
CKPT_ARGS=(
   --hf-checkpoint /root/Qwen3-4B     # HF 权重（用于 vLLM 初始化和 tokenizer）
   --ref-load /root/Qwen3-4B_torch_dist  # 参考模型检查点
   --load /root/Qwen3-4B_vime/        # Actor 加载路径（不存在则从 ref-load 加载）
   --save /root/Qwen3-4B_vime/        # Actor 保存路径
   --save-interval 20                 # 每 20 步保存一次
)
Rollout 参数
ROLLOUT_ARGS=(
   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --rm-type deepscaler              # 奖励模型类型
   --num-rollout 3000                # 总训练轮次
   --rollout-batch-size 32           # 每轮 Prompt 数量
   --n-samples-per-prompt 8          # 每 Prompt 采样数（GRPO）
   --rollout-max-response-len 8192   # 最大响应长度
   --rollout-temperature 1           # 采样温度
   --global-batch-size 256           # 全局批大小
   --balance-data                    # 训练数据负载均衡
)

关键约束：

rollout_batch_size × n_samples_per_prompt = global_batch_size × num_steps_per_rollout

即每轮"产出"与"消耗"必须相等。若设置了 --num-steps-per-rollout，--global-batch-size 未设置则自动计算。

算法参数
GRPO_ARGS=(
   --advantage-estimator grpo      # 使用 GRPO 算法
   --use-kl-loss                   # 启用 KL 散度监控
   --kl-loss-coef 0.00             # KL 系数（0 = 仅监控，不参与 loss）
   --kl-loss-type low_var_kl
   --entropy-coef 0.00             # 熵系数
   --eps-clip 0.2                  # PPO clip 范围（下界）
   --eps-clip-high 0.28            # PPO clip 范围（上界）
)
性能与并行参数
PERF_ARGS=(
   --tensor-model-parallel-size 2  # TP 并行度
   --sequence-parallel              # SP（建议启用）
   --pipeline-model-parallel-size 1 # PP 并行度
   --context-parallel-size 1       # CP（Ring Attention）
   --recompute-granularity full    # 重计算粒度
   --recompute-method uniform
   --recompute-num-layers 1
   --use-dynamic-batch-size        # 动态批处理（强烈推荐）
   --max-tokens-per-gpu 9216       # 每卡最大 token 数
)

提示：vime 总是通过 data packing 方法训练模型，并严格保证 per sample / per token loss 正确。开启动态批处理不影响 loss 计算，强烈推荐开启。

优化器参数
OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-6
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98
)
vLLM 参数
VLLM_ARGS=(
   --rollout-num-gpus-per-engine 2     # vLLM 的 tp_size
   --vllm-gpu-memory-utilization 0.7   # vLLM 显存占用（共享模式需调低）
)

vLLM 参数通过添加 --vllm- 前缀透传，例如 vLLM 的 --max-model-len 对应 vime 的 --vllm-max-model-len。

7. 参数详解

7.1 集群资源分配

参数 说明 示例
--actor-num-nodes 训练节点数 1
--actor-num-gpus-per-node 每节点训练 GPU 数 8
--rollout-num-gpus 推理总 GPU 数 8
--rollout-num-gpus-per-engine 每推理引擎 GPU 数（= vLLM tp_size） 2
--colocate 启用训推一体模式 （flag，无需赋值）
--critic-num-nodes Critic 节点数（仅 PPO） 1
--critic-num-gpus-per-node Critic 每节点 GPU 数（仅 PPO） 4

GPU 总数计算：

  • GRPO（训推分离）：actor_gpus + rollout_gpus
  • GRPO（训推一体）：actor_gpus（训练和推理共享）
  • PPO（训推分离）：actor_gpus + critic_gpus + rollout_gpus

7.2 Megatron 并行策略

参数 说明 推荐值
--tensor-model-parallel-size 张量并行 (TP) 根据模型大小调整
--sequence-parallel 序列并行（TP 的优化） 使用 TP 时建议开启
--pipeline-model-parallel-size 流水线并行 (PP) 1（按需调整）
--context-parallel-size 上下文并行 (CP，Ring Attention) 长序列时开启
--expert-model-parallel-size 专家并行 (EP，用于 MoE) MoE 模型按需设置
--expert-tensor-parallel-size 专家张量并行 (ETP) MoE 模型按需设置

7.3 重计算配置

参数 说明 可选值
--recompute-granularity 重计算粒度 full / selective / 不设置
--recompute-method 重计算方法 uniform
--recompute-num-layers 每组重计算层数 1

7.4 检查点参数

参数 说明
--ref-load 参考模型 Megatron 检查点路径
--load Actor 检查点路径（不存在则从 --ref-load 初始化）
--save Actor 检查点保存路径
--save-interval 保存间隔（步数）
--ckpt-format 检查点格式（torch / torch_dist，推荐 torch_dist）
--hf-checkpoint HuggingFace 检查点路径（用于 vLLM 和 tokenizer）

7.5 Rollout 采样参数

参数 说明 默认值
--rollout-batch-size 每轮 Prompt 数量 —
--n-samples-per-prompt 每 Prompt 采样数（GRPO） —
--rollout-max-response-len 最大响应 token 数 —
--rollout-temperature 采样温度 —
--rollout-top-p Top-p 采样 —
--global-batch-size 全局批大小 自动计算
--num-steps-per-rollout 每轮训练更新步数 1
--num-rollout 总训练轮次 —

8. RL 算法选择与配置

8.1 算法对比

算法 命令 是否需要 Critic 特点
GRPO --advantage-estimator grpo 否 资源效率高，组内相对比较
PPO --advantage-estimator ppo 是 经典算法，需要额外 Critic GPU
GSPO --advantage-estimator gspo 否 GRPO 改进变体
Reinforce++ --advantage-estimator reinforce_plus_plus 否 Reinforce 带基线改进
Reinforce++ Baseline --advantage-estimator reinforce_plus_plus_baseline 否 带基线的 Reinforce++

8.2 GRPO 配置

GRPO (Group Relative Policy Optimization) 对同一 Prompt 采样多个响应，通过组内相对奖励计算优势，无需 Critic 模型。

--advantage-estimator grpo
--n-samples-per-prompt 8      # 组内采样数（建议 4-16）
--eps-clip 0.2                # clip 范围
--normalize-advantages        # 优势归一化（可选）

相关参数：

  • --use-kl-loss — 启用 KL 散度监控（结合 --kl-loss-coef 控制是否计入 loss）
  • --entropy-coef — 熵正则系数
  • --calculate-per-token-loss — 切换 per-token loss（默认 per-sample loss）

8.3 PPO 配置

PPO 需要额外的 Critic 模型来估计价值函数：

--advantage-estimator ppo
--critic-num-nodes 1           # Critic 节点数
--critic-num-gpus-per-node 4   # Critic GPU 数（与 Actor 并行）
--critic-load /path/to/ckpt    # Critic 检查点
--critic-save /path/to/save    # Critic 保存路径
--critic-lr 1e-5               # Critic 学习率
--eps-clip 0.2                 # 策略 clip 范围
--value-clip 0.2               # 价值 clip 范围
--kl-coef 0.01                 # KL 惩罚系数

注意：PPO 中 Actor 和 Critic 各占一套独立的 GPU 资源。总 GPU 数 = Actor GPUs + Critic GPUs + Rollout GPUs。

8.4 Megatron 配置覆盖（PPO）

PPO 中可通过 --megatron-config-path 分别覆盖 Actor 和 Critic 的 Megatron 参数：

# megatron_config.yaml
megatron:
  - name: default
    role: actor
    overrides:
      lr: 1e-6
  - name: default
    role: critic
    overrides:
      lr: 1e-5

当前限制：此功能仅支持 PPO，且 Actor 和 Critic 必须使用相同的并行拓扑。

9. 训推分离与训推一体

9.1 训推分离（默认）

训练和推理使用独立的 GPU 资源，可并行运行：

--actor-num-nodes 1
--actor-num-gpus-per-node 4
--rollout-num-gpus 4

上述配置：Actor 使用 4 卡，Rollout 使用 4 卡，共需 8 卡。两者可同时运行。

9.2 训推一体（Colocate）

训练和推理共享同一组 GPU，交替执行（Megatron offload → vLLM 加载 → 推理 → vLLM offload → Megatron 加载 → 训练）：

--actor-num-nodes 1
--actor-num-gpus-per-node 8
--colocate

此时共需 8 卡，训练和推理共享。

注意：

  • 共享模式下需调整 --vllm-gpu-memory-utilization 降低 vLLM 显存占用，通常建议 0.8
  • torch_memory_saver 的部分优化仅在训推一体模式可用
  • 训推一体模式的 GPU 利用率通常低于分离模式，但节省总卡数

10. 评估配置

评估过程继承大部分 Rollout 参数，但可通过专用参数覆盖：

EVAL_ARGS=(
   --eval-interval 20                    # 每 20 轮 Rollout 评估一次
   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl  # 评估数据集
   --n-samples-per-eval-prompt 16        # 评估时每 Prompt 采样数
   --eval-max-response-len 16384         # 评估最大响应长度
   --eval-top-p 1                        # 评估采样 top-p
)

多数据集评估

vime 支持同时评估多个数据集：

--eval-prompt-data aime /root/aime-2024/aime-2024.jsonl \
--eval-prompt-data math /root/math-500/math-500.jsonl

评估结果会按数据集分别汇报。

11. 自定义扩展

vime 的核心设计理念是通过 --*-path 参数注入自定义逻辑，无需修改框架代码。

11.1 扩展接口总览

接口参数 用途
--rollout-function-path 替换整个 Rollout 生成逻辑
--custom-generate-function-path 仅替换生成步骤（保留框架 Rollout 循环）
--custom-rm-path 自定义奖励计算
--dynamic-sampling-filter-path 动态采样过滤
--buffer-filter-path 缓冲区数据过滤
--rollout-sample-filter-path 单样本过滤（是否参与 loss）
--rollout-data-postprocess-path Rollout 数据后处理
--custom-loss-function-path 自定义训练损失
--custom-reward-post-process-path 奖励后处理
--custom-convert-samples-to-train-data-path 样本转训练数据格式
--data-source-path 自定义 Prompt 数据源
--eval-function-path 自定义评估函数
--custom-rollout-log-function-path 自定义训练 Rollout 日志
--custom-eval-rollout-log-function-path 自定义评估 Rollout 日志
--custom-megatron-init-path Megatron 初始化钩子
--custom-megatron-before-log-prob-hook-path Log prob 计算前钩子
--custom-megatron-before-train-step-hook-path 训练步前钩子

11.2 自定义奖励函数

签名（单样本模式）：

async def custom_rm(args, sample: Sample) -> float:
    """计算单个样本的奖励"""
    ...
    return reward_value

签名（批量模式，需配合 --group-rm）：

async def batched_custom_rm(args, samples: list[Sample]) -> list[float]:
    """批量计算样本奖励"""
    ...
    return [reward_1, reward_2, ...]

使用方式：

--custom-rm-path my_module.reward:custom_rm

内置奖励类型（--rm-type）：

类型 说明
math 数学答案验证
dapo DAPO 风格评分
deepscaler DeepScaler 规则奖励
f1 F1 分数
gpqa GPQA 奖励
ifbench IFBench 奖励
remote_rm 远程奖励服务（需 --rm-url）

11.3 自定义生成函数

适用于需要工具调用、RAG、多轮交互等场景：

async def generate(args, sample: Sample, sampling_params) -> Sample:
    """自定义生成逻辑"""
    # 1. 发送请求到 vLLM
    # 2. 处理响应
    # 3. 填充 Sample 字段
    sample.tokens = prompt_token_ids + response_token_ids
    sample.response_length = len(response_token_ids)
    sample.status = Sample.Status.COMPLETED  # 或 TRUNCATED / ABORTED
    sample.response = tokenizer.decode(response_token_ids)
    return sample

--custom-generate-function-path my_module.generate:generate

11.4 自定义 Rollout 函数

完全替换框架的 Rollout 逻辑：

def generate_rollout(args, rollout_id, data_source, evaluation=False):
    """
    Args:
        args: 完整参数
        rollout_id: 当前 Rollout 轮次 ID
        data_source: 数据源对象
        evaluation: 是否为评估模式
    Returns:
        RolloutFnTrainOutput | RolloutFnEvalOutput
    """
    ...

--rollout-function-path my_module.rollout:generate_rollout

推荐：大多数场景用 --custom-generate-function-path 即可，仅当需要完全替换 Rollout 编排时才使用 --rollout-function-path。

11.5 动态采样

实现 DAPO 风格的动态采样，过滤奖励标准差为零的样本组：

--over-sampling-batch-size 64 \
--dynamic-sampling-filter-path \
  vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std

其中 over_sampling_batch_size 须大于 rollout_batch_size。系统会过采样 → 过滤 → 不足时自动触发新一轮过采样。

11.6 Partial Rollout

在动态采样中被中止的样本可缓存起来，在下一轮 Rollout 继续生成：

--partial-rollout
--buffer-filter-path my_module:pop_first   # 可选，默认 pop_first

12. 多轮交互（Agent）适配

vime 支持复杂的多轮交互场景（工具调用、环境反馈等），通过自定义函数实现。

12.1 适配步骤

  1. 数据准备：将多轮交互数据映射为 vime 的 Sample 格式，附加信息存入 metadata
  2. 自定义生成函数：实现"模型生成 → 执行工具 → 拼接观察"的交互循环
  3. 自定义奖励函数：评估完整交互轨迹

12.2 数据映射

在 JSONL 中将额外信息聚合到 metadata 字段：

{
  "prompt": [...],
  "label": "answer",
  "metadata": "{\"session_id\": \"sess_123\", \"tool_code\": \"code_A\"}"
}

训练脚本中：

--metadata-key metadata

自定义函数中通过 sample.metadata['session_id'] 访问。

12.3 自定义生成函数示例

async def generate(args, sample: Sample, sampling_params) -> Sample:
    prompt, full_response, loss_masks = sample.prompt, "", []

    for _ in range(max_turns):
        # 1. 模型生成动作
        model_output = await call_vllm(prompt + full_response, ...)
        loss_masks += [1] * len(model_tokens)     # 模型输出 → loss_mask = 1
        full_response += model_output

        # 2. 解析并执行动作
        action, content = parse_action(model_output)
        if action == "search":
            tool_output = await google_search(content)
            loss_masks += [0] * len(tool_tokens)   # 工具输出 → loss_mask = 0
            full_response += tool_output
        elif action == "answer":
            break

    # 3. 填充 Sample
    sample.response = full_response
    sample.tokens = ...
    sample.loss_mask = loss_masks
    return sample

12.4 配置

--custom-generate-function-path your_module.multiturn:generate
--custom-rm-path your_module.multiturn:reward_func

完整示例参见 examples/multi_agent/ 和 examples/fully_async/。

13. 多机分布式训练

13.1 启动 Ray 集群

# Node0（HEAD 节点）
ray start --head --node-ip-address ${MASTER_ADDR} \
  --num-gpus 8 --disable-usage-stats

# 其他 Worker 节点
ray start --address=${MASTER_ADDR}:6379 --num-gpus 8

13.2 提交训练任务

ray job submit --address="http://127.0.0.1:8265" \
   --runtime-env-json='{
     "env_vars": {
        "PYTHONPATH": "/root/Megatron-LM/",
        "CUDA_DEVICE_MAX_CONNECTIONS": "1"
     }
   }' \
   -- python3 train_async.py \
   --train-backend megatron \
   --actor-num-nodes 2 \
   --actor-num-gpus-per-node 8 \
   --rollout-num-gpus 8 \
   ...（其他参数）

13.3 单机启动（开发/调试）

# 启动本地 Ray
ray start --head --node-ip-address 127.0.0.1 --num-gpus 8 --disable-usage-stats

# 提交任务
ray job submit --address="http://127.0.0.1:8265" \
   --runtime-env-json='{"env_vars": {"PYTHONPATH": "/root/Megatron-LM/"}}' \
   -- python3 train.py \
   --colocate \
   ...

14. 高级特性

14.1 PD 分离（Prefill/Decode Disaggregation）

将推理的 Prefill 和 Decode 阶段分离到不同服务器：

--prefill-num-servers 2   # Prefill 服务器数量

14.2 推测解码（Speculative Decoding）

使用轻量级 draft model 加速推理：

MTP 层（适用于 GLM-4.7、DeepSeek-V3/R1）：

--vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}'

外部 Draft Model：

--vllm-speculative-config '{"method":"eagle","num_speculative_tokens":3,"model":"/your/draft/model/path"}'

在线训练 MTP 层（防止 draft model 与 target model 漂移）：

--mtp-num-layers 1
--enable-mtp-training
--mtp-loss-scaling-factor 0.2

注意：MTP 训练需要包含 MTP 权重的检查点，权重转换时也需加 --mtp-num-layers 1。

14.3 多模型部署

通过 --vllm-config YAML 文件配置多模型同时部署：

vllm:
  - name: actor
    update_weights: true
    server_groups:
      - worker_type: regular
        num_gpus: 8
        num_gpus_per_engine: 4
  - name: ref
    model_path: /path/to/ref_model
    update_weights: false
    server_groups:
      - worker_type: regular
        num_gpus: 4
        num_gpus_per_engine: 2

每个模型拥有独立的 router，通过 args.vllm_model_routers 访问。自定义函数中可使用 get_model_url(args, "ref") 路由请求。

14.4 MoE Routing Replay

稳定 MoE RL 训练，重放专家路由决策确保一致性：

参数 说明
--use-routing-replay 前向-反向路由一致性
--use-rollout-routing-replay R3：从 Rollout 重放路由到训练

14.5 动态批处理

--use-dynamic-batch-size      # 启用动态批处理
--max-tokens-per-gpu 9216     # 每卡最大 token 数

启用后，系统智能地将长短不一的样本打包，使每个 micro-batch 的总 token 数接近此限制。CP 模式下，N 张 CP 卡共享 N × max_tokens_per_gpu 的总长度。

14.6 vLLM Router

vime 使用 vllm-router 管理多个 vLLM 引擎：

--vllm-router-ip 127.0.0.1    # Router IP（不设置则自动启动）
--vllm-router-port 8000       # Router 端口

Router 支持负载均衡和 OpenAI 兼容 API。默认使用 consistent_hash 路由。可通过 --router-* 前缀传递 router 参数，例如 --router-balance-abs-threshold 0 强制均衡分布。

15. 容错与调试

15.1 开启容错

--use-fault-tolerance \
--rollout-health-check-first-wait 600 \
--rollout-health-check-interval 10 \
--rollout-health-check-timeout 5

当前覆盖范围：

  • vLLM Engine 健康检查（定期 heartbeat）
  • 超时后自动重启 Engine
  • 重启后正确更新参数
  • 保存 Debug Rollout Dump

15.2 Debug 路径

参数 说明
--debug-rollout-only 仅跑 Rollout 并保存数据，不训练
--save-debug-rollout-data /path/to/rollout_{rollout_id}.pt 保存 Rollout 数据
--load-debug-rollout-data /path/to/rollout_{rollout_id}.pt 加载已保存 Rollout 数据
--debug-train-only 仅跑训练侧逻辑，跳过 Rollout

15.3 推荐生产模式

  1. 开启 --use-fault-tolerance
  2. 通过 --save-interval 定期保存 checkpoint
  3. 对新 workload 保存 Rollout debug dump
  4. 使用 Trace Viewer 检查 long-tail samples
  5. 使用 Profiling 区分 Rollout 与 Training 瓶颈

15.4 需要关注的信号

  • 大 MoE 模型启动阶段 health check 失败 → 增大 --rollout-health-check-first-wait
  • 短暂负载高峰导致误判 → 增大 --rollout-health-check-timeout
  • Engine 在 weight sync 后反复重启 → 检查 vLLM 日志和 debug dump
  • Trainer 侧失败 → 从 checkpoint 恢复，用 debug replay 确认数据有效性

16. 常见问题

Q1: --hf-checkpoint 和 --ref-load 有什么区别？

  • --hf-checkpoint：提供 HuggingFace 格式的检查点，用于 vLLM 初始化和加载 tokenizer。首个训练步前，vime 会将 Megatron 参数同步到 vLLM，因此不需要包含最新训练参数。
  • --ref-load：提供 Megatron torch_dist 格式的参考模型检查点，用于计算 KL 散度等。

Q2: 训推一体模式下显存不足怎么办？

  1. 降低 --vllm-gpu-memory-utilization（建议 0.7-0.8）
  2. 减小 --max-tokens-per-gpu
  3. 增大 --recompute-num-layers 或开启 --recompute-granularity full
  4. 增大 --tensor-model-parallel-size

Q3: 如何从训练中断处恢复？

vime 支持从检查点恢复训练。只需确保 --load 指向包含有效检查点的目录，vime 会自动加载最新的迭代步。无需修改 --hf-checkpoint。

Q4: 如何选择 GRPO 的 n_samples_per_prompt？

  • 过小（如 2-4）：组内方差估计不准确
  • 过大（如 32+）：计算资源消耗大，边际收益递减
  • 建议范围：4-16，常用 8

Q5: --calculate-per-token-loss 和默认 loss 有何区别？

  • 默认（per sample）：mean(sum(sample_i) / len(sample_i)) — 每个样本对 loss 贡献相同
  • per token：sum(sum(sample_i)) / sum(len(sample_i)) — 长样本贡献更大

Q6: 如何监控训练？

# WandDB
--use-wandb
--wandb-project your-project
--wandb-group your-group
--wandb-key ${WANDB_KEY}

# TensorBoard
--tensorboard-dir /path/to/tb_logs

Q7: 权重转换时 embedding 形状不匹配怎么办？

在转换时显式设置 --vocab-size，因为 Megatron 会对 embedding 做 padding。

参考文档

文档 路径
快速开始 docs/zh/get_started/quick_start.md
参数详解 docs/zh/get_started/usage.md
自定义指南 docs/en/get_started/customization.md
常见问题 docs/zh/get_started/qa.md
PD 分离 docs/zh/advanced/pd-disaggregation.md
推测解码 docs/zh/advanced/speculative-decoding.md
容错 docs/zh/advanced/fault-tolerance.md
Megatron 配置 docs/zh/advanced/megatron-config.md
vLLM 配置 docs/zh/advanced/vllm-config.md
调试 docs/zh/developer_guide/debug.md
Profiling docs/zh/developer_guide/profiling.md
Trace docs/zh/developer_guide/trace.md
示例：Qwen3-4B docs/zh/examples/qwen3-4B.md
示例：Qwen3-30B-A3B docs/zh/examples/qwen3-30B-A3B.md
# 人工智能 # 深度学习 # 语言模型
Logo
加入AMD AI开发者计划！

免费领 200 小时云算力，进群参与显卡、AI PC 幸运抽奖

更多推荐

用 vLLM 在 Instinct 显卡上抗高并发，显存优化与吞吐实测数据分享

本文分享在 Instinct 显卡上利用 vLLM 实现高并发推理的实战经验。通过 PagedAttention 优化显存碎片、FlashAttention 提升吞吐及动态批处理调优，显著降低延迟。实测数据显示，ROCm 7.x 环境下系统性能大幅提升，为企业级大模型部署提供高效解决方案。

avatar AMD开发者中国社区

ROCm 用户福音，TileLang 跨平台算子开发实录

本文记录 TileLang 在 ROCm 环境的跨平台算子开发实战。通过详解环境配置、依赖陷阱及 GEMM 算子适配，验证其“一次编写，多处运行”能力。文章分享编译排错经验，助力 AMD 显卡用户高效进行高性能算子开发，打破硬件生态壁垒。

avatar AMD开发者中国社区

大模型长文本推理显存不够用，ROCm 7.x 量化与重计算实战技巧

本文详解 ROCm 7.x 环境下大模型长文本推理的显存优化方案。通过 FP8 量化与激活值重计算实战技巧，有效解决 HIP out of memory 报错，在单卡资源有限时实现上下文翻倍，助力开发者低成本部署百亿参数模型。

avatar AMD开发者中国社区