Repository: pytorch
Issue: [autograd engine && compiled autograd] support for partial backward compilation
Labels: module: autograd,triaged,oncall: pt2,module: compiled autograd

### 🚀 The feature, motivation and pitch

Hello, in our task, we aim to optimize some specific layers during backpropagation while keeping other layers' backpropagation and the complete forward pass in eager mode. As PyTorch doesn't currently support partial backward compilation directly, we have attempted to achieve this by decomposing the full backward pass into blocked ones via `torch.autograd.backward(..., inputs=[blocked_parameters, boundry_activation])`: 
```python
# backward-block2: eager
torch.auto.backward(loss, inputs=[model_block_2.parameters, activation_1])

# backward-block1: compile
with with torch._dynamo.compiled_autograd.enable(torch.compile(backend=my_compiler, fullgraph=False))::
    torch.auto.backward(activation_1, activation_1.grad, inputs=[model_block_1.parameters, activation_0])

# backward-block0: eager
torch.auto.backward(activation_0, activation_0.grad, inputs=[model_block_0.parameters])
```
During the process, we encountered some problems and made some attempts as follows:
- In "backward-block2", we must set `retain_graph=True`, otherwise the variables related to `activation_1` will be released and the subsequent backward block can not be launched successfully. However, `retain_graph=True` leads to unignorable memory overhead and `activation_1`'s grad_fn is not necessary for `activation_1`'s gradient calculation. Is there any method to prevent the inclusion of `activation_1`'s grad_fn in the GraphTask? 
- In "backward-block1", when we pass a non-leaf activation into `inputs`, the compiled autograd throws the error: `Runtime error: retains_grad_hooks not implemented for compiled autograd`. We temporarily comment the `TORCH_CHECK` for `retain_grad_hook` and access `activation_0`'s gradient via a `param_hook` that writes the gradient into a python global variable, which, however, is very cumbersome. We are quite curious if  there are plans to support `retain_grad_hook` in compiled autograd?
- `torch.autograd.backward` encounters an issue when the `inputs` includes non-leaf activations and weight that is shared in different layers, such as the `tied_weight` in higgingface models.  The error is `Runtime error: could not compute gradients for some functions`. It seems that the autograd engine attempts to incorporate all compute operations related to the weight gradient into the GraphTask, but in doing so, some operations are overlooked, and certain dependencies cannot be executed as expected. However, in our specific task, we would prefer the autograd engine not to include all dependencies associated with the shared weight, as it hinders the execution of blocked backwards. We are wondering if there is any interest in implementing something like `partial gradient mode `that generates only the minimal subgraph to reach the target in `inputs`, potentially allowing for the omission of certain input branches of the `AccumulateGrad` node?
- When we try to optimize the backward pass by compiled autograd, some unexpected `torch.clone` appears in the traced fx graph, which increases the performance overhead a lot. For more detailed information, you can refer to the relevant GitHub issue at https://github.com/pytorch/pytorch/issues/139862. Is there any method to prevent these `torch.clone` being generated in the traced fx graph?

 @bdhirsh @Chillee @xmfan @ezyang 

### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan @chauhang @penguinwu @yf225