Training RNNs as Fast as CNNs

Hi!

I was happily using SRUs with Pytorch native AMP, however I started experimenting with training using Microsoft DeepSpeed and bumped in to an issue.

Basically the issues is that I observed that FP16 training using DeepSpeed doesn't work for both GRUs and SRUs. However when using Nvidia APEX AMP, DeepSpeed training using GRUs does work.

So, based on the tips in one of your issues, I started looking in to how I could enable Pytorch native AMP and Nvidia APEX AMP for SRUs, so I could train models based on SRUs using DeepSpeed.

That is why I created this pull request. Basically, I found that by making the code simpler, I can make SRUs work with both methods of AMP.

Now amp_recurrence_fp16 can be used for both types of AMP. When amp_recurrence_fp16=True, the tensor's are cast to float16, otherwise nothing special happens. So, I also removed the torch.cuda.amp.autocast(enabled=False) region; I might be wrong, but it seems that we don't need it.

I did some tests with my own code and it works in the different scenarios of interest:

• Using PyTorch native AMP, not using DeepSpeed
• Not using PyTorch native AMP, not using DeepSpeed
• Using Nvidia APEX AMP, using DeepSpeed
• Not using Nvidia APEX AMP, using DeepSpeed

It would be beneficial if we can test this with an official SRU repo test, maybe repurposing the language_model/train_lm.py?

When I convert my model, which using this SRU unit, into float16 enabled one, it fails. Is this SRU not implemented to use in float16 environment, or is it hard to fix it?

This is on 3.0.0-dev branch for now

A non-trivial PR to support GPU inference in torchscript

• Load CUDA kernels as non-python modules; this is needed for torchscript compilation
• Refactored CUDA APIs as functions that return output as tensors, instead of procedures that modify some passed-in tensors.
• Added a workaround in case TS tries to locate and compile CUDA methods on machines that don't have CUDA / GPUs

The refactored code has passed the forward() & backward() test. I also checked the outputs are the same for the non-torchscript and torchscript versions of the same model.

Hello @taolei87 , After updating to the latest version, my code broke. It works great on the previous 2.3.5 version and with nn.LSTM.

File "C:\xxx\lib\site-packages\torch\nn\modules\module.py", line 722, in _call_impl
  result = self.forward(*input, **kwargs)
File "C:\xxx\lib\site-packages\sru\modules.py", line 576, in forward
  mask_pad = (mask_pad >= batch_sizes.view(length, 1)).contiguous()
RuntimeError: shape '[393, 1]' is invalid for input of size 384

I can see that in the previous version the unpacking code on forward was different:

        input_packed = isinstance(input, nn.utils.rnn.PackedSequence)
        if input_packed:
            input, lengths = nn.utils.rnn.pad_packed_sequence(input)
            max_length = lengths.max().item()
            mask_pad = torch.ByteTensor([[0] * l + [1] * (max_length - l) for l in lengths.tolist()])
            mask_pad = mask_pad.to(input.device).transpose(0, 1).contiguous()

Now is:

        orig_input = input
        if isinstance(orig_input, PackedSequence):
            input, batch_sizes, sorted_indices, unsorted_indices = input
            length = input.size(0)
            batch_size = input.size(1)
            mask_pad = torch.arange(batch_size,
                                    device=batch_sizes.device).expand(length, batch_size)
            mask_pad = (mask_pad >= batch_sizes.view(length, 1)).contiguous()

I'm trying to implement a model that includes a SRUCell. This are my specs:

Tesla M60 GPU torch.version: 0.4.1.post2 torch.cuda.version: 9.0.176

Although its training, every epoch the memory usage in the GPU increases until it fills it. I made a toy example where this error occurs:

import torch
from torch.autograd import Variable
from sru import SRUCell


batch_size = 5
seq_len = 60
epochs = 1000
cuda = torch.cuda.is_available()

model = SRUCell(100, 100)

if cuda:
    model.cuda(0)

optimizer = torch.optim.Adam([
        {'params':model.parameters()}], lr=1e-3)

loss_function = torch.nn.MSELoss()
    
seq = Variable(torch.rand(batch_size,seq_len,100))
y = Variable(torch.rand(batch_size,100))


if cuda:
    seq = seq.cuda(0)
    y = y.cuda(0)


model.train()

for e in range(epochs):
    model.zero_grad()
    
    h = Variable(torch.zeros(batch_size, 100))
    c = Variable(torch.zeros(batch_size, 100))
    
    if cuda:
        h = h.cuda(0)
        c = c.cuda(0)
    
    for i in range(seq_len):
        x = seq[:,i,:]
        h, c = model(x, c)
    loss = loss_function(h, y)
    loss.backward()
    optimizer.step()
    print('Epoch: {} - Loss: {}'.format(e, loss))

In vanilla it work like this

rnn = nn.LSTMCell(10, 20)
input = torch.randn(6, 3, 10)
hx = torch.randn(3, 20)
cx = torch.randn(3, 20)
output = []
for i in range(6):
    hx, cx = rnn(input[i], (hx, cx))
    output.append(hx)

How can i do same for sru cell?

Firstly i ran download.sh, and it succesfully downloaded glove and train/dev jsons for SQuAD. However, python prepro.py gave me this:

Traceback (most recent call last):
  File "prepro.py", line 243, in <module>
    vocab_tag = list(nlp.tagger.tag_names)
AttributeError: 'Tagger' object has no attribute 'tag_names'

My Spacy version is 2.0.3, and it seems like something broke in update from 1.x that is written in requirements, and I didn't succeed in fixing it myself. Any suggests?

I'm not sure how, but I'm seeing this error when I try to compute the backwards function. Don't know if you've come across this during your debug?

Traceback (most recent call last):
  File "gan_language.py", line 341, in <module>
    G.backward(one)
  File "/usr/local/lib/python2.7/dist-packages/torch/autograd/variable.py", line 156, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, retain_variables)
  File "/usr/local/lib/python2.7/dist-packages/torch/autograd/__init__.py", line 98, in backward
    variables, grad_variables, retain_graph)
  File "/home/nick/wgan-gp/sru/cuda_functional.py", line 417, in backward
    stream=SRU_STREAM
  File "cupy/cuda/function.pyx", line 129, in cupy.cuda.function.Function.__call__ (cupy/cuda/function.cpp:4010)  File "cupy/cuda/function.pyx", line 111, in cupy.cuda.function._launch (cupy/cuda/function.cpp:3647)
  File "cupy/cuda/driver.pyx", line 127, in cupy.cuda.driver.launchKernel (cupy/cuda/driver.cpp:2541)
  File "cupy/cuda/driver.pyx", line 62, in cupy.cuda.driver.check_status (cupy/cuda/driver.cpp:1446)
cupy.cuda.driver.CUDADriverError: CUDA_ERROR_INVALID_HANDLE: invalid resource handle

The data loading was inefficient and was found to be the bottleneck of BILLION WORD training. This PR rewrote the sharding (which data goes to a certain GPU / training process), and improved the training speed significantly.

The figure compares a previous run and a new test run. We see 40% reduction on training time.

This means our reported training efficiency will be much stronger from 59 GPU days to 36 GPU days, and 4x more efficient than FairSeq Transformer results.

Hi, I'm trying to implement a naive version of this paper in Keras, and was wondering how is the case that - n_in != n_out handled.

I went through the code a few times, and couldn't understand the element wise multiplication of (1 - r_t) with x_t, if x_t is of a different shape than r_t.

This PR works for master branch, v2.5 and v2.6 release

A non-trivial PR to support GPU inference in torchscript

• Load CUDA kernels as non-python modules; this is needed for torchscript compilation
• Refactored CUDA APIs as functions that return output as tensors, instead of procedures that modify some passed-in tensors.
• Added a workaround in case TS tries to locate and compile CUDA methods on machines that don't have CUDA / GPUs
• The refactored code has passed the forward() & backward() test.
• I also checked the outputs are the same for the non-torchscript and torchscript versions of the same model.

Hi,

first of all I want to thank you for your great work. I'm using SRUs for speech enhancement, they do very well on a reasonable computational cost.

I would like to know if there is a possibility to train SRUs in mixed precision mode? I tried to enable it, by setting precision=16 in the pytorch lightning trainer, but that didn't do the trick.

Kind of regards, Zadagu

Hello, I've read and really appreciated your team's wonderful works on SRU++. I want to implement this architecture in other tasks, but i'm having problem finding the documentation on SRU++, as how I can use SRU++ the same way as SRU (calling directly from sru library after installing by pip install sru). I have looked into the dev-3.0.0 branch, which seems like the latest updated branch, but I still have no clues how to call and integrate sru++ modules into my custom defined pytorch modules. Could you help me ?

when i run example, i meet this issue:FAILED: sru_cuda_kernel.cuda.o ,and in the end, it report ninja: build stopped: subcommand failed. what should i do to slove this problem?

We noticed the package initialization for sru is eagerly triggering the initialization because of the following stack of module imports sru.modules -> sru.ops -> cuda_functional and this last module is executing the function load of torch.utils.cpp_extension.

This was detected because of issues caused when running with the server framework in SUBPROCESS_MODE, that is forking a new process for it to run the model. We got an error complaining that CUDA had been already initialized in the parent process, which was not necessary because it is not meant to run the inference in the model.

This PR changes this loading to be more lazy, more concretely we changed the code in sru.modules to avoid the eager import of sru.ops and instead postpone it to the instantiation of a first SRUCell.

The changes in this PR have been tested doing a checkout of this branch in an AWS instance with GPU and running pytest -sv test which resulted in 141 passed, 161 warnings and no failures. So we understand this is working as expected for both CPU and GPU settings.

When using a bidirectional SRU, regular usage seems to be fine, and compilation to torchscript proceeds without error, but upon trying to infer with the compiled torchscript I get:

Unknown builtin op: sru_cuda::sru_bi_forward_simple.

Using pytorch 1.10, sru 2.6.0, cuda 11.3