feat: Add AMD ROCm GPU support alongside CUDA#1
Open
alexhegit wants to merge 1 commit into
Open
Conversation
Add explicit ROCm 7.2 support via uv extras while preserving the existing CUDA workflow (uv sync picks up PyPI CUDA wheels on Linux). Changes: - pyproject.toml: Add rocm extra with ROCm 7.2 index routing, conflict resolution (rocm + end2end are mutually exclusive), remove torch-backend=auto in favor of explicit [tool.uv.sources] index routing, relax torch upper bound to allow ROCm 2.12+ wheels - AGENTS.md: Add ROCm support documentation with what works / what doesn't Verified on AMD Radeon PRO W7900 (48GB VRAM) + ROCm 7.2: - torch==2.12.0+rocm7.2 detects GPU correctly - ptv3_vanilla backbone forward pass runs on GPU - 93/96 tests pass (3 skipped = integration tests needing live server) - pointnet2_ops gracefully degrades (warning only) - All demo scripts, ZMQ server/client, gripper wizard work What works on ROCm: - Base inference (point clouds, meshes, scenes) - All demo scripts (demo_object_pc, demo_scene_pc, demo_object_mesh) - ZMQ server/client (uv sync --extra rocm --extra serve) - Gripper config wizard (browser GUI) - Tests (uv run pytest -m "not end2end") What does NOT work on ROCm: - End-to-end pipeline (nvidia-curobo, newton, warp are CUDA-only) - Docker training (NGC base image is NVIDIA-only) - pointnet2_ops PointNet++ backbone (gracefully degrades, use ptv3_vanilla) Usage: uv sync # CUDA (unchanged) uv sync --extra rocm # AMD ROCm Signed-off-by: AMD <amd@amd.com> Signed-off-by: AlexHe99 <alehe@amd.com>
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
2 participants
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.
Summary
Add explicit AMD ROCm 7.2 GPU support via uv extras while preserving the existing CUDA workflow.
Changes
rocmextra with ROCm 7.2 index routing, conflict resolution (rocm+end2endare mutually exclusive), removetorch-backend=autoin favor of explicit[tool.uv.sources]index routing, relax torch upper bound to allow ROCm 2.12+ wheelsVerified on AMD Radeon PRO W7900 (48GB VRAM) + ROCm 7.2
What Works on ROCm
uv sync --extra rocmuv run python scripts/demo_*.pyuv sync --extra rocm --extra serveuv run python scripts/gripper_config_wizard.pyuv run pytest -m "not end2end"What Does NOT Work on ROCm
nvidia-curobo,newton,warpare CUDA-onlyptv3_vanillabackboneUsage
Conflicts
The
rocmandend2endextras are mutually exclusive — uv enforces this automatically.