Decomposition to Aten IR

Jerome_Ku: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a></p>
<p>Trying to understand how the decomposition of <code>torch.nn.functional.linear</code> is traced.</p>
<p>Example code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch._inductor.decomposition import select_decomp_table
from torch.fx import GraphModule
from torch.fx.experimental.proxy_tensor import make_fx

device = "cuda"
x = torch.randn(10, 10).to(device)
y = torch.randn(10, 10).to(device)

g: GraphModule = make_fx(torch.nn.functional.linear, select_decomp_table(), tracing_mode="real")(x, y)
g.print_readable()
</code></pre>
<p>This prints</p>
<pre><code class="lang-auto">class &lt;lambda&gt;(torch.nn.Module):
    def forward(self, arg0_1: "f32[10, 10]", arg1_1: "f32[10, 10]"):
        # No stacktrace found for following nodes
        permute: "f32[10, 10]" = torch.ops.aten.permute.default(arg1_1, [1, 0]);  arg1_1 = None
        mm: "f32[10, 10]" = torch.ops.aten.mm.default(arg0_1, permute);  arg0_1 = permute = None
        return mm
</code></pre>
<p>I’ve tried to step through the code using <code>pdb</code> but still can’t figure out how the decomposition of <code>linear</code> is happening:</p>
<ul>
<li>where is the mapping of <code>linear</code> to <code>permute</code> and <code>mm</code>?</li>
<li>if there is no direct mapping, how can I determine where along the trace these registered ops are being called?</li>
</ul>
<p>Noticed that something similar is called during the pregrad passes that <code>inductor</code> runs to normalize the IR into something more amenable for downstream passes.</p>
<p>Also related are the <a href="https://github.com/pytorch/workshops/tree/master/ASPLOS_2024/aot_autograd.pdf" rel="noopener nofollow ugc">AOTAutograd</a> slides from ASPLOS 2024.  On slide 23, it is mentioned that AOTAutograd handles “operator decomposition” among other Pytorch framework related functions that occur at the C++ level, and it does this through <code>__torch_dispatch__</code>.</p>
<p>Where / how does this call back to Python happen after the dispatching, autograd, and decomposition work is handled in C++?</p>
<p>Thanks!</p>
<p>-Jerome</p>

bdhirsh: <p>Hey <a class="mention" href="/u/jerome_ku">@Jerome_Ku</a>. The short answer is:</p>
<p>We have a handful of operator decompositions that always run by-default, inside of the dispatcher (in C++), before making it into <code>__torch_dispatch__</code>.</p>
<p>linear/matmul is by far the most common</p>
<ul>
<li>
<p><code>aten.linear</code> decomposes into <code>aten.matmul</code> + <code>aten.transpose</code> here: <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/Linear.cpp#L111" class="inline-onebox" rel="noopener nofollow ugc">pytorch/aten/src/ATen/native/Linear.cpp at main · pytorch/pytorch · GitHub</a></p>
</li>
<li>
<p><code>aten.matmul</code> decomposes into a few ops (either <code>aten.mm</code>, <code>aten.baddbmm</code> or a few others) here: <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/LinearAlgebra.cpp#L1995" class="inline-onebox" rel="noopener nofollow ugc">pytorch/aten/src/ATen/native/LinearAlgebra.cpp at main · pytorch/pytorch · GitHub</a></p>
</li>
</ul>
<p>If you’re wondering why, the historical answer is that there are a number of ops that we don’t have dedicated derivative formulas for (e.g. linear), and so rather than writing a brand new formula, we just have the autograd engine decompose the op into more primitive ops that it does have formulas for (e.g. <code>aten.mm</code> and transpose).</p>
<p>If you are interested in export and only care about inference, we actually recently made it so that exporting for inference can preserve all ATen ops, including these special ops like linear, in the graph:</p>
<pre><code class="lang-auto">m = torch.nn.Linear(...)
graph_module = torch.export.export(m, args).run_decompositions().module()
# you should see aten.linear, as long as you didn't manually specify that you want it decomposed
print(graph_module)
</code></pre>

Jerome_Ku: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a></p>
<p>Thanks for the response!</p>
<p>Trying to get a better mental picture of how the PT2 stack fits together in the context of more complex use cases such as tensor subclasses (quantized types, <code>DTensor</code>) and distributed training (<code>FSDP</code>).</p>
<p>Hoping you can help shed some light on the following or provide some pointers where to dig deeper (I’ve reviewed the relevant presos from the most recent Pytorch conf as well as the ASPLOS tutorial):</p>
<ul>
<li>How is <code>AOTAutograd</code> able to trace through torch framework code that lives in C++ (decomp, functionalization, etc) that dynamo can’t (since it only “sees” Python bytecode)?</li>
<li>Why is there a need for <code>compiled_autograd</code> in addition to <code>AOTAutograd</code>?  I.e., why aren’t <code>AccumulateGrad</code> nodes able to be included in the compiled backwards graph?</li>
<li>How do each of these components –  <code>dynamo</code>, <code>AOTAutograd</code>, <code>compiled_autograd</code>, and other parts of the PT2 stack – interact to enable tracing and compilation of larger, more representative graphs?</li>
<li>What are the limitations of each component that necessitates the others and at what level are they operating (Python, C++)?</li>
<li>For example, when compiling <code>FSDP</code>-wrapped models, how do each of these parts fit together?</li>
</ul>
<p>Many thanks!</p>

bdhirsh: <blockquote>
<ul>
<li>H  ow is <code>AOTAutograd</code> able to trace through torch framework code that lives in C++ (decomp, functionalization, etc) that dynamo can’t (since it only “sees” Python bytecode)?</li>
</ul>
</blockquote>
<p>The short answer is <code>__torch_dispatch__</code>. If you want some code pointers: we have another API <code>make_fx</code>, which is effectively a tracer that uses <code>__torch_dispatch__</code> to capture every ATen op that gets executed into a graph (including ATen ops that are executed in C++, e.g. through the autograd engine). <code>make_fx</code> lives <a href="https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/proxy_tensor.py#L2163" rel="noopener nofollow ugc">here</a>, called by AOTAutograd <a href="https://github.com/pytorch/pytorch/blob/main/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py#L55" rel="noopener nofollow ugc">here</a>).</p>
<p>If you want to naively capture a graph of the forward and backward (without all of the bells and whistles of functionalizing the graph, or safety from dynamo guards), <code>make_fx</code> is a relatively light-weight tracer that you can use to do that:</p>
<pre><code class="lang-auto"># captures and prints a graph of the combined forward and backward, tracing through the C++ autograd engine
import torch
from torch.fx.experimental.proxy_tensor import make_fx

def f(x, w):
    return torch.matmul(x, w).sum().backward()

x = torch.randn(64, 64)
w = torch.randn(64, 64, requires_grad=True)
fx_graph = make_fx(f)(x, w)
fx_graph.print_readable()
</code></pre>
<blockquote>
<ul>
<li>Why is there a need for <code>compiled_autograd</code> in addition to <code>AOTAutograd</code>? I.e., why aren’t <code>AccumulateGrad</code> nodes able to be included in the compiled backwards graph?</li>
</ul>
</blockquote>
<p>The answer here (at least my interpretration) is something like: <code>torch.compile</code> captures <em>almost</em> everything in the backward that is worth compiling, but there are a few distinct things in the backward that are difficult to capture, where having dynamo run on the backward graph is helpful:</p>
<ul>
<li>user backward hooks. in particular, backward hooks involve tracing through user code. This code might involve specializing on globals, or performing python side effects that we can’t safely graph capture, and so we need dynamo to (somehow) be in the loop.</li>
<li>AccumulateGrad. There are a few things that make capturing AccumulateGrad difficult  (naively, without compiled autograd) , mainly around handling all of the edge cases: handling post accumulate grad hooks, or ensuring that we properly guard on any behavior that AccumulateGrad specializes on</li>
<li>slightly better graph break behavior. Normally if there are two graphs in the forward, compile generates two graphs in the backward. Compiled autograd can instead ensure that we get a single graph for the entire backward, even if there are graph breaks in the forward.</li>
</ul>
<p>Compiled autograd morally works by taking the existing autograd tape that lives in the C++ autograd engine, and retracing all of it into a python FX graph that we can then run dynamo on and compile, which handles all of those cases nicely.</p>
<blockquote>
<ul>
<li>How do each of these components – <code>dynamo</code>, <code>AOTAutograd</code>, <code>compiled_autograd</code>, and other parts of the PT2 stack – interact to enable tracing and compilation of larger, more representative graphs?.. What are the limitations of each component that necessitates the others and at what level are they operating (Python, C++)?</li>
</ul>
</blockquote>
<p>I might be missing part of your question. But the short answer is that they are all (complementary) tracing systems.</p>
<p>Dynamo has the advantage that it can handle arbitrary user python code, guarding on any implicitly specialized code, and falling back to cpython when necessary, and emits a graph full of <code>torch.*</code> IR.  AOTAutograd is a “dumb” tracer that can trace through arbitrary code, and also further trace through the C++ dispatcher to emit a graph of <code>aten.*</code> IR.</p>
<p>Dynamo first generates a graph of torch ops (falling back to cpython and guarding on any specializations along the way), and after proving safety it sends this torch IR graph to AOTAutograd to be lowered further, also generating a backward graph. That backward graph gets inserted as a <code>CompiledFunctionBackward</code> node into the eager autograd engine’s tape.</p>
<p>Compiled autograd is effectively an extra layer that runs at the time a user calls <code>.backward()</code>. Instead of letting the eager autograd engine execute each node  in the tape eagerly, it traces the autograd tape itself back into an FX graph, running dynamo (and the rest of the stack) on the resulting FX graph.</p>
<blockquote>
<ul>
<li>For example, when compiling <code>FSDP</code>-wrapped models, how do each of these parts fit together?</li>
</ul>
</blockquote>
<p>Hmm i’m not sure if I’m giving you a full answer. But FSDP is effectively a piece of python framework code, so the main differences in the PT2 stack around FSDP handling are mostly in dynamo. Dynamo will graph break on bits of FSDP that are difficult to capture.  On thing worth noting though is that FSDP is implemented partially with autograd hooks, and so compiled autograd is one of the tools that we need to capture them properly</p>

Jerome_Ku: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a></p>
<p>This is a fantastic explanation – much appreciated.</p>
<p>A quick follow-up: how does <code>AOTAutograd</code> use <code>__torch_dispatch__</code> to trace the backwards graph?  Does it unroll the autograd tape that is recorded from the forward graph in something like FakeTensorMode and generate an fx graph from the aten ops that <code>__torch_dispatch__</code> sees (at a Python level)?</p>

bdhirsh: <p>It’s actually even simpler than that - we are effectively <strong>running</strong> the autograd engine as-is, except every operation that gets called ends up hitting our <code>__torch_dispatch__</code> handler before we need to run its implementation.</p>
<p>A simple example might make things more clear. Below is a simple torch_dispatch mode that just prints every op it sees.</p>
<pre><code class="lang-auto">import torch
from torch.utils._python_dispatch import TorchDispatchMode

class LoggingMode(TorchDispatchMode):
    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        rs = func(*args, **kwargs)
        print(f'func={str(func)}')
        return rs

def f(x):
    out = x.sin()
    out.sum().backward()

x = torch.randn(4, requires_grad=True)
with LoggingMode():
    f(x)
</code></pre>
<p>If you run it, you’ll see that both <code>aten.sin</code> and <code>aten.cos</code> are printed. This maps to the fact that autograd engine has a <code>SinBackward</code> node saved somewhere in the autograd tape - and upon running the backward, autograd will execute that node, which eventually calls <code>aten.cos</code> (the derivative), which we intercept:</p>
<pre><code class="lang-auto">func=aten.sin.default
func=aten.sum.default
func=aten.ones_like.default
func=aten.expand.default
func=aten.cos.default
func=aten.mul.Tensor
func=aten.detach.default
func=aten.detach.default
</code></pre>
<p>in the <code>make_fx</code> case, our <code>__torch_dispatch__</code> code is effectively intercepting each op and adding it to an FX graph (code <a href="https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/proxy_tensor.py#L901" rel="noopener nofollow ugc">here</a>)</p>

Jerome_Ku: <p><a class="mention" href="/u/bdhirsh">@bdhirsh</a></p>
<p>Awesome – thanks for the great explanation.  All clear now!</p>