Apply suggestions from code review

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

Author: anijain2305

recipes_source/recipes/regional_compilation.py

@@ -2,14 +2,17 @@
 Reducing torch.compile cold start compilation time with regional compilation
 ============================================================================
 
-Introduction
-------------
+**Author:** `Animesh Jain <https://github.com/anijain2305>`_
 As deep learning models get larger, the compilation time of these models also
-increase. This increase in compilation time can lead to a large startup time in
-inference services or wasted resources in large scale training. This recipe
+increases. This extended compilation time can result in a large startup time in
+inference services or wasted resources in large-scale training. This recipe
 shows an example of how to reduce the cold start compilation time by choosing to
 compile a repeated region of the model instead of the entire model.
 
+Prerequisites
+----------------
+
+* Pytorch 2.5 or later
 Setup
 -----
 Before we begin, we need to install ``torch`` if it is not already
@@ -19,6 +22,10 @@
 
    pip install torch
 
+.. note:: 
+   This feature is available starting with the 2.5 release. If you are using version 2.4,
+   you can enable the configuration flag ``torch._dynamo.config.inline_inbuilt_nn_modules=True``
+   to prevent recompilations during regional compilation. In version 2.5, this flag is enabled by default.
 """
 
 
@@ -27,28 +34,30 @@
 # Steps
 # -----
 # 
-# 1. Import all necessary libraries
+# In this recipe, we will follow these steps:
+#
+# 1. Import all necessary libraries.
 # 2. Define and initialize a neural network with repeated regions.
 # 3. Understand the difference between the full model and the regional compilation.
 # 4. Measure the compilation time of the full model and the regional compilation.
 # 
-# 1. Import necessary libraries for loading our data
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# First, let's import the necessary libraries for loading our data: 
+# 
 # 
 # 
 
 import torch
 import torch.nn as nn
 from time import perf_counter
 
+##########################################################
+# Next, let's define and initialize a neural network with repeated regions.
 # 
-# 2. Define and initialize a neural network with repeated regions.
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-# Typically neural networks are composed of repeated layers. For example, a
+# Typically, neural networks are composed of repeated layers. For example, a
 # large language model is composed of many Transformer blocks. In this recipe,
-# we will create a `Layer` `nn.Module` class as a proxy for a repeated region.
-# We will then create a `Model` which is composed of 64 instances of this
-# `Layer` class.
+# we will create a ``Layer`` using the ``nn.Module`` class as a proxy for a repeated region.
+# We will then create a ``Model`` which is composed of 64 instances of this
+# ``Layer`` class.
 # 
 class Layer(torch.nn.Module):
     def __init__(self):
@@ -83,53 +92,51 @@ def forward(self, x):
             x = layer(x)
         return x
 
-#
-# 3. Understand the difference between the full model and the regional compilation.
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+####################################################
+# Next, let's review the difference between the full model and the regional compilation.
 # 
-# In full model compilation, the full model is compiled as a whole. This is how
-# most users use torch.compile. In this example, we can apply torch.compile to
-# the `model` object. This will effectively inline the 64 layers, producing a
+# In full model compilation, the entire model is compiled as a whole. This is the common approach 
+# most users take with ``torch.compile``. In this example, we apply ``torch.compile`` to
+# the ``Model`` object. This will effectively inline the 64 layers, producing a
 # large graph to compile. You can look at the full graph by running this recipe
-# with `TORCH_LOGS=graph_code`.
+# with ``TORCH_LOGS=graph_code``.
 # 
 #
 
 model = Model(apply_regional_compilation=False).cuda()
 full_compiled_model = torch.compile(model)
 
 
-# 
+###################################################
 # The regional compilation, on the other hand, compiles a region of the model.
-# By wisely choosing to compile a repeated region of the model, we can compile a
-# much smaller graph and then reuse the compiled graph for all the regions. We
-# can apply regional compilation in the example as follows. `torch.compile` is
-# applied only to the `layers` and not the full model.
+# By strategically choosing to compile a repeated region of the model, we can compile a
+# much smaller graph and then reuse the compiled graph for all the regions.
+# In the example, ``torch.compile`` is applied only to the ``layers`` and not the full model.
 # 
 
 regional_compiled_model = Model(apply_regional_compilation=True).cuda()
 
 # Applying compilation to a repeated region, instead of full model, leads to
 # large savings in compile time. Here, we will just compile a layer instance and
-# then reuse it 64 times in the `model` object.
+# then reuse it 64 times in the ``Model`` object.
 # 
 # Note that with repeated regions, some part of the model might not be compiled.
-# For example, the `self.linear` in the `Model` is outside of the scope of
+# For example, the ``self.linear`` in the ``Model`` is outside of the scope of
 # regional compilation.
 # 
 # Also, note that there is a tradeoff between performance speedup and compile
-# time.  The full model compilation has larger graph and therefore,
-# theoretically, has more scope for optimizations. However for practical
+# time. Full model compilation involves a larger graph and,
+# theoretically, offers more scope for optimizations. However, for practical
 # purposes and depending on the model, we have observed many cases with minimal
 # speedup differences between the full model and regional compilation.
 
 
+###################################################
+# Next, let's measure the compilation time of the full model and the regional compilation.
 #
-# 4. Measure the compilation time of the full model and the regional compilation.
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-# `torch.compile` is a JIT compiler, i.e., it compiles on the first invocation.
-# Here, we measure the total time spent in the first invocation. This is not
-# precise, but it gives a good idea because the majority of time is spent in
+# ``torch.compile`` is a JIT compiler, which means that it compiles on the first invocation.
+# In the code below, we measure the total time spent in the first invocation. While this method is not 
+# precise, it provides a good estimate since the majority of the time is spent in
 # compilation.
 
 def measure_latency(fn, input):
@@ -150,11 +157,13 @@ def measure_latency(fn, input):
 print(f"Regional compilation time = {regional_compilation_latency:.2f} seconds")
 
 ############################################################################
+# Conclusion
+# -----------
+#
 # This recipe shows how to control the cold start compilation time if your model
-# has repeated regions. This requires user changes to apply `torch.compile` to
+# has repeated regions. This approach requires user modifications to apply `torch.compile` to
 # the repeated regions instead of more commonly used full model compilation. We
-# are continually working on reducing cold start compilation time. So, please
-# stay tuned for our next tutorials.
+# are continually working on reducing cold start compilation time.
 # 
 # This feature is available with 2.5 release. If you are on 2.4, you can use a
 # config flag - `torch._dynamo.config.inline_inbuilt_nn_modules=True` to avoid