# PR #6717 完整报告

- 仓库：`verl-project/verl`
- 标题：[worker] feat: add tinker training worker primitives
- 合并时间：2026-06-15 14:54
- 原文链接：http://prhub.com.cn/verl-project/verl/pull/6717

---

# 执行摘要

- 一句话：添加 Tinker 分步训练 Worker 原语
- 推荐动作：值得精读，特别是其设计模式：通过类属性控制 worker 子类选择、通过上下文管理器参数控制梯度清零，展示了一种低侵入的向后兼容扩展方式。

# 功能与动机

当前 TrainingWorker.train_batch() 是一个原子操作，无法在调用前单独控制梯度清零和优化器步进。此变更为需要更精细训练调度的场景（如梯度累积、阶段性数据加载）提供可选的分步原语，同时保持向后兼容。

# 实现拆解

1. 在 `verl/workers/engine/base.py` 的 `BaseEngineCtx.__init__` 中增加 `zero_grad_on_exit` 参数（默认 True），供子类按需控制退出时是否调用 `optimizer_zero_grad`。
2. 在 automodel、fsdp、megatron、torchtitan、veomni 五个后端的 `TrainModeCtx.__exit__` 中，将无条件 `optimizer_zero_grad()` 改为条件判断：`if self.zero_grad_on_exit or exc_type is not None`。这样当 `zero_grad_on_exit=False` 且正常退出时，梯度会被保留。
3. 新建 `verl/workers/engine_workers_tinker.py`，定义 `TinkerTrainingWorker`（继承 `TrainingWorker`）和 `TinkerActorRolloutRefWorker`（继承 `ActorRolloutRefWorker`）。前者暴露三个拆分的 RPC 方法，后者通过 `actor_worker_cls = TinkerTrainingWorker` 路由到 actor 的拆分原语。
4. 修改 `verl/workers/engine_workers.py`：为 `ActorRolloutRefWorker` 添加 `actor_worker_cls` 和 `ref_worker_cls` 类属性（默认为 `TrainingWorker`），并在 `init_model` 中使用它们创建 worker，使子类能简单替换。
5. 在 `tests/models/test_engine.py` 中添加契约测试和 FSDP 后端梯度累积端到端测试，验证三个拆分原语的存在和正确性。

关键文件：
- `verl/workers/engine_workers_tinker.py`（模块 工作器；类别 source；类型 dependency-wiring；符号 TinkerTrainingWorker, optimizer_zero_grad, forward_backward, optimizer_step）: 新增文件，包含 TinkerTrainingWorker 和 TinkerActorRolloutRefWorker 实现，是本 PR 的核心。
- `tests/models/test_engine.py`（模块 引擎测试；类别 test；类型 test-coverage；符号 test_tinker_workers_expose_split_training_primitives, _local_tensor, _snapshot_trainable_params, _param_delta_norm）: 添加了契约测试和 FSDP 梯度累积端到端测试，验证新原语行为。
- `verl/workers/engine_workers.py`（模块 工作器；类别 source；类型 core-logic；符号 ActorRolloutRefWorker, ActorRolloutRefWorker.actor_worker_cls, ActorRolloutRefWorker.ref_worker_cls, ActorRolloutRefWorker.__init__）: 修改 ActorRolloutRefWorker，增加类属性并改用类属性实例化 worker，使子类化扩展成为可能。
- `verl/workers/engine/base.py`（模块 引擎；类别 source；类型 core-logic；符号 BaseEngineCtx.__init__）: BaseEngineCtx 添加 zero_grad_on_exit 参数，作为所有 TrainModeCtx 条件清零的基础。
- `verl/workers/engine/automodel/transformer_impl.py`（模块 引擎；类别 source；类型 core-logic；符号 AutomodelTrainModeCtx.__exit__）: 修改 AutomodelTrainModeCtx.__exit__，条件执行 optimizer_zero_grad，支持保留梯度。
- `verl/workers/engine/fsdp/transformer_impl.py`（模块 引擎；类别 source；类型 core-logic；符号 FSDPTrainModeCtx.__exit__）: 类似 automodel；修改 FSDP 后端的 TrainModeCtx。
- `verl/workers/engine/megatron/transformer_impl.py`（模块 引擎；类别 source；类型 core-logic；符号 MegatronTrainModeCtx.__exit__）: 类似 automodel；修改 Megatron 后端的 TrainModeCtx。
- `verl/workers/engine/torchtitan/transformer_impl.py`（模块 引擎；类别 source；类型 core-logic；符号 TorchTitanTrainModeCtx.__exit__）: 类似 automodel；修改 TorchTitan 后端的 TrainModeCtx。
- `verl/workers/engine/veomni/transformer_impl.py`（模块 引擎；类别 source；类型 core-logic；符号 VeOmniTrainModeCtx.__exit__）: 类似 automodel；修改 VeOmni 后端的 TrainModeCtx。

关键符号：TinkerTrainingWorker.optimizer_zero_grad, TinkerTrainingWorker.forward_backward, TinkerTrainingWorker.optimizer_step, TinkerActorRolloutRefWorker.optimizer_zero_grad, TinkerActorRolloutRefWorker.forward_backward, TinkerActorRolloutRefWorker.optimizer_step, ActorRolloutRefWorker.actor_worker_cls, ActorRolloutRefWorker.ref_worker_cls, BaseEngineCtx.__init__, AutomodelTrainModeCtx.__exit__, FSDPTrainModeCtx.__exit__, MegatronTrainModeCtx.__exit__, TorchTitanTrainModeCtx.__exit__, VeOmniTrainModeCtx.__exit__, test_tinker_workers_expose_split_training_primitives, _snapshot_trainable_params, _param_delta_norm, _grad_norm, _has_any_grad, _make_split_step_batch, _split_training_primitives_fsdp_worker

## 关键源码片段

### `tests/models/test_engine.py`

添加了契约测试和 FSDP 梯度累积端到端测试，验证新原语行为。

```python
def test_tinker_workers_expose_split_training_primitives():
    # 契约测试：验证 Tinker 子类继承正确，且 split 方法不在父类 __dict__ 中
    assert issubclass(TinkerTrainingWorker, TrainingWorker)
    assert issubclass(TinkerActorRolloutRefWorker, ActorRolloutRefWorker)

    for name in ("optimizer_zero_grad", "forward_backward", "optimizer_step"):
        assert name not in TrainingWorker.__dict__
        assert name in TinkerTrainingWorker.__dict__
        assert name not in ActorRolloutRefWorker.__dict__
        assert name in TinkerActorRolloutRefWorker.__dict__


def _snapshot_trainable_params(module):
    # 获取当前可训练参数快照（用于后续计算 delta norm）
    return [_local_tensor(param).detach().float().clone() for param in module.parameters() if param.requires_grad]


def _param_delta_norm(module, before) -> torch.Tensor:
    # 计算参数更新量的全局 L2 范数（跨卡 all_reduce）
    device = torch.device("cuda", torch.cuda.current_device())
    total = torch.zeros((), device=device)
    for param, param_before in zip((p for p in module.parameters() if p.requires_grad), before, strict=True):
        param_local = _local_tensor(param).detach().float()
        total += (param_local - param_before.to(param_local.device)).pow(2).sum()
    dist.all_reduce(total, op=dist.ReduceOp.SUM)
    return total.sqrt()

```

# 评论区精华

PR 无 review 评论，自动审查无实质反馈，由 wuxibin89 直接 approve。未发现设计争议。

- 暂无高价值评论线程

# 风险与影响

- 风险：修改涉及 automodel、fsdp、megatron、torchtitan、veomni 共五个 engine 后端的 TrainModeCtx，若 `zero_grad_on_exit` 默认值（True）在边缘场景（如异常退出）行为差异导致梯度未清零，可能影响后续训练状态。新 worker 类的 RPC 注册与原 train_batch 并行，需确保 ray 调度无冲突。测试覆盖了 FSDP 后端，但其他后端的集成测试缺失。
- 影响：用户层面：提供可选的分步训练 API，无需修改现有代码即可使用 split 原语。系统层面：新增代码量较小，但多个 engine 后端均需同步修改。团队层面：为后续实现累积梯度、交错数据等训练实验提供基础。
- 风险标记：核心路径变更（5 引擎后端）, 多后端集成测试缺失

# 关联脉络

- 暂无明显关联 PR