Neuron distributed training validation

[michaelbenayoun]

What does this PR do?

The goal of this PR is to experiment and validate the new distributed training API on Neuron hardware. We check that:

1. The save/reload process works fine
2. The model is learning well

FSDP v2

Training run:

Save correctness check passed: 311 parameters match the unsharded checkpoint exactly.

{'loss': '1.248', 'grad_norm': '0', 'learning_rate': '0.00041', 'epoch': '10'}
{'loss': '0.02439', 'grad_norm': '0', 'learning_rate': '0.00031', 'epoch': '20'}
{'loss': '0.005542', 'grad_norm': '0', 'learning_rate': '0.00021', 'epoch': '30'}
{'loss': '0.004281', 'grad_norm': '0', 'learning_rate': '0.00011', 'epoch': '40'}
{'loss': '0.004143', 'grad_norm': '0', 'learning_rate': '1e-05', 'epoch': '50'}

{'train_runtime': '272.6', 'train_samples_per_second': '2.935', 'train_steps_per_second': '0.183', 'train_loss': '0.2572', 'epoch': '50'}

It works fine.

Tensor Parallelism

Training run:

Save correctness check passed: 311 parameters match the unsharded checkpoint exactly.

{'loss': '3.076', 'grad_norm': '0', 'learning_rate': '0.00041', 'epoch': '2.5'}
{'loss': '0.7326', 'grad_norm': '0', 'learning_rate': '0.00031', 'epoch': '5'}
{'loss': '0.08189', 'grad_norm': '0', 'learning_rate': '0.00021', 'epoch': '7.5'}
{'loss': '0.02404', 'grad_norm': '0', 'learning_rate': '0.00011', 'epoch': '10'}
{'loss': '0.01477', 'grad_norm': '0', 'learning_rate': '1e-05', 'epoch': '12.5'}
{'train_runtime': '93.97', 'train_samples_per_second': '2.128', 'train_steps_per_second': '0.532', 'train_loss': '0.7859', 'epoch': '12.5'}

It works fine.

Tensor Parallelism & LoRA

Currently integrating the new API in huggingface/peft#3614 & #48485.

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

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33526605282:1
Result: failure | Jobs: 16 | Tests: 185,106 | Failures: 7 | Duration: 15h 16m