Adjust torch.compile() best practices (#3336)

* Adjust torch.compile() best practices

1. Add best practice to prefer `mod.compile` over `torch.compile(mod)`, which avoids `_orig_` naming problems.
Repro steps:
- opt_mod = torch.compile(mod)
- train opt_mod
- save checkpoint
In another script, potentially on a machine that does NOT support `torch.compile`: load checkpoint.
This fails with an error, because the checkpoint on `opt_mod` got its params renamed by `torch.compile`:
```
RuntimeError: Error(s) in loading state_dict for VQVAE:
	Missing key(s) in state_dict: "embedding.weight", "encoder.encoder.net.0.weight", "encoder.encoder.net.0.bias", ...
	Unexpected key(s) in state_dict: "_orig_mod.embedding.weight", "_orig_mod.encoder.encoder.net.0.weight", "_orig_mod.encoder.encoder.net.0.bias", ...
```
- Add best practice to use, or at least try, `fullgraph=True`. This doesn't always work, but we should encourage it.

---------

Co-authored-by: Svetlana Karslioglu <[email protected]>

Author: punkeel

intermediate_source/torch_compile_tutorial.py

@@ -101,8 +101,11 @@ def forward(self, x):
         return torch.nn.functional.relu(self.lin(x))
 
 mod = MyModule()
-opt_mod = torch.compile(mod)
-print(opt_mod(t))
+mod.compile()
+print(mod(t))
+## or:
+# opt_mod = torch.compile(mod)
+# print(opt_mod(t))
 
 ######################################################################
 # torch.compile and Nested Calls
@@ -135,8 +138,8 @@ def forward(self, x):
         return torch.nn.functional.relu(self.outer_lin(x))
 
 outer_mod = OuterModule()
-opt_outer_mod = torch.compile(outer_mod)
-print(opt_outer_mod(t))
+outer_mod.compile()
+print(outer_mod(t))
 
 ######################################################################
 # We can also disable some functions from being compiled by using
@@ -197,6 +200,12 @@ def outer_function():
 # 4. **Compile Leaf Functions First:** In complex models with multiple nested
 # functions and modules, start by compiling the leaf functions or modules first.
 # For more information see `TorchDynamo APIs for fine-grained tracing <https://pytorch.org/docs/stable/torch.compiler_fine_grain_apis.html>`__.
+#
+# 5. **Prefer ``mod.compile()`` over ``torch.compile(mod)``:** Avoids ``_orig_`` prefix issues in ``state_dict``.
+#
+# 6. **Use ``fullgraph=True`` to catch graph breaks:** Helps ensure end-to-end compilation, maximizing speedup
+# and compatibility with ``torch.export``.
+
 
 ######################################################################
 # Demonstrating Speedups