pytorch_rocm_gtt

Python package to allow ROCm to overcome the reserved iGPU memory limits.

Based on https://github.com/pomoke/torch-apu-helper/tree/main, after discussion here: https://github.com/ROCm/ROCm/issues/2014

Install it from PyPI

pip install pytorch_rocm_gtt

Usage

Just call this before starting pytorch allocations (model or torch):

import pytorch_rocm_gtt

pytorch_rocm_gtt.patch()

hipcc command should be in your $PATH.

After that, just allocate GPU memory as you would with cuda:

import torch

torch.rand(1000).to("cuda")

Compatibility

In order to use this package, your APU must be compatible with ROCm in the first place.

Check AMD documentation on how to install ROCm for your distribution.

Development

Read the CONTRIBUTING.md file.

How to release

Update pyproject.toml file with the desired version, and run make release to create the new tag.

After that, the github action will publish to pypi.

Once it is published, run the docker_build_and_publish.sh <version-number> script to update the docker images.