Repository: pytorch
Issue: [RFC] Dynamo Single Step Graph
Labels: feature,module: autograd,triaged,oncall: pt2

### 🚀 The feature, motivation and pitch


## Context

### Background

_Torch Dynamo is the graph capture mechanism of PyTorch introduced in the PyTorch 2.0 release. Torch Dynamo will take a PyTorch Model and emit the TrochFX IR representing the computation. Currently TorchDynamo will emit separate graphs for the model forward and the backward which is inefficient for compilers like XLA since it will eliminate some fusion opportunities and incur higher execution overhead. This inefficiency is the biggest blocker to fully migrate PyTorch/XLA training stack to TorchDynamo._


## Design


### Overview

_The high level idea of this design is to use CompiledAutograd to generate the TorchFX IR representing the backward and then inline the TorchFX using Torch Dynamo. This way we will have a single graph representing both forward and the backward of the model. Proof concept PR in [here](https://github.com/pytorch/pytorch/pull/112296), I was able to trace MNIST and ResNet without issue. I used a couple hacks to bypass some major issues which I will describe in the following section. _

_In my opinion the backward is difficult to capture because the autograd engine is implemented in C++ and dynamo can only understand the python code. CompiledAutograd is very handy in this case because it can convert the C++ backward into a Torch FX IR. Torch FX IR can be executed as normal python code so dynamo can directly inlining through it._


#### Current Flow
![image](https://github.com/pytorch/pytorch/assets/59073027/e8b9db87-3258-4926-b8cb-f863c7f04fc7)

#### Proposed Flow
![image](https://github.com/pytorch/pytorch/assets/59073027/5ae8faf6-143a-42a7-bfa3-f1dd120f338d)


#### High Level Changes Breakdown



1. Inline all module’s forward method
2. Inline all non-core-aten-ops using `make_fx`
3. Inline loss functions similar to `cross_entropy_loss`
4. Inline a dummy version of the `accumulate_grad_`
5. Create a mapping from fake tensor to tensor Variables
6. Inline the `tensor.backward` until hitting the `torch._C._EngineBase` pybind
7. Inline the `torch._C._EngineBase.run_backward`
    1. Capture the backward as FX graph by executing the backward with fake tensors(converted from proxy tensor) under a `compiled_autograd` context manager.
    2. Inline the FX graph with tensorVariables(map input fake tensors back to Tensor variables) inputs.


### Detail Design

Note that I will use a different order to go over the changes I made to better explain why certain changes are needed. 


#### Inline `tensor.backward` 

Currently dynamo will graph break upon encountering `tensor.backward`
```
torch/_dynamo/variables/[tensor.py](http://tensor.py/)

.....
        elif name in ("backward", "data_ptr"):
            unimplemented(f"Tensor.{name}")
```

We need to force dynamo to continue tracing and inlining all python functions until hitting `torch._C._EngineBase.run_backward`.

```
torch/_dynamo/variables/[tensor.py](http://tensor.py/)

......
        elif config.use_single_step_graph and name == "backward":
......
            backward_variable = TorchInGraphFunctionVariable(torch.autograd.backward)
            result = backward_variable.call_function(
                tx,
                [
                    self,  # tensors
                    grad_tensors,  # grad_tensors
                    none_variable,  # retain_graph
                    false_variable,  # create_graph
                    none_variable,  # grad_variables
                    none_variable,  # inputs
                ],
                {},
            )
            return ConstantVariable.create(None)
```

Dynamo will eventually hit the `UserDefinedObject.callfunction`

```
torch/_dynamo/variables/user_defined.py

......
        if (
            self.value.__str__()
            == "<method 'run_backward' of 'torch._C._EngineBase' objects>"
            and name == "__call__"
        ):
......
            tensors_example = tuple(
                [proxy.node.meta.get("example_value") for proxy in args[0].as_proxy()]
            )
            grad_tensors_example = tuple(
                proxy.node.meta.get("example_value") for proxy in args[1].as_proxy()
            )
            retain_graph = args[2].as_python_constant()
            create_graph = args[3].as_python_constant()
            # ignore the inputs for backward for now
            inputs = ()
            with compiled_autograd.enable(get_real_backward_fn):
                Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
                    tensors_example,
                    grad_tensors_example,
                    retain_graph,
                    create_graph,
                    inputs,
                    allow_unreachable=True,
                    accumulate_grad=True,
                )

```

Note that a couple things happened here



1. We convert both `tensor` and `grad` from `TensorVariable` to `fake_tensor`(by accessing the example value)
2. We run the backward function on fake tensors(example values) under the compiled autograd.

Dynamo will capture all inputs as `VariableTracker` but in order to run the compiled autograd we need a tensor object that’s tracked by autograd. Luckily we can use the `example value` associated with the `TensorVariable`.


##### Compiled autograd usage

Compiled Autograd typically works as
![image](https://github.com/pytorch/pytorch/assets/59073027/6a03fd49-5732-4f7a-bfae-fc83b740c9dc)


However we used it slightly differently 
![image](https://github.com/pytorch/pytorch/assets/59073027/ae5e0e78-8d5e-4531-84db-a257accf4535)


In this process we need to convert the `tensor` and `grad` from `TensorVariable` to `fake_tensor` to run backward inside `compiled_autograd`. After backward FX is generated, we need to convert all `fake_tensor` inputs back to `TensorVariables`. Note that `fake_tensors` inputs to the FX graph are not the same tensor we passed to the compiled autograd.

For a simple MNIST
```
def train_model(model, data, target):
  loss_fn = torch.nn.CrossEntropyLoss()
  pred = model(data)
  loss = loss_fn(pred, target)
  loss.backward()
  return pred

mnist = MNIST().to(device)
dynamo_train_model = torch.compile(train_model, backend='eager')
optimizer = optim.SGD(mnist.parameters(), lr=0.001, momentum=0.9)

for data, target in loader:
    optimizer.zero_grad()
    output = dynamo_train_model(mnist, data, target)

```

Backward fx looks like this
```
def forward(self, inputs, sizes, hooks):
    # No stacktrace found for following nodes
    getitem: "f32[]" = inputs[0]
    getitem_1: "f32[4, 10]" = inputs[1]
......
    getitem_23: "f32[10]" = inputs[23];  inputs = None
    nll_loss_backward: "f32[4, 10]" = torch.ops.aten.nll_loss_backward.default(getitem, getitem_1, getitem_2, None, 1, -100, getitem_3);  getitem = getitem_2 = getitem_3 = None
    _log_softmax_backward_data: "f32[4, 10]" = torch.ops.aten._log_softmax_backward_data.default(nll_loss_backward, getitem_1, 1, torch.float32);  nll_loss_backward = getitem_1 = None
    _log_softmax_backward_data_1: "f32[4, 10]" = torch.ops.aten._log_softmax_backward_data.default(_log_softmax_backward_data, getitem_4, 1, torch.float32);  _log_softmax_backward_data = getitem_4 = None
    t: "f32[10, 50]" = torch.ops.aten.t.default(getitem_6);  getitem_6 = None
    mm: "f32[4, 50]" = torch.ops.aten.mm.default(_log_softmax_backward_data_1, t);  t = None
.......
    t_7: "f32[50, 320]" = torch.ops.aten.t.default(t_6);  t_6 = None
.......
    accumulate_grad__7 = torch.ops.inductor.accumulate_grad_.default(getitem_23, view);  getitem_23 = view = None
    return []
```

It takes 24 input which includes all of the intermediate tensors we saved in the forward to compute the gradients of model parameters. This brings up another issue for design which is for each of these fake tensor inputs to the backward, dynamo needs to be able to convert it back to the TensorVariable. This means the dynamo needs to see all of these intermediate saved tensors in the forward. This part is actually surprisingly non-trival. 


#### Create FakeTensor -> TensorVariable mapping

We can already map `TensorVariable` to `FakeTensor` by `var.as_proxy.proxy.node.meta.get("example_value")` so we only need to implement the reverse mapping. I used a dummy approach which is to maintain a python dict of `FakeTensor` to `TensorVariable` whenever we create a new `TensorVaribale`
```
class TensorVariable(VariableTracker):

    def __init__(
......

        fake_tensor_to_tensor_variable_dict[
            self.proxy.node.meta.get("example_value")
        ] = self

```


#### Inline CompositeImplicitAutograd Ops


##### Why

Let’s take a look at MNIST \
```
class MNIST(nn.Module):

  def __init__(self):
    super(MNIST, self).__init__()
    self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
    self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
    self.fc1 = nn.Linear(320, 50)
    self.fc2 = nn.Linear(50, 10)

  def forward(self, x):
    x = F.relu(F.max_pool2d(self.conv1(x), 2))
    x = F.relu(F.max_pool2d(self.conv2(x), 2))
    x = x.view(-1, 320)
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return F.log_softmax(x, dim=1)
```


Without any change, forward FX captured looks like this
```
def forward(self, L_data_ : torch.Tensor, L_target_ : torch.Tensor):
    l_data_ = L_data_
    l_target_ = L_target_
    
    # File: xla/mnist.py:23, code: x = F.relu(F.max_pool2d(self.conv1(x), 2))
    l__model___conv1 = self.L__model___conv1(l_data_);  l_data_ = None
    max_pool2d = torch.nn.functional.max_pool2d(l__model___conv1, 2);  l__model___conv1 = None
    x = torch.nn.functional.relu(max_pool2d);  max_pool2d = None
    
    # File: xla/mnist.py:24, code: x = F.relu(F.max_pool2d(self.conv2(x), 2))
    l__model___conv2 = self.L__model___conv2(x);  x = None
    max_pool2d_1 = torch.nn.functional.max_pool2d(l__model___conv2, 2);  l__model___conv2 = None
    x_1 = torch.nn.functional.relu(max_pool2d_1);  max_pool2d_1 = None
    
    # File: xla/mnist.py:25, code: x = x.view(-1, 320)
    x_2 = x_1.view(-1, 320);  x_1 = None
    
    # File: xla/mnist.py:26, code: x = F.relu(self.fc1(x))
    l__model___fc1 = self.L__model___fc1(x_2);  x_2 = None
    x_3 = torch.nn.functional.relu(l__model___fc1);  l__model___fc1 = None
    
    # File: xla/mnist.py:27, code: x = self.fc2(x)
    x_4 = self.L__model___fc2(x_3);  x_3 = None
    
    # File: xla/mnist.py:28, code: return F.log_softmax(x, dim=1)
    pred = torch.nn.functional.log_softmax(x_4, dim = 1);  x_4 = None
    
    # File: xla/mnist.py:33, code: loss = loss_fn(pred, target)
    loss = torch.nn.functional.cross_entropy(pred, l_target_, None, None, -100, None, 'mean', 0.0);  l_target_ = None
    return (loss, pred)
```

This forward seems legit at first glance, but if we try to run the backward graph generated by the compiled autograd dynamo will complain that it doesn’t know how to map some of the `FakeTensor` needed by the backward to `TensorVariable`. This is because dynamo by default does not try to decompose torch ops. 

Let’s take `max_pool2d` as an example. `max_pool2d` will get decomposed in to `max_pool2d_with_indices` and the underlying functions looks something like

```
def forward(self, arg0_1, arg1_1):
    max_pool2d_with_indices = torch.ops.aten.max_pool2d_with_indices.default(arg0_1, [2, 2])
    getitem = max_pool2d_with_indices[0]
    getitem_1 = max_pool2d_with_indices[1];  max_pool2d_with_indices = None
    return getitem
```

Note that `max_pool2d_with_indices` has 2 outputs which are result and indices. This indices tensor will be saved for the backward, but this tensor is invisible to the dynamo by default. We have to inline all of the torch functions to the very bottom to make sure dynamo can capture all of the tensors that might be saved for the backward.


##### How

There are two parts of part



1. Which `op`/`torch function` we should decompose
2. How do we inline

We need to figure out the first one because otherwise we will be stuck in an infinite loop. My strategy was trying to inline all `torch functions` until we hit a torch op that’s not CompositeImplicitAutograd. Note that not all `torch` functions are ops. For example <code>[torch.nn.functional.max_pool2d](torch.nn.functional.max_pool2d)</code> is not a <code>op</code> from python class perspective, but it will get decomposed into <code>torch.ops.aten.max_pool2d_with_indices.default</code>.

I also kept a black list for some functions that we don’t need to decompose.

```
torch/_dynamo/variables/[torch.py](http://torch.py/)

def should_decompose_torch_op(fn):
    from torch._dynamo import compiled_autograd

    # TODO([JackCaoG](https://who.corp.google.com/JackCaoG)): we need a better way to tell if a torch function should we decompose
    allowed_torch_fn = fn.__name__ != "_make_grads"
    definanilly_not_composite_kernel = type(
        fn
    ) == torch._ops.OpOverload and not torch._C._dispatch_has_kernel_for_dispatch_key(
        fn.name(), torch._C.DispatchKey.CompositeImplicitAutograd
    )

    # only decompoization torch ops for forward
    in_compiled_backward = compiled_autograd.compiled_autograd_enabled
    return (
        config.use_single_step_graph
        and allowed_torch_fn
        and not definanilly_not_composite_kernel
        and not in_compiled_backward
    )
```

One thing to note is that we only need to perform this decomposition in the forward. For the FX graph generated by the compiled autograd we don’t need to inline them.

Now knowing which function we should decompose we can use `make_fx`(inspired by how `aot_autograd` performs decomp) to generate the FX graph representing the decomposed op. Similar to how to inline the FX graph generated by the compiled autograd, we can inline the FX graph generated by `make_fx` here.

```
torch/_dynamo/variables/[torch.py](http://torch.py/)

    with enable_python_dispatcher():
        fx_g = make_fx(op, pre_dispatch=True)(*complete_fake_value_args)

    user_fn_variable = SourcelessBuilder()(tx, dummy_user_function_to_inline_gm)
    gm_variable = SourcelessBuilder()(tx, fx_g)
.....
    res = tx.inline_user_function_return(
    user_fn_variable, (gm_variable, input_list_variable), {}
    )
```


#### Inline loss functions

Loss functions are slightly different torch ops because it returns another python function that can be called with inputs.

For the purpose of this POC, I used `CrossEntropyLoss` which actually has a special handling in dynamo. `CrossEntropyLoss` decomposes into `logsoftmax` and other ops, so we just need to make sure we force the decomposition and inline the returned FX similar to the regular torch function above.


#### Accumulated\_grad\_


accumulated_grad_ takes a tensor and a grad then updates the tensor's gradient in place. It is currently being captured as a torch op in the graph returned by the compield_autograd

```
    accumulate_grad__7 = torch.ops.inductor.accumulate_grad_.default(getitem_23, view);  getitem_23 = view = None
```


There are some issue with this op that it works perfectly on CPU and inductor but on XLA device it becomes a no op. In order to unblock myself I have to implement a dummy version of the `accumulated_grad_` and inline it as a normal python function.


```
def dummy_accumulate_grad_(t1, t2):
   if t1.grad is None:
       t1.grad = t2
   else:
       t1.grad += t2
```

## Known Issues


### Missing TensorVariable due to `detach` in autograd

Dynamo needs to see all tensors in the backward to make inlining backward fx graph works. However in the cases that output of an op(`output = relu(a)` for example) needs to be saved for the backward, autograd will save the `output.tensor_data` and call `detach` to create the input used for the backward. This is to avoid the reference cycle since `output` will point to the `grad_fn` object so `grad_fn.saved_xxx` should not point back to the `output`. Currently my workaround is to save the `output` directly in the autograd. More details in https://github.com/pytorch/pytorch/pull/112296#issuecomment-1853155214. 


### `accumulate_grad_` op can not be executed properly in XLA

`Accumulate_grad_` is an op that will show up in the backward graph generated by the `compiled_autograd`. However this op currently appears to be registered under inductor. It works for the CPU, but when running with XLA I hit a weird error. I am wondering if this op can be something that’s lowerable or decomposeable. Other options might be to implement `accumulate_grad_` in python and let dynamo inlining it. Currently my workaround is to implement a dummy `accumulate\_grad\_` in python


```
def dummy_accumulate_grad_(t1, t2):
   if t1.grad is None:
       t1.grad = t2
   else:
       t1.grad += t2
```



### `tensor.grad` is not captured in the fwd but might be used in the backward

This issue might be an artifact of the previous issue. When the grad of the tensor is not empty, `accumulated_grad_` will take `tensor.grad` as an input. However `tensor.grad` is not captured in the fwd, so we will have a missing tensor variable. I don’t have a workaround for this issue yet.


### Lacking a good rule of thumb to decide when to decompose a <code>[torch.xxx](http://torch.xxx)</code> functions

I haven’t found a reliable way to decide when to decompose a torch ops. It is easier for ops like `torch.cos`, `torch.sum` since I can use helpers similar to `torch._C._dispatch_has_kernel_for_dispatch_key` to find out if it is already a core aten ops. However for functions like <code>[torch.nn.functional.log_softmax](http://torch.nn.functional.log_softmax)</code> which should be decomposed but themselves are not a torch op. My current workaround is to try to decompose all <code>[torch.xxx](http://torch.xxx)</code> functions and blacklist the ops that will run into issues like `make\_grads`.


## Next Steps

On 01/11/2024’s composability sync I presented this doc and most people in the meeting agree that current design is directly correct. In my opinion there are couple thing we need to do next



1. Fix a couple reaming issues including version counter issue when `relu\_` is used.
2. Works with PyTorch maintainers to tackle the known issues
3. Break this POC in pieces and refine the implementation of each subcomponent and land them in pieces.
4. Try to benchmark the performance of single step graph vs current dynamo implementation and understand the performance differences on smaller models(ResNet for example).

### Alternatives

Directly use `make_fx` to generate the backward and stich the fwd and bwd together. However this way we won't be able to handle the backward hook which is use heavily for FSDP.

### Additional context

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @msaroufim @bdhirsh @anijain2305 @chauhang @wconstab @voznesenskym @jansel