Multi-node training support

[michaelbenayoun]

This PR adds support for multi-node training.

• Check training TP + PP on multiple nodes with llama-7b
• Check training TP + PP on multiple nodes with llama-70b
• Check that files are pushed to the remote cache properly (with and without custom remote cache)
• Test evaluation with PP on multiple nodes
• Check that the logged loss and metrics are correct

In coming PRs:

• Add better support for gradient checkpointing (full vs feature to specify which layers to use)
• Handle GQA better with the custom layer for neuronx_distributed

[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.

[JingyaHuang]

Multi node is huge, thanks for the work 🎊 . The PR looks good to me, just left some minor nits. Wonder if @michaelbenayoun if it's possible to test the least functionality with CIs or other paths(I know setting up for multiple nodes is tricky...

[JingyaHuang]

No need for quote neither

[michaelbenayoun]

Here it is needed since if neuronx_distributed is not available it will not be imported.

[michaelbenayoun]

We test settings with TP + PP, but testing the multi-node setting is more tricky, it requires to create a node cluster, etc. For now let's keep it like that.

[dacorvo]

I only looked at the part that is related to the cache: I would rather move the import inside the utils to avoid multiple imports.

[michaelbenayoun]

Merging, tested on custom instances.

[shamanez]

@michaelbenayoun can we conduct multi-node training with Mistral?

[michaelbenayoun]

Yes it should be!
We haven't run many experiments for now, but support is here.