Repository: pytorch
Issue: torch.multiprocessing.spawn raise PicklingError inside a decorator
Labels: module: multiprocessing,module: serialization,triaged

### 🐛 Describe the bug

I wrote a decorator to simplify the process of launching multiprocessing which takes a function as an argument and calls torch.multiprocessing.spawn to spawn multiple processes that runs the input function. 

However, I met a quite odd situation. When calling the decorator directly as a function factory (`print_rank_B` in the script below), `torch.multiprocessing.spawn` works as expected and the input function indeed runs multiple times. But when I try to run a function that is decorated by this decorator (`print_rank_C` in the script below), a `PicklingError` is raised.

I understand that this `PicklingError` sometimes can be caused by inconsistent names of functions (i.e. `__name__`) so I used `wraps` from `functools` library, but the error still exists.

I'm using torch@1.10.0 with Python 3.8.13.

Thanks in advance!

The script to reproduce this error:
```python
import random
from functools import wraps
import torch.multiprocessing
import torch.distributed as dist

# from .utils.distributed import get_rank


def get_rank():
    if dist.is_initialized():
        return dist.get_rank()
    else:
        return 0

def _distributed_init(id, world_size, init_method, fn, args):
    dist.init_process_group(backend='nccl', init_method=init_method, world_size=world_size, rank=id)
    fn(*args)


def distributed(world_size):
    if world_size > 1:
        def decorator(fn):
            @wraps(_distributed_init)
            def wrapper(*args):
                port = random.randint(10000, 20000)
                init_method = f'tcp://localhost:{port}'
                torch.multiprocessing.spawn(fn=_distributed_init, args=(world_size, init_method, fn, args), nprocs=world_size)
            return wrapper
    else:
        def decorator(fn):
            return fn
    return decorator


def run_distributed(world_size, fn, *args):
    port = random.randint(10000, 20000)
    init_method = f'tcp://localhost:{port}'
    if world_size > 1:
        torch.multiprocessing.spawn(fn=_distributed_init, args=(world_size, init_method, fn, args), nprocs=world_size)
    else:
        fn(*args)


def print_rank():
    print(get_rank())


def print_rank_A():
    run_distributed(4, print_rank)


def print_rank_B():
    distributed(4)(print_rank)()


@distributed(4)
def print_rank_C():
    print('hi')


if __name__ == '__main__':
    print('A:')
    print_rank_A()
    print('B:')
    print_rank_B()
    print('C:')
    print_rank_C()
```

The output of the script:
```shell
A:
0
3
1
2
B:
0
1
2
3
C:
Traceback (most recent call last):
  File "/home/tb5zhh/SUField/sufield/lib/distributed_launch.py", line 68, in <module>
    print_rank_C()
  File "/home/tb5zhh/SUField/sufield/lib/distributed_launch.py", line 28, in wrapper
    torch.multiprocessing.spawn(fn=_distributed_init, args=(world_size, init_method, fn, args), nprocs=world_size)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 179, in start_processes
    process.start()
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/home/tb5zhh/.conda/envs/sufield/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <function print_rank_C at 0x7f432e3f6700>: it's not the same object as __main__.print_rank_C
```


### Versions

Collecting environment information...
PyTorch version: 1.10.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.13.0-39-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.10.0
[pip3] torchaudio==0.10.0
[pip3] torchvision==0.11.0
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               11.3.1               h2bc3f7f_2    defaults
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640    defaults
[conda] mkl-service               2.4.0            py38h7f8727e_0    defaults
[conda] mkl_fft                   1.3.1            py38hd3c417c_0    defaults
[conda] mkl_random                1.2.2            py38h51133e4_0    defaults
[conda] numpy                     1.21.5           py38he7a7128_2    defaults
[conda] numpy-base                1.21.5           py38hf524024_2    defaults
[conda] pytorch                   1.10.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.10.0               py38_cu113    pytorch
[conda] torchvision               0.11.0               py38_cu113    pytorch

cc @VitalyFedyunin @mruberry