Merge branch 'main' into svekars-patch-54

Author: sekyondaMeta

beginner_source/introyt/modelsyt_tutorial.py

@@ -144,7 +144,7 @@ def forward(self, x):
 # We saw convolutional layers in action in LeNet5 in an earlier video:
 # 
 
-import torch.functional as F
+import torch.nn.functional as F
 
 
 class LeNet(torch.nn.Module):

intermediate_source/torch_compile_full_example.py

@@ -235,6 +235,6 @@ def train(mod, data):
 #
 # For a gentle introduction to ``torch.compile``, please check out `the introduction to torch.compile tutorial <https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html>`__.
 #
-# To troubleshoot issues and to gain a deeper understanding of how to apply ``torch.compile`` to your code, check out `the torch.compile programming model <https://docs.pytorch.org/docs/main/compile/programming_model.html>`__.
+# To troubleshoot issues and to gain a deeper understanding of how to apply ``torch.compile`` to your code, check out `the torch.compile programming model <https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/compile/programming_model.html>`__.
 #
 # We hope that you will give ``torch.compile`` a try!

recipes_source/distributed_async_checkpoint_recipe.rst

@@ -1,9 +1,9 @@
 Asynchronous Saving with Distributed Checkpoint (DCP)
 =====================================================
 
-**Author:** `Lucas Pasqualin <https://github.com/lucasllc>`__, `Iris Zhang <https://github.com/wz337>`__, `Rodrigo Kumpera <https://github.com/kumpera>`__, `Chien-Chin Huang <https://github.com/fegin>`__
+**Author:** `Lucas Pasqualin <https://github.com/lucasllc>`__, `Iris Zhang <https://github.com/wz337>`__, `Rodrigo Kumpera <https://github.com/kumpera>`__, `Chien-Chin Huang <https://github.com/fegin>`__, `Yunsheng Ni <https://github.com/niyunsheng>`__
 
-Checkpointing is often a bottle-neck in the critical path for distributed training workloads, incurring larger and larger costs as both model and world sizes grow.
+Checkpointing is often a bottleneck in the critical path for distributed training workloads, incurring larger and larger costs as both model and world sizes grow.
 One excellent strategy for offsetting this cost is to checkpoint in parallel, asynchronously. Below, we expand the save example
 from the `Getting Started with Distributed Checkpoint Tutorial <https://github.com/pytorch/tutorials/blob/main/recipes_source/distributed_checkpoint_recipe.rst>`__
 to show how this can be integrated quite easily with ``torch.distributed.checkpoint.async_save``.
@@ -26,7 +26,7 @@ to show how this can be integrated quite easily with ``torch.distributed.checkpo
 
 Asynchronous Checkpointing Overview
 ------------------------------------
-Before getting started with Asynchronous Checkpointing, it's important to understand it's differences and limitations as compared to synchronous checkpointing.
+Before getting started with Asynchronous Checkpointing, it's important to understand its differences and limitations as compared to synchronous checkpointing.
 Specifically:
 
 * Memory requirements - Asynchronous checkpointing works by first copying models into internal CPU-buffers.
@@ -35,8 +35,8 @@ Specifically:
     the memory constraints of their systems. Specifically, pinned memory implies the usage of ``page-lock`` memory, which can be scarce as compared to
     ``pageable`` memory.
 
-* Checkpoint Management - Since checkpointing is asynchronous, it is up to the user to manage concurrently run checkpoints. In general, users can
-    employ their own management strategies by handling the future object returned form ``async_save``. For most users, we recommend limiting
+* Checkpoint Management - Since checkpointing is asynchronous, it is up to the user to manage concurrently run checkpoints.
+    In general, users can employ their own management strategies by handling the future object returned form ``async_save``. For most users, we recommend limiting
     checkpoints to one asynchronous request at a time, avoiding additional memory pressure per request.
 
 
@@ -72,7 +72,7 @@ Specifically:
             self.optimizer = optimizer
 
         def state_dict(self):
-            # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
+            # this line automatically manages FSDP FQNs, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
             model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
             return {
                 "model": model_state_dict,
@@ -196,7 +196,7 @@ checkpoint requests users can take advantage of direct memory access to speed up
             self.optimizer = optimizer
 
         def state_dict(self):
-            # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
+            # this line automatically manages FSDP FQNs, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
             model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
             return {
                 "model": model_state_dict,
@@ -279,6 +279,154 @@ checkpoint requests users can take advantage of direct memory access to speed up
         )
 
 
+Fully Asynchronous Staging with DefaultStager
+--------------------------------------------
+
+.. versionadded:: 2.9
+   The ``async_stager`` argument and ``DefaultStager`` class were introduced in PyTorch 2.9.
+
+While ``async_save`` handles the disk write asynchronously, the process of copying data from GPU to CPU (known as "staging") typically happens on the main thread. Even with Pinned Memory, this Device-to-Host (D2H) copy can block the training loop for large models.
+
+To achieve maximum overlap between computation and checkpointing, we can use the ``DefaultStager``. This component offloads the state dictionary creation and the D2H copy to a background thread.
+
+**Timeline Comparison:**
+
+* **Standard async_save:** ``[GPU Compute] -> [CPU Copy (Blocking)] -> [Disk Write (Async)]``
+* **With AsyncStager:** ``[GPU Compute] || [CPU Copy (Async)] -> [Disk Write (Async)]``
+
+.. note::
+    Using ``AsyncStager`` introduces a background thread that consumes CPU resources. Ensure your environment has sufficient CPU cores to handle this without impacting the main training process.
+
+.. code-block:: python
+
+    import os
+
+    import torch
+    import torch.distributed as dist
+    import torch.distributed.checkpoint as dcp
+    import torch.multiprocessing as mp
+    import torch.nn as nn
+
+    from torch.distributed.fsdp import fully_shard
+    from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
+    from torch.distributed.checkpoint.stateful import Stateful
+    from torch.distributed.checkpoint.staging import DefaultStager
+    from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
+
+    CHECKPOINT_DIR = "checkpoint"
+
+
+    class AppState(Stateful):
+        """This is a useful wrapper for checkpointing the Application State. Since this object is compliant
+        with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the
+        dcp.save/load APIs.
+
+        Note: We take advantage of this wrapper to hande calling distributed state dict methods on the model
+        and optimizer.
+        """
+
+        def __init__(self, model, optimizer=None):
+            self.model = model
+            self.optimizer = optimizer
+
+        def state_dict(self):
+            # this line automatically manages FSDP FQNs, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT
+            model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
+            return {
+                "model": model_state_dict,
+                "optim": optimizer_state_dict
+            }
+
+        def load_state_dict(self, state_dict):
+            # sets our state dicts on the model and optimizer, now that we've loaded
+            set_state_dict(
+                self.model,
+                self.optimizer,
+                model_state_dict=state_dict["model"],
+                optim_state_dict=state_dict["optim"]
+            )
+
+    class ToyModel(nn.Module):
+        def __init__(self):
+            super(ToyModel, self).__init__()
+            self.net1 = nn.Linear(16, 16)
+            self.relu = nn.ReLU()
+            self.net2 = nn.Linear(16, 8)
+
+        def forward(self, x):
+            return self.net2(self.relu(self.net1(x)))
+
+
+    def setup(rank, world_size):
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = "12355 "
+
+        # initialize the process group
+        dist.init_process_group("gloo", rank=rank, world_size=world_size)
+        torch.cuda.set_device(rank)
+
+
+    def cleanup():
+        dist.destroy_process_group()
+
+
+    def run_fsdp_checkpoint_save_example(rank, world_size):
+        print(f"Running basic FSDP checkpoint saving example on rank {rank}.")
+        setup(rank, world_size)
+
+        # create a model and move it to GPU with id rank
+        model = ToyModel().to(rank)
+        model = fully_shard(model)
+
+        loss_fn = nn.MSELoss()
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
+
+        checkpoint_future = None
+        for step in range(10):
+            print(f"Step {step} starting...")
+            optimizer.zero_grad()
+            model(torch.rand(8, 16, device="cuda")).sum().backward()
+
+            # Critical: We must ensure the previous checkpoint's D2H copy (staging)
+            # is complete before the optimizer modifies the model parameters.
+            # Placing this await AFTER the backward pass allows us to overlap
+            # the D2H copy with the current step's Forward and Backward computation.
+            if checkpoint_future is not None:
+                checkpoint_future.staging_completion.result()
+            optimizer.step()
+
+            # waits for checkpointing to finish if one exists, avoiding queuing more then one checkpoint request at a time
+            if checkpoint_future is not None:
+                checkpoint_future.upload_completion.result()
+
+            state_dict = { "app": AppState(model, optimizer) }
+
+            # Pass the DefaultStager to enable fully asynchronous staging.
+            # This offloads the state_dict creation and GPU-to-CPU copy to a background thread.
+            # The return object (AsyncSaveResponse) exposes distinct futures for staging and upload.
+            checkpoint_future = dcp.async_save(
+                state_dict,
+                checkpoint_id=f"{CHECKPOINT_DIR}_step{step}",
+                async_stager=DefaultStager(),
+            )
+
+        # Ensure the last checkpoint completes
+        if checkpoint_future:
+            checkpoint_future.upload_completion.result()
+
+        cleanup()
+
+
+    if __name__ == "__main__":
+        world_size = torch.cuda.device_count()
+        print(f"Running async checkpoint example on {world_size} devices.")
+        mp.spawn(
+            run_fsdp_checkpoint_save_example,
+            args=(world_size,),
+            nprocs=world_size,
+            join=True,
+        )
+
 Conclusion
 ----------
 In conclusion, we have learned how to use DCP's :func:`async_save` API to generate checkpoints off the critical training path. We've also learned about the

recipes_source/torch_export_challenges_solutions.rst

@@ -33,7 +33,7 @@ handle the unsupported code, and then resumes capturing the graph. This break in
 One of the key differences between ``torch.export`` and ``torch.compile`` is that ``torch.export`` doesn’t support graph breaks
 which means that the entire model or part of the model that you are exporting needs to be a single graph. This is because handling graph breaks
 involves interpreting the unsupported operation with default Python evaluation, which is incompatible with what ``torch.export`` is
-designed for. You can read details about the differences between the various PyTorch frameworks in this `link <https://pytorch.org/docs/main/export.html#existing-frameworks>`__
+designed for. You can read details about the differences between the various PyTorch frameworks in this `link <https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/export.html#existing-frameworks>`__
 
 You can identify graph breaks in your program by using the following command:
 

unstable_source/openvino_quantizer.rst

@@ -24,7 +24,7 @@ OpenVINO backend compiles the FX Graph generated by TorchDynamo into an optimize
 
 The quantization flow mainly includes four steps:
 
-- Step 1: Capture the FX Graph from the eager Model based on the `torch export mechanism <https://pytorch.org/docs/main/export.html>`_.
+- Step 1: Capture the FX Graph from the eager Model based on the `torch export mechanism <https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/export.html>`_.
 - Step 2: Apply the PyTorch 2 Export Quantization flow with OpenVINOQuantizer based on the captured FX Graph.
 - Step 3: Lower the quantized model into OpenVINO representation with the `torch.compile <https://docs.openvino.ai/2024/openvino-workflow/torch-compile.html>`_ API.
 - Optional step 4: : Improve quantized model metrics via `quantize_pt2e <https://openvinotoolkit.github.io/nncf/autoapi/nncf/experimental/torch/fx/index.html#nncf.experimental.torch.fx.quantize_pt2e>`_ method.