Neuron integration

[michaelbenayoun]

What does this PR do?

Add support for AWS Trainium chips (Neuron Cores).

Missing integrations:

• src/accelerate/hooks.py

[HuggingFaceDocBuilderDev]

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.

[SunMarc]

Thanks for the integration ! Can you fix the merge conflict ?

[michaelbenayoun]

Done

[BrownianNotion]

@michaelbenayoun awesome work! I'm not sure I understand the relationship between accelerate's Trainium support and optimum-neuron, would you be able to help me out?

Is the idea that optimum-neuron is a higher-level wrapper, and is the trainium drop-in replacement for TRL, while accelerate only handles parallelism and allows finer-grained control of the training loop? What made trainium sufficiently different such that it received its own library, as opposed to being part of TRL? Are the current plans to keep the two separate or will they be unified in future?

Thank you so much!

[michaelbenayoun]

Trainium now has native support in PyTorch, which was not the case before when it was built on top of torch-xla. It is the main motivation behind moving from a standalone library (optimum-neuron) to integration inside the various libraries we have at HF.