Repository: pytorch/pytorch

Issue:
cuDNN RNNs silently break weight tying in __init__
While the fast path for cuDNN RNNs introduced in https://github.com/pytorch/pytorch/pull/2179 includes a check intended to fall back to copying weights if weights are moved or aliased, this only works if the aliasing happens after the weights are first compacted. But the most natural place to tie weights is inside a module's `__init__` method, which runs before `.cuda()`, the typical time when `flatten_parameters` is first invoked.

Here's an MWE. The three runs (CPU, weight tying before `.cuda()`, and weight tying after `.cuda()`) should all give the same results, and they do before #2179, but they have not since:

```python
import torch

for cuda in (False, True):
    print('CUDA:', cuda)
    for when in ('before', 'after') if cuda else ('before',):
        if cuda: print('tying', when, '.cuda()')
        torch.manual_seed(1)
        rnn = torch.nn.LSTM(1, 1, num_layers=1, bidirectional=True, dropout=0)
        if when == 'before': rnn.bias_ih_l0_reverse = rnn.bias_ih_l0
        inputs = torch.autograd.Variable(torch.rand(1, 1, 1))
        if cuda:
            inputs = inputs.cuda(0)
            rnn.cuda(0)
        if when == 'after': rnn.bias_ih_l0_reverse = rnn.bias_ih_l0
        opt = torch.optim.SGD(rnn.parameters(), lr=0.1)

        opt.zero_grad()
        output = rnn(inputs)[0]
        print('iter 1 output:', output.data[0, 0, 0])
        output.sum().backward()

        opt.step()
        output = rnn(inputs)[0]
        print('iter 2 output:', output.data[0, 0, 0])
```


Discussion:

Resolving patch:
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -66,6 +66,15 @@ def flatten_parameters(self):
             self._data_ptrs = []
             return
 
+        # If any parameters alias, we fall back to the slower, copying code path. This is
+        # a sufficient check, because overlapping parameter buffers that don't completely
+        # alias would break the assumptions of the uniqueness check in
+        # Module.named_parameters().
+        unique_data_ptrs = set(p.data_ptr() for l in self.all_weights for p in l)
+        if len(unique_data_ptrs) != sum(len(l) for l in self.all_weights):
+            self._data_ptrs = []
+            return
+
         with torch.cuda.device_of(any_param):
             # This is quite ugly, but it allows us to reuse the cuDNN code without larger
             # modifications. It's really a low-level API that doesn't belong in here, but
@@ -98,8 +107,8 @@ def flatten_parameters(self):
             fn.w_desc = rnn.init_weight_descriptor(fn, fn.weight_buf)
 
             # Slice off views into weight_buf
-            params = rnn.get_parameters(fn, handle, fn.weight_buf)
             all_weights = [[p.data for p in l] for l in self.all_weights]
+            params = rnn.get_parameters(fn, handle, fn.weight_buf)
 
             # Copy weights and update their storage
             rnn._copyParams(all_weights, params)


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -2254,6 +2254,40 @@ def test_cudnn_weight_format(self):
             weight_data[:] = 4
             self.assertEqual(weight_data, all_vars[4].data)
 
+    @unittest.skipIf(not TEST_CUDNN, 'CUDNN not available')
+    def test_cudnn_weight_tying(self):
+        rnns = [
+            nn.LSTM(10, 20, batch_first=True, bidirectional=True),
+            nn.GRU(10, 20, batch_first=True, bidirectional=True),
+            nn.RNN(10, 20, batch_first=True, bidirectional=True)
+        ]
+        for rnn in rnns:
+            rnn.bias_ih_l0_reverse = rnn.bias_ih_l0
+            rnn.cuda()
+            input = Variable(torch.randn(5, 4, 10).cuda(), requires_grad=True)
+            hx = Variable(torch.randn(2, 5, 20).cuda(), requires_grad=True)
+            all_vars = [input, hx] + list(rnn.parameters())
+            opt = torch.optim.SGD(rnn.parameters(), lr=0.1)
+            opt.zero_grad()
+            if isinstance(rnn, nn.LSTM):
+                cx = Variable(torch.randn(2, 5, 20).cuda(), requires_grad=True)
+                all_vars[2:2] = [cx]
+                hx = (hx, cx)
+
+            with warnings.catch_warnings(record=True) as w:
+                output = rnn(input, hx)
+            self.assertEqual(len(w), 1)
+            self.assertIn('weights are not part of single contiguous chunk of memory', w[0].message.args[0])
+            output[0].sum().backward()
+
+            opt.step()
+            with warnings.catch_warnings(record=True) as w:
+                output_cuda = rnn(input, hx)
+            rnn.cpu()
+            hx = (hx[0].cpu(), hx[1].cpu()) if isinstance(rnn, nn.LSTM) else hx.cpu()
+            output_cpu = rnn(input.cpu(), hx)
+            self.assertEqual(output_cuda, output_cpu)
+
     @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
     def test_cuda_rnn_fused(self):
         def copy_rnn(rnn1, rnn2):