Google TPU support in verl with Ray

[andrewsykim]

Feature request

I'm interested to see if we can add TPU support in verl. I saw an issue from last year indicating it's not supported.

Ray has made some significant strides in the past year to improve TPU support. I'd like to understand what the gaps are to bring TPU support into verl with Ray. I'm also happy to help contribute but need some guidance from maintainers.

Motivation

Ray is improving TPU support and would be great if verl can also support TPUs when using with Ray.

Your contribution

Happy to submit or review PRs.

[wuxibin89]

@andrewsykim To support TPU in verl, we need support TPU both in model engine and rollout engine:

• model engine: adapt existed backend: fsdp, torchtitan to TPU, while I'm not sure how usable PyTorch for TPU now. It's better to add a new model engine with native TPU support, e.g maxtext
• rollout engine: integrate vllm-tpu, sglang-jax with trainer-rollout weight synchronization.

[wuxibin89]

Oh, besides model engine and rollout engine, we heavily rely on pytorch ecosystem, e.g transformers, TensorDict, so I think it's better to use Pytorch for TPU.