Fix training

[michaelbenayoun]

What does this PR do?

• Updates the NeuronTrainer to match Transformers version
• Fixes the way mixed-precision training is handled. It was breaking multiple trainings such as Llama. We do not use XLA_USE_BF16 and XLA_DOWNCAST_BF16 flags anymore following instructions from here
• Fixes support for gradient clipping: now it always happens between the reduction of the gradients across the devices and the optimizer step. The gradient norm is also now always reported when logging.
• Adds support for ignore_index in the parallel_cross_entropy_loss. There was a big issue in training when using TP, the model was not learning. After investigating, it was linked to the input being padded and the vanilla parallel_cross_entropy from neuronx_distributed not supporting ignore_index:
  • First, loss.mean() does not work in this case because the loss for the ignored tokens is not zeroed.
  • Second, the ignored tokens contributed to the gradient, which effectively destroys training.

For now DP + TP can lead to compilation issues with SDK 2.20, but they seem to be gone with SDK 2.21.

Tests performed

• Llama (HuggingFaceTB/SmolLM2-135M-Instruct) can overfit with dp=1 tp=1
• Llama (HuggingFaceTB/SmolLM2-135M-Instruct) + LoRA can overfit with dp=1 tp=1
• Llama (meta-llama/Llama-3.2-1B) can overfit with dp=1 tp=2
• Llama + LoRA (meta-llama/Llama-3.2-1B) can overfit with dp=1 tp=2
• Llama (meta-llama/Llama-3.2-1B) can overfit with dp=4 tp=2 (Only tested on SDK 2.21, otherwise compiler error)
• Actual training of Llama (meta-llama/Llama-3.2-1B) dp=4 tp=2 on SDK 2.21 and compared to GPUs

To be done in following PRs

• Add AdamW_FP32Params for a more stable training in mixed-precision

[tengomucho]

I find difficult to understand what your changes have accomplished, can you give more details about that please?
Also, would you mind pointing to a test (or example) that now works after your changes?

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

[dacorvo]

Several failures in the CI:

• style check,
• errors in the training code directly related to some of the changes (_PARALLEL_CROSS_ENTROPY_SHOULD_PRESERVE_INPUT is not found).

Maybe this should be rebased on the 2.21.1 branch once it is merged.

[JingyaHuang]

Not too much to comment about, kudos for finding the fix!!

(will approve when the CIs pass)

[JingyaHuang]

What about raising an error when the version>= 0.10.0 (to not forget removing it).

[michaelbenayoun]

Several failures in the CI:

• style check,
• errors in the training code directly related to some of the changes (_PARALLEL_CROSS_ENTROPY_SHOULD_PRESERVE_INPUT is not found).

Maybe this should be rebased on the 2.21.1 branch once it is merged.

I am working on fixing the CI failures.
If you want you can merge the 2.21.1 branch, then I will rebase and adapt: some changes in the PR are linked to 2.20, so if we are going to move to 2.21 very soon it does not make sense to add them.

[dacorvo]

I bumped the SDK version to 2.21.1: you can now rebase your branch and drop the 2.20 specifics.