Prhub

#2249 Remove --train-memory-margin-bytes

原始 PR 作者 zhuzilin 合并时间 2026-08-03 19:54 文件变更 4 提交数 1 评论 0 代码增减 +0 / -17

执行摘要

移除 --train-memory-margin-bytes 参数及相关内存 margin 配置

PR body 为空,但从变更内容推断,维护者认为 --train-memory-margin-bytes 这个自定义内存 margin 已不再必要,可能被 torch_memory_saver 的默认行为接管,或该参数在实践中未发挥预期作用。移除后参数体系更简洁,也减少了 debug_rollout_only 分支中特殊处理 margin 的额外逻辑。需要说明的是,由于缺少 body 和 issue,动机无法完全确认。

该 PR 变更简单直接,不建议作为重点精读对象;但它展示了如何清理已废弃或不再适用的训练参数,以及参数移除时连带处理校验、actor 初始化和示例脚本的完整配套流程。关注点在于:确认 torch_memory_saver 默认 margin 是否满足 Slime 的 offload_train 需求,以及是否已在文档中移除该参数说明。

讨论亮点

该 PR 没有 review 评论和讨论线程,作者直接合并,未记录设计权衡或备选方案。

实现拆解

  1. 在 slime/utils/arguments.py 的 add_train_arguments 中删除 --train-memory-margin-bytes 参数注册(原默认 1GB),该参数此前用于后续设置 torch_memory_saver.memory_margin_bytes。
  2. 在 slime/utils/arguments.py 的 slime_validate_args 中删除 debug_rollout_only 分支里对 train_memory_margin_bytes 的告警与强制清零逻辑,因为参数已不存在,该兜底判断也随之失效。
  3. 在 slime/backends/megatron_utils/actor.py 的 init 中删除 offload_train 分支下读取 train_memory_margin_bytes 并写入 torch_memory_saver.memory_margin_bytes 的 5 行代码,内存 margin 改由 torch_memory_saver 库自身默认值控制。
  4. 同步清理配套:tests/test_megatron_argument_validation.py 的 make_slime_validate_args 中移除 train_memory_margin_bytes=0 字段;scripts/run-gpt-oss-20B.sh 中删除 --train-memory-margin-bytes 268435456 参数,避免用户在脚本中继续传入已被删除的参数。
文件 模块 状态 重要度
slime/utils/arguments.py 参数解析 modified 5.08
slime/backends/megatron_utils/actor.py 训练执行 modified 4.8
tests/test_megatron_argument_validation.py 参数校验 modified 2.88
scripts/run-gpt-oss-20B.sh 示例脚本 modified 1.94

关键符号

add_train_arguments slime_validate_args init

关键源码片段

slime/utils/arguments.py configuration

核心变更文件:删除 --train-memory-margin-bytes 的参数注册(9 行)以及 slime_validate_args 中 debug_rollout_only 分支的强制置零逻辑(3 行),是本次移除的源头。

def add_train_arguments(parser):
    # --train-backend 会由 _pre_parse_mode() 提前解析,稍后合并进来。
    parser.add_argument(
        "--qwen-gdn-backend",
        type=str,
        choices=["fla", "flashqla"],
        default="fla",
        help="GDN implementation backend for Qwen linear-attention layers.",
    )
    parser.add_argument(
        "--train-env-vars",
        type=json.loads,
        default="{}",
        help="Extra environment variables for training process, e.g. PyTorch memory management ones.",
    )
    # 已移除 --train-memory-margin-bytes:此前用于在 offload_train 模式下给
    # torch_memory_saver.memory_margin_bytes 追加约 1GB margin,现改由库默认行为接管。
    parser.add_argument(
        "--megatron-to-hf-mode",
        choices=["raw", "bridge"],
        default="raw",
        help="The method to convert megatron weights to hugging face weights for SGLang.",
    )
    # Delta weight sync.
    parser.add_argument(
        "--update-weight-mode",
        choices=["full", "delta"],
        default="full",
        help=(
            "Weight sync strategy. 'full' (default) broadcasts every parameter "
            "every sync. 'delta' diffs each sync against a pinned-CPU snapshot of the "
            "previous one and ships only the changed bytes (disk transport only)."
        ),
    )
    parser.add_argument(
        "--update-weight-transport",
        choices=["nccl", "disk"],
        default="nccl",
        help=(
            "Carrier for weight sync. In full mode, 'nccl' broadcasts chunks and "
            "'disk' writes a complete HF checkpoint under --update-weight-disk-dir "
            "before engines reload it. Delta mode is 'disk' only: each host applies the "
            "published deltas into its local checkpoint and reloads via update_weights_from_disk."
        ),
    )

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 兼容性风险:--train-memory-margin-bytes 属于用户可见命令行参数,删除后所有仍传入该参数的脚本或启动命令会因未知参数立即报错,属于潜在 breaking change。
  2. 内存行为变化:删除后 torch_memory_saver.memory_margin_bytes 不再由 Slime 显式设置为默认 1GB,而是使用 torch_memory_saver 库的默认值;若库默认值小于 1GB,在 offload_train 场景下可能增加 OOM 概率;若大于 1GB,可能增加不必要的 CPU↔GPU 交换。
  3. 文档同步风险:本次变更未更新任何 docs 文件,文档中若仍提及 --train-memory-margin-bytes(如 usage.md 等),会导致用户读到已不存在的参数。
  4. 行为回归风险:scripts/run-gpt-oss-20B.sh 中原注释提到该参数是为了“减少 colocate 模式下 CPU↔GPU 过度交换”,移除后该脚本的内存表现可能与注释描述不一致。

影响范围集中在训练参数解析与 actor 初始化路径:使用 slime_validate_args 校验的所有训练启动都会少一个可选参数;依赖 torch_memory_saver 的 offload_train / colocate 场景不再由 Slime 显式配置内存 margin;examples 和 scripts 中若沿用旧参数会启动失败。对普通用户和系统整体影响较小,但需要确保文档与启动脚本同步更新。

breaking change:命令行参数移除 内存行为依赖库默认值 文档可能未同步

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论