Skip to content

note AMD ROCm support (runs unchanged on AMD Instinct MI300X +) #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ZJLi2013 wants to merge 1 commit into from
Closed

note AMD ROCm support (runs unchanged on AMD Instinct MI300X +) #2

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,9 @@ Config files are packaged under `lingbot_vision/configs/` and selected automatic

## 🔧 Installation

**Requirements**: Python ≥ 3.10 · PyTorch ≥ 2.0 · CUDA-capable GPU (recommended for large-model inference)
**Requirements**: Python ≥ 3.10 · PyTorch ≥ 2.0 · a CUDA- or ROCm-capable GPU (recommended for large-model inference)

> **AMD GPUs (ROCm)**: LingBot-Vision runs unchanged on AMD GPUs via ROCm — see [Running on AMD ROCm](#-running-on-amd-rocm).

**1. Clone the repository**

Expand All @@ -83,6 +85,27 @@ python -m pip install -r requirements.txt
python -m pip install -e .
```

## 🟥 Running on AMD ROCm

LingBot-Vision runs **out of the box on AMD GPUs via ROCm with zero source changes**. The default attention path (`torch.nn.functional.scaled_dot_product_attention`, backend `sdpa`) is auto-dispatched to PyTorch's AOTriton Flash Attention kernel on ROCm, so `lingbot_vision/attention.py` works unchanged. The pure-PyTorch stack (`torch/torchvision/numpy/opencv-python-headless/pillow/omegaconf/huggingface_hub`, no custom CUDA/C++ extension, no `xformers`/`flash-attn`) has nothing CUDA-specific to build.

Verified on an AMD Instinct MI300X (gfx942) using the official ROCm PyTorch image `rocm/pytorch:rocm6.4.3_ubuntu24.04_py3.12_pytorch_release_2.6.0` (torch 2.6.0, HIP 6.4), where `torch.cuda.is_available() == True`:

```bash
# The ROCm PyTorch image already provides torch/torchvision — install the rest.
python -m pip install numpy opencv-python-headless pillow omegaconf huggingface_hub
python -m pip install -e .

# Recommended ViT-L backbone (weights auto-download from Hugging Face).
python -m lingbot_vision.pca_demo \
--config-file lingbot_vision/configs/lbot_vision_vitl.yaml \
--ckpt /path/to/model.pt \
--input examples/example.png \
--out outputs/pca_demo --size 512 --mode square --dtype bf16
```

Both the recommended ViT-L and the flagship 1.1B ViT-g backbones load with 0 missing / 0 unexpected keys and produce the expected PCA-of-patch-token maps. On ROCm, PyTorch may print `Using AOTriton backend for Flash Attention forward` — this is expected and requires no changes.

## 🚀 Quick Start

### Load LingBot-Vision Pretrains with PyTorch
Expand Down