Why is custom RNN implementation so slow?

omnomc: <p>Hello, I am currently creating a generic custom RNN implementation in PyTorch, however it takes about 60x longer to train than if I were to just use nn.RNN.</p>
<pre><code class="lang-auto">def forward(self, idx, targets=None):
        emb = self.embedding(idx)

        B, T, C = emb.shape

        h = torch.zeros((B, C)).to(emb.device)

        logits = []

        for t in range(T):
            token = emb[:, t, :]
            h = F.tanh(self.ih(token) + self.hh(h))
            logits.append(self.head(h))

        logits = torch.stack(logits, dim=1)

        logits = self.head(out)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss
</code></pre>

ptrblck: <p>The <code>nn.RNN</code> module dispatches to cuDNN by default (you could verify it by profiling the use case and check the kernel names), which is uses an optimized algorithm. Your custom implementation launches a lot of kernels internally and would thus suffer from this overhead. You could try to use <code>torch.compile</code> to see if it can improve the performance.</p>

omnomc: <p>Thanks, I have tried torch.compile however it gives minimal performance boost. Is it possible to get near cuDNN performance with optimizations? And what optimizations would I need to do?</p>

J_Johnson: <p>You could go to the GitHub repo for PyTorch’s RNNBase class and see what calls are used to optimize internal functions:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/rnn.py">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/nn/modules/rnn.py</a></h4>


      <pre><code class="lang-py"># mypy: allow-untyped-decorators
# mypy: allow-untyped-defs
import math
import numbers
import warnings
import weakref
from typing import List, Optional, overload, Tuple
from typing_extensions import deprecated

import torch
from torch import _VF, Tensor
from torch.nn import init
from torch.nn.parameter import Parameter
from torch.nn.utils.rnn import PackedSequence

from .module import Module


__all__ = [
    "RNNBase",
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/rnn.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Or modify that file in your torch install to what you need.</p>

omnomc: <p>Thanks! Where is the actual forward loop? It seems to lead to _VF.rnn_tanh but I checked files that are called _VF but I cant find anything to do with rnns.</p>

J_Johnson: <p><code>RNNBase</code> is a parent class called below in various RNN implementations in the same file such as RNN, GRU, LSTM, etc.</p>
<p>For example, go to line 467 where the RNN class is defined. And that has a forward method defined. Within that method, as you noticed, are called <code>_VF.rnn_tanh</code>, etc.</p>
<p>These are handled in C++. That sends you here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/_VF.py">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/_VF.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/_VF.py" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/_VF.py</a></h4>


      <pre><code class="lang-py">"""
This makes the functions in torch._C._VariableFunctions available as
    torch._VF.&lt;funcname&gt;
without mypy being able to find them.

A subset of those functions are mapped to ATen functions in
torch/jit/_builtins.py

See https://github.com/pytorch/pytorch/issues/21478 for the reason for
introducing torch._VF

"""

import sys
import types

import torch


class VFModule(types.ModuleType):
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/main/torch/_VF.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Which is just a wrapper for files in ATen/native, such as for cudnn here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cudnn/RNN.cpp">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cudnn/RNN.cpp" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cudnn/RNN.cpp" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/aten/src/ATen/native/cudnn/RNN.cpp</a></h4>


      <pre><code class="lang-cpp">#define TORCH_ASSERT_ONLY_METHOD_OPERATORS
#include &lt;ATen/Config.h&gt;
#include &lt;ATen/MatrixRef.h&gt;
#include &lt;ATen/TensorUtils.h&gt;
#include &lt;ATen/core/Tensor.h&gt;
#include &lt;ATen/cuda/CUDAConfig.h&gt;
#include &lt;ATen/cuda/CUDAEvent.h&gt;
#include &lt;ATen/cuda/Exceptions.h&gt;
#include &lt;ATen/native/RNN.h&gt;
#include &lt;c10/util/Exception.h&gt;
#include &lt;c10/util/accumulate.h&gt;
#include &lt;c10/util/irange.h&gt;
#include &lt;ATen/cuda/CUDAGraphsUtils.cuh&gt;

#ifndef AT_PER_OPERATOR_HEADERS
#include &lt;ATen/Functions.h&gt;
#include &lt;ATen/NativeFunctions.h&gt;
#else
#include &lt;ATen/ops/_cudnn_init_dropout_state.h&gt;
#include &lt;ATen/ops/_cudnn_init_dropout_state_native.h&gt;
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cudnn/RNN.cpp" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Note, PyTorch backend functions are optimized for each device (i.e. CPU, Cuda). Also note that some functions may be called directly from the cudnn library, which is closed source.</p>