Any interest in training with Rocm???

[safe-sentinel-co]

I modified the code to train with Rocm bc I have an 8x Mi 300X. It shards well but the training is 4x slower compared to a NVIDIA H200. If there is any interest from others, I can tidy it up a bit and push it out :)

Base image I used is: rocm/jax-community:rocm6.3.2-jax0.5.0-py3.12.8

[gauravpradeep]

@PRB-R-NINE-T Hi, our team working on finetuning pi0 on different embodiments and would appreciate if you could share the code to train on rocm; i dont have too much experience working with jax so this would be really helpful-thanks!