Prhub

#2296 fix(train): skip optimizer and scheduler for eval-only

原始 PR 作者 EazyReal 合并时间 2026-08-21 10:50 文件变更 2 提交数 3 评论 0 代码增减 +10 / -6

执行摘要

eval-only 跳过优化器与调度器构建

PR body 明确指出 eval-only 是预期功能:train.py 已有特殊分支在 --num-rollout 0 时调用 rollout_manager.eval,train loop 是 range(..., num_rollout) 所以实际零训练步。但 bring-up 阶段 setup_model_and_optimizer 总是构造 Megatron 的 optimizer 与 OptimizerParamScheduler,默认 --lr-decay-iters 下 slime 计算 train_iters=0 → lr_decay_steps=0,触发 Megatron 断言 lr_decay_steps > 0,导致永远走不到 eval 分支。作者还解释了为何不用 #2109 的 train_iters=1 方案:那是对调度器撒谎,且 lr_warmup_steps >= lr_decay_steps 时仍会崩溃。

该 PR 值得精读,尤其是“为什么不用 train_iters=1”的设计权衡。建议在合入后补一条 CPU 侧测试:断言 num_rollout == 0 时 setup_model_and_optimizer 返回 (model, None, None)、create_training_models 不分配 critic,并验证正常训练路径仍构造 optimizer 与 scheduler。

讨论亮点

仓库中未发现针对本 PR 的 review 评论。最有价值的讨论集中在 PR body 对 #2109 方案的反驳:

  • 作者指出 train_iters=1 是“对调度器撒谎”,会让 train_iters 失去“实际训练步数”的语义,且 lr_warmup_steps >= lr_decay_steps 时仍然崩溃。
  • 结论:eval-only 不应构造训练栈,直接跳过 optimizer 与 scheduler 才是正确方向。
  • 该决策最终被合入,说明团队认可“非训练模式不建训练栈”的设计。

实现拆解

变更入口是 slime/backends/megatron_utils/model.py 的 setup_model_and_optimizer 和 slime/ray/placement_group.py 的 create_training_models。

  1. setup_model_and_optimizer 在 get_model 之后新增 early return:当 args.num_rollout == 0 时,设置 args.no_load_optim = True 并返回 (model, None, None)。这一步绕过了 OptimizerConfig 组装、Adam 补丁、get_megatron_optimizer 和 get_optimizer_param_scheduler 等全部训练栈构造,同时通过 no_load_optim 让 checkpoint 加载跳过 optimizer 状态,与 ref/teacher 使用 optimizer=None 的既有路径保持一致。
  2. 同步更新 setup_model_and_optimizer 和 initialize_model_and_optimizer 的返回类型标注,将 optimizer 与 scheduler 类型改为 MegatronOptimizer | None、OptimizerParamScheduler | None,明确新契约。
  3. create_training_models 将 critic 创建条件从 args.use_critic 收紧为 args.use_critic and args.num_rollout != 0,避免 eval-only 时白白分配 critic 训练组;随后 start_rollout_ids 的来源判断从 args.use_critic 改为 critic_model is not None,避免 use_critic 与 num_rollout 组合时引用未定义的 critic_start_rollout_ids。
  4. 测试配套:PR 最终未包含自动化测试。作者在第二笔 commit 中曾加入针对 train.py 控制流、critic 分配和 setup_model_and_optimizer 的 CPU 测试,但在第三笔 commit 中将其删除,理由是那只是本地验证脚本而非 slime 应保留的测试。
文件 模块 状态 重要度
slime/backends/megatron_utils/model.py 后端训练 modified 6.18
slime/ray/placement_group.py 训练编排 modified 5.07

关键符号

setup_model_and_optimizer initialize_model_and_optimizer create_training_models

关键源码片段

slime/backends/megatron_utils/model.py data-contract

核心修复点:setup_model_and_optimizer 在 num_rollout == 0 时跳过 optimizer 与 scheduler 构造,避免 Megatron 因 lr_decay_steps == 0 断言失败;返回类型同步变为 Optional。

def setup_model_and_optimizer(
    args: Namespace, role: str = "actor"
) -> tuple[list[DDP], MegatronOptimizer | None, OptimizerParamScheduler | None]:
    # eval-only 模式(--num-rollout 0)不训练,不应构造优化器与调度器。
    # 若照常构造,Megatron 会因 lr_decay_steps == 0 触发 assert 失败。
    assert not args.moe_use_upcycling
    assert args.load is not None or args.pretrained_checkpoint is not None
​
    model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder)
​
    if args.num_rollout == 0:
        # 标记不加载 optimizer 状态,与 ref / teacher 的 checkpoint 加载路径保持一致。
        args.no_load_optim = True
        # optimizer 与 scheduler 均为 None,调用方按“无优化器”路径处理。
        return model, None, None
​
    # 正常训练路径:组装 OptimizerConfig,再构造 optimizer 与 scheduler。
    kwargs = {}
    for field in dataclasses.fields(OptimizerConfig):
        if hasattr(args, field.name):
            kwargs[field.name] = getattr(args, field.name)
    config = OptimizerConfig(**kwargs)
    config.timers = None
​
    # stateless Adam 需要替换 Adam 类并禁用分布式优化器状态初始化。
    optimizer_context = (
        _patch_megatron_adam(StatelessAdam) if args.use_stateless_adam else nullcontext()
    )
    with optimizer_context:
        optimizer = get_megatron_optimizer(
            config=config,
            model_chunks=model,
            use_gloo_process_groups=args.enable_gloo_process_groups,
        )
    if args.use_stateless_adam:
        _disable_distributed_optimizer_state_initialization(optimizer)
    opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer)
    return model, optimizer, opt_param_scheduler
slime/ray/placement_group.py core-logic

create_training_models 在 num_rollout == 0 时跳过 critic 分配,并将 start_rollout_ids 的判断依据从 args.use_critic 改为 critic_model is not None,避免 eval-only 组合时引用未定义变量。

def create_training_models(args, pgs, rollout_manager, actor_cls=None):
    actor_model, actor_start_rollout_ids = create_actor_model(args, pgs, rollout_manager, actor_cls=actor_cls)
​
    critic_model = None
    # eval-only 模式不训练,无需分配 critic 训练组;仅当实际训练时创建。
    if args.use_critic and args.num_rollout != 0:
        from slime.utils.arguments import parse_megatron_role_args
​
        critic_args = (
            parse_megatron_role_args(args, args.megatron_config_path, role="critic")
            if args.megatron_config_path is not None
            else copy.deepcopy(args)
        )
        if args.megatron_config_path is None:
            critic_args.disable_param_buffers_cpu_backup = False
​
        critic_model = allocate_train_group(
            args=critic_args,
            num_nodes=args.critic_num_nodes,
            num_gpus_per_node=args.critic_num_gpus_per_node,
            pg=pgs["critic"],
            role="critic",
        )
        critic_start_rollout_ids = critic_model.create(rollout_manager=rollout_manager)
​
    # 以 critic_model 是否实际创建为判断依据,避免 use_critic 与 num_rollout 组合时的误判。
    if critic_model is not None:
        start_rollout_ids = critic_start_rollout_ids
    else:
        start_rollout_ids = actor_start_rollout_ids
​
    assert len(set(start_rollout_ids)) == 1
​
    if args.start_rollout_id is None:
        args.start_rollout_id = start_rollout_ids[0]
​
    if args.rollout_global_dataset:
        ray.get(rollout_manager.load.remote(args.start_rollout_id - 1))
​
    return actor_model, critic_model

评论区精华

为何不用 train_iters=1 绕过 scheduler 构造 设计

PR body 指出 #2109 通过设置 train_iters=1 让 scheduler 构造通过,但这是对调度器撒谎:train_iters 不再表示真实训练步数,且 lr_warmup_steps >= lr_decay_steps 时依然崩溃。

结论:不采用 train_iters=1,改为 eval-only 直接不构造优化器与调度器。 · 已解决

风险与影响

  • 返回契约变更:setup_model_and_optimizer 与 initialize_model_and_optimizer 的返回值从固定三元/四元组变为可含 None。虽然 ref/teacher 加载路径已有 optimizer=None 的先例,但本 PR 未逐一排查其余调用方,若某处直接使用返回的 optimizer 属性(如 optimizer.param_groups)会触发 AttributeError。
  • 缺少测试覆盖:最终 PR 无自动化测试,--num-rollout 0 的 eval-only 启动与正常训练回归都依赖人工验证,存在回归风险。
  • critic 跳过逻辑:若配置 --use-critic 且 --num-rollout 0,create_training_models 不再创建 critic,start_rollout_ids 将退回 actor 的 id;若上层逻辑仍按 use_critic 期待 critic 的 rollout id,可能产生不一致。
  • 影响是正面的:eval-only 不再创建 optimizer、scheduler 和 critic,显存与启动开销下降;正常训练 num_rollout > 0 分支未改动,风险可控。
  • 对用户:修复了 --num-rollout 0 --eval-interval 1 纯评测场景的启动崩溃,用户无需再通过 hack train_iters 绕过断言。
  • 对系统:eval-only 启动路径不再构造训练栈,内存占用更低,启动更快。
  • 对团队:明确了 train_iters 的语义只属于真实训练;同时暴露了创建训练模型与评估模式组合时缺少自动化回归的问题。
缺少测试覆盖 返回 Optional 契约变更 eval-only 启动路径

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论