Repository: pytorch
Issue: How to calculate second derivative using PyTorch with GPU (cuda)
Labels: module: double backwards,module: cudnn,module: autograd,module: rnn,triaged,module: functorch

### 🚀 The feature, motivation and pitch

I have a python code segment related to a deep RL algorithm where it calculates the second order optimization and second derivative with Hessian matrix and fisher information matrix. Normally I run the whole code on GPU (cuda), but since I got a computational issue to calculate second derivative in cuda,
```
NotImplementedError: the derivative for '_cudnn_rnn_backward' is not implemented. Double backwards is not supported for CuDNN RNNs due to limitations in the CuDNN API. To run double backwards, please disable the CuDNN backend temporarily while running the forward pass of your RNN. For example: 
with torch.backends.cudnn.flags(enabled=False):
    output = model(inputs)
```
I had to move to CPU for this code segment, and now the code is executing sequentially instead of in parallel, which takes a long time to run:
```
grads = torch.autograd.grad(policy_loss, self.policy.Actor.parameters(), retain_graph=True)
loss_grad = torch.cat([grad.view(-1) for grad in grads])

def Fvp_fim(v = -loss_grad):
    with torch.backends.cudnn.flags(enabled=False):
        M, mu, info = self.policy.Actor.get_fim(states_batch)
        #pdb.set_trace()
        mu = mu.view(-1)
        filter_input_ids = set([info['std_id']])

        t = torch.ones(mu.size(), requires_grad=True, device=mu.device)
        mu_t = (mu * t).sum()
        Jt = compute_flat_grad(mu_t, self.policy.Actor.parameters(), filter_input_ids=filter_input_ids, create_graph=True)
        Jtv = (Jt * v).sum()
        Jv = torch.autograd.grad(Jtv, t)[0]
        MJv = M * Jv.detach()
        mu_MJv = (MJv * mu).sum()
        JTMJv = compute_flat_grad(mu_MJv, self.policy.Actor.parameters(), filter_input_ids=filter_input_ids, create_graph=True).detach()
        JTMJv /= states_batch.shape[0]
        std_index = info['std_index']
        JTMJv[std_index: std_index + M.shape[0]] += 2 * v[std_index: std_index + M.shape[0]]
        return JTMJv + v * self.damping
```
Above is the main function, where it calculates the second derivative. below are the supportive functions and relevant classes it has used.
```
def compute_flat_grad(output, inputs, filter_input_ids=set(), retain_graph=True, create_graph=False):
    if create_graph:
        retain_graph = True

    inputs = list(inputs)
    params = []
    for i, param in enumerate(inputs):
        if i not in filter_input_ids:
            params.append(param)

    grads = torch.autograd.grad(output, params, retain_graph=retain_graph, create_graph=create_graph, allow_unused=True)

    j = 0
    out_grads = []
    for i, param in enumerate(inputs):
        if (i in filter_input_ids):
            out_grads.append(torch.zeros(param.view(-1).shape, device=param.device, dtype=param.dtype))
        else:
            if (grads[j] == None):
                out_grads.append(torch.zeros(param.view(-1).shape, device=param.device, dtype=param.dtype))
            else:
                out_grads.append(grads[j].view(-1))
            j += 1
    grads = torch.cat(out_grads)

    for param in params:
        param.grad = None
    return grads

------

import torch
import torch.nn as nn


from agents.models.feature_extracter import LSTMFeatureExtractor
from agents.models.policy import PolicyModule
from agents.models.value import ValueModule


class ActorNetwork(nn.Module):
    def __init__(self, args):
        super(ActorNetwork, self).__init__()
        self.FeatureExtractor = LSTMFeatureExtractor(args)
        self.PolicyModule = PolicyModule(args)

    def forward(self, s):
        lstmOut = self.FeatureExtractor.forward(s)
        mu, sigma, action, log_prob = self.PolicyModule.forward(lstmOut)
        return mu, sigma, action, log_prob
    
    def get_fim(self, x):
        mu, sigma, _, _ = self.forward(x)

        if sigma.dim() == 1:
            sigma = sigma.unsqueeze(0)

        cov_inv = sigma.pow(-2).repeat(x.size(0), 1)

        param_count = 0
        std_index = 0
        id = 0
        std_id = id
        for name, param in self.named_parameters():
            if name == "sigma.weight":
                std_id = id
                std_index = param_count
            param_count += param.view(-1).shape[0]
            id += 1

        return cov_inv.detach(), mu, {'std_id': std_id, 'std_index': std_index}
```
In the bigger picture there are large amounts of batches going through this function, since all of 'em have to go sequentially through this function, it highly increases the total running time. Is there a possible way to calculate the second derivative with Pytorch while running on cuda/GPU?

### Alternatives

_No response_

### Additional context

_No response_

cc @csarofeen @ptrblck @xwang233 @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan @mikaylagawarecki @zou3519 @Chillee @samdow @kshitij12345