Repository: pytorch/pytorch

Issue:
Multi-GPU tarining with DataParallel must has GPU0
The loss will never decrease when I train a model using multi-gpus but without gpu0. This is an example I modified from the [official mnist example](https://github.com/pytorch/examples/tree/master/mnist) : https://gist.github.com/longcw/e5efdb0e5a580a6b75604561b6e4f613

```python
gpus = [1, 2]
model = torch.nn.DataParallel(model, device_ids=gpus)
```

If `gpus = [1, 2]` the result after one epoch is 
```
...
Train Epoch: 1 [58240/60000 (97%)]      Loss: 2.363280
Train Epoch: 1 [58880/60000 (98%)]      Loss: 2.373188
Train Epoch: 1 [59520/60000 (99%)]      Loss: 2.351889

Test set: Average loss: 2.3285, Accuracy: 685/10000 (7%)
```

If `gpus = [0, 2]` the result after one epoch is 
```
...
Train Epoch: 1 [58240/60000 (97%)]	Loss: 0.268342
Train Epoch: 1 [58880/60000 (98%)]	Loss: 0.286652
Train Epoch: 1 [59520/60000 (99%)]	Loss: 0.514245

Test set: Average loss: 0.1972, Accuracy: 9436/10000 (94%)
```



Discussion:
`device_ids` should always contain `0`. If you want to specify which GPUs to use, you can set the environment variable `CUDA_VISIBLE_DEVICES`, either by running your script like this:
```
CUDA_VISIBLE_DEVICES=1,2 python your_script.py
```
or set it programmatically in your script like this:
```
os.environ['CUDA_VISIBLE_DEVICES'] = '1, 2'
```
So that device_ids 0, 1 will map to GPUs 1, 2, respectively.
(Note that for the second method, you should put that line before any cuda relevant function call.)

@allenye0119 Thank you.

Resolving patch:
diff --git a/torch/cuda/comm.py b/torch/cuda/comm.py
--- a/torch/cuda/comm.py
+++ b/torch/cuda/comm.py
@@ -74,22 +74,26 @@ def reduce_add(inputs, destination=None):
     """
     # TODO: try to find an input on another gpu, copy it,
     # and accumulate into the copy
+    if destination is None:
+        destination = torch.cuda.current_device()
     input_size = inputs[0].size()
+    nccl_root = None
     for i, inp in enumerate(inputs):
         assert inp.is_cuda, "reduce_add expects all inputs to be on GPUs"
+        if inp.get_device() == destination:
+            nccl_root = i
         if inp.size() != input_size:
             got = 'x'.join(str(x) for x in inp.size())
             expected = 'x'.join(str(x) for x in input_size)
             raise ValueError("input {} has invalid size: got {}, but expected "
                              "{}".format(i, got, expected))
-    if destination is None:
-        destination = torch.cuda.current_device()
+    assert nccl_root is not None, "reduce_add expects destination to be on the same GPU with one of the tensors"
     with torch.cuda.device(destination):
         result = type(inp)(input_size).zero_()
 
     if nccl.is_available(inputs) and inputs[0].get_device() == destination:
         outputs = [result] + [t.new(t.size()) for t in inputs[1:]]
-        nccl.reduce(inputs, outputs, root=destination)
+        nccl.reduce(inputs, outputs, root=nccl_root)
         return result
 
     for inp in inputs:


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -1426,13 +1426,27 @@ def test_data_parallel(self):
         l = nn.Linear(10, 5).float().cuda()
         i = Variable(torch.randn(20, 10).float().cuda(1))
         l.cuda(1)
-        expected_out = l(i).data
-        l.cuda(0)
-        out = dp.data_parallel(l, i, (0, 1))
-        self.assertEqual(out.get_device(), 0)
-        self.assertEqual(out.data, expected_out)
+        expected_out = l(i)
+        loss = expected_out.sum()
+        loss.backward()
+        expected_grads = []
+        for param in l.parameters():
+            expected_grads.append(param.grad.clone())
+        dev_ids_list = [(0, 1), (1, 0)]
+        for dev_id in dev_ids_list:
+            with torch.cuda.device(dev_id[0]):
+                l.cuda()
+                l.zero_grad()
+                out = dp.data_parallel(l, i, dev_id)
+                loss = out.sum()
+                loss.backward()
+                self.assertEqual(out.get_device(), dev_id[0])
+                self.assertEqual(out.data, expected_out.data)
+                for expected, param in zip(expected_grads, l.parameters()):
+                    self.assertEqual(param.grad.data, expected.data)
 
         # Check for None device_ids
+        l = l.cuda()
         out = dp.data_parallel(l, i)
 
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")