Add AWS Neuron torchrun support
Author: jeffhataws
Pull request: yes

# What does this PR do?

This PR adds support for torchrun for AWS Neuron SDK.

Existing [HF tutorial](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/tutorials/training/finetune_hftrainer.html) for Neuron SDK requires users to modify the HF example script (ie run_glue.py). This change will help minimize the changes required.

This change will require future AWS Neuron PyTorch 1.13 support.

This is an update to https://github.com/huggingface/transformers/pull/19907 .

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


## 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.

@sgugger

Comments:
_The documentation is not available anymore as the PR was closed or merged._

@jeffhataws could you maybe please explain a bit more about how users would benefit from that? I quickly checked the [HF tutorial](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/tutorials/training/finetune_hftrainer.html) and with the change you propose users would still need to modify the scripts, e.g., for 
```python
# Fixup to enable distributed training with XLA
from packaging import version
from transformers import __version__
if version.parse(__version__) < version.parse("4.20.0"):
    Trainer._wrap_model = lambda self, model, training=True: model
else:
    Trainer._wrap_model = lambda self, model, training=True, dataloader=None: model

# Workaround for NaNs seen with transformers version >= 4.21.0
# https://github.com/aws-neuron/aws-neuron-sdk/issues/593
if os.environ.get("XLA_USE_BF16") or os.environ.get("XLA_DOWNCAST_BF16"):
    transformers.modeling_utils.get_parameter_dtype = lambda x: torch.bfloat16
```


> Thanks for adding this new integration. The test won't be run on our CI since `torch_neuroncore` is not installed. Is it possible to install it in regular images or do we need to be on an AWS instance>

Yes for this test we will need Trainium instance. Over time, once https://github.com/pytorch/xla/pull/3609 is released, we can make it more generic for GPU/XLA. For now, Neuron team will test this. Test is currently passing on Trainium instance.

> @jeffhataws could you maybe please explain a bit more about how users would benefit from that? I quickly checked the [HF tutorial](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/tutorials/training/finetune_hftrainer.html) and with the change you propose users would still need to modify the scripts, e.g., for
> 
> ```python
> # Fixup to enable distributed training with XLA
> from packaging import version
> from transformers import __version__
> if version.parse(__version__) < version.parse("4.20.0"):
>     Trainer._wrap_model = lambda self, model, training=True: model
> else:
>     Trainer._wrap_model = lambda self, model, training=True, dataloader=None: model
> 
> # Workaround for NaNs seen with transformers version >= 4.21.0
> # https://github.com/aws-neuron/aws-neuron-sdk/issues/593
> if os.environ.get("XLA_USE_BF16") or os.environ.get("XLA_DOWNCAST_BF16"):
>     transformers.modeling_utils.get_parameter_dtype = lambda x: torch.bfloat16
> ```

The first workaround is for missing DDP support which will be available in Neuron's PyTorch-XLA version 1.13 (future release). The second workaround is already fixed in transformers==4.25.1 by https://github.com/huggingface/transformers/pull/20562.

Thanks for the precisions. Let's wait until the release of Neuron's PyTorch-XLA version 1.13 to merge this, then?

> Thanks for the precisions. Let's wait until the release of Neuron's PyTorch-XLA version 1.13 to merge this, then?

@sgugger  since we already have a workaround for DDP wrapper by overwriting the _wrap_model function, we can actually merge this first. The reason is that 1) we want it in for next transformer release ahead of 1.13, and 2) I will need this change to post another PR for the default compiler flag for transformer model type. Let me know if this is acceptable.

Thanks for your patience on this.