执行摘要

• 一句话：添加 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 梯度累积端到端测试，验证新原语行为。

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