huggingface / accelerate

🚀 A simple way to launch, train, and use PyTorch models on almost any device and distributed configuration, automatic mixed precision (including fp8), and easy-to-configure FSDP and DeepSpeed support
https://huggingface.co/docs/accelerate
Apache License 2.0
7.32k stars 872 forks source link

Add XLA Dynamo backends for training and inference #2892

Open johnsutor opened 1 week ago

johnsutor commented 1 week ago

What does this PR do?

This PR introduces the TorchDynamo integration with PyTorch XLA. Though I haven't thorougly benchmarked the code, it works on the CLM example from the Transformers library, when spawned using the accelerate launch command on a v5e TPU VM with 8 TPUs running on Ubuntu 22.04. It also passes all tests I ran on the TPU.

Fixes #2870

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

@SunMarc

HuggingFaceDocBuilderDev commented 6 days ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.