Repository: pytorch
Issue: PicklingError when saving a ddp module with torch.save()
Labels: oncall: distributed

## 🐛 Bug

when saving a ddp module with torch.save(), unexpected picking errors occured.
my model uses encoder-decoder framework, and the encoder contains a BertModel from transformers(Huggingface)

## To Reproduce
Steps to reproduce the behavior:

1. using pytorch1.9, installed with command:
```bash
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia
```
2. using apex with O1 optimization level

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

3. The code logic is as follows：
```python
from torch.nn.parallel import DistributedDataParallel as DDP
my_model = MyModel()
ddp_model = DDP(my_model)
# after one epoch training
model_to_save= ddp_model.module
torch.save(model_to_save)
```

and then, the picking error occurs like:
```
  ......(error stack in my code)
  ......
  File "path_to_env/.conda/envs/pt19/lib/python3.8/site-packages/fastNLP/core/dist_trainer.py", line 426, in save_check_point
    torch.save(model_to_save, path)
  File "path_to_env/.conda/envs/pt19/lib/python3.8/site-packages/torch/serialization.py", line 379, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "path_to_env/.conda/envs/pt19/lib/python3.8/site-packages/torch/serialization.py", line 484, in _save
    pickler.dump(obj)
_pickle.PicklingError: Can't pickle <function gelu at 0x2abf3711a1f0>: it's not the same object as torch.nn.functional.gelu
```
gelu is using in BertModel.

<!-- A clear and concise description of what you expected to happen. -->

i want to save the model as expected without errors.

## Environment

```python
due to some reasons in server, the script collect_env.py cannot work correctly. But i think, it may not matters for this bug.

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.1
[pip3] torchaudio==0.9.0a0+a85b239
[pip3] torchvision==0.10.1
[conda] blas                      1.0                         mkl    
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] mkl                       2021.3.0           h06a4308_520    
[conda] mkl-service               2.4.0            py38h7f8727e_0   
[conda] mkl_fft                   1.3.0            py38h42c9631_2   
[conda] mkl_random                1.2.2            py38h51133e4_0    
[conda] numpy                     1.20.3           py38hf144106_0    
[conda] numpy-base                1.20.3           py38h74d4b33_0    
[conda] pytorch                   1.9.1           py3.8_cuda11.1_cudnn8.0.5_0    pytorch
[conda] torchaudio                0.9.1                      py38    pytorch
[conda] torchvision               0.10.1               py38_cu111    pytorch
```


 - PyTorch Version: 1.9.1
 - OS: CentOS 7
 - How you installed PyTorch (`conda`, `pip`, source):
```bash
   conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia
``` 
 - Python version: 3.8
 - CUDA/cuDNN version: CUDA/11.1.0
 - GPU models and configuration: Tesla V100(16GB)
 - Any other relevant information: using apex 0.1(fully build) 

## Additional information:
if not using ddp, just run and save the model in single GPU, torch.save() can work correctly.
i checked other issues with picking error, this  #37703 may be relevant, but solution is unavailable in my case.(https://github.com/pytorch/pytorch/issues/37703#issuecomment-693551398)


cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang