Add compiled autograd tutorial

Author: xmfan

intermediate_source/compiled_autograd_tutorial.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+
+"""
+Compiled Autograd: Capturing a larger backward graph for ``torch.compile``
+==========================================================================
+
+"""
+
+######################################################################
+# Compiled Autograd is a torch.compile extension introduced in PyTorch 2.4
+# that allows the capture of a larger backward graph. It is highly recommended
+# to familiarize yourself with `torch.compile <https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html>`_.
+# 
+
+######################################################################
+# Doesn't torch.compile already capture the backward graph?
+# ------------
+# Partially. AOTAutograd captures the backward graph ahead-of-time, but with certain limitations:
+# - Graph breaks in the forward lead to graph breaks in the backward
+# - `Backward hooks <https://pytorch.org/docs/stable/notes/autograd.html#backward-hooks-execution>`_ are not captured
+# 
+# Compiled Autograd addresses these limitations by directly integrating with the autograd engine, allowing
+# it to capture the full backward graph at runtime. Models with these two characteristics should try
+# Compiled Autograd, and potentially observe better performance.
+#
+
+######################################################################
+# Basic Usage
+# ------------
+#
+
+import torch
+
+class Model(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(10, 10)
+
+    def forward(self, x):
+        return self.linear(x)
+
+torch._dynamo.config.compiled_autograd = True
+
+@torch.compile
+def train(model, x):
+    loss = model(x).sum()
+    loss.backward()
+
+######################################################################
+# TODO: add an image of the graph
+# Note: In a future release, we will prevent the graph break caused by loss.backward()
+# 
+
+######################################################################
+# Compiling the forward and backward pass using different flags
+# ------------
+# 
+
+def train(model, x):
+    model = torch.compile(model)
+    loss = model(x).sum()
+    torch.compile(lambda: loss.backward(), fullgraph=True)(loss)
+
+######################################################################
+# Appendix: Compatibility
+# ------------
+#
+# Compiled Autograd is not yet compatible with all existing PyTorch features.
+# Below is a list of known incompatibilities.