Install PyTorch for ROCm

Install PyTorch for ROCm#

Refer to this section for the recommended PyTorch via PIP installation method.

PCIe atomics

ROCm is an extension of HSA platform architecture, and shares queuing model, memory model, signaling and synchronization protocols.

Platform atomics are integral to perform queuing and signaling memory operations, where there may be multiple-writers across CPU and GPU agents.

For more details, see How ROCm uses PCIe atomics.

Install methods#

AMD recommends the PIP install method to create a PyTorch environment when working with ROCm™ for machine learning development.

PyTorch via PIP installation#

AMD recommends the PIP install method to create a PyTorch environment when working with ROCm™ for machine learning development.

Check Pytorch.org for latest PIP install instructions and availability. See Compatibility matrices for support information.

Note
To install the following wheels, Python 3.12 must be set up.

  1. Enter the following command to unpack and begin set up.

    sudo apt install python3-pip -y

  2. Enter this command to update the pip wheel.

    pip3 install --upgrade pip wheel

  3. Select the applicable Ubuntu version and enter the commands to install Torch and Torchvision for ROCm AMD GPU support.

    This may take several minutes.

    Important! AMD recommends proceeding with ROCm WHLs available at repo.radeon.com. The ROCm WHLs available at PyTorch.org are not tested extensively by AMD as the WHLs change regularly when the nightly builds are updated.

    Important! When manually downloading WHLs from repo.radeon, ensure to select the compatible WHLs for specific Python versions.

    See Compatibility matrices for support information.

    wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchaudio-2.9.0%2Brocm7.1.1.gite3c6ee2b-cp312-cp312-linux_x86_64.whl
pip3 uninstall torch torchvision triton torchaudio
pip3 install torch-2.9.1+rocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl torchvision-0.24.0+rocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl torchaudio-2.9.0+rocm7.1.1.gite3c6ee2b-cp312-cp312-linux_x86_64.whl triton-3.5.1+rocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl

    Note

    The --break-system-packages flag must be added when installing wheels for Python 3.12 in a non-virtual environment.

    wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp310-cp310-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp310-cp310-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp310-cp310-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchaudio-2.9.0%2Brocm7.1.1.gite3c6ee2b-cp310-cp310-linux_x86_64.whl
pip3 uninstall torch torchvision triton torchaudio
pip3 install torch-2.9.1+rocm7.1.1.lw.git351ff442-cp310-cp310-linux_x86_64.whl torchvision-0.24.0+rocm7.1.1.gitb919bd0c-cp310-cp310-linux_x86_64.whl torchaudio-2.9.0+rocm7.1.1.gite3c6ee2b-cp310-cp310-linux_x86_64.whl triton-3.5.1+rocm7.1.1.gita272dfa8-cp310-cp310-linux_x86_64.whl

    wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchaudio-2.9.0%2Brocm7.1.1.gite3c6ee2b-cp312-cp312-linux_x86_64.whl
pip3 uninstall torch torchvision triton torchaudio
pip3 install torch-2.9.1+rocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl torchvision-0.24.0+rocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl torchaudio-2.9.0+rocm7.1.1.gite3c6ee2b-cp312-cp312-linux_x86_64.whl triton-3.5.1+rocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl

    Note

    The --break-system-packages flag must be added when installing wheels for Python 3.12 in a non-virtual environment.

Next, verify your PyTorch installation.

Verify PyTorch installation#

Confirm if PyTorch is correctly installed.

  1. Verify if Pytorch is installed and detecting the GPU compute device.

    python3 -c 'import torch' 2> /dev/null && echo 'Success' || echo 'Failure'

    Expected result:

    Success

  2. Enter command to test if the GPU is available.

    python3 -c 'import torch; print(torch.cuda.is_available())'

    Expected result:

    True

  3. Enter command to display installed GPU device name.

    python3 -c "import torch; print(f'device name [0]:', torch.cuda.get_device_name(0))"

    Expected result: Example: device name [0]: Radeon RX 7900 XTX

    device name [0]: <Supported AMD GPU>

  4. Enter command to display component information within the current PyTorch environment.

    python3 -m torch.utils.collect_env

    Expected result:

    PyTorch version
 
ROCM used to build PyTorch
 
OS
 
Is CUDA available
 
GPU model and configuration
 
HIP runtime version
 
MIOpen runtime version

Environment set-up is complete, and the system is ready for use with PyTorch to work with machine learning models, and algorithms.

Contents