# PR #7272 完整报告

- 仓库：`verl-project/verl`
- 标题：[fsdp,veomni] feat: support pad_to_length to reduce jit compile time
- 合并时间：2026-08-06 09:01
- 原文链接：http://prhub.com.cn/verl-project/verl/pull/7272

---

# 执行摘要

- 一句话：FSDP/VeOmni 新增 pad_to_length 减少 JIT 重复编译
- 推荐动作：值得精读的 perf 型 PR。重点关注三处设计：一是“按 attention workload 均衡分 micro-batch 导致 token 数越界，因此必须用细 bucket 而非固定预算对齐”的量化论证；二是 `prepare_model_inputs` 中补位时机（roll 之后、SP 切分之前）与 `_gather_and_unpad_packed` 统一裁剪的配合；三是 router replay 在补位下的掩码处理——补位 token 无 RECORD 记录所以必须 gate 掉走原生路由，而真实 token 打 gate 才会破坏 R2 bit-equal 保证。测试文件对这两个引擎用 `MagicMock` 桩掉 `veomni.*` 依赖的做法也值得复用。

# 功能与动机

PR body 明确指出痛点：With `use_remove_padding=True` + `use_dynamic_bsz=True`, every packed micro-batch has a different token count, so shape-specialized kernels (torch.compile, Triton/DeepGEMM autotune) recompile or re-autotune continuously throughout training。并且解释了为什么必须用 bucket 对齐而非固定目标长度：`rearrange_micro_batches` 按 token 预算推导微批数量，却按 attention 负载（`Σ 24576·seqlen + seqlen²`）做 Karmarkar-Karp 均衡分配样本，导致微批 token 数既非恒定也不受预算约束，直接对齐预算会留下大量“完全动态”的超预算微批。

# 实现拆解

实现按 5 步展开，全部围绕「让 packed 微批形状稳定」这一目标：

1. **新增张量级补位工具 `pad_packed_inputs`**（`verl/workers/engine/utils.py`）：对 `(1, total_nnz)` 的 packed 序列右补 `pad_size` 个 token。关键约定是追加的 `position_ids` 从 0 重新开始，使补位 token 在 `cu_seqlens` 中自成一个尾部 varlen 段，而非把最后一条真实序列拉长；同时兼容 mRoPE 的 `(rope_dim, 1, total_nnz)` 三维布局，并支持自定义 `pad_value`（temperature 补位必须用 1 而不是 0，因为它会除进 logits）。

2. **FSDP 引擎接入 **（`verl/workers/engine/fsdp/transformer_impl.py`）：`FSDPEngine.__init__` 从配置读取 `pad_to_length` / `pad_to_length_bucket`；新增 `_get_packed_pad_size` 用 `ceildiv` 上取整到 bucket 倍数；`prepare_model_inputs` 在 `torch.roll` 之后、SP 切分之前计算 `static_pad_size` 并对 `input_ids_rmpad`、`position_ids_rmpad`、`input_ids_rmpad_rolled`、`temperature_rmpad` 统一补位，随后把 `output_args["pad_size"]` 从“仅 SP 对齐”扩展为“静态 + SP 对齐”的总和；`prepare_model_outputs` 中原来散落的 `if self.use_ulysses_sp: gather_outputs_and_unpad(...)` 分支统一收敛到新增的 `_gather_and_unpad_packed`，保证 log_probs、entropy、sum_pi_squared、distillation 辅助量全部按总 pad 裁剪。

3. **VeOmni 引擎与 router replay 适配 **（`verl/workers/engine/veomni/transformer_impl.py`）：`VeOmniEngineWithLMHead` 同样读取 pad 配置（继承 `_get_packed_pad_size`）；`_maybe_push_router_replay_state` 在 REPLAY 分支先记录 `real_nnz = flat.size(0)`，再按 `output_args["pad_size"]` 把路由目标 `flat` 扩展 pad_size 行、把掩码尾部置 False 后重新走 `slice_microbatch_replay_targets` / `slice_microbatch_replay_mask`。这里的正确性论证是：补位 token 没有 RECORD 记录，必须被 gate 掉走原生路由，而对真实 token 打 gate 才会破坏 R2 的 bit-equal 保证。

4. **配置与入口联动**：`verl/workers/config/engine.py` 给 `FSDPEngineConfig` 和 `VeOmniEngineConfig` 各加 `pad_to_length: bool = False`、`pad_to_length_bucket: int = 1024` 两个字段及详细 docstring；`verl/workers/config/actor.py` 为 `FSDPActorConfig`、`VeOmniActorConfig` 增加转发字段；同步更新 `fsdp.yaml`、`veomni.yaml`、`veomni_actor.yaml`、`dp_actor.yaml`、`dp_ref.yaml`、`veomni_ref.yaml` 和两个 `_generated_*.yaml`；`run_deepseek_v4_veomni.sh`、`run_deepseek_v4_flash_megatron.sh` 两个示例脚本按新配置路径修正参数。

5. **测试配套**：新增 `tests/workers/test_packed_pad_to_length_on_cpu.py`（186 行），用 `MagicMock` 桩掉 `veomni.*` 模块，覆盖 pad_size=0 无操作、补位自成一 varlen 段、mRoPE 逐 rope_dim 补位、pad 值与 dtype 保持、bucket 取整性质、超预算微批仍落 bucket、以及 top-K 蒸馏守卫；`test_router_replay_engine_helpers_on_cpu.py` 新增 R2/R3 在 pad_to_length 下的 2 个测试；`test_distillation_topk_symmetry_on_cpu.py` 增加 2 行守卫断言。

关键文件：
- `verl/workers/engine/fsdp/transformer_impl.py`（模块 FSDP 引擎；类别 source；类型 core-logic；符号 _get_packed_pad_size, _gather_and_unpad_packed, prepare_model_inputs, prepare_model_outputs）: 核心实现所在：新增 _get_packed_pad_size / _gather_and_unpad_packed，prepare_model_inputs 完成静态补位，prepare_model_outputs 统一按总 pad 裁剪，并加 top-K 蒸馏守卫。
- `verl/workers/engine/utils.py`（模块 引擎工具；类别 source；类型 core-logic；符号 pad_packed_inputs）: 新增 pad_packed_inputs 核心工具函数，定义补位 token 的 position_ids 约定（自成 varlen 段）与 mRoPE 兼容逻辑。
- `verl/workers/engine/veomni/transformer_impl.py`（模块 VeOmni 引擎；类别 source；类型 core-logic；符号 _maybe_push_router_replay_state, VeOmniEngineWithLMHead）: VeOmni 引擎接入同一配置，并完成 router replay（R2/R3）在补位下的目标扩展与掩码 gate，是本 PR 后半部分正确性核心。
- `tests/workers/test_packed_pad_to_length_on_cpu.py`（模块 单元测试；类别 test；类型 test-coverage；符号 _make_engine, TestPadPackedInputs, test_pad_size_zero_is_a_noop, test_pad_tokens_form_their_own_varlen_segment）: 新增 186 行 CPU 单测，覆盖补位工具、bucket 取整性质、超预算微批、蒸馏守卫与解补裁剪，并用 MagicMock 桩掉 veomni 依赖。
- `verl/workers/config/engine.py`（模块 引擎配置；类别 source；类型 configuration；符号 pad_to_length, pad_to_length_bucket）: FSDPEngineConfig 与 VeOmniEngineConfig 新增 pad_to_length / pad_to_length_bucket 字段与详尽文档，是整个特性的配置契约入口。
- `tests/workers/test_router_replay_engine_helpers_on_cpu.py`（模块 单元测试；类别 test；类型 test-coverage；符号 test_r2_pad_to_length_extends_targets_and_gates_the_pad_tail, test_r3_pad_to_length_gates_both_prompt_and_pad）: 新增 R2/R3 在 pad_to_length 下目标扩展与补位尾部 gate 的测试，锁住 router replay 与补位的交互契约。

关键符号：pad_packed_inputs, _get_packed_pad_size, _gather_and_unpad_packed, prepare_model_inputs, prepare_model_outputs, _maybe_push_router_replay_state

## 关键源码片段

### `verl/workers/engine/utils.py`

新增 pad_packed_inputs 核心工具函数，定义补位 token 的 position_ids 约定（自成 varlen 段）与 mRoPE 兼容逻辑。

```python
# verl/workers/engine/utils.py
def pad_packed_inputs(
    input_ids_rmpad: torch.Tensor,
    position_ids_rmpad: torch.Tensor | None,
    pad_size: int,
    pad_value: float = 0,
):
    """右补 packed 序列 (1, total_nnz) 尾部 pad_size 个 token。

    与 verl.utils.ulysses.ulysses_pad 补到序列并行倍数的方式保持一致：
    追加的 position_ids 从 0 重新开始，因此补位 token 在 cu_seqlens 里
    自成一个尾部 varlen 段，而不是把最后一条真实序列“拉长”。
    """
    if pad_size <= 0:
        return input_ids_rmpad, position_ids_rmpad

    input_ids_rmpad = torch.nn.functional.pad(input_ids_rmpad, (0, pad_size), value=pad_value)
    if position_ids_rmpad is not None:
        # 单段 position_ids 形状为 (1, pad_size)；mRoPE 下为 (rope_dim, 1, pad_size)
        pad_position_ids = torch.arange(
            pad_size, dtype=position_ids_rmpad.dtype, device=position_ids_rmpad.device
        ).unsqueeze(0)
        if position_ids_rmpad.dim() == 3:  # (rope_dim, 1, total_nnz) mRoPE 布局
            pad_position_ids = pad_position_ids.unsqueeze(0).repeat(position_ids_rmpad.size(0), 1, 1)
        position_ids_rmpad = torch.cat((position_ids_rmpad, pad_position_ids), dim=-1)
    return input_ids_rmpad, position_ids_rmpad

```

### `verl/workers/engine/veomni/transformer_impl.py`

VeOmni 引擎接入同一配置，并完成 router replay（R2/R3）在补位下的目标扩展与掩码 gate，是本 PR 后半部分正确性核心。

```python
# verl/workers/engine/veomni/transformer_impl.py
# _maybe_push_router_replay_state 的 REPLAY 分支（节选）
# flat 是 [mb_nnz, L, topk] 的压平路由目标；pad_to_length 会往 input_ids
# 尾部追加补位 token，这些位置在 RECORD 阶段没有记录路由，因此要把目标
# 扩展到补位后的长度，并给尾部打 False gate 让它们走原生路由。
# 对真实 token 打 gate 才会破坏 R2 的 bit-equal 前向保证，而补位 token
# 的输出在进入 loss 前就被丢弃，gate 掉是安全的。
real_nnz = flat.size(0)

# R3：response 位置才有真实记录，prompt 位置是占位零值，需先构建逐 token
# 的 mask_flat（response_mask 相关校验与构造逻辑此处省略）
mask_flat = None
if self._router_replay_mode == "R3":
    ...

pad_size = int(output_args.get("pad_size", 0))
if pad_size:
    if mask_flat is None:
        # R2 默认全 True：真实 token 全部回放，保证 bit-equal 前向
        mask_flat = torch.ones(real_nnz, dtype=torch.bool, device=flat.device)
    mask_flat = torch.cat([mask_flat, mask_flat.new_zeros(pad_size)])
    flat = torch.cat([flat, flat.new_zeros((pad_size, *flat.shape[1:]))])

per_layer = rr.slice_microbatch_replay_targets(flat)
# 与 routed_experts 使用同一套 pad + slice 规则切分 mask
replay_mask = rr.slice_microbatch_replay_mask(mask_flat) if mask_flat is not None else None
rr.set_microbatch_targets(per_layer, replay_mask=replay_mask)

```

# 评论区精华

该 PR 无任何 review 评论，reviewer `Luosuu` 直接 APPROVED。真正的设计讨论沉淀在 PR body：作者用一组量化数据对比三种补位方案——“对齐预算、超预算保持动态”会带来 +24.8% token 开销且形状无界；“上取整到预算倍数”形状收敛但开销高达 +49.5%；而“按 1024 bucket 对齐”（本 PR 方案）仅 +1.7% 开销且形状收敛到 `{6144, 10240, 12288}`。另外 4 个 commit（`[fsdp,veomni] feat → simplify → fix ci → pad to bucket`）显示 bucket 对齐是最后一轮才定型的方向，前两轮是更简单的固定目标补位，后来才针对 `rearrange_micro_batches` 的 workload 均衡特性改成细粒度 bucket。

- bucket 对齐 vs 固定目标长度对齐的设计权衡 (design): 采用细粒度 bucket 方案；提交历史中 'pad to bucket' 是最后一个 commit，说明该方案在迭代末期定型。无 reviewer 异议，Luosuu 直接 approve。
- router replay 在补位下如何保证正确性 (correctness): 实现中对 real token 保持原掩码语义，仅对补位尾部追加 False gate，并断言 flat 扩展前后长度一致；新增 R2/R3 两个测试锁住该行为。

# 风险与影响

- 风险：具体风险点如下：
 - **核心训练路径变更**：`prepare_model_inputs` / `prepare_model_outputs` 是 FSDP/VeOmni 共用热路径，`output_args["pad_size"]` 语义从“仅 SP 对齐”扩展为“SP + 静态”总和，任何仍只读 SP 部分的消费方都会漏裁导致序列错位。代码通过收敛到 `_gather_and_unpad_packed` 单点裁剪来降低此风险。
 - **补位 token 的隐形开销**：补位 token 仍走完整前向（含 EP all-to-all），bucket 太细则形状集合大、收益下降，太粗则浪费算力；默认 1024 是经验值，文档建议按 `max_token_len_per_gpu * sp` 配置。
 - **补位 token 的语义边界**：`input_ids` 补位值默认 0 可能与真实 token id 0 相同，但因补位 token 自成 varlen 段且输出在进入 loss 前被丢弃，不会污染训练。
 - **router replay 交互**：R2 此前从不带 mask，现在会带“真实 token 全 True + 补位尾部 False”的掩码，pad 尾部走原生路由会略微改变 EP 负载分布；控制器 `set_microbatch_targets` 的掩码语义已同步放宽并加测试锁住。
 - **功能边界**：top-K distillation 与 `pad_to_length` 不兼容（teacher 张量只按 Ulysses 规则切分、不知晓静态补位），已用 RuntimeError 显式拒绝，避免静默错位。
 - 影响：对用户：开启 `pad_to_length=true` 后，动态 batch 训练的 packed 长度收敛到有限形状集合，torch.compile / Triton / DeepGEMM 的编译与 autotune 从“每步一次”变为“训练早期一次”，以约 +1.7% token 前向开销换取大幅减少的 JIT 卡顿；默认关闭，未开启用户无任何行为变化。对系统：改动覆盖 FSDP 与 VeOmni 两个引擎的输入输出管线及配置契约，新增 2 个配置键并同步到全部生成配置；VeOmni router replay 的 R2/R3 掩码逻辑是本 PR 后半部分新增的复杂度。对团队：后续若把该特性推广到 Torchtitan/Megatron 引擎，需要复制同样的补位与裁剪管线；`pad_to_length` 已成为 verl 侧的引擎级配置，与 VeOmni 上游 `train.pad_to_length` 语义对齐。
 - 风险标记：核心训练路径变更 , 补位 token 增加前向开销 , Ulysses SP 补位语义合并 , router replay 掩码交互

# 关联脉络

- PR #7242 [veomni] feat: add DeepSeek V4 support: 同样改动 verl/workers/engine/veomni/transformer_impl.py 与示例 run_deepseek_v4_veomni.sh；本 PR 是 verl 侧对 VeOmni 上游 train.pad_to_length 的引擎级移植，DeepSeek V4 VeOmni 示例是主要采纳场景。
- PR #7224 [vllm] feat: enhance DeepSeek V4 fp8/fp4 linear and moe weight refit: 同样触碰 run_deepseek_v4_flash_megatron.sh；两个 PR 都把 DeepSeek V4 示例脚本作为配置契约的联调面，说明该示例线是 pad_to_length 等新配置的主要落地路径。
- PR #7225 [algo] fix: micro-batch normalization for distillation loss: 蒸馏损失功能领域相关：本 PR 显式拒绝 pad_to_length 与 top-K distillation 组合（RuntimeError 守卫），与蒸馏相关的前向布局改动同属需保持张量流对齐的区域。