Repository: pytorch
Issue: Retain Subgraph or Save Intermediate Grad support?
Labels: module: autograd,triaged

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Hi, I'd like to know if it possible to support *retain **subgraph*** for backward calculation.

Or some ways to save **intermediate gradients** that I need exactly for later use. The gradient hook is not precise enough.

## Motivation
I've tried current available **retrain_graph=True**, it works fine when model is small, but causes **out of memory error when model gets bigger**. Actually, I only need some intermediate gradients values that related to 1/50 subgraph or even smaller to be retained. Not the structure of the subgraph.

I also tried hooks, and hooks are just not precise enough to get exactly the gradient tensors I need especially when multi-head model is used. The hook always sums them up, while I need exactly values of each gradient.
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

## Pitch
### Two examples, one for subgraph, the other for save intermediate grad. Both should work.
#### Subgraph example:

```python
from torch.autograd import grad

input_x = torch.tensor([1.0, 2.0]) # input data
groundtruths = torch.tensor([...]) # some ground truth labels
model = build_model(...) # This model might be a very big multi-head model

preds = model(input_x) # multi-head prediction

# This is $dx$ of $\frac{dy}{dx}$
x1 = model.head_a.fc1.weight
x2 = model.head_a.fc1.bias

 # multi-head loss, assume 3 heads, loss1 is $dy$ of $\frac{dy}{dx}$
loss1, loss2, loss3 = loss_eval(preds, groundtruths)

loss = sum([loss1, loss2, loss3])
# only retain exactly the needed part, not whole graph
loss.backward(retain_graph=True, retain_subgraph_y=[loss1], retain_subgraph_x=[x1, x2])

# The following gradients from part of the graph are what I need, 
# but retain whole graph uses too much memory
# Here retain_graph only retain the subgraph, others released for memory efficiency.
gradient1 = grad([loss1], [x1], allow_unused=True, retain_graph=True)
gradient2 = grad([loss1], [x2], allow_unused=True, retain_graph=True)
```

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

#### Example of Save Specific Intermediate Grad: 
I think this is **much more *simple* and *efficient***. The definition of model /input/groundtruth/intermediate variables are the same as above.

```python
# Before backward, use a dict to record the gradients need to be saved
model.save_grad_during_backward(name='g1', y=[loss1], x=[x1])
model.save_grad_during_backward(name='g2', y=[loss1], x=[x2])
# No need to save any part of the graph, only save specific gradients during backpropagation 
loss.backward(retain_graph=False)

# Get exactly values of gradients
gradient1 = model.saved_grad['g1']
gradient2 = model.saved_grad['g2']
```

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->


cc @ezyang @SsnL @albanD @zou3519 @gqchen