Skip to content

[docker] migrate ROCm installation to use prebuilt deb packages - #86

Open
lauri9 wants to merge 4 commits into
mainfrom
rocm-deb-migration
Open

lauri9 wants to merge 4 commits into
mainfrom
rocm-deb-migration

Conversation

@lauri9

@lauri9 lauri9 commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Install ROCm in the CI image from a pinned nightly deb snapshot, and keep optional a lean runtime option.

Background

docker/Dockerfile.ci built ROCm from TheRock source, including ccache, git mirrors, and a full artifact install into /opt/rocm. That compile dominated CI image builds, and the image always carried the developer tree.

This PR installs ROCm from a pinned TheRock nightly deb snapshot on the legacy multi-arch native index. ROCM_RELEASE_ID and ROCM_DEB_SERIES select the snapshot. ROCM_GFX_TARGETS selects the architecture shards and is independent of PYTORCH_ROCM_ARCH. The nightly version in this PR is chosen so that it closely matches the current version installed via THEROCK_COMMIT on main. We will move to the current multi-arch native index later with the next ROCm update.

The default stages keep developer tools. ROCM_PARENT=rocm_base builds the same stages without LLVM static archives outside lib/clang and without the developer-tools packages. This manual path enables leaner "runtime" style images. The nightly image no longer has /opt/rocm/share/therock, so CI records installed amdrocm-* package versions instead.

Goals

Replace the in-image ROCm source build with a pinned nightly deb install, and make the developer and lean runtime layouts explicit.

  • Install the ROCm libraries PyTorch and the AITER JIT need from pinned nightly debs.
  • Keep the default image on the developer layout, and provide a lean runtime parent.
  • Record the installed deb identity in CI, and document how to move the snapshot pin.

Tasks

  • Replace the TheRock ROCm source build with docker/install_rocm_debs.sh, and drop the artifact manifests.
  • Add self-retiring shims for amd-smi bindings, the classic llvm/bin/ld.lld path, and the origami CMake package.
  • Split rocm_tree, rocm_base, and rocm_devel, with ROCM_PARENT selecting the parent of core, deps, and final.
  • Record installed ROCm deb versions from the working image instead of copying TheRock metadata.
  • Document the snapshot pins, package targets, and shim retirement in docker/README.md and CONTRIBUTING.md.
  • Treat ROCM_RELEASE_ID and ROCM_DEB_SERIES as release-note pins, and stop diffing rocm-libraries and rocm-systems commit overrides.

Test plan

  • Test that lean image builds and works in E2E (local test)
  • Test that CI builds - build time on cold cache ~40 minutes CI run
  • Test that CI build layer caching continues to work CI run
  • Test E2E performance on supported workloads and architectures on full CI build path

Performance tests

Compared this PR to v26.10 release candidate. Performance looks good on MI350X.

model ROCm-debs (this PR) v26.10 RC Ratio
flux.usp 0.77 0.77 100.2%
flux.single_gpu 2.23 2.22 99.2%
flux.usp_2k 2.19 2.20 100.5%
flux.single_gpu_2k 11.30 11.33 100.3%
flux2.default 4.35 4.39 100.8%
flux2.quantgemm.gfx950 2.39 2.40 100.4%
flux2.quantgemm_quantattn.gfx950 2.17 2.19 100.8%
flux2.ti2i_2k 12.83 12.83 100.0%
flux2.t2i_1k 2.86 2.87 100.4%
flux2.t2i_2k 10.92 10.96 100.4%
flux2.t-multi-i2i_1k 5.85 5.84 99.9%
flux2.t-multi-i2i_2k 15.67 15.56 99.3%
flux2_klein.default 0.49 0.49 99.9%
flux_kontext.default 1.48 1.48 99.9%
hunyuanvideo.default.gfx950 88.91 89.15 100.3%
hunyuanvideo_1_5.default.gfx950 114.45 114.97 100.5%
ltx2_5.default 6.07 6.17 101.6%
minimax_h3.default 28.92 29.04 100.4%
qwen_image.default 4.83 4.83 100.1%
qwen_image_edit.default 14.82 14.84 100.2%
stablediffusion_3_5.default 0.87 0.87 100.0%
wan2_1.default.gfx950 73.53 73.87 100.5%
wan2_2.default.gfx950 75.11 75.01 99.9%
wan2_2.quantgemm.gfx950 65.48 65.60 100.2%
wan2_2.quantgemm_quantattn.gfx950 36.30 36.33 100.1%
wan2_2.quantgemm_sparseattn.gfx950 28.85 28.82 99.9%
z_image.default 2.59 2.59 100.1%

Other

  • ROCM_LIBRARIES_COMMIT_OVERRIDE and ROCM_SYSTEMS_COMMIT_OVERRIDE are removed. A change that is not in the pinned nightly has to be published in a snapshot before this image can use it.
  • Image size gets a nice reduction, amdsiloai/pytorch-xdit-staging:0539bf1-temp clocks in at 17.2GB as reported by docker vs. 21.5GB for v26.10

@lauri9
lauri9 force-pushed the rocm-deb-migration branch 2 times, most recently from 1b59e14 to 8e738cd Compare October 5, 2026 13:09
lauri9 and others added 4 commits October 6, 2026 02:21
Replace the TheRock source build with runtime, JIT, and devel deb layers so CI no longer compiles ROCm. PyTorch, Triton, and the TheRock checkout pin stay on main; TheRock is cloned only for the PyTorch build scripts. Self-retiring shims cover nightly packaging gaps for amdsmi, ld.lld, and origami.

Co-authored-by: Cursor <cursoragent@cursor.com>
Release notes now treat ROCM_RELEASE_ID and ROCM_DEB_SERIES as snapshot pins rather than TheRock commit overrides.

Co-authored-by: Cursor <cursoragent@cursor.com>
The nightly image has no /opt/rocm/share/therock tree, so both benchmark workflows record installed package versions from the image. The ITT workflow passes its runner list into the architecture matrix.

Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the amdrocm-core metapackage with the libraries PyTorch and AITER JIT need. rocm_devel keeps LLVM static archives and developer tools; ROCM_PARENT=rocm_base builds the product stages without them.

Co-authored-by: Cursor <cursoragent@cursor.com>
@lauri9
lauri9 force-pushed the rocm-deb-migration branch from 8e738cd to a025e21 Compare October 6, 2026 07:21
@lauri9
lauri9 marked this pull request as ready for review October 6, 2026 09:40
@lauri9
lauri9 requested review from a team as code owners October 6, 2026 09:40
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.

1 participant