PyTorch ROCm Installation Alternative for AMD GPUs Using uv: A Complete Guide

You can install PyTorch with ROCm support for AMD GPUs using uv by configuring a ROCm optional dependency in pyproject.toml, pointing to the ROCm wheel index, and running uv sync --extra rocm5.7.

The baonguyen6742/uv-install-torch repository demonstrates a streamlined approach to PyTorch ROCm installation for AMD GPUs using uv, a fast Python package manager and installer. While the original configuration targets NVIDIA CUDA, the same pattern applies to AMD hardware by substituting the ROCm wheel index and optional dependencies.

Why Use uv for PyTorch ROCm Installation on AMD GPUs

uv provides deterministic, fast environment resolution that outperforms traditional pip workflows. Unlike pip, uv supports multiple package indexes and conditional sources based on optional dependencies. This allows a single pyproject.toml to support CPU, CUDA, and ROCm configurations without conflicts, making it ideal for PyTorch ROCm installation on AMD GPUs.

Repository Structure and Key Files

The baonguyen6742/uv-install-torch repository contains three critical files that demonstrate the uv workflow:

• pyproject.toml: Declares dependencies, optional extras, and uv-specific index configuration in the [tool.uv.sources] table.
• main.py: Validates the installation by importing libraries and checking GPU availability.
• README.md: Documents the setup process and troubleshooting steps for CUDA-based setups, which translates directly to ROCm workflows.

Configuring pyproject.toml for ROCm Support

To adapt the repository for AMD GPUs, you must add a ROCm optional dependency and configure the uv sources to point to the ROCm wheel index.

Declaring Optional Dependencies

Add a rocm5.7 extra alongside existing CPU or CUDA extras in pyproject.toml:

[project.optional-dependencies]
cpu = ["torch==2.4.1", "torchvision==0.19.1", "torchaudio==2.4.1"]
rocm5.7 = ["torch==2.4.1", "torchvision==0.19.1", "torchaudio==2.4.1"]

Configuring uv Sources and Index

Map the ROCm extra to the AMD wheel index using the [tool.uv.sources] table:

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-rocm5.7", extra = "rocm5.7" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-rocm5.7", extra = "rocm5.7" },
]
torchaudio = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-rocm5.7", extra = "rocm5.7" },
]

[[tool.uv.index]]
name = "pytorch-rocm5.7"
url = "https://download.pytorch.org/whl/rocm5.7"
explicit = true

Installing PyTorch with ROCm Using uv

After configuring pyproject.toml, install the ROCm build with a single command:

uv sync --extra rocm5.7

This command resolves dependencies using the ROCm wheel index and installs the AMD-compatible versions of torch, torchvision, and torchaudio.

Verifying the ROCm Installation

The repository's main.py script validates the installation. For AMD GPUs, verify ROCm availability by checking the HIP runtime:

import torch
import torchvision
import torchaudio

print("torch version:", torch.__version__)
print("torchvision version:", torchvision.__version__)
print("torchaudio version:", torchaudio.__version__)

# Check ROCm availability

print("ROCm available:", torch.version.hip is not None)
print("GPU count:", torch.cuda.device_count())

Run the verification:

uv run main.py

When ROCm is functional, torch.version.hip returns a version string, while torch.cuda.is_available() typically returns False on pure ROCm systems.

Troubleshooting Cache Issues

If you encounter stale wheel errors when switching from CUDA to ROCm, clear the uv cache to force fresh downloads:

uv cache prune          # Remove unused cache entries

uv cache clean torch    # Clear torch-related cache

Summary

• The baonguyen6742/uv-install-torch repository provides a template for PyTorch ROCm installation using uv.
• Configure pyproject.toml with a rocm5.7 optional dependency and point to the ROCm wheel index at https://download.pytorch.org/whl/rocm5.7.
• Use uv sync --extra rocm5.7 to install AMD-compatible PyTorch builds.
• Verify installation by checking torch.version.hip in the main.py script.
• Clear uv cache with uv cache clean torch if switching between CUDA and ROCm builds.

Frequently Asked Questions

Can I use the same pyproject.toml for both CUDA and ROCm GPUs?

Yes. Define separate optional dependencies for each platform (e.g., cu124 and rocm5.7) and use the [tool.uv.sources] table to map each extra to its respective wheel index. Install with uv sync --extra rocm5.7 or uv sync --extra cu124 depending on your hardware. You can also add a conflicts entry in [tool.uv] to prevent simultaneous installation of incompatible extras.

What ROCm version should I specify in the wheel index?

The configuration uses rocm5.7 as an example, but you should match your system's ROCm installation. Check your AMD GPU compatibility against the ROCm documentation and install the corresponding ROCm stack. Then use the matching wheel URL such as https://download.pytorch.org/whl/rocm5.7 or the latest available ROCm version supported by PyTorch.

How do I verify that PyTorch is using my AMD GPU instead of CPU?

After installation, run python -c "import torch; print(torch.version.hip)". If it returns a version string (e.g., 5.7), PyTorch is using the ROCm backend. Note that torch.cuda.is_available() typically returns False on ROCm systems, so checking torch.version.hip is the reliable method for AMD GPUs. You can also run torch.cuda.device_count() to see detected ROCm devices.

Does uv support mutual exclusion between CUDA and ROCm extras?

Yes. Add a conflicts entry in the [tool.uv] section of pyproject.toml to prevent simultaneous installation of incompatible extras. This ensures that uv will not allow both cu124 and rocm5.7 to be active in the same environment, avoiding runtime conflicts between NVIDIA and AMD libraries. This feature is particularly useful when sharing a single configuration file across different hardware platforms.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}