Skip to content

feat: add ROCm GPU support alongside CUDA#104

Open
alexhegit wants to merge 2 commits into
johnsutor:mainfrom
alexhegit:feat/rocm-support
Open

feat: add ROCm GPU support alongside CUDA#104
alexhegit wants to merge 2 commits into
johnsutor:mainfrom
alexhegit:feat/rocm-support

Conversation

@alexhegit

Copy link
Copy Markdown

Summary

Add ROCm 7.2 GPU support so AMD GPU users can develop and test without conflicting with the default CUDA environment.

Changes

  • scripts/setup-rocm.sh: Creates .venv-rocm with ROCm 7.2 PyTorch using uv pip install --torch-backend rocm7.2 --force-reinstall to prevent CUDA torch from being pulled in as a transitive dependency.
  • Makefile: Added rocm-* targets (rocm-sync, rocm-test, rocm-test-warp, rocm-format, rocm-lint, rocm-typecheck) that run against .venv-rocm.

Usage

# CUDA (default)
uv sync --extra train --extra warp --group dev
make test

# ROCm
make rocm-sync
make rocm-test

Two venvs (.venv and .venv-rocm) coexist independently.

amd and others added 2 commits July 1, 2026 15:25
Add scripts/setup-rocm.sh to create a separate .venv-rocm with
ROCm 7.2 PyTorch, and Makefile targets (rocm-*) for test/lint/format
in the ROCm environment. The two venvs coexist independently.
@johnsutor

Copy link
Copy Markdown
Owner

Hi @alexhegit thank you for contributing. I realize I don't have a contribution guide since this project is still very early on, but I look forward to integrating your proposed changes. One thought I have off the bat is, why not bundle this as an extra to be installed (to the effect of so101-nexus[rocm])?

Also, I would prefer to avoid install scripts. I use uv in the first place so that it is simple to run teleop without even having to install this environment. I will create an issue to track this and figure out how to better integrate ROCm support, if you are open to this.

@johnsutor

Copy link
Copy Markdown
Owner

@alexhegit I did some work on this branch to support ROCm. Would you mind checking it out to see if that is sufficient?

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.

2 participants