trainer fails when fsdp = full_shard auto_wrap
Author: chijames
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.20.0.dev0
- Platform: Linux-5.4.0-1072-aws-x86_64-with-debian-buster-sid
- Python version: 3.7.12
- Huggingface_hub version: 0.7.0
- PyTorch version (GPU?): 1.13.0.dev20220610 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes
```


### Who can help?

@sgugger @patric

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```bash
torchrun --nproc_per_node=4 \
    run_summarization.py \
    --model_name_or_path google/pegasus-large \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size=1 \
    --per_device_eval_batch_size=1 \
    --overwrite_output_dir \
    --predict_with_generate \
    --fsdp "full_shard auto_wrap" \
    --fsdp_min_num_params 20000
```

Running the above script will generate the following error:
`File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 1242, in _wrap_model
    from torch.distributed.fsdp.wrap import default_auto_wrap_policy
ImportError: cannot import name 'default_auto_wrap_policy' from 'torch.distributed.fsdp.wrap'`

A little bit digging into the torch/distributed/fsdp/wrap.py shows default_auto_wrap_policy is no longer in the file. I tried to change it to size_based_auto_wrap_policy as it seems to have the same function signature. Unfortunately, another error pops up:
`File "run_summarization.py", line 734, in <module>                                                                                                         
    main()
  File "run_summarization.py", line 653, in main
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 1610, in _inner_training_loop
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 1372, in train
    tr_loss_step = self.training_step(model, inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 2301, in training_step
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 1610, in _inner_training_loop
    loss = self.compute_loss(model, inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 2333, in compute_loss
    tr_loss_step = self.training_step(model, inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 2301, in training_step
    outputs = model(**inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    loss = self.compute_loss(model, inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/trainer.py", line 2333, in compute_loss
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2303, in forward
    outputs = model(**inputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    outputs = self._fsdp_wrapped_module(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2303, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    return self.module(*inputs, **kwinputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    outputs = self._fsdp_wrapped_module(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 1414, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    return self.module(*inputs, **kwinputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return_dict=return_dict,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 1414, in forward
return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    return self.module(*inputs, **kwinputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return_dict=return_dict,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 1414, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 1245, in forward
    return_dict=return_dict,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return_dict=return_dict,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
        return forward_call(*input, **kwargs)return forward_call(*input, **kwargs)

  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 1245, in forward
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2303, in forward
    return_dict=return_dict,
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    outputs = self._fsdp_wrapped_module(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2303, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    return self.module(*inputs, **kwinputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    outputs = self._fsdp_wrapped_module(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 761, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 476, in forward
    inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return self.module(*inputs, **kwinputs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/sparse.py", line 160, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/functional.py", line 2156, in embedding
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/transformers/models/pegasus/modeling_pegasus.py", line 761, in forward
    inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
  File "/home/ubuntu/anaconda3/envs/transformers/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Output 0 of ViewBackward0 is a view and its base or another view of its base has been modified inplace. This view is the output of a function that retur
ns multiple views. Such functions do not allow the output views to be modified inplace. You should replace the inplace operation by an out-of-place one.`

This time I have no idea how the problem should be solved.

Any help is greatly appreciated! Thanks.

### Expected behavior

```shell
The script should run without errors when fsdp is enabled.
```


Comments:
cc @pacman100 

@pacman100 Thanks for looking into this! Would love to provide any additional information!

Hello @chijames, thanks for letting us know that `default_auto_wrap_policy` is no more, will be fixing it shortly. Regarding the subsequent error, it is unrelated to the integration and I have opened issue in PyTorch repo for the same and mentioned this issue in that issue as seen above.   


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.

> Hello @chijames, thanks for letting us know that `default_auto_wrap_policy` is no more, will be fixing it shortly. Regarding the subsequent error, it is unrelated to the integration and I have opened issue in PyTorch repo for the same and mentioned this issue in that issue as seen above.

One way to get the **default_auto_wrap_policy** is to get Nvidia's docker nvcr.io/nvidia/pytorch:**22.05-py3**

The definition of **default_auto_wrap_policy** is in /opt/pytorch/pytorch/torch/distributed/fsdp/wrap.py:31