Repository: pytorch/PiPPy

Issue:
local_test_ddp + CUDA failure
Log [here](https://gist.github.com/kwen2501/1dc2bee00e5e3bcc56ec9c28174b5d1d)
```
RuntimeError: DDP Pickling/Unpickling are only supported when using DDP with the default process group. That is, when you have called init_process_group and have not passed process_group argument to DDP constructor
```
cc @mrshenli 


Discussion:

Resolving patch:
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -399,6 +399,13 @@ def __init__(self, stage_id, mod, rank_worker, dp_pg_cb, pp_rank):
 
         self.peer_executors : Dict[int, torch._C._distributed_rpc.PyRRef] = None
 
+    def __getstate__(self):
+        # Adding an empty __getstate__ function here to work around the DDP pickling issue (#153) that occurs when the
+        # PipelineDiver asks PipeStageExecutors to install_peer_executor(a list of RRefs)
+        # More elegant solution is needed in CUDAFuture or RPC to avoid pickling when users do not need to transfer
+        # tensors
+        pass
+
     def install_peer_executors(self, peer_executors):
         assert self.peer_executors is None
         self.peer_executors = peer_executors


Test patch:
diff --git a/test/local_test_ddp.py b/test/local_test_ddp.py
--- a/test/local_test_ddp.py
+++ b/test/local_test_ddp.py
@@ -60,7 +60,6 @@ def run_master(args, pp_ranks):
     bs = 503
     CHUNKS = 5
     DEBUG_MASK_MINIBATCHES = True
-    REPLICATE = os.environ.get('REPLICATE', '0') != '0'
     MULTI_USE_PARAM_CONFIG = MultiUseParameterConfig.REPLICATE if args.replicate else MultiUseParameterConfig.TRANSMIT
     print(f'REPLICATE config: {args.replicate} -> {MULTI_USE_PARAM_CONFIG}')
 
@@ -100,7 +99,8 @@ def forward(self, x):
             return x
 
     ec = ExampleCode()
-    ec(torch.randn(bs, d_hid))
+    ec.to(args.device)
+    ec(torch.randn(bs, d_hid, device=args.device))
     ec.train()
 
     # TODO: works with sum, need to define semantics for e.g. mean
@@ -121,8 +121,8 @@ def forward(self, x):
     print(f'Rank {args.rank} Instantiated pipe with ranks {pp_ranks}')
 
     torch.manual_seed(args.rank)
-    input = torch.randn(bs, d_hid)
-    target = torch.randn(bs, d_hid)
+    input = torch.randn(bs, d_hid, device=args.device)
+    target = torch.randn(bs, d_hid, device=args.device)
 
     # TODO: distributed optimizer
     out = pipe_driver.run(CHUNKS, input, target)
@@ -166,13 +166,31 @@ def forward(self, x):
     if len(not_close_grads):
         raise AssertionError(f'Gradients not close: {not_close_grads}')
 
+    print('Gradient equivalence test passed')
+
 
 def run_worker(rank, world_size, args):
-    print(f"rank = {rank} host/pid = {socket.gethostname()}/{os.getpid()}")
     os.environ['MASTER_ADDR'] = args.master_addr
     os.environ['MASTER_PORT'] = args.master_port
-    torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
-    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256)
+    # Exclude IB for metadata transport due to lack of EFA support on AWS
+    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256,
+                                              _transports=["shm", "uv"])
+    if args.cuda:
+        n_devs = torch.cuda.device_count()
+        if n_devs > 0:
+            dev_id = rank % n_devs
+            for i in range(world_size):
+                options.set_device_map(f"worker{i}", {dev_id: i % n_devs})
+        else:
+            args.cuda = 0
+    args.device = f'cuda:{dev_id}' if args.cuda else 'cpu'
+    print(f"rank = {rank} host/pid/device = "
+          f"{socket.gethostname()}/{os.getpid()}/{args.device}")
+
+    # Init DDP process group
+    backend = "nccl" if args.cuda else "gloo"
+    torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size)
+
     rpc.init_rpc(
         f"worker{rank}",
         rank=rank,
@@ -218,6 +236,7 @@ def run_worker(rank, world_size, args):
     parser.add_argument('--master_port', type=str, default=os.getenv('MASTER_PORT', '29500'))
     parser.add_argument('-s', '--schedule', type=str, default=list(schedules.keys())[0], choices=schedules.keys())
     parser.add_argument('--replicate', type=int, default=int(os.getenv("REPLICATE", '0')))
+    parser.add_argument('--cuda', type=int, default=int(torch.cuda.is_available()))
     args = parser.parse_args()
 
     assert args.dp_group_size * args.pp_group_size == args.world_size