feat: add ROCm support for AMD training via cuda/rocm extras - #25
Open
wlgys8 wants to merge 1 commit into
Open
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
The ROCm extra currently lacks OS/arch markers (risking incorrect resolution on unsupported platforms) and the new make install docs need an explicit non-make fallback or prerequisite note for Windows.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds an ROCm install profile alongside the existing CUDA flow so AMD GPUs can install the correct PyTorch wheels, while keeping a single UV workspace + lockfile. It also extends runtime GPU telemetry to support AMD SMI and updates docs/CI to use the new make install* workflow.
Changes:
- Introduce mutually exclusive
cuda/rocmextras and route torch wheels to vendor-specific PyTorch indexes viatool.uv. - Add a Makefile-based install workflow (
make install,make install-dev) and update CI/docs to use activated-venvpythoninvocations instead ofuv run. - Add AMD SMI GPU utilization/VRAM sampling (with tests) and refresh documentation + architecture diagrams to reflect dual-vendor GPU support.
File summaries
| File | Description |
|---|---|
| README.zh-CN.md | Update install/run commands to make install + activated venv; mention ROCm/CUDA. |
| README.md | Update install/run commands to make install + activated venv; mention ROCm/CUDA. |
| pyproject.toml | Add cuda/rocm extras, uv conflict rules, and ROCm index/source routing. |
| motrix_rl/tests/test_system_metrics.py | Add unit tests covering AMD SMI GPU samplers. |
| motrix_rl/src/motrix_rl/system_metrics.py | Add AMD SMI + sysfs fallback GPU samplers alongside NVML. |
| motrix_rl/README.md | Update training instructions to activated-venv python + make install BACKEND=.... |
| motrix_rl/pyproject.toml | Broaden torch requirement to allow profile-selected versions via root extras. |
| motrix_deploy/README.md | Switch setup and commands to make install-dev + direct CLI usage. |
| motrix_deploy_unitree/README.md | Switch setup and commands to make install-dev + direct CLI usage. |
| motrix_deploy_tasks/README.md | Switch command examples away from uv run. |
| Makefile | Add install/install-dev targets with GPU-profile autodetection. |
| docs/source/zh_CN/user_guide/tutorial/training_environment_config.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/tutorial/training_and_result.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/tutorial/runs_and_checkpoints.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/tutorial/motrix_deploy.md | Switch setup and commands to make install-dev + direct CLI usage. |
| docs/source/zh_CN/user_guide/tutorial/export_onnx.md | Update install/export instructions to the new extras/profile workflow. |
| docs/source/zh_CN/user_guide/tutorial/custom_training_backend.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/tutorial/basic_frame.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/robots.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/getting_started/installation.md | Make make install/make install-dev the primary install path. |
| docs/source/zh_CN/user_guide/getting_started/hello_motrixlab.md | Update tutorial commands to activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/whole_body_tracking/motion_format.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/whole_body_tracking/adding_wbt_task.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/quadruped_velocity_tracking/adding_robot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/quadruped_velocity_tracking.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/quadruped_locomotion/go1_rough_terrain.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/shadow_hand_repose.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/rm65_open_cabinet.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/rm65_insert_peg.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/franka_open_cabinet.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/franka_lift_cube.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/manipulation/bounce_ball.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/index.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/humanoid_velocity_tracking/adding_robot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/humanoid_velocity_tracking.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_walker.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_reacher.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_quadruped.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_point_mass.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_lqr.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_humanoid.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_hopper.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_finger.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/dm_cheetah.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/dm_control/bring_ball.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/basic/stewart.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/basic/pendulum.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/basic/cartpole.md | Replace uv run examples with activated-venv python usage. |
| docs/source/zh_CN/user_guide/envs/basic/acrobot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/tutorial/training_environment_config.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/tutorial/training_and_result.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/tutorial/runs_and_checkpoints.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/tutorial/motrix_deploy.md | Switch setup and commands to make install-dev + direct CLI usage. |
| docs/source/en/user_guide/tutorial/export_onnx.md | Update install/export instructions to the new extras/profile workflow. |
| docs/source/en/user_guide/tutorial/custom_training_backend.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/tutorial/basic_frame.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/robots.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/getting_started/installation.md | Make make install/make install-dev the primary install path. |
| docs/source/en/user_guide/getting_started/hello_motrixlab.md | Update tutorial commands to activated-venv python usage. |
| docs/source/en/user_guide/envs/whole_body_tracking/motion_format.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/whole_body_tracking/index.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/whole_body_tracking/adding_wbt_task.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/quadruped_velocity_tracking/adding_robot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/quadruped_velocity_tracking.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/quadruped_locomotion/go1_rough_terrain.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/shadow_hand_repose.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/rm65_open_cabinet.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/rm65_insert_peg.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/franka_open_cabinet.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/franka_lift_cube.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/manipulation/bounce_ball.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/index.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/humanoid_velocity_tracking/adding_robot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/humanoid_velocity_tracking.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_walker.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_reacher.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_quadruped.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_point_mass.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_lqr.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_humanoid.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_hopper.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_finger.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/dm_cheetah.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/dm_control/bring_ball.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/basic/stewart.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/basic/pendulum.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/basic/cartpole.md | Replace uv run examples with activated-venv python usage. |
| docs/source/en/user_guide/envs/basic/acrobot.md | Replace uv run examples with activated-venv python usage. |
| docs/source/_static/images/architecture-light.svg | Refresh architecture diagram to include NVIDIA/AMD GPU backends. |
| docs/source/_static/images/architecture-dark.svg | Refresh architecture diagram to include NVIDIA/AMD GPU backends. |
| CONTRIBUTING.md | Update contributor workflow to make install* + activated-venv commands. |
| configs/task/microduck-walk-flat/motrix.fastsac.yaml | Adjust FastSAC config (batch size + strict UTD mode). |
| AGENTS.md | Update repo agent/dev environment guidance to the new install flow. |
| .github/workflows/tests.yml | Update CI to install via make install-dev and run pytest via venv python. |
| .github/PULL_REQUEST_TEMPLATE.md | Update validation checklist command to python -m pytest -q. |
Review details
Suppressed comments (1)
docs/source/zh_CN/user_guide/getting_started/installation.md:80
- 这里将
make install作为默认安装命令，但未说明需要 GNU Make；且 Windows 原生环境通常没有make。建议在此处补充无make的uv sync兜底命令或明确前置要求。
- Files reviewed: 94/97 changed files
- Comments generated: 6
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
wlgys8
force-pushed
the
feat/rocm-fastsac-amd
branch
from
September 9, 2026 15:28
1effc55 to
5f32f4c
Compare
wlgys8
force-pushed
the
feat/rocm-fastsac-amd
branch
from
September 9, 2026 15:41
5f32f4c to
62399c5
Compare
There was a problem hiding this comment.
🟡 Changes recommended
The new install.sh has a few verified user-facing correctness/documentation issues (misleading --all help text, inaccurate WSL mention, and missing --gpu value validation) that should be fixed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
install.sh:22
- The
--allhelp text says it installs “all packages, groups and extras”, but in the ROCm branch it installs a curated subset (e.g., it does not include theskrl-jaxextra). Consider clarifying the help text to match what the script actually does.
- Files reviewed: 96/99 changed files
- Comments generated: 2
- Review effort level: Lite
Comment on lines
+99
to
+103
| if [ -z "$GPU" ]; then | ||
| if lsmod 2>/dev/null | grep -q '^nvidia '; then | ||
| GPU=cuda | ||
| elif lsmod 2>/dev/null | grep -q '^amdgpu '; then | ||
| GPU=rocm |
Comment on lines
+8
to
+10
| # GPU wheel flavor (--gpu): cuda or rocm. When not given explicitly, detect | ||
| # from the loaded kernel driver (nvidia -> cuda, amdgpu -> rocm); CPU-only | ||
| # and WSL2 hosts fall back to cuda. |
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
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.
Closes #24
Summary
cuda/rocmextras in the rootpyproject.toml:cudapins torch 2.7.0 from the pytorch-cu128 index,rocmpins torch 2.11.0 + triton-rocm + amdsmi from the new pytorch-rocm72 index; declared conflicting under[tool.uv]so uv keeps both branches forked in a singleuv.lockMakefile:make installauto-detects the GPU vendor from the loaded kernel module (nvidia-> cuda,amdgpu-> rocm; CPU-only and WSL2 fall back to cuda), overridable withPROFILE=cuda|rocm make install;BACKEND=skrl-torch|skrl-jax|rslrlselects the training backend extra (defaultskrl-torch);make install-devbootstraps the full dev environment (rocm hosts get an explicit extras list since--all-extrascannot combine both profiles)motrix_rl.system_metricsGPU samplers now talk to AMD SMI on ROCm hosts and NVML on NVIDIA hosts, so the FastSAC console panel reports VRAM/utilization on both vendors; unit tests added for the AMD pathtests.yml) installs throughmake install-devinstead of rawuv syncmicroduck-walk-flatfastsac config:batch_size: 2048and strictutd_modeNotes
make installon hosts with neither driver loaded falls back to the cuda profile (CPU torch wheels from the cu128 index)