Thanks for this! I've been trying to get local training working on Fedora with ROCm for a while now, and this is the first project that ran end-to-end on my RX 7600 XT 16GB

The whole reason for me buying the card was to play around with deep learning and models, so its awesome to have a stepping off point now!

No CUDA meant ROCm suprisingly just worked with JAX. Given recent events at NVDA its also very timely :)

Gotchas for anyone with a 7600 / 7600 XT:

1. ROCm 10.1 supports gfx1102, but JAX's ROCm plugin rejects it ("no supported devices found"), so we need to set HSA_OVERRIDE_GFX_VERSION=11.0.0

2. XLA's Triton GEMM fusions exceed Navi 33's dispatch limits at the default config and hangs the GPU, so set XLA_FLAGS=--xla_gpu_enable_triton_gemm=false

3. JAX ≥0.5 changed the default PRNG, so set JAX_THREEFRY_PARTITIONABLE=false if you are looking to replicate precisely

My setup: Fedora 45 with the in-kernel amdgpu driver, rootless podman, and AMD's rocm/jax image (ROCm 10.1, JAX 0.11.1). Nothing ROCm is installed on the host, which is nice

PR with the container setup, a comparison script, and plot: https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax/pull/1