Repository navigation
Add ROCm support for AMD Instinct GPUs - #9153
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (3)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughAdds a configurable ROCm container recipe for MONAI. Extension detection accepts PyTorch’s Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Severity of issue fixed: Medium Merge Risk: ⚪ Minimal · up to The ROCm support is mergeable after normal checks; the supplied evidence identifies no outstanding defect. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation 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.)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
ericspod
left a comment
ericspod
left a comment
There was a problem hiding this comment.
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.
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>
- Add ARG PYTHON_VERSION; the rocm-sdk site-packages paths were hardcoded to python3.12 and broke silently on a different BASE_IMAGE. - Build the extensions with BUILD_MONAI=1 FORCE_CUDA=1 (the build host has no GPU, so setup.py would otherwise skip them) and set ENV BUILD_MONAI=1, which deviceconfig needs at runtime to gate USE_COMPILED. - Drop the test_densenet skipTest that duplicated its @skipIf decorator. Assisted-by: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Patel, Nilaykumar K <NilaykumarKantibhai.Patel@amd.com>
1360ca8 to
bc09a1e
Compare
Thanks @ericspod , I have addressed the comment. For ongoing maintenance, I have to check, There is something called amd developer cloud(https://devcloud.amd.com/login) but not sure how to integrate here(if you mean that by having mean to test). I also am planning to have one more non-invasive PR for clean install of MONAI for rocm, will raise PR soon. |
|
Hi @nilapate we're good to merge now I think. I would want a means of automatically testing with AMD hardware if we could so I'm not sure how Developer Cloud fits, but we may want something for manual testing either way. |
2 participants
Fixes #9152.
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_HOMEis unset on a ROCm build, where the toolkit is located viaROCM_HOME, soBUILD_CUDAevaluatedFalseand the C++ extensions were silently skipped. Accept either.CUDAExtensionhandles the source translation, so no kernel source changes are needed.monai/_extensions/loader.py— the JIT build cache key includedtorch.version.cuda, unset on ROCm, so every ROCm toolkit version collided on one cache entry. Fall back totorch.version.hip.tests/networks/nets/test_densenet.py—test_pretrain_consistencycompares 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 defaultDockerfileuses 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 thecucimmodule 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.pypasses (5 passed), which exercises theloader.pychange through a real JIT extension build.Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.