find_unused_parameters is not passed from Trainer to Sagemaker DistributedModel
Author: marcuscollins
Pull request: no

### System Info

When trying to launch a model on SageMaker with the Huggingface Estimator and the `transformers.Trainer` class, I discovered the the trainer argument `ddp_find_unused_parameters` is not passed to the Sagemaker DistributedModel, see https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/trainer.py#L1336

It is possible to work around this by wrapping the model with DistributedModel before passing it to the Trainer, so that I can pass any arguments I want, but really the trainer argument ought to just work.

I'm running fine tuning of a MaskedLM model, using a stripped down version of the example run_mlm.py script.

### Who can help?

@sgugger

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Pass the argument `ddp_find_unused_parameters = True` to the Trainer class in a SageMaker Model Parallel environment, when training one of the ESM models (e.g., Facebook/esm2_t33_650M_UR50D) loaded using on sequence data, e.g., agemagician/uniref30. 
1. In the training script, e.g., `run_mlm.py`, load a masked LM model, e.g. `AutoModelForMaskedLM("Facebook/esm2_t33_650M_UR50D")`
2. Use a dataset like agemagician/uniref30, by downloading the files to Sagemaker via its "data channels". Load it in the training script using `datasets.load_dataset` with the data_files argument.

Then use a Huggingface estimator in a SageMaker notebook, e.g. this example from AWS: https://github.com/PacktPublishing/Applied-Machine-Learning-and-High-Performance-Computing-on-AWS/blob/main/Chapter12/protein-secondary-structure-model-parallel.ipynb, the model training should crash indicating that there are unreduced parameters, telling you to use `find_unused_parameters`:

```
[1,mpirank:0,algo-1]<stderr>:RuntimeError: Expected to have finished reduction in the prior iteration before 
[1,mpirank:0,algo-1]<stderr>:starting a new one. This error indicates that your module has parameters that 
[1,mpirank:0,algo-1]<stderr>:were not used in producing loss. You can enable unused parameter detection by 
[1,mpirank:0,algo-1]<stderr>:passing the keyword argument `find_unused_parameters=True` to 
[1,mpirank:0,algo-1]<stderr>:`torch.nn.parallel.DistributedDataParallel`, and by 
[1,mpirank:0,algo-1]<stderr>:making sure all `forward` function outputs participate in calculating loss. 
[1,mpirank:0,algo-1]<stderr>:If you already have done the above, then the distributed data parallel module 
[1,mpirank:0,algo-1]<stderr>:wasn't able to locate the output tensors in the return value of your module's 
[1,mpirank:0,algo-1]<stderr>:`forward` function. Please include the loss function and the structure of the 
[1,mpirank:0,algo-1]<stderr>:return value of `forward` of your module when reporting this issue (e.g. list, 
[1,mpirank:0,algo-1]<stderr>:dict, iterable).
[1,mpirank:0,algo-1]<stderr>:Parameter indices which did not receive grad for rank 0: 1 132 133
[1,mpirank:0,algo-1]<stderr>: In addition, you can set the environment variable TORCH_DISTRIBUTED_DEBUG to 
[1,mpirank:0,algo-1]<stderr>:either INFO or DETAIL to print out information about which particular parameters
[1,mpirank:0,algo-1]<stderr>:did not receive gradient on this rank [1,mpirank:0,algo-1]<stderr>:as part of this error
```

3. Now, using a sagemaker Huggingface estimator, pass the hyper-parameter `{"ddp_find_unused_parameters": True,...}` to the estimator. If using the script `run_mlm.py`, this will be parsed and passed to the `Trainer` class as part of `training_args`. However, you should still see the same error. This is because, as noted above, the argument is not passed to `smp.DistributedModel`. 

### Expected behavior

The argument should be passed to `DistributedModel`, so that it can resolve the error noted above.

Comments:
Update: the workaround does not appear to solve the problem either. We need some way to pass this argument into SageMaker's version of parallelism (which I think still uses torch's DistributedDataParallel under the hood.)

Possibly @pacman100 might know about this? 

Are we sure this SageMaker class actually supports this argument with the same name as PyTorch?

I just wanted to add that this is a problem with [run.plm](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_plm.py) as well, but the example for [question-answer](https://github.com/huggingface/notebooks/blob/main/sagemaker/03_distributed_training_data_parallelism/sagemaker-notebook.ipynb) does work.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.