I've been working on migrating PyTorch/XLA from our legacy XRT runtime to PJRT. We have detailed documentation on the differences and changes here: https://github.com/pytorch/xla/blob/master/docs/pjrt.md

This PR collects all of the fixes that I've made so far to support XRT and PJRT interchangeably through Accelerate. In order of significance:

1. Update implementation of synchronize_rng_types to use xm.collective_broadcast to broadcast the RNG state tensor. Collective operations in general should be called from the main thread of each process to avoid unpredictable behaviors in XLA. But, our implementation of MpDeviceLoader calls accelerate.DataLoaderShard's __iter__ (which synchronizes the RNG) in all of the preloading threads. To ensure that the RNG is synchronized once from the main thread, call it in the MpDeviceLoaderWrapper's __iter__ instead.
2. In general, you should call xm.mark_step to finish any remaining steps before checkpointing. MpDeviceLoader is responsible for calling xm.mark_step at the beginning of each new step and at the end of the dataset iterator. If you checkpoint in the middle of iteration, replica 0 won't reach the mark_step at the beginning of the next iteration before it tries to checkpoint. To avoid making the user call mark_step themselves, call it for them on replica 0 before checkpointing in save_state.
3. Use XLA's all_gather in gather instead of mesh_reduce + torch.cat. With PJRT, rendezvous and mesh_reduce both use XLA collective ops to broadcast pickled data (docs) and move it back to the CPU. Batching all of the recursive all_gather calls together and calling xm.mark_step once reduces the number of transfers between host and device.
4. Don't override XLA_USE_BF16 when mixed_precision is not set. This isn't really mixed precision as much as it is a mechanism to silently convert all torch.floats to BF16 on the TPU. Real mixed precision support in PT/XLA is still a WIP, and we can update back here when it's stable.
5. Remove obsolete/unused _mp_fn in test script.

Tested:

• accelerate test on v4-8 with XRT and PJRT
• diffusers Stable Diffusion fine-tuning example on TPU v4-8

Accelerate will not work on TPU v2 and v3 with this PR, because both of them use use multithreading due to TPU design constraints. I'll follow up with the remaining fixes for TPU v2 and v3 in #1385.

cc @sgugger @JackCaoG