[AMD] docker: retire the rocm700 variants, add a rocm10-mi35x variant - #2854
Merged
guapisolo merged 8 commits intoSep 9, 2026
Merged
Conversation
… the ROCm 10 base Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
XinyuJiangCMU
requested review from
guapisolo,
yueming-yuan and
yushengsu-thu
as code owners
September 1, 2026 04:33
The wheel-SDK base ships the devel tree's runtime libs as byte-identical copies of the core/libraries ones (ROCm/TheRock#7818), and devel/lib puts unversioned *.so aliases on LD_LIBRARY_PATH, so an unversioned dlopen (torch: libamd_smi.so; triton/proton: libamdhip64.so, libhsa-runtime64.so, libhipblaslt.so) pulls a second copy of an already-loaded runtime into the process. Two comgr copies abort at startup with LLVM's "Option registered more than once"; the base image already collapses the libamd_smi pair, where the second copy silently breaks device enumeration (sgl-project/sglang#36434). Collapse every byte-identical top-level devel duplicate back onto its original.
The ROCm 7.2 base ships Apex, but the ROCm 10 wheel-SDK base does not, and Megatron's gradient_accumulation_fusion path imports apex's fused_weight_gradient_mlp_cuda, so the Megatron CI suites fail on import without it. Install a prebuilt cp312/gfx950 wheel from the variant's wheels release (ROCm/apex@40608ba, built on this base with GPU access and APEX_BUILD_{CPP,CUDA}_OPS=1), mirroring how the CUDA images consume Apex from a wheel shelf. Gated by APEX_USE_PREBUILT so the ROCm 7.2 variant is unchanged.
Megatron's async_strategy=nvrx hard-imports nvidia-resiliency-ext, which only publishes manylinux_2_39 wheels: the ROCm 7.2 base (Ubuntu 22.04, glibc 2.35) cannot install it, so it stays excluded there, but the ROCm 10 base (Ubuntu 24.04, glibc 2.39) installs 0.6.0 cleanly and the real NVRx async checkpoint path runs end-to-end on MI355X. Gated by NVRX_INSTALL so the ROCm 7.2 variant is unchanged.
fzyzcjy/torch_memory_saver#109 (HIP VMM allocation-granularity rounding) is still in review; without it every Megatron colocate CI test crashes at init on ROCm 10. Pin the fork commit under review and switch back to upstream once radixark#109 merges.
torch_memory_saver PR 109 (HIP VMM allocation-granularity rounding) merged upstream on 2026-09-04, so drop the temporary fork pin and point at the merge commit 06caa534822c5b980b61733a5d79ac731f9a5f9c. The merge carries no changes of its own beyond the merged branch, so this is byte-identical to the fork commit the rocm10 images were built and tested against. The pin is still 48 commits ahead of the f05a8754 that docker/Dockerfile uses for the CUDA images (Intel XPU backend, CPU-backup RSS reclaim, the paused-memory double-release fix, retained CPU backups); the bump stays scoped to the ROCm images.
…60901 # Conflicts: # docker/Dockerfile.rocm
guapisolo
approved these changes
Sep 9, 2026
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Upstream SGLang ROCm 7.0 bases are no longer built, while ROCm 10 MI35X bases now build daily with Python 3.12 and the wheel-based ROCm SDK. ROCm 10 also no longer reproduces the VMM-pause or IPC-export memory issues seen on ROCm 7.2, so the old ROCr and legacy-IPC workarounds are not needed.
Modifications
rocm700-*withrocm10-mi35xand pin each ROCm variant to its wheels release.rocm10-gfx950-v0.5.18cp312 wheels: TE 2.17.0, flash-attn 2.8.3, router, and gateway.numpy<2.torch_memory_saverwith--no-build-isolationfor the ROCm 10 wheel SDK.tile_kernelspatch path dynamically for py3.10/py3.12.rocm-sdk inithardlinks can become duplicate shared libraries across Docker layers ROCm/TheRock#7818).nvidia-resiliency-exton the ROCm 10 variant (glibc 2.39) for NVRx async checkpointing.Dependencies
hipMemCreatecrash that breaks all Megatron colocate CI on ROCm 10) merged upstream on 2026-09-04. The ROCm pin here is that merge commit,06caa534, which is byte-identical to the branch this PR was validated against. No blockers remain.torch_memory_saver48 commits forward fromf05a8754for the ROCm images only;docker/Dockerfilekeepsf05a8754for CUDA.