Merge branch 'main' into patch-1

Author: Svetlana Karslioglu

.jenkins/validate_tutorials_built.py

@@ -22,10 +22,11 @@
     "beginner_source/former_torchies/tensor_tutorial_old",
     "beginner_source/examples_autograd/polynomial_autograd",
     "beginner_source/examples_autograd/polynomial_custom_function",
-    "beginner_source/t5_tutorial", # re-enable after this is fixed: https://github.com/pytorch/text/issues/1756 
+    "beginner_source/t5_tutorial", # re-enable after this is fixed: https://github.com/pytorch/text/issues/1756
     "intermediate_source/parametrizations",
     "intermediate_source/mnist_train_nas",  # used by ax_multiobjective_nas_tutorial.py
     "intermediate_source/fx_conv_bn_fuser",
+    "intermediate_source/_torch_export_nightly_tutorial",  # does not work on release
     "advanced_source/super_resolution_with_onnxruntime",
     "advanced_source/ddp_pipeline",  # requires 4 gpus
     "prototype_source/fx_graph_mode_ptq_dynamic",

conf.py

@@ -106,7 +106,8 @@ def reset_seeds(gallery_conf, fname):
     'first_notebook_cell': ("# For tips on running notebooks in Google Colab, see\n"
                             "# https://pytorch.org/tutorials/beginner/colab\n"
                             "%matplotlib inline"),
-    'reset_modules': (reset_seeds)
+    'reset_modules': (reset_seeds),
+    'ignore_pattern': r'_torch_export_nightly_tutorial.py'
 }
 
 if os.getenv('GALLERY_PATTERN'):

en-wordlist.txt

@@ -52,6 +52,7 @@ CPUs
 CPython
 CUDA
 Caffe
+callable's
 Captum
 Captum's
 CartPole
@@ -72,6 +73,8 @@ DQN
 DataPipe
 DataPipes
 DataLoaders
+Decompositions
+decompositions
 DeepMind
 DeiT
 DenseNet
@@ -160,6 +163,8 @@ ONNX Runtime
 ONNX Script
 OpenAI
 OpenMP
+Opset
+opset
 Ornstein
 OU
 PIL

intermediate_source/FSDP_tutorial.rst

@@ -8,7 +8,7 @@ Getting Started with Fully Sharded Data Parallel(FSDP)
 
 Training AI models at a large scale is a challenging task that requires a lot of compute power and resources. 
 It also comes with considerable engineering complexity to handle the training of these very large models.
-`Pytorch FSDP <https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/>`__, released in PyTorch 1.11 makes this easier.
+`PyTorch FSDP <https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/>`__, released in PyTorch 1.11 makes this easier.
 
 In this tutorial, we show how to use `FSDP APIs <https://pytorch.org/docs/1.11/fsdp.html>`__, for simple MNIST models that can be extended to other larger models such as `HuggingFace BERT models <https://huggingface.co/blog/zero-deepspeed-fairscale>`__, 
 `GPT 3 models up to 1T parameters <https://pytorch.medium.com/training-a-1-trillion-parameter-model-with-pytorch-fully-sharded-data-parallel-on-aws-3ac13aa96cff>`__ . The sample DDP MNIST code has been borrowed from `here <https://github.com/yqhu/mnist_examples>`__. 
@@ -18,7 +18,7 @@ How FSDP works
 --------------
 In `DistributedDataParallel <https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html>`__, (DDP) training, each process/ worker owns a replica of the model and processes a batch of data, finally it uses all-reduce to sum up gradients over different workers. In DDP the model weights and optimizer states are replicated across all workers. FSDP is a type of data parallelism that shards model parameters, optimizer states and gradients across DDP ranks. 
 
-FSDP GPU memory footprint would be smaller than DDP across all workers. This makes the training of some very large models feasible and helps to fit larger models or batch sizes for our training job. This would come with the cost of increased communication volume. The communication overhead is reduced by internal optimizations like communication and computation overlapping.
+When training with FSDP, the GPU memory footprint is smaller than when training with DDP across all workers. This makes the training of some very large models feasible by allowing larger models or batch sizes to fit on device. This comes with the cost of increased communication volume. The communication overhead is reduced by internal optimizations like overlapping communication and computation.
 
 .. figure:: /_static/img/distributed/fsdp_workflow.png
    :width: 100%
@@ -27,7 +27,7 @@ FSDP GPU memory footprint would be smaller than DDP across all workers. This mak
 
    FSDP Workflow
 
-At high level FSDP works as follow:
+At a high level FSDP works as follow:
 
 *In constructor*
 
@@ -48,11 +48,11 @@ At high level FSDP works as follow:
 
 How to use FSDP
 --------------
-Here we use a toy model to run training on MNIST dataset for demonstration purposes. Similarly the APIs and logic can be applied to larger models for training. 
+Here we use a toy model to run training on the MNIST dataset for demonstration purposes. The APIs and logic can be applied to training larger models as well. 
 
 *Setup*
 
-1.1 Install Pytorch along with Torchvision
+1.1 Install PyTorch along with Torchvision
 
 .. code-block:: bash 
 
@@ -139,7 +139,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
             output = F.log_softmax(x, dim=1)
             return output
 
-2.2 define a train function 
+2.2 Define a train function 
 
 .. code-block:: python
 
@@ -189,7 +189,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
 
 2.4 Define a distributed train function that wraps the model in FSDP
 
-**Note: to save the FSDP model, we need to call the state_dict on each rank then on Rank 0 save the overall states. This is only available in Pytorch nightlies, current Pytorch release is 1.11 at the moment.**
+**Note: to save the FSDP model, we need to call the state_dict on each rank then on Rank 0 save the overall states.**
 
 .. code-block:: python
 
@@ -250,7 +250,6 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
         if args.save_model:
             # use a barrier to make sure training is done on all ranks
             dist.barrier()
-            # state_dict for FSDP model is only available on Nightlies for now
             states = model.state_dict()
             if rank == 0:
                 torch.save(states, "mnist_cnn.pt")
@@ -259,7 +258,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
 
 
 
-2.5 Finally parsing the arguments and setting the main function
+2.5 Finally parse the arguments and set the main function
 
 .. code-block:: python
 
@@ -319,7 +318,7 @@ Alternatively, we will look at adding the fsdp_auto_wrap_policy next and will di
     )
  )
 
-Following is the peak memory usage from FSDP MNIST training on g4dn.12.xlarge AWS EC2 instance with 4 gpus captured from Pytorch Profiler. 
+The following is the peak memory usage from FSDP MNIST training on g4dn.12.xlarge AWS EC2 instance with 4 GPUs captured from PyTorch Profiler. 
 
 
 .. figure:: /_static/img/distributed/FSDP_memory.gif
@@ -329,7 +328,7 @@ Following is the peak memory usage from FSDP MNIST training on g4dn.12.xlarge AW
 
    FSDP Peak Memory Usage
 
-*Applying fsdp_auto_wrap_policy* in FSDP otherwise, FSDP will put the entire model in one FSDP unit, which will reduce computation efficiency and memory efficiency. 
+Applying *fsdp_auto_wrap_policy* in FSDP otherwise, FSDP will put the entire model in one FSDP unit, which will reduce computation efficiency and memory efficiency. 
 The way it works is that, suppose your model contains 100 Linear layers. If you do FSDP(model), there will only be one FSDP unit which wraps the entire model. 
 In that case, the allgather would collect the full parameters for all 100 linear layers, and hence won't save CUDA memory for parameter sharding.
 Also, there is only one blocking allgather call for the all 100 linear layers, there will not be communication and computation overlapping between layers. 
@@ -354,7 +353,7 @@ Finding an optimal auto wrap policy is challenging, PyTorch will add auto tuning
     model = FSDP(model,
         fsdp_auto_wrap_policy=my_auto_wrap_policy)
 
-Applying the FSDP_auto_wrap_policy, the model would be as follows:
+Applying the fsdp_auto_wrap_policy, the model would be as follows:
 
 .. code-block:: bash
 
@@ -381,7 +380,7 @@ Applying the FSDP_auto_wrap_policy, the model would be as follows:
 
     CUDA event elapsed time on training loop 41.89130859375sec
 
-Following is the peak memory usage from FSDP with auto_wrap policy of MNIST training on g4dn.12.xlarge AWS EC2 instance with 4 gpus captured from Pytorch Profiler. 
+The following is the peak memory usage from FSDP with auto_wrap policy of MNIST training on a g4dn.12.xlarge AWS EC2 instance with 4 GPUs captured from PyTorch Profiler. 
 It can be observed that the peak memory usage on each device is smaller compared to FSDP without auto wrap policy applied, from ~75 MB to 66 MB.
 
 .. figure:: /_static/img/distributed/FSDP_autowrap.gif
@@ -391,11 +390,11 @@ It can be observed that the peak memory usage on each device is smaller compared
 
    FSDP Peak Memory Usage using Auto_wrap policy
 
-*CPU Off-loading*: In case the model is very large that even with FSDP wouldn't fit into gpus, then CPU offload can be helpful here. 
+*CPU Off-loading*: In case the model is very large that even with FSDP wouldn't fit into GPUs, then CPU offload can be helpful here. 
 
 Currently, only parameter and gradient CPU offload is supported. It can be enabled via passing in cpu_offload=CPUOffload(offload_params=True).
 
-Note that this currently implicitly enables gradient offloading to CPU in order for params and grads to be on the same device to work with the optimizer. This API is subject to change. Default is None in which case there will be no offloading.
+Note that this currently implicitly enables gradient offloading to CPU in order for params and grads to be on the same device to work with the optimizer. This API is subject to change. The default is None in which case there will be no offloading.
 
 Using this feature may slow down the training considerably, due to frequent copying of tensors from host to device, but it could help improve memory efficiency and train larger scale models. 
 
@@ -409,7 +408,7 @@ In 2.4 we just add it to the FSDP wrapper
         cpu_offload=CPUOffload(offload_params=True))
 
 
-Compare it with DDP, if in 2.4 we just normally wrap the model in ddp, saving the changes in “DDP_mnist.py”.
+Compare it with DDP, if in 2.4 we just normally wrap the model in DPP, saving the changes in “DDP_mnist.py”.
 
 .. code-block:: python
 
@@ -423,7 +422,7 @@ Compare it with DDP, if in 2.4 we just normally wrap the model in ddp, saving th
 
     CUDA event elapsed time on training loop 39.77766015625sec
 
-Following is the peak memory usage from DDP MNIST training on g4dn.12.xlarge AWS EC2 instance with 4 gpus captured from Pytorch profiler. 
+The following is the peak memory usage from DDP MNIST training on g4dn.12.xlarge AWS EC2 instance with 4 GPUs captured from PyTorch profiler. 
 
 .. figure:: /_static/img/distributed/DDP_memory.gif
    :width: 100%
@@ -434,8 +433,8 @@ Following is the peak memory usage from DDP MNIST training on g4dn.12.xlarge AWS
 
 
 Considering the toy example and tiny MNIST model we defined here, we can observe the difference between peak memory usage of DDP and FSDP. 
-In DDP each process holds a replica of the model, so the memory footprint is higher compared to FSDP that shards the model parameter, optimizer states and gradients over DDP ranks.
+In DDP each process holds a replica of the model, so the memory footprint is higher compared to FSDP which shards the model parameters, optimizer states and gradients over DDP ranks.
 The peak memory usage using FSDP with auto_wrap policy is the lowest followed by FSDP and DDP. 
 
-Also, looking at timings, considering the small model and running the training on a single machine, FSDP with/out auto_wrap policy performed almost as fast as DDP.
+Also, looking at timings, considering the small model and running the training on a single machine, FSDP with and without auto_wrap policy performed almost as fast as DDP.
 This example does not represent most of the real applications, for detailed analysis and comparison between DDP and FSDP please refer to this `blog post  <https://pytorch.medium.com/6c8da2be180d>`__ .