Repository: pytorch
Issue: RFC: Add flag for RNN decomposition to all RNN modules
Labels: feature,module: rnn,triaged

**tl;dr** The basic proposal here is to add a flag to RNN (and subclasses like GRU or LSTM) where instead of running the RNN kernel, it will run the linear, dropout, etc. calls that create an equivalent decomposition. Without this, the monolithic rnn functions and buffers returned from the _cudnn_rnn function make it difficult to extend RNNs in the cases of extending RNNs, computing per sample gradients, and AOTAutograd. The proposed API adds a flag to the RNN class that determines whether or not to use the decomposed version, which will be defaulted to False in order to not incur perf penalties

## Problem
The basic problem with the RNN kernel in particular is that the cuda versions pass around buffers that are used during the backward computation. This is particularly problematic when someone wants to use a custom derivative for the RNN since CUDA doesn't have any stability guarantees for what is being passed back in the buffers. Therefore, a developer cannot even try to recompute the intermediate values and pass those to CUDA's RNN backwards function and hope to produce correct results.

## Use Cases
### RNN Experimentation
For a long time (even since [issue 1932](https://github.com/pytorch/pytorch/issues/1932)), people have been asking for ways to adapt RNNs. Some of the asks include [using layer norm](https://github.com/pytorch/pytorch/issues/7032) as the activation to [having different hidden sizes per layer](https://github.com/pytorch/pytorch/issues/55910). Although RNNs have somewhat fallen out of style with the rise of transformers, new research on them is [still hitting ICML](https://icml.cc/virtual/2021/poster/10541). Right now, everything exists in monolithic kernels (like rnn_tanh and rnn_relu) that are performant but make it difficult to understand what's happening. Although a user could write the same decomposition we plan to in Python, there's so many flags to an RNN that make it difficult to know if you've implemented the decomposition correctly. Having a deomposed Python version that we know works correctly will let users experiment with these new versions easily

### Expanded Weights and per sample gradients
These kernels are also problematic for Expanded Weights, our new system for computing per sample gradients. The mechanism behind this uses torch function and autograd.Function since we need to change the autograd behavior. In doing this, we also need to recompute the batched gradient with respect to the input. So, we would need to decide which backwards to use and then pass the correct byffers if we're using _cudnn_rnn_backward. As mentioned, this won't work because NVIDIA doesn't guarantee that the values in the buffers will be consistent between versions.

To work around this, libraries that want to support RNNs while computing per sample gradients like Opacus have hacky solutions that we shouldn't copy upstream. Specifically, they implement RNNs as two linear layers, which gets them the correct behavior. However, in order to make it exportable, they reset the names so that it looks like it's an RNN module. More concretely, a vanilla Pytorch RNN may have a weight with a "weight_hh_l0". An Opacus version of this would be decomposed into multiple Linear layers where the equivalent parameter should have the name "l0.hh.weight". In order to make their models save and loadable, they patch it to have the same name as the vanilla PyTorch RNN. However, we should not be copying this hack upstream since it breaks mechanisms like make_stateless that assumes that the name of the weights follows the structure of the nn.Module.

### AOTAutograd
AOTAutograd has mentioned that they've noticed these functions show up in traces. Although they are able to support the current mechanism, having a decomposition can help support backends that don't have an RNN kernel and allow for custom optimizations for different backends.

This would also fix https://github.com/pytorch/functorch/issues/586, which is an issue that stems from LSTMs not properly forwarding the `requires_grad_`-ness of its weights through to the `_cudnn_rnn` kernel

## Proposed API
Our proposal is to add a flag to the RNN module that determines whether to use the decomposed version or the RNN kernel like before. By keeping this flag off by default, users should not see any changes from the original behavior. User should be able to determine this while building the RNN while also toggle the flag without rebuilding their RNN, similar to a training flag. Unlike the training flag, we should be able to set this on a layer-by-layer basis instead of only at a whole model level.

`net = RNN(input_size, hidden_size, num_layers=4, use_decomposition=True)`
`net.set_decomposition_(False)`

## Perf Concerns
The decompositions written in Python will have worse perf than the custom C++ kernels. First, these decompositions will be necesssary for backends that don't have a custom RNN kernel, as noted in the AOTAutograd section. Additionally, systems like Opacus that currently require this decomposition to do their per sample gradient computation currently pay this cost. So, we will not be worsening their perf metrics.

Finally, with incoming systems like torchdynamo, we can hope to recover some of this performance for backends that have custom implemented kernels. Until then, by leaving the flag as the default of `use_decomposition=False`, we should not see any performance hits for users that do not use this flag.

## Alternatives

### `__torch_dispatch__` Decomposition
One alternate would be to implement the decomposition at the torch dispatch level, like the decompositions that AOTAutograd use. Since users will still be able to use the undecomposed version, AOTAutograd will still need a torch dispatch decomposition and we will probably want to even use the same decomposition. The only issue is if we only have a torch dispatch decomposition. Since Expanded Weights needs to be at the torch function level in order to extend autograd, it won't work to have the decomposition only exist at the torch dispatch level

### Add a `__torch_function__` for RNN
Currently, RNN doesn't have a functional version nor a torch function intercept. One argument may be to add these and then have a user intercept the RNN at that level and decompose if necessary. Given how monolithic the RNN kernels are, we won't be able to decompose it much between the module's forward call and this call. So if we just added this extension point, we end up with the same issue where a user could decompose the function themselves but runs a lot of risk of implementing it incorrectly. Additionally, this is BC breaking since we're intercepting torch function calls where we weren't before.

### Additional context

_No response_

cc @zou3519