FSDP1 deprecation msg

Author: weifengpy

index.rst

@@ -766,14 +766,14 @@ Welcome to PyTorch Tutorials
    :tags: Parallel-and-Distributed-Training
 
 .. customcarditem::
-   :header: Getting Started with Fully Sharded Data Parallel(FSDP)
-   :card_description: Learn how to train models with Fully Sharded Data Parallel package.
+   :header: Getting Started with Fully Sharded Data Parallel (FSDP2)
+   :card_description: Learn how to train models with Fully Sharded Data Parallel (fully_shard) package.
    :image: _static/img/thumbnails/cropped/Getting-Started-with-FSDP.png
    :link: intermediate/FSDP_tutorial.html
    :tags: Parallel-and-Distributed-Training
 
 .. customcarditem::
-   :header: Advanced Model Training with Fully Sharded Data Parallel (FSDP)
+   :header: Advanced Model Training with Fully Sharded Data Parallel (FSDP1)
    :card_description: Explore advanced model training with Fully Sharded Data Parallel package.
    :image: _static/img/thumbnails/cropped/Getting-Started-with-FSDP.png
    :link: intermediate/FSDP_advanced_tutorial.html

intermediate_source/FSDP1_tutorial.rst

@@ -4,19 +4,19 @@ Getting Started with Fully Sharded Data Parallel(FSDP)
 **Author**: `Hamid Shojanazeri <https://github.com/HamidShojanazeri>`__, `Yanli Zhao <https://github.com/zhaojuanmao>`__, `Shen Li <https://mrshenli.github.io/>`__
 
 .. note::
-   |edit| View and edit this tutorial in `github <https://github.com/pytorch/tutorials/blob/main/intermediate_source/FSDP_tutorial.rst>`__.
+   |edit| FSDP1 is deprecated. Please check out `FSDP2 tutorial <https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html>`_.
 
-Training AI models at a large scale is a challenging task that requires a lot of compute power and resources. 
+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.
 
-In this tutorial, we show how to use `FSDP APIs <https://pytorch.org/docs/stable/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 courtesy of `Patrick Hu <https://github.com/yqhu/>`_. 
+In this tutorial, we show how to use `FSDP APIs <https://pytorch.org/docs/stable/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 courtesy of `Patrick Hu <https://github.com/yqhu/>`_.
 
 
 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. 
+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.
 
 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.
 
@@ -44,7 +44,7 @@ At a high level FSDP works as follow:
 * Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit
 * Run backward computation
 * Run reduce_scatter to sync gradients
-* Discard parameters. 
+* Discard parameters.
 
 One way to view FSDP's sharding is to decompose the DDP gradient all-reduce into reduce-scatter and all-gather. Specifically, during the backward pass, FSDP reduces and scatters gradients, ensuring that each rank possesses a shard of the gradients. Then it updates the corresponding shard of the parameters in the optimizer step. Finally, in the subsequent forward pass, it performs an all-gather operation to collect and combine the updated parameter shards.
 
@@ -57,7 +57,7 @@ One way to view FSDP's sharding is to decompose the DDP gradient all-reduce into
 
 How to use FSDP
 ---------------
-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. 
+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*
 
@@ -116,7 +116,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
     def cleanup():
         dist.destroy_process_group()
 
-2.1  Define our toy model for handwritten digit classification. 
+2.1  Define our toy model for handwritten digit classification.
 
 .. code-block:: python
 
@@ -131,7 +131,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
             self.fc2 = nn.Linear(128, 10)
 
         def forward(self, x):
-        
+
             x = self.conv1(x)
             x = F.relu(x)
             x = self.conv2(x)
@@ -146,7 +146,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
 
@@ -169,7 +169,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
         if rank == 0:
             print('Train Epoch: {} \tLoss: {:.6f}'.format(epoch, ddp_loss[0] / ddp_loss[1]))
 
-2.3 Define a validation function 
+2.3 Define a validation function
 
 .. code-block:: python
 
@@ -230,8 +230,8 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
             size_based_auto_wrap_policy, min_num_params=100
         )
         torch.cuda.set_device(rank)
-        
-        
+
+
         init_start_event = torch.cuda.Event(enable_timing=True)
         init_end_event = torch.cuda.Event(enable_timing=True)
 
@@ -261,7 +261,7 @@ We add the following code snippets to a python script “FSDP_mnist.py”.
             states = model.state_dict()
             if rank == 0:
                 torch.save(states, "mnist_cnn.pt")
-        
+
         cleanup()
 
 
@@ -309,7 +309,7 @@ We have recorded cuda events to measure the time of FSDP model specifics. The CU
     CUDA event elapsed time on training loop 40.67462890625sec
 
 Wrapping the model with FSDP, the model will look as follows, we can see the model has been wrapped in one FSDP unit.
-Alternatively, we will look at adding the auto_wrap_policy next and will discuss the differences. 
+Alternatively, we will look at adding the auto_wrap_policy next and will discuss the differences.
 
 .. code-block:: bash
 
@@ -326,7 +326,7 @@ Alternatively, we will look at adding the auto_wrap_policy next and will discuss
     )
  )
 
-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. 
+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
@@ -336,18 +336,18 @@ The following is the peak memory usage from FSDP MNIST training on g4dn.12.xlarg
 
    FSDP Peak Memory Usage
 
-Applying *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. 
+Applying *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. 
+Also, there is only one blocking allgather call for the all 100 linear layers, there will not be communication and computation overlapping between layers.
 
 To avoid that, you can pass in an auto_wrap_policy, which will seal the current FSDP unit and start a new one automatically when the specified condition is met (e.g., size limit).
 In that way you will have multiple FSDP units, and only one FSDP unit needs to collect full parameters at a time. E.g., suppose you have 5 FSDP units, and each wraps 20 linear layers.
 Then, in the forward, the 1st FSDP unit will allgather parameters for the first 20 linear layers, do computation, discard the parameters and then move on to the next 20 linear layers. So, at any point in time, each rank only materializes parameters/grads for 20 linear layers instead of 100.
 
 
 To do so in 2.4 we define the auto_wrap_policy and pass it to FSDP wrapper, in the following example, my_auto_wrap_policy defines that a layer could be wrapped or sharded by FSDP if the number of parameters in this layer is larger than 100.
-If the number of parameters in this layer is smaller than 100, it will be wrapped with other small layers together by FSDP. 
+If the number of parameters in this layer is smaller than 100, it will be wrapped with other small layers together by FSDP.
 Finding an optimal auto wrap policy is challenging, PyTorch will add auto tuning for this config in the future. Without an auto tuning tool, it is good to profile your workflow using different auto wrap policies experimentally and find the optimal one.
 
 .. code-block:: python
@@ -388,7 +388,7 @@ Applying the auto_wrap_policy, the model would be as follows:
 
     CUDA event elapsed time on training loop 41.89130859375sec
 
-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. 
+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
@@ -398,13 +398,13 @@ 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. 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. 
+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.
 
 In 2.4 we just add it to the FSDP wrapper
 
@@ -430,7 +430,7 @@ Compare it with DDP, if in 2.4 we just normally wrap the model in DPP, saving th
 
     CUDA event elapsed time on training loop 39.77766015625sec
 
-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. 
+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%
@@ -440,9 +440,9 @@ The following is the peak memory usage from DDP MNIST training on g4dn.12.xlarge
    DDP Peak Memory Usage using Auto_wrap policy
 
 
-Considering the toy example and tiny MNIST model we defined here, we can observe the difference between peak memory usage of DDP and FSDP. 
+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 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. 
+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 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>`__ .

intermediate_source/FSDP_tutorial.rst

@@ -1,10 +1,10 @@
-Getting Started with Fully Sharded Data Parallel(FSDP)
+Getting Started with Fully Sharded Data Parallel (FSDP2)
 ======================================================
 
 **Author**: `Wei Feng <https://github.com/weifengpy>`__, `Will Constable <https://github.com/wconstab>`__, `Yifan Mao <https://github.com/mori360>`__
 
 .. note::
-   |edit| Check out the code in this tutorial from `pytorch/examples <https://github.com/pytorch/examples/tree/main/distributed/FSDP2>`__.
+   |edit| Check out the code in this tutorial from `pytorch/examples <https://github.com/pytorch/examples/tree/main/distributed/FSDP2>`_. FSDP1 will be deprecated. The old tutorial can be found `here <https://docs.pytorch.org/tutorials/intermediate/FSDP1_tutorial.html>`_.
 
 How FSDP2 works
 --------------
@@ -166,7 +166,7 @@ Explicit prefetching works well in following situation:
 Enabling Mixed Precision
 ~~~~~~~~~~~~~~~
 
-FSDP2 offers a flexible `mixed precision policy <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.MixedPrecisionPolicy>`_ to speed up training. One typical use case are
+FSDP2 offers a flexible `mixed precision policy <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.MixedPrecisionPolicy>`_ to speed up training. One typical use case is
 
 * Casting float32 parameters to bfloat16 for forward/backward computation, see ``param_dtype=torch.bfloat16``
 * Upcasting gradients to float32 for reduce-scatter to preserve accuracy, see ``reduce_dtype=torch.float32``
@@ -399,13 +399,13 @@ sync_module_states=True/False: Moved to DCP. User can broadcast state dicts from
 
 forward_prefetch: Manual control over prefetching is possible with
 
-* Manually call ``fsdp_module.unshard()``
-* Use these APIs to control automatic prefetching, ``set_modules_to_forward_prefetch`` and ``set_modules_to_backward_prefetch``
+* Manually call `fsdp_module.unshard() <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.FSDPModule.unshard>`_
+* Use these APIs to control automatic prefetching, `set_modules_to_forward_prefetch <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.FSDPModule.set_modules_to_forward_prefetch>`_ and `set_modules_to_backward_prefetch <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.FSDPModule.set_modules_to_backward_prefetch>`_
 
 limit_all_gathers: No longer needed, because ``fully_shard`` removed cpu synchronization
 
 use_orig_params: Original params are always used (no more flat parameter)
 
 no_sync(): `set_requires_gradient_sync <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.FSDPModule.set_requires_gradient_sync>`_
 
-ignored_params and ignored_states: ignored_params
+ignored_params and ignored_states: `ignored_params <https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html#torch.distributed.fsdp.fully_shard>`_