Prhub

#2768 [AMD] fix(docker): drop the obsolete ROCm Megatron fused-kernels patch

原始 PR 作者 michaelzhang-ai 合并时间 2026-08-28 06:12 文件变更 4 提交数 1 评论 3 代码增减 +3 / -40

infra docker amd megatron bugfix

执行摘要

删除过时的 ROCm Megatron patch，修复镜像构建

PR 描述指出，ROCm 镜像构建在 Dockerfile.rocm 第 17 步失败，因为 radixark/Megatron-LM@miles-main 已包含上游 NVIDIA 提交 3aa73966 (Guard non-core imports, NVIDIA/Megatron-LM#3993)，该提交删除了 megatron/legacy/fused_kernels 包和 _compile_dependencies() 中的 fused_kernels.load(args) 调用，导致 git apply 报错 'No such file or directory'。此问题导致 8 月 25 日起 nightly ROCm 镜像发布全部失败。由于 Dockerfile.rocm 未提供 MEGATRON_COMMIT 参数，无法通过调用工作流固定补丁版本，因此需在此处修复。

该 PR 值得精读，尤其是了解如何处理上游依赖变更导致的构建失败。关键设计决策是删除死代码而不试图修复补丁，这基于对当前 tree 的检查。对于维护者，应关注后续是否有替代方案处理 Megatron 内部行为变化。

讨论亮点

该 PR 的讨论较少，主要是作者请求评审和两位维护者的 LGTM 评论。自动化审查因是 fork PR 而未执行。没有实质性的审查意见或争议。

实现拆解

变更分为四部分：

  1. 删除补丁文件: docker/amd_patch/latest/megatron.patch 被整体删除。该补丁原本为 Megatron 的 legacy fused-kernels 加载器添加 ROCm 守卫（if not torch.version.cuda: return），但现在 miles-main 已不再加载这些内核，属于死代码。

  2. 更新 Dockerfile: docker/Dockerfile.rocm 中删除了 COPY docker/amd_patch/latest/megatron.patch /tmp/amd_patch/megatron.patch 和 git apply /tmp/amd_patch/megatron.patch 两个步骤，构建阶段少了一步。

  3. 简化 CI 工作流: .github/workflows/_run-ci-rocm.yml 中，此前在 MEGATRON_PR 覆盖时应用的 git apply --check / git apply / git apply --reverse --check / 错误处理的整段逻辑被删除，现在无需打补丁。

  4. 调整测试: tests/ci/test/test_run_suite.py 中将原来的 test_megatron_override_preserves_rocm_patch 改为 test_megatron_override_installs_the_checked_out_ref_unpatched，断言不再检查 amd_patch 相关步骤，仅验证 git checkout -f FETCH_HEAD 在 pip install -e . 之前执行。

配套说明：此次未改动 docker/image_inputs.py，因此构建缓存跳过标签不受影响；未执行的 docker push 步骤不在变更范围内。

文件 模块 状态 重要度
docker/amd_patch/latest/megatron.patch Megatron removed 5.93
tests/ci/test/test_run_suite.py 测试套件 modified 5.38
.github/workflows/_run-ci-rocm.yml CI 工作流 modified 3.73
docker/Dockerfile.rocm Docker 构建 modified 2.78

关键符号

load test_megatron_override_installs_the_checked_out_ref_unpatched

关键源码片段

tests/ci/test/test_run_suite.py test-coverage

测试逻辑需与 CI 工作流变更保持一致，确保覆盖新的无补丁流程。

# tests/ci/test/test_run_suite.py (head 版本摘录 )
class TestRunSuite:
    def test_megatron_override_installs_the_checked_out_ref_unpatched(self):
        # 读取 _run-ci-rocm.yml 中 MEGATRON_PR 覆盖逻辑
        reusable = (Path(__file__).resolve().parents[3] / ".github" / "workflows" / "_run-ci-rocm.yml").read_text()
        override = reusable.split('if [ -n "$MEGATRON_PR" ]; then', 1)[1].split("          cd $GITHUB_WORKSPACE", 1)[0]
​
        # 验证 checkout 发生在 install 之前
        checkout = override.index("git checkout -f FETCH_HEAD")
        install = override.index("pip install -e . --no-deps --break-system-packages")
​
        assert checkout < install
        # 不再期望 amd_patch 相关步骤
        assert "amd_patch" not in override
.github/workflows/_run-ci-rocm.yml infrastructure

CI 工作流中 Megatron PR 覆盖逻辑删除了补丁验证和应用，直接安装未修补的 ref。

# .github/workflows/_run-ci-rocm.yml (head 版本节选 )
- name: Install Megatron
  run: |
    if [ -n "$MEGATRON_PR" ]; then
      git fetch origin "$MEGATRON_PR"
      git checkout -f FETCH_HEAD
      git log --oneline -1
      # 已移除的补丁应用步骤：
      # if git apply --check ...; then git apply ...; fi
      pip install -e . --no-deps --break-system-packages
    fi

评论区精华

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

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

风险与影响

风险较低，但需注意：

  • megatron 兼容性: 假设当前 miles-main 不再依赖 legacy fused kernels，但若未来 Megatron 恢复或引入新的 legacy 依赖，该补丁删除可能导致 ROCm 下缺失守卫。已通过验证当前 tree 中无 megatron.legacy 引用 (除 fp16_deprecated shim) 来缓解。

  • CI 验证缺失: 由于未执行 docker push 步骤，未验证镜像实际发布环节，可能遗漏推送阶段的新问题。不过构建已成功，发布失败风险较小。

  • 重复补丁: docker/amd_patch/sglv0.5.10/megatron.patch 是同一补丁的拷贝，虽未被引用但可能造成维护上的混淆。作者已说明保留以保持变更范围最小，未来可考虑删除。

该变更影响 ROCm 镜像构建流程，对使用 ROCm 镜像的用户有直接影响，修复了构建失败问题，使 nightly 发布恢复。对系统构建时间略有减少（少一个阶段）。对团队的维护成本降低，消除了死代码和失效的补丁。影响范围限于 ROCm 镜像构建相关流程，不涉及核心训练逻辑。

构建流程变更 上游依赖变化 发布环节未验证

关联 Issue

#3993 Guard non-core imports

完整报告

Markdown PDF

参与讨论