Revamp TPU internals to be more efficient

[muellerzr]

Revamp TPU internals

What does this add?

• Make prepare_model use MpModelWrapper to distribute the model across all devices efficiently when used
• Changes prepare_dataloader to create an MpDeviceLoader allowing for the dataloaders to be more efficient.
• Changes DataLoaderShard to no longer do xm.mark_step, MpDeviceLoader will handle this for us. Instead if on TPU we set the device as None in prepare_dataloader letting MpDeviceLoader take over with the device when needed.
• Allow for FP16 and BF16 precision types on the TPU via AcceleratorState since they are now supported

Who is it for?

• Users of TPUs

Why is it needed?

We currently have a number of "bad practices" in TPU handling due to improvements in the xla API. Here are some benchmarks I ran on a high-memory colab TPU instance with the nlp example script:

Baseline:

• Post warmup: 56 seconds

W/ MpModelWrapper and MpDeviceLoader:

• Post warmup: 48 seconds

W/ previous and default_tensor_type change:

• Post warmup: 34 to 36 seconds

Roughly a 40% boost to speed once this is all done. I also saw some speed increases on the initial launch as well anecdotally, but I did not time them.

I also saw a 2x speedup when using the new DataLoader class vs the old one

About default_tensor_type:

Though we do set the device to bf16 which helps with automatically converting the tensors to the right types, if we add torch.set_default_tensor_type('torch.FloatTensor') there is a considerable speedup when it comes to training TPUs.

I'm a bit unsure as to where it would be best to put this, as either something hidden when we initialize the Accelerator or as a util that should get called when you are training on TPUs, open to ideas

Anticipated maintenance burden? (What will happen in say, 3 months if something changes)

This is pretty stable and considered as "good practices" when running on TPUs w/ XLA, so it's unlikely these will change. However after this PR a subsequent PR will be opened to change the nlp_example and cv_example notebook, as another best practice is to declare the model outside of xm.spawn. This includes the internals to make the model as memory efficient as we can, so that will be the last stage needed.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[tmabraham]

Thanks for fixing up the TPU support with best practices! Looks great to me! 🔥

[sgugger]

Thanks for your PR and running the benchmarks in Colab. Could you link to the documentation that shows those are best practices recommended by the torch XLA team? I haven't seen anything personally.

I'd like to make sure those changes are not speeding up the experience in Colab at the cost of the experience on TPU machines, so I'd like the same benchmarks to be run on a TPU VM before fully approving :-)

[muellerzr]

@sgugger this PR also modifies the TPU selection in AcceleratorState to actually use mixed precision types.

• FP16 Support
• BF16 Support

[muellerzr]

We're passing 😄

[sgugger]

Thanks for revamping this!