AMD Official Support: Add ROCm Docker support for PyTorch training on AMD GPUs

README.md

@@ -31,6 +31,8 @@ To run the models in this repository, you will need an NVIDIA GPU with at least
 
 The repo has been tested with Ubuntu 22.04, we do not currently support other operating systems.
 
+AMD GPUs are supported for PyTorch training through a separate ROCm Docker image; see [Training on AMD GPUs (ROCm)](docs/rocm.md).
+
 ## Installation
 
 When cloning this repo, make sure to update submodules:
@@ -51,7 +53,7 @@ GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
 
 NOTE: `GIT_LFS_SKIP_SMUDGE=1` is needed to pull LeRobot as a dependency.
 
-**Docker**: As an alternative to uv installation, we provide instructions for installing openpi using Docker. If you encounter issues with your system setup, consider using Docker to simplify installation. See [Docker Setup](docs/docker.md) for more details.
+**Docker**: As an alternative to uv installation, we provide instructions for installing openpi using Docker. If you encounter issues with your system setup, consider using Docker to simplify installation. See [Docker Setup](docs/docker.md) for more details. On AMD GPUs, Docker is the only supported path — see [Training on AMD GPUs (ROCm)](docs/rocm.md).
 
 
 

docs/docker.md

@@ -7,6 +7,7 @@ All of the examples in this repo provide instructions for being run normally, an
 - To use your GPU you must also install the [NVIDIA container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
 - The version of docker installed with `snap` is incompatible with the NVIDIA container toolkit, preventing it from accessing `libnvidia-ml.so` ([issue](https://github.com/NVIDIA/nvidia-container-toolkit/issues/154)). The snap version can be uninstalled with `sudo snap remove docker`.
 - Docker Desktop is also incompatible with the NVIDIA runtime ([issue](https://github.com/NVIDIA/nvidia-container-toolkit/issues/229)). Docker Desktop can be uninstalled with `sudo apt remove docker-desktop`.
+- These instructions assume an NVIDIA GPU. For AMD GPUs, see [Training on AMD GPUs (ROCm)](rocm.md) instead, which uses a separate ROCm image and does not need a container toolkit.
 
 
 If starting from scratch and your host machine is Ubuntu 22.04, you can use accomplish all of the above with the convenience scripts `scripts/docker/install_docker_ubuntu22.sh` and `scripts/docker/install_nvidia_container_toolkit.sh`.

docs/rocm.md

@@ -0,0 +1,64 @@
+# Training on AMD GPUs (ROCm)
+
+The default setup targets NVIDIA GPUs: `pyproject.toml` pins the CUDA builds of `torch` and `jax`, so `uv sync` cannot be used on a ROCm host. Instead, `scripts/docker/train_rocm.Dockerfile` builds on top of a ROCm base image that already ships a matching PyTorch, and installs the remaining openpi dependencies without disturbing it.
+This covers PyTorch training only. Everything under [PyTorch Support](../README.md#pytorch-support) in the main README applies, with the same feature limitations. JAX training is not supported here. Tested on 8×MI308X (gfx942) with ROCm 7.14 and PyTorch 2.12. Other gfx942 parts (MI300X, MI325X) use the same base image and should work; older generations have not been tested.
+
+## Requirements
+
+- A ROCm-capable GPU and the `amdgpu` kernel driver on the host. Check with `rocm-smi`.
+- Docker. Unlike NVIDIA, there is no container toolkit to install — GPU access comes from passing `/dev/kfd` and `/dev/dri` into the container and adding the `video` and `render` groups, which `compose_rocm.yml` already does.
+
+## Quick start
+
+```bash
+docker compose -f scripts/docker/compose_rocm.yml up -d --build
+docker compose -f scripts/docker/compose_rocm.yml exec openpi_rocm bash
+```
+
+The repo is mounted at `/app`, so local edits take effect without rebuilding. Checkpoints and datasets are large, so point `OPENPI_DATA_HOME` at a local disk rather than a network share — on NFS the data loader becomes the bottleneck:
+
+```bash
+OPENPI_DATA_HOME=/mnt/local/openpi docker compose -f scripts/docker/compose_rocm.yml up -d
+```
+
+Verify that the GPUs are visible from inside the container:
+
+```bash
+python3 -c "import torch; print(torch.cuda.device_count(), torch.cuda.get_device_name(0))"
+```
+
+## Running training
+
+The workflow is the same as [Finetuning with PyTorch](../README.md#finetuning-with-pytorch), except that commands run directly rather than through `uv run`. Convert a base checkpoint, compute norm stats, then train:
+
+```bash
+python3 examples/convert_jax_model_to_pytorch.py \
+    --config_name <config name> \
+    --checkpoint_dir /path/to/jax/base/model \
+    --output_path /path/to/pytorch/base/model
+
+python3 scripts/compute_norm_stats.py --config-name <config name>
+
+# single-node training
+torchrun --standalone --nnodes=1 --nproc_per_node=<gpus_per_node> scripts/train_pytorch.py <config_name> --exp_name <run_name> --pytorch-weight-path /path/to/pytorch/base/model --num-workers 8
+
+# multi-node training
+torchrun --nnodes=<num_nodes> --nproc_per_node=<gpus_per_node> --node_rank=<rank_of_node> --master_addr=<master_ip> --master_port=<port> \
+    scripts/train_pytorch.py <config_name> --exp_name=<run_name> --save_interval <interval> --pytorch-weight-path /path/to/pytorch/base/model
+```
+
+## How the image differs from the uv setup
+
+Three deviations from `pyproject.toml` are needed to keep the ROCm stack intact. They are applied in the Dockerfile and explained here because they are easy to undo by accident.
+
+**`uv sync` is not used.** Installing openpi's pinned `torch==2.7.1` and `jax[cuda12]` would replace the ROCm builds with CUDA wheels. The Dockerfile instead generates a constraints file from the versions already present in the base image and passes it to every `pip install`. Packages that would otherwise pull in torch — openpi, openpi-client and lerobot — are installed with `--no-deps`.
+
+**numpy stays on 2.x** rather than the `<2.0.0` that `pyproject.toml` requires. The ROCm extensions in the base image are compiled against the numpy 2 ABI and break if numpy is downgraded. jax 0.5.3, flax, orbax and the openpi data pipeline all work under numpy 2.
+
+**JAX is installed as the CPU build** and `JAX_PLATFORMS=cpu` is set. openpi imports jax for its configs, data transforms and orbax checkpoint loading, but all GPU work goes through PyTorch, so jax should not claim GPU memory.
+
+As a result `pip` prints dependency conflict warnings during the build, and `pip check` reports the numpy and torch pins above alongside a handful of unrelated packages that the base image ships and openpi never imports. These are expected. The final build layer imports the model module and instantiates a training config, so the build still fails if the environment is genuinely broken.
+
+## Known issues
+
+**`--overwrite` is racy under multi-GPU training.** `scripts/train_pytorch.py` calls `shutil.rmtree(config.checkpoint_dir)` without guarding on the main process, so every rank deletes the same directory concurrently and they remove files out from under each other. This surfaces as `FileNotFoundError` on `norm_stats.json` or on the checkpoint directory itself. It is not specific to ROCm. Use a fresh `--exp_name` for each run instead of `--overwrite`.

scripts/docker/compose_rocm.yml

@@ -0,0 +1,41 @@
+# Training on AMD GPUs (ROCm). See docs/rocm.md.
+#
+# Start with:
+# docker compose -f scripts/docker/compose_rocm.yml up -d --build
+#
+# Then open a shell in the container:
+# docker compose -f scripts/docker/compose_rocm.yml exec openpi_rocm bash
+services:
+  openpi_rocm:
+    image: openpi_rocm
+    build:
+      context: ../..
+      dockerfile: scripts/docker/train_rocm.Dockerfile
+    init: true
+    tty: true
+    network_mode: host
+    # Keep the container up so that training can be launched with `exec`.
+    command: sleep infinity
+    volumes:
+      - $PWD:/app
+      - ${OPENPI_DATA_HOME:-~/.cache/openpi}:/openpi_assets
+    environment:
+      - OPENPI_DATA_HOME=/openpi_assets
+      - IS_DOCKER=true
+
+    # ROCm needs both device nodes plus membership in the video and render
+    # groups; unlike NVIDIA there is no container toolkit to inject them.
+    devices:
+      - /dev/kfd
+      - /dev/dri
+    group_add:
+      - video
+      - render
+    security_opt:
+      - seccomp:unconfined
+    cap_add:
+      - SYS_PTRACE
+    # Multi-GPU training runs one DataLoader worker set per rank, which needs
+    # far more shared memory than the 64 MB default.
+    ipc: host
+    shm_size: 64gb

scripts/docker/train_rocm.Dockerfile

@@ -0,0 +1,87 @@
+# Dockerfile for training openpi on AMD GPUs (ROCm).
+#
+# The base image ships ROCm 7.14 and a matching PyTorch 2.12 build for gfx942
+# (MI300 series). openpi's pyproject pins the CUDA builds of torch and jax, so
+# this image does not use `uv sync`; it installs the dependencies under a
+# constraints file that keeps the pre-installed ROCm stack in place.
+#
+# Build the container:
+# docker build . -t openpi_rocm -f scripts/docker/train_rocm.Dockerfile
+
+# Run the container:
+# docker compose -f scripts/docker/compose_rocm.yml up -d
+
+# See docs/rocm.md for hardware requirements and known issues.
+
+FROM rocm/primus:v26.4@sha256:8c8ecc6fe14b5061423cc82517831e3515f6eb15647ee6d746fd6a0c1963da24
+
+# Kept in sync with pyproject.toml.
+ARG LEROBOT_REF=0cf864870cf29f4738d3ade893e6fd13fbd7cdb5
+ARG TRANSFORMERS_VER=4.53.2
+
+ENV PIP_ROOT_USER_ACTION=ignore \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    GIT_LFS_SKIP_SMUDGE=1 \
+    ROCPROFILER_LOG_LEVEL=fatal
+
+WORKDIR /app
+
+# Pin the ROCm stack. pyproject asks for torch==2.7.1 (CUDA), jax[cuda12] and
+# numpy<2, all of which would replace the ROCm builds that the base image ships.
+# Generate a constraints file from the versions already installed and pass it to
+# every install below. numpy stays on 2.x: downgrading breaks the ROCm
+# extensions, which are compiled against the numpy 2 ABI.
+RUN set -eux; \
+    python3 -c "import importlib.metadata as m; \
+names='torch torchvision triton numpy pillow tokenizers huggingface_hub safetensors av'.split(); \
+open('/etc/openpi-constraints.txt','w').write(''.join(n+'=='+m.version(n)+chr(10) for n in names))"; \
+    cat /etc/openpi-constraints.txt
+
+# JAX on CPU. openpi's configs, data transforms and orbax checkpoint loading
+# import jax, but all GPU work goes through PyTorch.
+RUN pip install -c /etc/openpi-constraints.txt \
+    "jax==0.5.3" "jaxlib==0.5.3" "flax==0.10.2" "orbax-checkpoint==0.11.13" \
+    "ml-dtypes==0.4.1" "tensorstore==0.1.74" "chex==0.1.90" "beartype==0.19.0" \
+    "jaxtyping==0.2.36" "equinox>=0.11.8" "augmax>=0.3.4" "dm-tree>=0.1.8" \
+    "flatbuffers>=24.3.25" "ml_collections==1.0.0" "treescope>=0.1.7"
+
+# Remaining runtime dependencies.
+RUN pip install -c /etc/openpi-constraints.txt \
+    "transformers==${TRANSFORMERS_VER}" \
+    "imageio[ffmpeg]>=2.36.1" "numpydantic>=1.6.6" "opencv-python-headless>=4.10.0.84" \
+    "polars>=1.30.0" "rich>=14.0.0" "tqdm-loggable>=0.2" "typing-extensions>=4.12.2" \
+    "tyro>=0.9.5" "wandb>=0.19.1" "filelock>=3.16.1" "sentencepiece>=0.2.0" \
+    "fsspec[gcs]>=2024.6.0" \
+    jsonlines deepdiff draccus termcolor
+
+# LeRobot, for its datasets module. --no-deps because it pulls CUDA torch; its
+# transitive dependencies are covered above.
+RUN pip install --no-deps -c /etc/openpi-constraints.txt \
+    "lerobot @ git+https://github.com/huggingface/lerobot@${LEROBOT_REF}"
+
+# Copied after the dependencies so that editing the repo does not invalidate the
+# layers above.
+COPY . /app
+
+# Editable installs so that mounting the repo over /app at runtime picks up
+# local edits. Same transformers patch as the uv setup in the main README:
+# it adds AdaRMS, activation precision control and a read-only KV cache, without
+# which PI0Pytorch fails to initialize.
+RUN set -eux; \
+    pip install --no-deps -e ./packages/openpi-client -e .; \
+    cp -r src/openpi/models_pytorch/transformers_replace/* \
+        "$(python3 -c 'import os, transformers; print(os.path.dirname(transformers.__file__))')/"; \
+    python3 -c "from transformers.models.siglip import check; \
+assert check.check_whether_transformers_replace_is_installed_correctly()"; \
+    JAX_PLATFORMS=cpu python3 -c "\
+import jax, openpi.training.config as c, openpi.models_pytorch.pi0_pytorch, torch; \
+print('torch', torch.__version__, '| jax', jax.__version__, jax.devices()); \
+print('config ok:', c.get_config('pi05_libero').name)"
+
+# /openpi_assets is where compose_rocm.yml mounts OPENPI_DATA_HOME.
+ENV OPENPI_DATA_HOME=/openpi_assets \
+    HF_HOME=/openpi_assets/hf \
+    JAX_PLATFORMS=cpu
+
+CMD ["/bin/bash"]