Достаточно нишевая инфа про скейлинг ллмного рля на Jax для TPU но мб вам интересно https://adityamakkar000.github.io/posts/async-rl-jax.html