PT Pruning + amp + LSTM/RNN == `RuntimeError` [Bug Report]

Bug report

Summary

It seems like the ptr arithmetic in the following function amp.utils function breaks when the pytorch pruning module is used for LSTMs:
https://github.com/NVIDIA/apex/blob/43a6f9fe91c242170cbc5c8bf13f466eaccab2e4/apex/amp/utils.py#L194

In the function, the calculated offset can (sometimes) be < 0 which causes the following error (full traceback below):

RuntimeError: Tensor: invalid storage offset at /opt/conda/conda-bld/pytorch_1587428266983/work/aten/src/THC/THCTensor.cpp:166 

Minimal Example

import torch 
import torch.nn.utils.prune as prune
from apex import amp


input_size = 256
hidden_size = input_size
num_attempts = 10

def gen_args():
    """Helper function to generate LSTM input args"""
    batch = 2
    seq_len = 20
    x = torch.randn(seq_len, batch, input_size)
    x = x.cuda()
    return x

lstm = torch.nn.LSTM(input_size, hidden_size, 1, batch_first=False, bidirectional=False)
lstm = lstm.cuda()
optimizer = torch.optim.Adam(lstm.parameters())
lstm, optimizer = amp.initialize(lstm, optimizer, opt_level="O1")

# run model + update weights before pruning to demonstrate
# that the issue is introduced by pruning 
for i in range(num_attempts):
    x = gen_args()
    res = lstm(x)
    # can optionally call res.mean().backward() here with no change in result
print(f"\nSucessfully ran {num_attempts} attempts with no error\n")

# prune model
additional_sparsity = 0.2
prune.l1_unstructured(lstm, 'weight_ih_l0', additional_sparsity)
prune.l1_unstructured(lstm, 'weight_hh_l0', additional_sparsity)

# run model arbitrary number of times - one should throw error below
for i in range(num_attempts):
    print(f"Sparse attempt #{i}")
    x = gen_args()
    res = lstm(x)

outputs

Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic

Sucessfully ran 10 attempts with no error

Sparse attempt #0
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-57c5e5db0261> in <module>
     38     print(f"Sparse attempt #{i}")
     39     x = gen_args()
---> 40     res = lstm(x)

~/miniconda3/envs/ttsmyrtle/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

~/miniconda3/envs/ttsmyrtle/lib/python3.7/site-packages/torch/nn/modules/rnn.py in forward(self, input, hx)
    568         if batch_sizes is None:
    569             result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
--> 570                               self.dropout, self.training, self.bidirectional, self.batch_first)
    571         else:
    572             result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,

~/miniconda3/envs/ttsmyrtle/lib/python3.7/site-packages/apex/amp/wrap.py in wrapper(*args, **kwargs)
    255                                                     dtype=torch.half)
    256                 casted_weights = utils.new_synthesize_flattened_rnn_weights(
--> 257                     arg, fp16_weight_buf, fn, verbose)
    258                 new_args.append(casted_weights)
    259             elif utils.is_fp_tensor(arg):

~/miniconda3/envs/ttsmyrtle/lib/python3.7/site-packages/apex/amp/utils.py in new_synthesize_flattened_rnn_weights(fp32_weights, fp16_flat_tensor, rnn_fn, verbose)
    203         w_fp16.set_(fp16_flat_tensor.storage(),
    204                     offset,
--> 205                     w_fp32.shape)
    206         w_fp16.copy_(w_fp32)
    207         if verbose:

RuntimeError: Tensor: invalid storage offset at /opt/conda/conda-bld/pytorch_1587428266983/work/aten/src/THC/THCTensor.cpp:166

Non-determinism

The bug is highly non-deterministic as it relies on the layout of tensors in memory (so I couldn't make it reproducible
using seeds and torch.backends.cudnn.deterministic=True). I found that it occurred more often as hidden_size and input_size increased but it should also be possible to increase num_attempts to see the error.

Debug offset <0

To confirm that the error only occurs when offset < 0 in new_synthesize_flattened_rnn_weights function, you can vim into the apex file and print the value of offset before the error is raised.

@mcarilli @cbcase

Start with apex/amp/utils.py, especially new_synthesize_flattened_rnn_weights around the offset calculation and set_ call. Run the minimal LSTM, pruning, and amp example, then inspect the offset when the failure occurs. Done means the example no longer raises an invalid storage offset error.