VITRA VLA inference runs on AMD GPUs (ROCm) — verified on MI300X, zero source changes

[ZJLi2013]

Summary

VITRA's VLA inference path (build_vla + predict_action) runs on AMD Instinct GPUs via ROCm with no source changes beyond the dtype handling already documented in the README minimal example.

VITRA is pure PyTorch / transformers (no custom CUDA kernels), so it works on ROCm out of the box once the environment is set up correctly.

Tested on

• GPU: AMD Instinct MI300X (also verified on MI308X, gfx942)
• ROCm: 7.2.1
• PyTorch: 2.9.1 (ROCm build)
• Docker: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1
• transformers==4.47.1 (the repo's pin)
• Weights: VITRA-VLA/VITRA-VLA-3B + google/paligemma2-3b-mix-224

Result

• VITRA_Paligemma loads the checkpoint with 0 missing / 0 unexpected state-dict keys.
• predict_action returns a finite [1, 16, 192] action chunk (16-step × 192-dim unified action).
• Timings (warm cache): load ≈ 15.5s, infer ≈ 10.0s.

Setup notes for ROCm / AMD users

A few environment details that make it work — none require source changes:

1. Pin transformers==4.47.1 (the repo's pyproject pin). Newer transformers change
   get_image_features' return type → 'BaseModelOutputWithPooling' has no attribute 'numel',
   and the correct pin also makes the checkpoint load cleanly (0/0 keys).
2. Install without training-only CUDA deps for inference: a full pip install -e . stalls
   building flash-attn / deepspeed, which are only needed for training. Using
   pip install --no-deps -e . plus the explicit pure-Python deps keeps the Docker ROCm
   PyTorch intact (torch.version.hip unchanged) and inference runs fine without them.
   Would a pip install-friendly inference-only path (moving flash-attn/deepspeed to an
   optional extra) be welcome? Happy to help.
3. float32 inputs: state / fov built from NumPy / np.deg2rad default to float64
   (torch.Double) and raise mat1 and mat2 must have the same dtype. This is already handled
   by the README minimal example (which casts to float32), so no change needed — just noting
   it for other ROCm users following along.

Ask

Nothing blocking — this is mainly a compatibility report so AMD/ROCm users know VITRA works.
If useful, I'm happy to open a follow-up PR for the inference-only optional-extras packaging
in note (2).

[arnoldland]

Hi, thanks a lot for testing VITRA on ROCm and for sharing such a detailed compatibility report!

It's great to know that the inference pipeline works on AMD GPUs without any source changes. We also appreciate you documenting the required environment details and installation notes—they'll be very helpful for other ROCm users.

One thing that caught our attention is the reported inference time (~10 s per inference on MI300X), which seems much slower than we would expect. Do you happen to know whether this timing corresponds to the entire Python inference pipeline (including hand reconstruction and others), or just the model forward pass?

Thanks again for taking the time to test VITRA and share your findings. We're looking forward to your PR! I'm sure your documentation will be a great help to other AMD/ROCm users.