VITRA trains end-to-end on AMD Instinct MI300X (ROCm) — no fork, no hipify, no custom kernels

Follow-up to #42, which reported that inference runs on AMD ROCm. This issue covers the other half: end-to-end multi-GPU training.

Summary

I ran VITRA's pretraining loop end-to-end on 8× AMD Instinct MI300X (gfx942) under ROCm and it works with stock upstream code — no fork, no hipify, no custom CUDA/HIP kernels. The training path is pure PyTorch FSDP + RCCL, and neither DeepSpeed nor FlashAttention is needed (the DiT action head uses torch SDPA and the VLM uses the default HF attention).

Getting a clean run required only one runtime env var, one small (platform-agnostic) dataloader fix, and honoring the utils3d pin that's already in pyproject.toml. Details below in case it helps other AMD users, and I'm happy to send a small PR for the code part.

Environment

• GPUs: 8× AMD Instinct MI300X OAM, gfx942, 192 GB each
• Container: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1 (torch 2.9.1+rocm7.2.1, HIP 7.2)
• Launcher: torchrun --standalone --nproc_per_node=8 scripts/train.py
• Data: a 10k-episode crop of VITRA-1M SSv2 (9,403 webm / 213,835 frames), data_mix=ssv2
• Config: human_pretrain.json with short-run overrides for a smoke of the full loop:
  batch_size=8 (per GPU), total_batch_size=64 (→ grad_accum=1), max_steps=1000,
  save_steps=250, llm_freeze_step=200, FSDP shard-grad-op (ZeRO-2), bf16.

What worked out of the box

• pip install --no-deps -e . + the pure-Python training deps (transformers 4.47.1, diffusers, timm, peft, etc.). DeepSpeed and FlashAttention were intentionally omitted and are not required.
• Model build (PaliGemma2-3B backbone + DiT head), bf16 forward/backward/optimizer step.
• FSDP sharding, checkpoint save and resume (optimizer state included).
• apex fused_layer_norm JIT-compiled cleanly via hipcc --offload-arch=gfx942 ("unsupported CUDA calls: 0").

Three changes needed

Only #2 is an actual source change; the files themselves are fine (they decode in isolation at any thread count). Suggest also making _grab_window_images fail-fast after its retries instead of falling through to np.stack(imgs, ...).

Evidence (1000-step run)

What "training works" means here — each of these was exercised on gfx942 with no source changes:

• Full step: forward → backward → optimizer.step() in bf16, finite loss (no NaN/Inf).
• FSDP sharding (shard-grad-op / ZeRO-2) across 8 ranks.
• RCCL collectives (init + all-gather / reduce-scatter) — the same path that fails without env fix #1.
• Checkpoint save at steps 250 / 500 / 750 / 1000.
• Checkpoint resume — model + optimizer state reloaded correctly and training continued.
• LLM freeze → unfreeze schedule (llm_freeze_step=200) so the backbone-trainable path is also exercised (backbone LR goes 0 → 1e-5).
• 8/8 GPUs at 100% utilization (rocm-smi), ~10 s/it, clean exit (torchrun exit code 0).

Raw per-step log (rank 0), loss stays finite and trends down (~0.29 → ~0.21):

=>> [Epoch 000] Global Step 000516 =>> Backbone LR :: 0.000000 - Loss :: 0.2664
=>> [Epoch 000] Global Step 000517 =>> Backbone LR :: 0.000000 - Loss :: 0.2943
=>> [Epoch 000] Global Step 000518 =>> Backbone LR :: 0.000000 - Loss :: 0.3123
...
=>> [Epoch 000] Global Step 000813 =>> Backbone LR :: 0.000010 - Loss :: 0.2362
...
=>> [Epoch 000] Global Step 000999 =>> Backbone LR :: 0.000010 - Loss :: 0.2145

(This is a short pipeline-validation run, not a convergence study — the point is the full ROCm training loop runs correctly, not the final metric.)

Offer

Happy to open a small PR with:

1. decord.VideoReader(name, num_threads=1) + fail-fast after retries in _grab_window_images, and
2. a short "Running on AMD ROCm (MI300X)" section in the README (the HSA_NO_SCRATCH_RECLAIM=1 note + the container tag above).

Thanks for open-sourcing VITRA — it's a great codebase to work with.

Start in video_utils.py at _grab_window_images and inspect the decord.VideoReader call and retry path. Reproduce the threaded-decoder failure with concurrent readers, then verify the decoder fix and fail-fast behavior. If covering the documentation part, update the README with the MI300X ROCm environment note; done means the smoke training run proceeds without the dataloader error.