Inquiry on Torch XLA Scalability and Performance Comparison with JAX on TPU

[Rayjryang]

Thank you very much for this project, it's really great. I would like to know if the improved Torch XLA has been tested for the maximum number of TPUs it can support when scaling on TPU. Because previously we found that on the original Torch XLA, some unknown errors occurred when the number of TPUs exceeded a certain amount. Also, has there been any testing of the training speed difference between Torch XLA and JAX under the same conditions on TPU?

[tengyifei]

@vlasenkoalexey @yaoshiang

[cjonesy20]

@Rayjryang I'm Chris, PM supporting PT/XLA. We're doing a lot of work right now on vastly improving our distributed apis to scale better but the short answer is that I don't have a specific answer to your very fair questions. Would you mind sharing what error messages did you run across and how many TPUs (and which TPUs) seemed to be the threshold where the problem started?