Repository: pytorch/pytorch

Issue:
nn.DataParallel ignores requires_grad setting when running
- PyTorch version: 0.3.0
- Python version: 3.5
- GPU models and configuration: 8x NVIDIA Titan

I found that requires_grad setting in module is ignored when module is wrapped with nn.DataParallel. For example,

```python
import torch
from torch.nn import DataParallel
from torchvision.models.resnet import resnet50

module_ = resnet50()
for name, param in module_.named_parameters():
    if name.startswith('conv1') or name.startswith('bn1'):
        param.requires_grad = False
    if name.startswith('layer1') or name.startswith('layer2'):
        param.requires_grad = False
    if name.startswith('layer3') or name.startswith('layer4'):
        param.requires_grad = False

module_ = DataParallel(module_).cuda()

x = torch.rand(32, 3, 1600, 1600)
x = torch.autograd.Variable(x)
x = module_(x)
print(x)
```
This code should be executed because very small part of network activation (only fc layer) is stored for backward computation, but it results to run-time error:

```
RuntimeError: cuda runtime error (2) : out of memory at /tmp/pip-8pfswvat-build/torch/lib/THC/generic/THCStorage.cu:58
```

Internally, I found a crack in ```replicate``` function which is in ```torch.nn.parallel.replicate```. In replicate function, it copies all parameter in module (# of replica times) with ```Broadcast.apply```. In broadcasting code, it just defines new ```torch.nn.Parameter``` with default constructor ```requires_grad``` parameter, which is always set to True.

I think there can be some choices to fix this issue.

1. It is intended behavior for DataParallel, so we should use volatile (or torch.no_grad) to implement transfer learning in DataParallel
2. Broadcast model should be fixed to handle requires_grad when copying parameters.
3. There should be synchronizing ```requires_grad``` in ```torch.nn.parallel.replicate```

I think any of solution doesn't need lots of effort to fix it.



Discussion:

Resolving patch:
diff --git a/torch/nn/parallel/_functions.py b/torch/nn/parallel/_functions.py
--- a/torch/nn/parallel/_functions.py
+++ b/torch/nn/parallel/_functions.py
@@ -15,6 +15,12 @@ def forward(ctx, target_gpus, *inputs):
         ctx.num_inputs = len(inputs)
         ctx.input_device = inputs[0].get_device()
         outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
+        non_differentiables = []
+        for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):
+            if not input_requires_grad:
+                for output in outputs:
+                    non_differentiables.append(output[idx])
+        ctx.mark_non_differentiable(*non_differentiables)
         return tuple([t for tensors in outputs for t in tensors])
 
     @staticmethod
diff --git a/torch/nn/parallel/data_parallel.py b/torch/nn/parallel/data_parallel.py
--- a/torch/nn/parallel/data_parallel.py
+++ b/torch/nn/parallel/data_parallel.py
@@ -26,6 +26,11 @@ class DataParallel(Module):
     other types will be a shallow copy and can be corrupted if written to in
     the model's forward pass.
 
+    .. warning::
+        Forward and backwrad hooks defined on :attr:`module` and its submodules
+        won't be invoked anymore, unless the hooks are initialized in the
+        :meth:`forward` method.
+
     Args:
         module: module to be parallelized
         device_ids: CUDA devices (default: all devices)
diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py
--- a/torch/nn/parallel/distributed.py
+++ b/torch/nn/parallel/distributed.py
@@ -78,6 +78,11 @@ class DistributedDataParallel(Module):
         (e.g. BatchNorm stats) are broadcast form the module in process of rank
         0, to all other replicas in the system in every iteration.
 
+    .. warning::
+        Forward and backwrad hooks defined on :attr:`module` and its submodules
+        won't be invoked anymore, unless the hooks are initialized in the
+        :meth:`forward` method.
+
     Args:
         module: module to be parallelized
         device_ids: CUDA devices (default: all devices)


Test patch:
diff --git a/test/common.py b/test/common.py
--- a/test/common.py
+++ b/test/common.py
@@ -237,6 +237,8 @@ def assertTensorsEqual(a, b):
             super(TestCase, self).assertEqual(len(x), len(y), message)
             for x_, y_ in zip(x, y):
                 self.assertEqual(x_, y_, prec, message)
+        elif isinstance(x, bool) and isinstance(y, bool):
+            super(TestCase, self).assertEqual(x, y, message)
         elif isinstance(x, Number) and isinstance(y, Number):
             if abs(x) == float('inf') or abs(y) == float('inf'):
                 if allow_inf:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -1596,6 +1596,20 @@ def test_broadcast_double_backwards_gpu(self):
                                               torch.randn(4, 4).cuda(),
                                               torch.randn(4, 4).cuda())
 
+    @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
+    def test_broadcast_not_requiring_grad(self):
+        variables = [
+            Variable(torch.randn(1, 2).cuda(), requires_grad=True),
+            Variable(torch.randn(1, 2).cuda(), requires_grad=False),
+            Variable(torch.randn(1, 2).cuda(), requires_grad=False),
+            Variable(torch.randn(1, 2).cuda(), requires_grad=True),
+            Variable(torch.randn(1, 2).cuda(), requires_grad=True),
+        ]
+        broadcasted_variables = Broadcast.apply((0, 1), *variables)
+        for output_idx, broadcasted_var in enumerate(broadcasted_variables):
+            input_var = variables[output_idx % len(variables)]
+            self.assertEqual(input_var.requires_grad, broadcasted_var.requires_grad)
+
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
     def test_replicate(self):
         module = nn.Linear(10, 5).float().cuda()