Skip to content

[AMD] docker: retire the rocm700 variants, add a rocm10-mi35x variant - #2854

Merged
guapisolo merged 8 commits into
radixark:mainfrom
XinyuJiangCMU:pr/rocm10-variant-20260901
Sep 9, 2026
Merged

guapisolo merged 8 commits into
radixark:mainfrom
XinyuJiangCMU:pr/rocm10-variant-20260901

Conversation

@XinyuJiangCMU

@XinyuJiangCMU XinyuJiangCMU commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

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

  • Replace rocm700-* with rocm10-mi35x and pin each ROCm variant to its wheels release.
  • Add rocm10-gfx950-v0.5.18 cp312 wheels: TE 2.17.0, flash-attn 2.8.3, router, and gateway.
  • Keep the base NumPy on py3.12 images instead of forcing numpy<2.
  • Build torch_memory_saver with --no-build-isolation for the ROCm 10 wheel SDK.
  • Resolve the tile_kernels patch path dynamically for py3.10/py3.12.
  • Remove dead ROCm 7.0 config and update build dispatch/docs.
  • Collapse duplicated wheel-SDK runtime libs that can be loaded twice in one process ([Issue]: [ROCm SDK] rocm-sdk init hardlinks can become duplicate shared libraries across Docker layers ROCm/TheRock#7818).
  • Install a prebuilt ROCm/apex wheel from the wheels shelf for Megatron's fused ops.
  • Install nvidia-resiliency-ext on the ROCm 10 variant (glibc 2.39) for NVRx async checkpointing.

Dependencies

… the ROCm 10 base

Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

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.
@guapisolo
guapisolo merged commit 8d9826e into radixark:main Sep 9, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants