Repository: pytorch
Issue: RuntimeError during converting Reformer model to TorchScript
Labels: triage review,oncall: jit,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. install reformer-pytorch from https://github.com/lucidrains/reformer-pytorch
2. install transformers
3. Run script 
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from reformer_pytorch import Reformer, ReformerLM
from transformers import BertTokenizer, AdamW

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
tokenizer.max_len = 40960
model = ReformerLM(
    num_tokens= tokenizer.vocab_size,
    dim = 768,
    depth = 12,
    max_seq_len = tokenizer.max_len,
    heads = 16,
    lsh_dropout = 0.1,
    emb_dim = 256,        # embedding factorization for further memory savings
    causal = False,        # auto-regressive or not
    bucket_size = 64,     # average size of qk per bucket, 64 was recommended in paper
    n_hashes = 4,         # 4 is permissible per author, 8 is the best but slower
    ff_chunks = 200,      # number of chunks for feedforward layer, make higher if there are memory issues
    weight_tie = False,   # tie parameters of each layer for no memory per additional depth
    attn_chunks = 8,        # process lsh attention in chunks, only way for memory to fit when scaling to 16k tokens
    num_mem_kv = 128,       # persistent learned memory key values, from all-attention paper
    twin_attention = False, # both branches of the reversible network will be attention
    use_full_attn = False,  # use full self attention, for comparison
    full_attn_thres = 128, # use full attention if context length is less than set value
    use_scale_norm = True,  # use scale norm from 'Transformers without tears' paper
    #fixed_position_emb = True
    axial_position_emb = True,
    axial_position_shape = (640, 64),
    axial_position_dims = (128, 128)
)
x = torch.randint(0, 30522, (2, 40960))
model.eval()
traced_cell = torch.jit.trace(model, x)
```
```
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-2cc80d9cd643> in <module>
     46 #print(datetime.datetime.now())
     47 #my_module = torch.jit.script(model)
---> 48 traced_cell = torch.jit.trace(model, x)
     49 print("ok")

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/jit/__init__.py in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, _force_outplace, _module_class, _compilation_unit)
    879         return trace_module(func, {'forward': example_inputs}, None,
    880                             check_trace, wrap_check_inputs(check_inputs),
--> 881                             check_tolerance, _force_outplace, _module_class)
    882 
    883     if (hasattr(func, '__self__') and isinstance(func.__self__, torch.nn.Module) and

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/jit/__init__.py in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, _force_outplace, _module_class, _compilation_unit)
   1041                 else:
   1042                     _check_trace([inputs], func, check_trace_method,
-> 1043                                  check_tolerance, _force_outplace, True, _module_class)
   1044     finally:
   1045         torch.jit._trace_module_map = old_module_map

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     13         def decorate_context(*args, **kwargs):
     14             with self:
---> 15                 return func(*args, **kwargs)
     16         return decorate_context
     17 

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/jit/__init__.py in _check_trace(check_inputs, func, traced_func, check_tolerance, force_outplace, is_trace_module, _module_class)
    673         traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, 'trace')
    674         fn_outs = run_mod_and_filter_tensor_outputs(func, inputs, 'Python function')
--> 675         if compare_outputs(traced_outs, fn_outs, 'Python function'):
    676             check_outs = run_mod_and_filter_tensor_outputs(check_mod_func, inputs, 'repeated trace')
    677             compare_outputs(traced_outs, check_outs, 'repeated trace')

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/jit/__init__.py in compare_outputs(original, reference, match_what)
    661                         ref = ref.dequantize()
    662                     torch.testing.assert_allclose(orig.double(), ref.double(), rtol=check_tolerance,
--> 663                                                   atol=torch.testing._get_default_tolerance(orig, ref)[1])
    664                 except AssertionError as e:
    665                     maybe_warn_nondeterministic()

/opt/anaconda/envs/torch-nigtly-reformer/lib/python3.6/site-packages/torch/testing/__init__.py in assert_allclose(actual, expected, rtol, atol, equal_nan)
     28         rtol, atol = _get_default_tolerance(actual, expected)
     29 
---> 30     close = torch.isclose(actual, expected, rtol, atol, equal_nan)
     31     if close.all():
     32         return

RuntimeError: [enforce fail at CPUAllocator.cpp:64] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 20002897920 bytes. Error code 12 (Cannot allocate memory)
```
## Expected behavior

Get model, converted to TorchScript

## Environment

PyTorch version: 1.5.0.dev20200305
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Quadro RTX 8000
GPU 1: Quadro RTX 8000

Nvidia driver version: 440.33.01
cuDNN version: Could not collect

Versions of relevant libraries:
[pip] numpy==1.18.1
[pip] reformer-pytorch==0.19.2
[pip] torch==1.5.0.dev20200305
[pip] torchvision==0.6.0a0+15aa7e2
[conda] blas                      1.0                         mkl  
[conda] mkl                       2019.5                      281    conda-forge
[conda] mkl-service               2.3.0            py36h516909a_0    conda-forge
[conda] mkl_fft                   1.1.0            py36hc1659b7_1    conda-forge
[conda] mkl_random                1.1.0            py36hb3f55d8_0    conda-forge
[conda] pytorch                   1.5.0.dev20200305 py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch-nightly
[conda] reformer-pytorch          0.19.2                   pypi_0    pypi
[conda] torchvision               0.6.0.dev20200305      py36_cu101    pytorch-nightly

 - PyTorch Version: 1.5.0.dev20200305
 - OS: Ubuntu 18.04 LTS
 - How you installed PyTorch: Conda
 - Python version: 3.6
 - CUDA/cuDNN version: 1.5.0.dev20200305/not collect
 - GPU models and configuration: RTX 8000 x 2

## Additional context

I trying to convert Reformer model, witch storing on cpu Into TorchScript. But, during converting model allocate all RAM(128 GB) and fail with RuntimeError. 
If I trying to infer model in torch, all going ok, and model allocate only 7GB RAM with batch size 2.


cc @suo