Add ROCm support for AMD Instinct GPUs by nilapate · Pull Request #9153 · Project-MONAI/MONAI

nilapate · 2026-10-05T11:14:56Z

Description

Minimal ROCm enablement for AMD Instinct GPUs. Most of MONAI is already portable; this fixes the few places that assume a single GPU toolkit, plus adds a container recipe.

setup.py — CUDA_HOME is unset on a ROCm build, where the toolkit is located via ROCM_HOME, so BUILD_CUDA evaluated False and the C++ extensions were silently skipped. Accept either. CUDAExtension handles the source translation, so no kernel source changes are needed.
monai/_extensions/loader.py — the JIT build cache key included torch.version.cuda, unset on ROCm, so every ROCm toolkit version collided on one cache entry. Fall back to torch.version.hip.
tests/networks/nets/test_densenet.py — test_pretrain_consistency compares two separately-constructed module graphs holding identical weights for bit-exactness. Convolution algorithm selection is free to differ between the graphs, so this was never a guaranteed property; skipped on ROCm for now.
Dockerfile.rocm — the default Dockerfile uses a vendor-specific base image, so ROCm gets its own recipe. Installs the ROCm SDK and a matching PyTorch, filters packages that don't apply on this platform (cucim-cu*, nvidia-ml-py, nni), and installs hipCIM to provide the cucim module the whole-slide-image paths need.

3 lines of library code changed; everything else is additive. Existing code paths are unchanged — both edits short-circuit on the current value, so CUDA and CPU-only behaviour is identical.

Validated on MI300X (gfx942); MI355X (gfx950) validated separately. tests/networks/layers/test_gmm.py passes (5 passed), which exercises the loader.py change through a real JIT extension build.

Types of changes

Non-breaking change (fix or new feature that would not break existing functionality).
New feature (non-breaking change which adds functionality).
Breaking change (fix or new feature that would cause existing functionality to change).
New tests added to cover the changes.
Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
In-line docstrings updated.
Documentation updated, tested make html command in the docs/ folder.

MONAI runs on ROCm unmodified for the most part, since a ROCm build of PyTorch presents itself as `cuda`. Three places assume a CUDA toolkit specifically, and the Docker image has no ROCm equivalent. setup.py: `CUDA_HOME` is None on a ROCm torch, where the toolkit is found via `ROCM_HOME` instead, so BUILD_CUDA evaluated False and the C++/HIP extensions were silently skipped. Accept either. `CUDAExtension` hipifies the .cu sources transparently, so no source changes are needed. monai/_extensions/loader.py: the JIT build cache key included `torch.version.cuda`, which is None on ROCm, so every ROCm toolkit version collided on one cache entry. Fall back to `torch.version.hip`. tests/networks/nets/test_densenet.py: test_pretrain_consistency compares two separately-constructed module graphs holding identical weights for bit-exactness. The backend may select different convolution algorithms per graph, so this was never guaranteed; skip it on ROCm. Dockerfile.rocm: the default Dockerfile builds on the NVIDIA PyTorch container, so ROCm gets its own recipe. It installs the ROCm SDK and a matching PyTorch from AMD's public index, filters the CUDA-only extras (cucim-cu*, nvidia-ml-py, nni), and installs hipCIM, which provides the `cucim` module that MONAI's whole-slide-image paths need on ROCm. Assisted-by: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Patel, Nilaykumar K <NilaykumarKantibhai.Patel@amd.com>

coderabbitai · 2026-10-05T11:15:21Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: Project-MONAI/MONAI/.coderabbit.yaml
Review profile: CHILL
Plan: Advanced
Run ID: 7dc3b1bd-975f-4eed-9744-2d70ed9fbcaf

📥 Commits

Reviewing files that changed from the base of the PR and between 4e89184 and 1360ca8.

📒 Files selected for processing (4)

Dockerfile.rocm
monai/_extensions/loader.py
setup.py
tests/networks/nets/test_densenet.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

Walkthrough

Adds a configurable ROCm container recipe for MONAI. Extension detection now accepts PyTorch’s ROCM_HOME, and extension cache keys use the HIP version when CUDA is unset. The DenseNet pretrain consistency test skips on ROCm.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 1360c

No confirmed issue currently prevents merging. The container build should still receive its normal validation.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 3 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding ROCm support for AMD Instinct GPUs.
Description check	✅ Passed	The description follows the template, identifies issue `#9152`, explains the changes and validation, and marks the applicable change types. It also states that the listed integration and quick test comm…
Linked Issues check	✅ Passed	For `#9152`, `setup.py` recognizes `ROCM_HOME`, so ROCm extensions are not silently skipped. `Dockerfile.rocm` adds the requested container recipe. The HIP cache-key fallback supports ROCm JIT builds. T…
Out of Scope Changes check	✅ Passed	All reported changes support `#9152`’s ROCm enablement: toolkit detection, extension caching, ROCm test behavior, and a ROCm container recipe. No unrelated changes are reported.

Full details: Docstring Coverage

Explanation

Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 3 files. (1 skipped: 1 unsupported.)

Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create a new PR

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

ericspod · 2026-10-05T12:38:58Z

+        "ROCm may select different conv algorithms per graph; bit-exactness not guaranteed.",
+    )
    def test_pretrain_consistency(self, model, input_param, input_shape):
+        if torch.version.hip is not None:


Is this needed with the decorator?

ericspod · 2026-10-05T13:07:29Z

+ENV ROCM_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_core"
+ENV ROCM_HOME="${ROCM_PATH}"
+ENV ROCM_DEVEL_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel"
+ENV ROCM_LIBRARIES_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_libraries"


Suggested change

ENV ROCM_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_core"

ENV ROCM_HOME="${ROCM_PATH}"

ENV ROCM_DEVEL_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel"

ENV ROCM_LIBRARIES_PATH="/opt/venv/lib/python3.12/site-packages/_rocm_sdk_libraries"

ARG PYTHON_VER=3.12

ENV ROCM_PATH="/opt/venv/lib/python${PYTHON_VER}/site-packages/_rocm_sdk_core"

ENV ROCM_HOME="${ROCM_PATH}"

ENV ROCM_DEVEL_PATH="/opt/venv/lib/python${PYTHON_VER}/site-packages/_rocm_sdk_devel"

ENV ROCM_LIBRARIES_PATH="/opt/venv/lib/python${PYTHON_VER}/site-packages/_rocm_sdk_libraries"

I'd suggest making the Python version an ARG here, if BASE_IMAGE is set to some other image then the Python that comes with it may be different from 3.12 so users would need an ARG to change it as well.

ericspod · 2026-10-05T13:09:17Z

+
+# Use print_dependencies.py rather than -e .[all,testing] to filter CUDA-only packages:
+#   cucim-cu* pulls in cuda-toolkit (~1.2 GB); nvidia-ml-py fails at import on ROCm; nni depends on it.
+# BUILD_MONAI is not set: docker build has no GPU, so extensions JIT-compile at container runtime.


BUILD_MONAI does need to be set for MONAI to actually use the compiled code at runtime. This is buried in the code unfortunately and is something we should update to make clearer or change the behaviour. If the argument causes a problem in this RUN it should be set afterward

ericspod

hi @nilapate thanks for this, we're happy to support ROCm in MONAI but we do have the issue of having no means to test it. The Dockerfile for example isn't something we can test without AMD hardware so we would have to rely on users to report any issues. I made a few comments about minor things but the changes look good to me as they are, if you can address things we should be good to merge once tests get through.

nilapate requested review from KumoLiu, Nic-Ma and ericspod as code owners October 5, 2026 11:15

ericspod reviewed Oct 5, 2026

View reviewed changes

ericspod requested changes Oct 5, 2026

View reviewed changes

Navigation Menu

Add ROCm support for AMD Instinct GPUs - #9153

Add ROCm support for AMD Instinct GPUs#9153
nilapate wants to merge 1 commit into
Project-MONAI:devfrom
AMD-Ecosystem:9152-rocm-support

nilapate commented Oct 5, 2026

Uh oh!

coderabbitai Bot commented Oct 5, 2026 •

edited

Loading

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

ericspod Oct 5, 2026

Uh oh!

ericspod Oct 5, 2026

Uh oh!

ericspod Oct 5, 2026

Uh oh!

ericspod left a comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

nilapate commented Oct 5, 2026

Description

Types of changes

Uh oh!

coderabbitai Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

ericspod Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ericspod Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ericspod Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ericspod left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

coderabbitai Bot commented Oct 5, 2026 •

edited

Loading