Skip to content

Update build directions for AMD GPUs #167

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JianzheXiao merged 1 commit into from
Apr 17, 2026
Merged

Update build directions for AMD GPUs #167

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 16 additions & 9 deletions docs/build.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
We recommend installation in [Nvidia PyTorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags).

#### if for AMD GPU:
- ROCM 6.3.0
- Torch 2.4.1 with ROCM support
- ROCM 7.1
- Torch 2.7.1 with ROCM support



Expand Down Expand Up @@ -139,7 +139,7 @@ bash .codebase/scripts/nvidia/run_tutorial_test.sh
See examples in the `tutorials` directory at the project root.

## To use Triton-distributed with the AMD backend:
Starting from the rocm/pytorch:rocm6.1_ubuntu22.04_py3.10_pytorch_2.4 Docker container
Starting from the rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1 Docker container
#### AMD Build Steps
1. Clone the repo
```sh
Expand All @@ -150,14 +150,21 @@ git clone https://github.com/ByteDance-Seed/Triton-distributed.git
cd Triton-distributed/
git submodule update --init --recursive
```
If you are updating an old repo, there may be issues if the rocshmem submodule is still present. Erase it if necessary:
```sh
rm -rf 3rdparty/rocshmem # only for updated repo
```
3. Install dependencies
```sh
sudo apt-get update -y
sudo apt install -y libopenmpi-dev
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3 --no-deps
bash ./shmem/rocshmem_bind/build.sh
python3 -m pip install -i https://test.pypi.org/simple hip-python>=6.3.0 # (or whatever Rocm version you have)
export TRITON_BUILD_WITH_CLANG_LLD=TRUE
export TRITON_USE_ASSERT_ENABLED_LLVM=TRUE
export TRITON_BUILD_PROTON=0
Copy link
Copy Markdown
Contributor

@drprajap drprajap Apr 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there any workaround for this, I see cuda header when compiling proton files, is that why you suggest disabling?

Copy link
Copy Markdown
Collaborator Author

@erieaton-amd erieaton-amd Apr 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Proton isn't supported on AMD hardware yet so this can cause a build error.

rm -f /usr/local/bin/cmake
apt-get update -y
apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake miopen-hip autoconf libtool flex ninja-build clang lld
python3 -m pip install -i https://test.pypi.org/simple hip-python>=7.1 # (or whatever Rocm version you have)
pip3 install pybind11
bash ./shmem/rocshmem_bind/build.sh
```
4. Build Triton-distributed
```sh
Expand All @@ -167,7 +174,7 @@ pip3 install -e python --verbose --no-build-isolation --use-pep517
#### GEMM ReduceScatter example on single node
```sh
bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ag_gemm_intra_node.py 8192 8192 29568
```
```
and see the following (reduced) output
```sh
✅ Triton and Torch match
Expand Down