m-yuasa/pytorch-mi300a-rocm72

PyTorch 2.13 for AMD Instinct MI300A (ROCm 7.2)

A pip-installable PyTorch 2.13 build with MI300A APU shared-memory support. Eligible CPU-to-GPU and GPU-to-CPU transfers alias the same physical memory instead of copying, so the transfer disappears rather than getting faster.

Disclosure: this build, its tests, and this page were prepared with the assistance of a generative AI tool (Claude Code). A human reviewed the change and the validation results before publishing.

Install

python3.12 -m venv mi300a && source mi300a/bin/activate
pip install https://huggingface.co/m-yuasa/pytorch-mi300a-rocm72/resolve/main/torch-2.13.0+mi300a.rocm72-cp312-cp312-linux_x86_64.whl

The ROCm 7.2 user-space libraries are bundled, so no ROCm module or /opt/rocm installation is needed at run time. A compatible AMDGPU/KFD kernel driver must already be present.

Requirements

• AMD Instinct MI300A (gfx942)
• CPython 3.12
• Linux x86_64 with x86-64-v2 or newer, glibc 2.34 or newer

TorchVision and TorchAudio

Do not run pip install torch torchvision. TorchVision pins an exact upstream torch version that this build's local version 2.13.0+mi300a.rocm72 does not satisfy, so pip silently replaces this wheel with the stock nightly and the APU support is lost. Install the companions with --no-deps:

pip install --no-deps --index-url https://download.pytorch.org/whl/nightly/rocm7.2 \
  torchvision==0.28.0+rocm7.2 torchaudio==2.11.0+rocm7.2
pip install --index-url https://download.pytorch.org/whl/nightly/rocm7.2 triton-rocm==3.7.1

pip check will report the TorchVision requirement as unmet. That is the version-string mismatch only; imports and native operators work.

Verify

import torch

cpu = torch.arange(1024, dtype=torch.float32)
gpu = cpu.to("cuda:0")
assert torch.cuda.apu.is_available(0)
assert torch.cuda.apu.is_shared(gpu)
torch.testing.assert_close(gpu.cpu(), cpu)
print("APU shared memory: OK")

What this changes

Measured against the stock 2.13.0+rocm7.2 build on the same MI300A with the same commands.

Compute speed is unchanged: resnet50 bf16 training at batch 256 reaches 1771 samples/sec here and 1774 samples/sec on the stock build, and GEMM throughput and copy bandwidth match. This build is not faster at arithmetic.

What changes is the memory a CPU-to-GPU handoff costs, because no copy is made:

Payload moved to the GPU	This build	Stock
256 MiB	0 MiB GPU, 1 MiB system	256 MiB GPU, 841 MiB system
1 GiB	0 MiB GPU, 5 MiB system	1024 MiB GPU, 1030 MiB system
4 GiB	0 MiB GPU, 9 MiB system	4096 MiB GPU, 4120 MiB system

The gain is real for workloads that stage large CPU-resident data onto the GPU. It is small for compute-bound training, where weights, activations, and gradients dominate: resnet50 at batch 256 peaks at 11059 MiB of GPU allocation here against 11307 MiB on the stock build.

Choosing a BLAS backend

The default hipBLASLt path collapses on some large shapes. A 16384 square bf16 matmul measures 15.7 TFLOPS with hipBLASLt and 617 TFLOPS with rocBLAS, while both agree at 8192. If large matmuls are slower than expected:

torch.backends.cuda.preferred_blas_library("cublas")  # rocBLAS

torch.cuda.tunable can also autotune GEMM kernels per shape. This comes from the ROCm libraries, not from the APU changes; the stock build behaves the same.

Validation

On a single MI300A APU: 20/20 APU shared-memory tests pass, pip check is clean, and TorchVision GPU NMS, TorchAudio, and torch.compile all work.

Source

The complete source for this build, including the APU shared-memory changes, is at:

• Branch: https://github.com/Mamoru-Yuasa/pytorch/tree/mi300a-apu-shared-memory-v2.13
• Tag: mi300a-apu-rocm7.2-torch213-v0.1.0
• Commit: f0fb8da574d2ccaa2772726cfd1f65857dfab2ce
• Release notes: https://github.com/Mamoru-Yuasa/pytorch/releases/tag/mi300a-apu-rocm7.2-torch213-v0.1.0

Checking out that tag reproduces exactly the tree this wheel was built from; torch.version.git_version records the same commit.

SHA-256 of this wheel: e3426c43a6b17ce174ef0ecb02306f3d97ec7ac2059d9dcb220ca7e9ab07e522

Licensing

This is an unofficial community build. It is not published by the PyTorch Foundation or by AMD, and it is not endorsed by either.

PyTorch and the components vendored into it are covered by the license files inside the wheel, under torch-*.dist-info/licenses/, and by the License-Expression field in the wheel metadata.

The bundled build additionally ships ROCm 7.2 user-space libraries and a small set of system libraries, added when the wheel was made self-contained. Their license texts are included under torch-*.dist-info/licenses/bundled/, alongside a README.txt mapping every bundled shared object to the component whose terms cover it:

• rocm/ covers MIOpen, rocBLAS, rocSOLVER, rocSPARSE, rocRAND, rocFFT, RCCL, hipBLAS, hipBLASLt, hipFFT, hipSOLVER, hipSPARSE, hipSPARSELt, HIP, amd_comgr, hsa-rocr, roctracer, rocprofiler, amd-smi, rocm-smi, rocm-core
• system/ covers elfutils, libomp, libquadmath, bzip2, xz, zstd, and carries notices for libdrm, libgfortran, libgomp, numactl, and ncurses
• external/ and vendor/ carry notices for AOTriton and for libhsa-amd-aqlprofile64.so, which AMD distributes under its own terms

Some of these are copyleft, notably elfutils (GPL-2.0-or-later OR LGPL-3.0-or-later) and numactl (LGPL-2.1 / GPL-2.0). Corresponding source for those components is available on request. Anyone redistributing this wheel further should keep the license directory intact and satisfy those obligations themselves.

The minimal wheel attached to the GitHub release does not bundle any of these; it links against the ROCm installation already present on the host.