Prhub

#2316 Remove megatron_patch for memory optimization

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

执行摘要

删除 megatron_patch 梯度合并补丁,简化 Megatron 适配层。

PR body 原话:'This is more or less like an optimization for megatron instead of slime, so remove for clarity of the code.' 作者希望把对上游 Megatron 内部实现的优化移出 Slime 代码库,降低维护面、减少 import 期副作用,并保持 Slime 与 Megatron 的边界清晰。无关联 Issue。

值得快速精读,重点在于理解 Slime 如何通过移除运行时 monkeypatch 来收敛对上游框架的侵入。合并前建议确认目标 Megatron 版本已具备等效的分块梯度同步机制,并在大模型 TP 场景做一次 OOM 回归验证;若无法确认,可先保留补丁并仅注释导入以避免直接行为回退。

讨论亮点

本 PR 没有任何 review 评论或讨论线程(comments_count=0,review_comments_count=0),由作者自行合并。可提炼的决策依据只有 PR body 中关于代码职责边界的判断:该补丁是 Megatron 的优化而非 Slime 的功能,因此从 Slime 仓库移除。

实现拆解

  1. 删除 slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py:该文件通过 try-import 在运行时替换 megatron.core.distributed.finalize_model_grads 中的 _allreduce_non_tensor_model_parallel_grads 及其旧别名 _allreduce_layernorm_grads,利用 _split_into_chunksSLIME_GRAD_COALESCE_CHUNK_BYTES(默认 1 GiB)分块执行 flatten + all_reduce,避免单个超大连续缓冲区在分配器碎片下 OOM。
  2. 删除 slime/backends/megatron_utils/megatron_patch/__init__.py:原文件只有一行 re-export 导入,删除后整个子包清空。
  3. 修改 slime/backends/megatron_utils/__init__.py:移除末尾 from . import megatron_patch # noqa: F401, E402,使加载 megatron_utils 不再自动触发 Megatron 运行时补丁;DeepEP Buffer.__init__ 的 TMS 区域包装逻辑保持不变。
  4. 调整 tests/test_deep_ep_tms_patch.py:删除 _load_megatron_utils_init 中向 sys.modules 预置假 megatron_patch 模块的 monkeypatch,因为被加载的 __init__.py 不再导入该子包;deep_ep 初始化与 TMS 区域恢复的断言本身不变。
文件 模块 状态 重要度
slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py 适配层 removed 8.71
slime/backends/megatron_utils/megatron_patch/__init__.py 适配层 removed 4.19
slime/backends/megatron_utils/__init__.py 适配层 modified 4.89
tests/test_deep_ep_tms_patch.py 测试配套 modified 3.11

关键符号

_grad_attr _fsdp_flag _split_into_chunks _allreduce_non_tensor_model_parallel_grads

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

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

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

风险与影响

  1. 回归风险:TP 梯度同步从分块 coalesce 回退到 Megatron 默认实现;如果目标 Megatron 版本尚未内建分块机制,大模型 + 大批次场景可能重现单个大连续缓冲区导致的 OOM。涉及文件即被删除的 megatron_chunked_grad_coalesce_patch.py
  2. 环境变量失效:SLIME_GRAD_COALESCE_CHUNK_BYTES 不再被读取,依赖该变量调整梯度同步内存峰值的作业需要重新评估。
  3. 版本兼容性收缩:原补丁显式兼容 core_v0.13.0 与 post-core_v0.15.0rc7 两套 Megatron API;删除后该兼容层消失,跨版本行为差异由上游决定。
  4. 测试缺口:没有新增针对去除补丁后 TP 梯度同步行为的回归测试,删除是否安全目前依赖作者对上游 Megatron 已知优化的判断。

对用户与系统:大 TP 并行、大模型训练脚本的显存峰值可能变化;不依赖该补丁的小规模模型基本无感知。对团队:megatron_utils 导入期副作用减少,启动日志中不再出现 'slime grad coalesce patch applied' 相关提示,也少了一个需要随 Megatron API 变化持续维护的兼容点。影响面集中在 Megatron 后端适配层,不涉及 rollout、agent、观察性等模块。

删除核心内存优化补丁 运行期 monkeypatch 移除 环境变量 SLIME_GRAD_COALESCE_CHUNK_BYTES 失效 缺少回归测试确认

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论