Fix #155018 (convert distributed rst to markdown) (#155528)

Used [rst2myst tool](https://rst-to-myst.readthedocs.io/en/latest/)

Fixes #155018

Docs comparison (check out the 'new' whenever docs build)

1. distributed.checkpoint ([old](https://docs.pytorch.org/docs/main/distributed.checkpoint.html) vs. [new](https://docs-preview.pytorch.org/pytorch/pytorch/155528/distributed.checkpoint.html))
2. distributed.elastic ([old](https://docs.pytorch.org/docs/main/distributed.elastic.html) vs. [new](https://docs-preview.pytorch.org/pytorch/pytorch/155528/distributed.elastic.html))
3. distributed.fsdp.fully_shard ([old](https://docs.pytorch.org/docs/main/distributed.fsdp.fully_shard.html) vs. [new](https://docs-preview.pytorch.org/pytorch/pytorch/155528/distributed.fsdp.fully_shard.html))
4. distributed.optim ([old](https://docs.pytorch.org/docs/main/distributed.optim.html) vs. [new](https://docs-preview.pytorch.org/pytorch/pytorch/155528/distributed.optim.html))
5. distributed.pipelining ([old](https://docs.pytorch.org/docs/main/distributed.pipelining.html) vs. [new](https://docs-preview.pytorch.org/pytorch/pytorch/155528/distributed.pipelining.html))

Pull Request resolved: https://github.com/pytorch/pytorch/pull/155528
Approved by: https://github.com/wz337, https://github.com/svekars

docs/source/distributed.checkpoint.md

@@ -1,100 +1,154 @@
+```{eval-rst}
 .. role:: hidden
     :class: hidden-section
+```
 
-Distributed Checkpoint - torch.distributed.checkpoint
+# Distributed Checkpoint - torch.distributed.checkpoint
-=====================================================
-
 
 Distributed Checkpoint (DCP) support loading and saving models from multiple ranks in parallel.
 It handles load-time resharding which enables saving in one cluster topology and loading into another.
 
 DCP is different than `torch.save` and `torch.load` in a few significant ways:
 
-* It produces multiple files per checkpoint, with at least one per rank.
+- It produces multiple files per checkpoint, with at least one per rank.
-* It operates in place, meaning that the model should allocate its data first and DCP uses that storage instead.
+- It operates in place, meaning that the model should allocate its data first and DCP uses that storage instead.
 
 The entrypoints to load and save a checkpoint are the following:
 
-Additional resources:
+## Additional resources:
----------------------
 
-* `Getting Started with Distributed Checkpoint (DCP) <https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html>`__
+- [Getting Started with Distributed Checkpoint (DCP)](https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html)
-* `Asynchronous Saving with Distributed Checkpoint (DCP) <https://pytorch.org/tutorials/recipes/distributed_async_checkpoint_recipe.html>`__
+- [Asynchronous Saving with Distributed Checkpoint (DCP)](https://pytorch.org/tutorials/recipes/distributed_async_checkpoint_recipe.html)
-* `TorchTitan Checkpointing Docs <https://github.com/pytorch/torchtitan/blob/main/docs/checkpoint.md>`__
+- [TorchTitan Checkpointing Docs](https://github.com/pytorch/torchtitan/blob/main/docs/checkpoint.md)
-* `TorchTitan DCP Implementation <https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpoint.py>`__
+- [TorchTitan DCP Implementation](https://github.com/pytorch/torchtitan/blob/main/torchtitan/components/checkpoint.py)
 
+```{eval-rst}
 .. automodule:: torch.distributed.checkpoint
+```
 
+```{eval-rst}
 .. currentmodule:: torch.distributed.checkpoint.state_dict_saver
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.state_dict_saver.AsyncCheckpointerType
   :members:
+```
 
+```{eval-rst}
 .. autofunction::  save
+```
+
+```{eval-rst}
 .. autofunction::  async_save
+```
+
+```{eval-rst}
 .. autofunction::  save_state_dict
+```
 
+```{eval-rst}
 .. currentmodule:: torch.distributed.checkpoint.state_dict_loader
+```
 
+```{eval-rst}
 .. autofunction::  load
+```
+
+```{eval-rst}
 .. autofunction::  load_state_dict
+```
 
 The following module is also useful for additional customization of the staging mechanisms used for asynchronous checkpointing (`torch.distributed.checkpoint.async_save`):
 
+```{eval-rst}
 .. automodule:: torch.distributed.checkpoint.staging
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.staging.AsyncStager
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.staging.BlockingAsyncStager
   :members:
+```
 
 In addition to the above entrypoints, `Stateful` objects, as described below, provide additional customization during saving/loading
-.. automodule:: torch.distributed.checkpoint.stateful
 
+```{eval-rst}
+.. automodule:: torch.distributed.checkpoint.stateful
+   :noindex:
+```
+
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.stateful.Stateful
   :members:
+```
 
-This `example <https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py>`_ shows how to use Pytorch Distributed Checkpoint to save a FSDP model.
+This [example](https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py) shows how to use Pytorch Distributed Checkpoint to save a FSDP model.
 
 The following types define the IO interface used during checkpoint:
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.StorageReader
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.StorageWriter
   :members:
+```
 
 The following types define the planner interface used during checkpoint:
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.LoadPlanner
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.LoadPlan
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.ReadItem
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.SavePlanner
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.SavePlan
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.planner.WriteItem
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.planner.BytesIOWriteData
   :members:
+```
 
 We provide a filesystem based storage layer:
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.FileSystemReader
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.FileSystemWriter
   :members:
+```
 
 We also provide other storage layers, including ones to interact with HuggingFace safetensors:
 
@@ -107,12 +161,16 @@ We also provide other storage layers, including ones to interact with HuggingFac
 We provide default implementations of `LoadPlanner` and `SavePlanner` that
 can handle all of torch.distributed constructs such as FSDP, DDP, ShardedTensor and DistributedTensor.
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.DefaultSavePlanner
   :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.DefaultLoadPlanner
   :members:
 
+```
 
 Due to legacy design decisions, the state dictionaries of `FSDP` and `DDP` may have different keys or fully qualified names (e.g., layer1.weight) even when the original unparallelized model is identical. Moreover, `FSDP` offers various types of model state dictionaries, such as full and sharded state dictionaries. Additionally, optimizer state dictionaries employ parameter IDs instead of fully qualified names to identify parameters, potentially causing issues when parallelisms are used (e.g., pipeline parallelism).
 
@@ -126,40 +184,71 @@ Note that `set_optimizer_state_dict()` can only be called before `backward()` or
 
 Note that this feature is experimental, and API signatures might change in the future.
 
-
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.get_state_dict
+```
 
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.get_model_state_dict
+```
 
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.get_optimizer_state_dict
+```
 
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.set_state_dict
+```
 
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.set_model_state_dict
+```
 
+```{eval-rst}
 .. autofunction:: torch.distributed.checkpoint.state_dict.set_optimizer_state_dict
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.state_dict.StateDictOptions
    :members:
+```
 
 For users which are used to using and sharing models in the `torch.save` format, the following methods are provided which provide offline utilities for converting betweeing formats.
 
+```{eval-rst}
 .. automodule:: torch.distributed.checkpoint.format_utils
+```
 
+```{eval-rst}
 .. currentmodule:: torch.distributed.checkpoint.format_utils
+```
 
+```{eval-rst}
 .. autofunction:: dcp_to_torch_save
+```
+
+```{eval-rst}
 .. autofunction:: torch_save_to_dcp
+```
 
 The following classes can also be utilized for online loading and resharding of models from the torch.save format.
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader
    :members:
+```
 
+```{eval-rst}
 .. autoclass:: torch.distributed.checkpoint.format_utils.DynamicMetaLoadPlanner
    :members:
+```
 
 The following experimental interfaces are provided for improved observability in production environments:
 
+```{eval-rst}
 .. py:module:: torch.distributed.checkpoint.logger
+```
+
+```{eval-rst}
 .. py:module:: torch.distributed.checkpoint.logging_handlers
+```

docs/source/distributed.elastic.md

@@ -0,0 +1,46 @@
+# Torch Distributed Elastic
+
+Makes distributed PyTorch fault-tolerant and elastic.
+
+## Get Started
+
+```{toctree}
+:caption: Usage
+:maxdepth: 1
+
+elastic/quickstart
+elastic/train_script
+elastic/examples
+```
+
+## Documentation
+
+```{toctree}
+:caption: API
+:maxdepth: 1
+
+elastic/run
+elastic/agent
+elastic/multiprocessing
+elastic/errors
+elastic/rendezvous
+elastic/timer
+elastic/metrics
+elastic/events
+elastic/subprocess_handler
+elastic/control_plane
+```
+
+```{toctree}
+:caption: Advanced
+:maxdepth: 1
+
+elastic/customization
+```
+
+```{toctree}
+:caption: Plugins
+:maxdepth: 1
+
+elastic/kubernetes
+```

docs/source/distributed.elastic.rst

@@ -1,44 +0,0 @@
-Torch Distributed Elastic
-============================
-
-Makes distributed PyTorch fault-tolerant and elastic.
-
-Get Started
----------------
-.. toctree::
-   :maxdepth: 1
-   :caption: Usage
-
-   elastic/quickstart
-   elastic/train_script
-   elastic/examples
-
-Documentation
----------------
-
-.. toctree::
-   :maxdepth: 1
-   :caption: API
-
-   elastic/run
-   elastic/agent
-   elastic/multiprocessing
-   elastic/errors
-   elastic/rendezvous
-   elastic/timer
-   elastic/metrics
-   elastic/events
-   elastic/subprocess_handler
-   elastic/control_plane
-
-.. toctree::
-   :maxdepth: 1
-   :caption: Advanced
-
-   elastic/customization
-
-.. toctree::
-   :maxdepth: 1
-   :caption: Plugins
-
-   elastic/kubernetes

docs/source/distributed.fsdp.fully_shard.md

@@ -1,85 +1,100 @@
-torch.distributed.fsdp.fully_shard
+# torch.distributed.fsdp.fully_shard
-==================================
 
-PyTorch FSDP2 (``fully_shard``)
+## PyTorch FSDP2 (`fully_shard`)
--------------------------------
 
 PyTorch FSDP2 provides a fully sharded data parallelism (FSDP) implementation
 targeting performant eager-mode while using per-parameter sharding for improved
 usability.
 
 - If you are new to FSDP, we recommend that you start with FSDP2 due to improved
-  usability. See `TorchTitan <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md>`_ for code examples.
+  usability. See [TorchTitan](https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md) for code examples.
 - If you are currently using FSDP1, consider evaluating the following
   differences to see if you should switch to FSDP2:
 
-Compared to PyTorch FSDP1 (``FullyShardedDataParallel``):
+Compared to PyTorch FSDP1 (`FullyShardedDataParallel`):
 
-- FSDP2 uses ``DTensor``-based dim-0 per-parameter sharding for a simpler
+- FSDP2 uses `DTensor`-based dim-0 per-parameter sharding for a simpler
   sharding representation compared to FSDP1's flat-parameter sharding, while
   preserving similar throughput performance. More specifically, FSDP2 chunks
   each parameter on dim-0 across the data parallel workers (using
-  ``torch.chunk(dim=0)``), whereas FSDP1 flattens, concatenates, and chunks a
+  `torch.chunk(dim=0)`), whereas FSDP1 flattens, concatenates, and chunks a
   group of tensors together, making reasoning about what data is present on
   each worker and resharding to different parallelisms complex. Per-parameter
   sharding provides a more intuitive user experience, relaxes constraints
   around frozen parameters, and allows for communication-free (sharded) state
   dicts, which otherwise require all-gathers in FSDP1.
 - FSDP2 implements a different memory management approach to handle the
-  multi-stream usages that avoids ``torch.Tensor.record_stream``. This ensures
+  multi-stream usages that avoids `torch.Tensor.record_stream`. This ensures
   deterministic and expected memory usage and does not require blocking the CPU
-  like in FSDP1's ``limit_all_gathers=True``.
+  like in FSDP1's `limit_all_gathers=True`.
 - FSDP2 exposes APIs for manual control over prefetching and collective
   scheduling, allowing power users more customization. See the methods on
-  ``FSDPModule`` below for details.
+  `FSDPModule` below for details.
 - FSDP2 simplifies some of the API surface: e.g. FSDP2 does not directly
   support full state dicts. Instead, users can reshard the sharded state dicts
-  containing ``DTensor`` s to full state dicts themselves using ``DTensor``
+  containing `DTensor` s to full state dicts themselves using `DTensor`
-  APIs like ``DTensor.full_tensor()`` or by using higher-level APIs like
+  APIs like `DTensor.full_tensor()` or by using higher-level APIs like
-  `PyTorch Distributed Checkpoint <https://pytorch.org/docs/stable/distributed.checkpoint.html>`_ 's
+  [PyTorch Distributed Checkpoint](https://pytorch.org/docs/stable/distributed.checkpoint.html) 's
   distributed state dict APIs. Also, some other args have been removed; see
-  `here <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md>`_ for
+  [here](https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md) for
   details.
 
 If you are onboarding FSDP for the first time or if any of the above appeals to
 your use case, we recommend that you consider using FSDP2.
 
-See `this RFC <https://github.com/pytorch/pytorch/issues/114299>`_ for details
+See [this RFC](https://github.com/pytorch/pytorch/issues/114299) for details
 on system design and implementation.
 
-.. note::
+:::{note}
-  ``torch.distributed.fsdp.fully_shard`` is currently in prototype state and
+`torch.distributed.fsdp.fully_shard` is currently in prototype state and
-  under development. The core API will likely not change, but we may make some
+under development. The core API will likely not change, but we may make some
-  API changes if necessary.
+API changes if necessary.
+:::
 
+```{eval-rst}
 .. currentmodule:: torch.distributed.fsdp
+```
 
-The frontend API is ``fully_shard`` that can be called on a ``module``:
+The frontend API is `fully_shard` that can be called on a `module`:
 
+```{eval-rst}
 .. autofunction:: fully_shard
+```
 
-Calling ``fully_shard(module)`` dynamically constructs a new class that
+Calling `fully_shard(module)` dynamically constructs a new class that
-subclasses ``type(module)`` and an FSDP class ``FSDPModule``. For example, if
+subclasses `type(module)` and an FSDP class `FSDPModule`. For example, if
-we call ``fully_shard(linear)`` on a module ``linear: nn.Linear``, then FSDP
+we call `fully_shard(linear)` on a module `linear: nn.Linear`, then FSDP
-constructs a new class ``FSDPLinear`` and changes ``linear`` 's type to this.
+constructs a new class `FSDPLinear` and changes `linear` 's type to this.
-Otherwise, ``fully_shard`` does not change the module structure and parameter
+Otherwise, `fully_shard` does not change the module structure and parameter
-fully-qualified names. The class ``FSDPModule`` allows providing some
+fully-qualified names. The class `FSDPModule` allows providing some
 FSDP-specific methods on the module.
 
+```{eval-rst}
 .. autoclass:: FSDPModule
     :members:
     :member-order: bysource
+```
 
+```{eval-rst}
 .. autoclass:: UnshardHandle
     :members:
+```
 
+```{eval-rst}
 .. autofunction:: register_fsdp_forward_method
+```
 
+```{eval-rst}
 .. autoclass:: MixedPrecisionPolicy
     :members:
+```
 
+```{eval-rst}
 .. autoclass:: OffloadPolicy
     :members:
+```
 
+```{eval-rst}
 .. autoclass:: CPUOffloadPolicy
     :members:
+```

docs/source/distributed.optim.md

@@ -1,11 +1,15 @@
+```{eval-rst}
 .. role:: hidden
     :class: hidden-section
+```
 
-Distributed Optimizers
+# Distributed Optimizers
-======================
 
-.. warning ::
+:::{warning}
-    Distributed optimizer is not currently supported when using CUDA tensors
+Distributed optimizer is not currently supported when using CUDA tensors
+:::
 
+```{eval-rst}
 .. automodule:: torch.distributed.optim
     :members: DistributedOptimizer, PostLocalSGDOptimizer, ZeroRedundancyOptimizer
+```

docs/source/distributed.pipelining.md

@@ -0,0 +1,515 @@
+```{eval-rst}
+.. role:: hidden
+    :class: hidden-section
+```
+
+# Pipeline Parallelism
+
+:::{note}
+`torch.distributed.pipelining` is currently in alpha state and under
+development. API changes may be possible. It was migrated from the [PiPPy](https://github.com/pytorch/PiPPy) project.
+:::
+
+## Why Pipeline Parallel?
+
+Pipeline Parallelism is one of the **primitive** parallelism for deep learning.
+It allows the **execution** of a model to be partitioned such that multiple
+**micro-batches** can execute different parts of the model code concurrently.
+Pipeline parallelism can be an effective technique for:
+
+- large-scale training
+- bandwidth-limited clusters
+- large model inference
+
+The above scenarios share a commonality that the computation per device cannot
+hide the communication of conventional parallelism, for example, the weight
+all-gather of FSDP.
+
+## What is `torch.distributed.pipelining`?
+
+While promising for scaling, pipelining is often difficult to implement because
+it needs to **partition the execution** of a model in addition to model weights.
+The partitioning of execution often requires intrusive code changes to your
+model. Another aspect of complexity comes from **scheduling micro-batches in a
+distributed environment**, with **data flow dependency** considered.
+
+The `pipelining` package provides a toolkit that does said things
+**automatically** which allows easy implementation of pipeline parallelism
+on **general** models.
+
+It consists of two parts: a
+**splitting frontend** and a **distributed runtime**.
+The splitting frontend takes your model code as-is, splits it up into "model
+partitions", and captures the data-flow relationship. The distributed runtime
+executes the pipeline stages on different devices in parallel, handling things
+like micro-batch splitting, scheduling, communication, and gradient propagation,
+etc.
+
+Overall, the `pipelining` package provides the following features:
+
+- Splitting of model code based on simple specification.
+- Rich support for pipeline schedules, including GPipe, 1F1B,
+  Interleaved 1F1B and Looped BFS, and providing the infrastructure for writing
+  customized schedules.
+- First-class support for cross-host pipeline parallelism, as this is where PP
+  is typically used (over slower interconnects).
+- Composability with other PyTorch parallel techniques such as data parallel
+  (DDP, FSDP) or tensor parallel. The [TorchTitan](https://github.com/pytorch/torchtitan) project demonstrates a "3D parallel"
+  application on the Llama model.
+
+## Step 1: build `PipelineStage`
+
+Before we can use a `PipelineSchedule`, we need to create `PipelineStage`
+objects that wrap the part of the model running in that stage. The
+`PipelineStage` is responsible for allocating communication buffers and
+creating send/recv ops to communicate with its peers. It manages intermediate
+buffers e.g. for the outputs of forward that have not been consumed yet, and it
+provides a utility for running the backwards for the stage model.
+
+A `PipelineStage` needs to know the input and output shapes for the stage
+model, so that it can correctly allocate communication buffers. The shapes must
+be static, e.g. at runtime the shapes can not change from step to step. A class
+`PipeliningShapeError` will be raised if runtime shapes do not match the
+expected shapes. When composing with other paralleisms or applying mixed
+precision, these techniques must be taken into account so the `PipelineStage`
+knows the correct shape (and dtype) for the output of the stage module at
+runtime.
+
+Users may construct a `PipelineStage` instance directly, by passing in an
+`nn.Module` representing the portion of the model that should run on the
+stage. This may require changes to the original model code. See the example
+in {ref}`option_1_manual`.
+
+Alternatively, the splitting frontend can use graph partitioning to split your
+model into a series of `nn.Module` automatically. This technique requires the
+model is traceable with `torch.Export`. Composability of the resulting
+`nn.Module` with other parallelism techniques is experimental, and may require
+some workarounds. Usage of this frontend may be more appealing if the user
+cannot easily change the model code. See {ref}`option_2_tracer` for more
+information.
+
+## Step 2: use `PipelineSchedule` for execution
+
+We can now attach the `PipelineStage` to a pipeline schedule, and run the
+schedule with input data. Here is a GPipe example:
+
+```python
+from torch.distributed.pipelining import ScheduleGPipe
+
+# Create a schedule
+schedule = ScheduleGPipe(stage, n_microbatches)
+
+# Input data (whole batch)
+x = torch.randn(batch_size, in_dim, device=device)
+
+# Run the pipeline with input `x`
+# `x` will be divided into microbatches automatically
+if rank == 0:
+    schedule.step(x)
+else:
+    output = schedule.step()
+```
+
+Note that the above code needs to be launched for each worker, thus we use a
+launcher service to launch multiple processes:
+
+```bash
+torchrun --nproc_per_node=2 example.py
+```
+
+## Options for Splitting a Model
+
+(option_1_manual)=
+
+### Option 1: splitting a model manually
+
+To directly construct a `PipelineStage`, the user is responsible for providing
+a single `nn.Module` instance that owns the relevant `nn.Parameters` and
+`nn.Buffers`, and defines a `forward()` method that executes the operations
+relevant for that stage. For example, a condensed version of the Transformer
+class defined in Torchtitan shows a pattern of building an easily partitionable
+model.
+
+```python
+class Transformer(nn.Module):
+    def __init__(self, model_args: ModelArgs):
+        super().__init__()
+
+        self.tok_embeddings = nn.Embedding(...)
+
+        # Using a ModuleDict lets us delete layers without affecting names,
+        # ensuring checkpoints will correctly save and load.
+        self.layers = torch.nn.ModuleDict()
+        for layer_id in range(model_args.n_layers):
+            self.layers[str(layer_id)] = TransformerBlock(...)
+
+        self.output = nn.Linear(...)
+
+    def forward(self, tokens: torch.Tensor):
+        # Handling layers being 'None' at runtime enables easy pipeline splitting
+        h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens
+
+        for layer in self.layers.values():
+            h = layer(h, self.freqs_cis)
+
+        h = self.norm(h) if self.norm else h
+        output = self.output(h).float() if self.output else h
+        return output
+```
+
+A model defined in this manner can be easily configured per stage by first
+initializing the whole model (using meta-device to avoid OOM errors), deleting
+undesired layers for that stage, and then creating a PipelineStage that wraps
+the model. For example:
+
+```python
+with torch.device("meta"):
+    assert num_stages == 2, "This is a simple 2-stage example"
+
+    # we construct the entire model, then delete the parts we do not need for this stage
+    # in practice, this can be done using a helper function that automatically divides up layers across stages.
+    model = Transformer()
+
+    if stage_index == 0:
+        # prepare the first stage model
+        del model.layers["1"]
+        model.norm = None
+        model.output = None
+
+    elif stage_index == 1:
+        # prepare the second stage model
+        model.tok_embeddings = None
+        del model.layers["0"]
+
+    from torch.distributed.pipelining import PipelineStage
+    stage = PipelineStage(
+        model,
+        stage_index,
+        num_stages,
+        device,
+    )
+```
+
+When composing with other Data or Model parallelism techniques, `output_args`
+may also be required, if the output shape/dtype of the model chunk will be
+affected.
+
+(option_2_tracer)=
+
+### Option 2: splitting a model automatically
+
+If you have a full model and do not want to spend time on modifying it into a
+sequence of "model partitions", the `pipeline` API is here to help.
+Here is a brief example:
+
+```python
+class Model(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.emb = torch.nn.Embedding(10, 3)
+        self.layers = torch.nn.ModuleList(
+            Layer() for _ in range(2)
+        )
+        self.lm = LMHead()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.emb(x)
+        for layer in self.layers:
+            x = layer(x)
+        x = self.lm(x)
+        return x
+```
+
+If we print the model, we can see multiple hierarchies, which makes it hard to split by hand:
+
+```python
+Model(
+  (emb): Embedding(10, 3)
+  (layers): ModuleList(
+    (0-1): 2 x Layer(
+      (lin): Linear(in_features=3, out_features=3, bias=True)
+    )
+  )
+  (lm): LMHead(
+    (proj): Linear(in_features=3, out_features=3, bias=True)
+  )
+)
+```
+
+Let us see how the `pipeline` API works:
+
+```python
+from torch.distributed.pipelining import pipeline, SplitPoint
+
+# An example micro-batch input
+x = torch.LongTensor([1, 2, 4, 5])
+
+pipe = pipeline(
+    module=mod,
+    mb_args=(x,),
+    split_spec={
+        "layers.1": SplitPoint.BEGINNING,
+    }
+)
+```
+
+The `pipeline` API splits your model given a `split_spec`, where
+`SplitPoint.BEGINNING` stands for adding a split point
+*before* execution of certain submodule in the `forward` function, and
+similarly, `SplitPoint.END` for split point *after* such.
+
+If we `print(pipe)`, we can see:
+
+```python
+GraphModule(
+  (submod_0): GraphModule(
+    (emb): InterpreterModule()
+    (layers): Module(
+      (0): InterpreterModule(
+        (lin): InterpreterModule()
+      )
+    )
+  )
+  (submod_1): GraphModule(
+    (layers): Module(
+      (1): InterpreterModule(
+        (lin): InterpreterModule()
+      )
+    )
+    (lm): InterpreterModule(
+      (proj): InterpreterModule()
+    )
+  )
+)
+
+def forward(self, x):
+    submod_0 = self.submod_0(x);  x = None
+    submod_1 = self.submod_1(submod_0);  submod_0 = None
+    return (submod_1,)
+```
+
+The "model partitions" are represented by submodules (`submod_0`,
+`submod_1`), each of which is reconstructed with original model operations, weights
+and hierarchies. In addition, a "root-level" `forward` function is
+reconstructed to capture the data flow between those partitions. Such data flow
+will be replayed by the pipeline runtime later, in a distributed fashion.
+
+The `Pipe` object provides a method for retrieving the "model partitions":
+
+```python
+stage_mod : nn.Module = pipe.get_stage_module(stage_idx)
+```
+
+The returned `stage_mod` is a `nn.Module`, with which you can create an
+optimizer, save or load checkpoints, or apply other parallelisms.
+
+`Pipe` also allows you to create a distributed stage runtime on a device given
+a `ProcessGroup`:
+
+```python
+stage = pipe.build_stage(stage_idx, device, group)
+```
+
+Alternatively, if you would like to build the stage runtime later after some
+modification to the `stage_mod`, you can use a functional version of the
+`build_stage` API. For example:
+
+```python
+from torch.distributed.pipelining import build_stage
+from torch.nn.parallel import DistributedDataParallel
+
+dp_mod = DistributedDataParallel(stage_mod)
+info = pipe.info()
+stage = build_stage(dp_mod, stage_idx, info, device, group)
+```
+
+:::{note}
+The `pipeline` frontend uses a tracer (`torch.export`) to capture your
+model into a single graph. If your model is not full-graph'able, you can use
+our manual frontend below.
+:::
+
+## Hugging Face Examples
+
+In the [PiPPy](https://github.com/pytorch/PiPPy) repo where this package was
+original created, we kept examples based on unmodified Hugging Face models.
+See the [examples/huggingface](https://github.com/pytorch/PiPPy/tree/main/examples/huggingface) directory.
+
+Examples include:
+
+- [GPT2](https://github.com/pytorch/PiPPy/tree/main/examples/huggingface/pippy_gpt2.py)
+- [Llama](https://github.com/pytorch/PiPPy/tree/main/examples/llama)
+
+## Technical Deep Dive
+
+### How does the `pipeline` API split a model?
+
+First, the `pipeline` API turns our model into a directed acyclic graph (DAG)
+by tracing the model. It traces the model using `torch.export` -- a PyTorch 2
+full-graph capturing tool.
+
+Then, it groups together the **operations and parameters** needed by a stage
+into a reconstructed submodule: `submod_0`, `submod_1`, ...
+
+Different from conventional submodule access methods like `Module.children()`,
+the `pipeline` API does not only cut the module structure of your model, but
+also the **forward** function of your model.
+
+This is necessary because model structure like `Module.children()` merely
+captures information during `Module.__init__()`, and does not capture any
+information about `Module.forward()`. Said differently, `Module.children()`
+lacks information about the following aspects key to pipelininig:
+
+- Execution order of child modules in `forward`
+- Activation flows between child modules
+- Whether there are any functional operators between child modules (for example,
+  `relu` or `add` operations will not be captured by `Module.children()`).
+
+The `pipeline` API, on the contrary, makes sure that the `forward` behavior
+is truly preserved. It also captures the activation flow between the partitions,
+helping the distributed runtime to make correct send/receive calls without human
+intervention.
+
+Another flexibility of the `pipeline` API is that split points can be at
+arbitrary levels within your model hierarchy. In the split partitions, the original model
+hierarchy related to that partition will be reconstructed at no cost to you.
+At a result, fully-qualified names (FQNs) pointing to a submodule or parameter
+would be still valid, and services that relies on FQNs (such as FSDP, TP or
+checkpointing) can still run with your partitioned modules with almost zero code
+change.
+
+## Implementing Your Own Schedule
+
+You can implement your own pipeline schedule by extending one of the following two class:
+
+- `PipelineScheduleSingle`
+- `PipelineScheduleMulti`
+
+`PipelineScheduleSingle` is for schedules that assigns *only one* stage per rank.
+`PipelineScheduleMulti` is for schedules that assigns multiple stages per rank.
+
+For example, `ScheduleGPipe` and `Schedule1F1B` are subclasses of `PipelineScheduleSingle`.
+Whereas, `ScheduleInterleaved1F1B`, `ScheduleLoopedBFS`, `ScheduleInterleavedZeroBubble`, and `ScheduleZBVZeroBubble`
+are subclasses of `PipelineScheduleMulti`.
+
+## Logging
+
+You can turn on additional logging using the `TORCH_LOGS` environment variable from [torch.\_logging](https://pytorch.org/docs/main/logging.html#module-torch._logging):
+
+- `TORCH_LOGS=+pp` will display `logging.DEBUG` messages and all levels above it.
+- `TORCH_LOGS=pp` will display `logging.INFO` messages and above.
+- `TORCH_LOGS=-pp` will display `logging.WARNING` messages and above.
+
+## API Reference
+
+```{eval-rst}
+.. automodule:: torch.distributed.pipelining
+```
+
+### Model Split APIs
+
+The following set of APIs transform your model into a pipeline representation.
+
+```{eval-rst}
+.. currentmodule:: torch.distributed.pipelining
+```
+
+```{eval-rst}
+.. autoclass:: SplitPoint
+```
+
+```{eval-rst}
+.. autofunction:: pipeline
+```
+
+```{eval-rst}
+.. autoclass:: Pipe
+```
+
+```{eval-rst}
+.. autofunction:: pipe_split
+```
+
+### Microbatch Utilities
+
+```{eval-rst}
+.. automodule:: torch.distributed.pipelining.microbatch
+```
+
+```{eval-rst}
+.. currentmodule:: torch.distributed.pipelining.microbatch
+```
+
+```{eval-rst}
+.. autoclass:: TensorChunkSpec
+```
+
+```{eval-rst}
+.. autofunction:: split_args_kwargs_into_chunks
+```
+
+```{eval-rst}
+.. autofunction:: merge_chunks
+```
+
+### Pipeline Stages
+
+```{eval-rst}
+.. automodule:: torch.distributed.pipelining.stage
+```
+
+```{eval-rst}
+.. currentmodule:: torch.distributed.pipelining.stage
+```
+
+```{eval-rst}
+.. autoclass:: PipelineStage
+```
+
+```{eval-rst}
+.. autofunction:: build_stage
+```
+
+### Pipeline Schedules
+
+```{eval-rst}
+.. automodule:: torch.distributed.pipelining.schedules
+```
+
+```{eval-rst}
+.. currentmodule:: torch.distributed.pipelining.schedules
+```
+
+```{eval-rst}
+.. autoclass:: ScheduleGPipe
+```
+
+```{eval-rst}
+.. autoclass:: Schedule1F1B
+```
+
+```{eval-rst}
+.. autoclass:: ScheduleInterleaved1F1B
+```
+
+```{eval-rst}
+.. autoclass:: ScheduleLoopedBFS
+```
+
+```{eval-rst}
+.. autoclass:: ScheduleInterleavedZeroBubble
+```
+
+```{eval-rst}
+.. autoclass:: ScheduleZBVZeroBubble
+```
+
+```{eval-rst}
+.. autoclass:: PipelineScheduleSingle
+  :members:
+```
+
+```{eval-rst}
+.. autoclass:: PipelineScheduleMulti
+  :members:
+```

docs/source/distributed.pipelining.rst

@@ -1,491 +0,0 @@
-.. role:: hidden
-    :class: hidden-section
-
-Pipeline Parallelism
-####################
-
-.. note::
-  ``torch.distributed.pipelining`` is currently in alpha state and under
-  development. API changes may be possible. It was migrated from the `PiPPy
-  <https://github.com/pytorch/PiPPy>`_ project.
-
-
-Why Pipeline Parallel?
-**********************
-
-Pipeline Parallelism is one of the **primitive** parallelism for deep learning.
-It allows the **execution** of a model to be partitioned such that multiple
-**micro-batches** can execute different parts of the model code concurrently.
-Pipeline parallelism can be an effective technique for:
-
-* large-scale training
-* bandwidth-limited clusters
-* large model inference
-
-The above scenarios share a commonality that the computation per device cannot
-hide the communication of conventional parallelism, for example, the weight
-all-gather of FSDP.
-
-
-What is ``torch.distributed.pipelining``?
-*****************************************
-
-While promising for scaling, pipelining is often difficult to implement because
-it needs to **partition the execution** of a model in addition to model weights.
-The partitioning of execution often requires intrusive code changes to your
-model. Another aspect of complexity comes from **scheduling micro-batches in a
-distributed environment**, with **data flow dependency** considered.
-
-The ``pipelining`` package provides a toolkit that does said things
-**automatically** which allows easy implementation of pipeline parallelism
-on **general** models.
-
-It consists of two parts: a
-**splitting frontend** and a **distributed runtime**.
-The splitting frontend takes your model code as-is, splits it up into "model
-partitions", and captures the data-flow relationship.  The distributed runtime
-executes the pipeline stages on different devices in parallel, handling things
-like micro-batch splitting, scheduling, communication, and gradient propagation,
-etc.
-
-Overall, the ``pipelining`` package provides the following features:
-
-* Splitting of model code based on simple specification.
-* Rich support for pipeline schedules, including GPipe, 1F1B,
-  Interleaved 1F1B and Looped BFS, and providing the infrastructure for writing
-  customized schedules.
-* First-class support for cross-host pipeline parallelism, as this is where PP
-  is typically used (over slower interconnects).
-* Composability with other PyTorch parallel techniques such as data parallel
-  (DDP, FSDP) or tensor  parallel. The `TorchTitan
-  <https://github.com/pytorch/torchtitan>`_ project demonstrates a "3D parallel"
-  application on the Llama model.
-
-
-Step 1: build ``PipelineStage``
-*******************************
-
-Before we can use a ``PipelineSchedule``, we need to create ``PipelineStage``
-objects that wrap the part of the model running in that stage.  The
-``PipelineStage`` is responsible for allocating communication buffers and
-creating send/recv ops to communicate with its peers.  It manages intermediate
-buffers e.g. for the outputs of forward that have not been consumed yet, and it
-provides a utility for running the backwards for the stage model.
-
-A ``PipelineStage`` needs to know the input and output shapes for the stage
-model, so that it can correctly allocate communication buffers.  The shapes must
-be static, e.g. at runtime the shapes can not change from step to step.  A class
-``PipeliningShapeError`` will be raised if runtime shapes do not match the
-expected shapes.  When composing with other paralleisms or applying mixed
-precision, these techniques must be taken into account so the ``PipelineStage``
-knows the correct shape (and dtype) for the output of the stage module at
-runtime.
-
-Users may construct a ``PipelineStage`` instance directly, by passing in an
-``nn.Module`` representing the portion of the model that should run on the
-stage.  This may require changes to the original model code.  See the example
-in :ref:`option_1_manual`.
-
-Alternatively, the splitting frontend can use graph partitioning to split your
-model into a series of ``nn.Module`` automatically.  This technique requires the
-model is traceable with ``torch.Export``. Composability of the resulting
-``nn.Module`` with other parallelism techniques is experimental, and may require
-some workarounds.  Usage of this frontend may be more appealing if the user
-cannot easily change the model code.  See :ref:`option_2_tracer` for more
-information.
-
-
-Step 2: use ``PipelineSchedule`` for execution
-**********************************************
-
-We can now attach the ``PipelineStage`` to a pipeline schedule, and run the
-schedule with input data. Here is a GPipe example:
-
-.. code-block:: python
-
-  from torch.distributed.pipelining import ScheduleGPipe
-
-  # Create a schedule
-  schedule = ScheduleGPipe(stage, n_microbatches)
-
-  # Input data (whole batch)
-  x = torch.randn(batch_size, in_dim, device=device)
-
-  # Run the pipeline with input `x`
-  # `x` will be divided into microbatches automatically
-  if rank == 0:
-      schedule.step(x)
-  else:
-      output = schedule.step()
-
-Note that the above code needs to be launched for each worker, thus we use a
-launcher service to launch multiple processes:
-
-.. code-block:: bash
-
-  torchrun --nproc_per_node=2 example.py
-
-
-Options for Splitting a Model
-*****************************
-
-.. _option_1_manual:
-
-Option 1: splitting a model manually
-====================================
-
-To directly construct a ``PipelineStage``, the user is responsible for providing
-a single ``nn.Module`` instance that owns the relevant ``nn.Parameters`` and
-``nn.Buffers``, and defines a ``forward()`` method that executes the operations
-relevant for that stage.  For example, a condensed version of the Transformer
-class defined in Torchtitan shows a pattern of building an easily partitionable
-model.
-
-.. code-block:: python
-
-  class Transformer(nn.Module):
-      def __init__(self, model_args: ModelArgs):
-          super().__init__()
-
-          self.tok_embeddings = nn.Embedding(...)
-
-          # Using a ModuleDict lets us delete layers without affecting names,
-          # ensuring checkpoints will correctly save and load.
-          self.layers = torch.nn.ModuleDict()
-          for layer_id in range(model_args.n_layers):
-              self.layers[str(layer_id)] = TransformerBlock(...)
-
-          self.output = nn.Linear(...)
-
-      def forward(self, tokens: torch.Tensor):
-          # Handling layers being 'None' at runtime enables easy pipeline splitting
-          h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens
-
-          for layer in self.layers.values():
-              h = layer(h, self.freqs_cis)
-
-          h = self.norm(h) if self.norm else h
-          output = self.output(h).float() if self.output else h
-          return output
-
-A model defined in this manner can be easily configured per stage by first
-initializing the whole model (using meta-device to avoid OOM errors), deleting
-undesired layers for that stage, and then creating a PipelineStage that wraps
-the model.  For example:
-
-.. code-block:: python
-
-  with torch.device("meta"):
-      assert num_stages == 2, "This is a simple 2-stage example"
-
-      # we construct the entire model, then delete the parts we do not need for this stage
-      # in practice, this can be done using a helper function that automatically divides up layers across stages.
-      model = Transformer()
-
-      if stage_index == 0:
-          # prepare the first stage model
-          del model.layers["1"]
-          model.norm = None
-          model.output = None
-
-      elif stage_index == 1:
-          # prepare the second stage model
-          model.tok_embeddings = None
-          del model.layers["0"]
-
-      from torch.distributed.pipelining import PipelineStage
-      stage = PipelineStage(
-          model,
-          stage_index,
-          num_stages,
-          device,
-      )
-
-When composing with other Data or Model parallelism techniques, ``output_args``
-may also be required, if the output shape/dtype of the model chunk will be
-affected.
-
-
-.. _option_2_tracer:
-
-Option 2: splitting a model automatically
-=========================================
-
-If you have a full model and do not want to spend time on modifying it into a
-sequence of "model partitions", the ``pipeline`` API is here to help.
-Here is a brief example:
-
-.. code-block:: python
-
-  class Model(torch.nn.Module):
-      def __init__(self) -> None:
-          super().__init__()
-          self.emb = torch.nn.Embedding(10, 3)
-          self.layers = torch.nn.ModuleList(
-              Layer() for _ in range(2)
-          )
-          self.lm = LMHead()
-
-      def forward(self, x: torch.Tensor) -> torch.Tensor:
-          x = self.emb(x)
-          for layer in self.layers:
-              x = layer(x)
-          x = self.lm(x)
-          return x
-
-
-If we print the model, we can see multiple hierarchies, which makes it hard to split by hand::
-
-  Model(
-    (emb): Embedding(10, 3)
-    (layers): ModuleList(
-      (0-1): 2 x Layer(
-        (lin): Linear(in_features=3, out_features=3, bias=True)
-      )
-    )
-    (lm): LMHead(
-      (proj): Linear(in_features=3, out_features=3, bias=True)
-    )
-  )
-
-Let us see how the ``pipeline`` API works:
-
-.. code-block:: python
-
-  from torch.distributed.pipelining import pipeline, SplitPoint
-
-  # An example micro-batch input
-  x = torch.LongTensor([1, 2, 4, 5])
-
-  pipe = pipeline(
-      module=mod,
-      mb_args=(x,),
-      split_spec={
-          "layers.1": SplitPoint.BEGINNING,
-      }
-  )
-
-The ``pipeline`` API splits your model given a ``split_spec``, where
-``SplitPoint.BEGINNING`` stands for adding a split point
-*before* execution of certain submodule in the ``forward`` function, and
-similarly, ``SplitPoint.END`` for split point *after* such.
-
-If we ``print(pipe)``, we can see::
-
-  GraphModule(
-    (submod_0): GraphModule(
-      (emb): InterpreterModule()
-      (layers): Module(
-        (0): InterpreterModule(
-          (lin): InterpreterModule()
-        )
-      )
-    )
-    (submod_1): GraphModule(
-      (layers): Module(
-        (1): InterpreterModule(
-          (lin): InterpreterModule()
-        )
-      )
-      (lm): InterpreterModule(
-        (proj): InterpreterModule()
-      )
-    )
-  )
-
-  def forward(self, x):
-      submod_0 = self.submod_0(x);  x = None
-      submod_1 = self.submod_1(submod_0);  submod_0 = None
-      return (submod_1,)
-
-
-The "model partitions" are represented by submodules (``submod_0``,
-``submod_1``), each of which is reconstructed with original model operations, weights
-and hierarchies.  In addition, a "root-level" ``forward`` function is
-reconstructed to capture the data flow between those partitions. Such data flow
-will be replayed by the pipeline runtime later, in a distributed fashion.
-
-The ``Pipe`` object provides a method for retrieving the "model partitions":
-
-.. code-block:: python
-
-  stage_mod : nn.Module = pipe.get_stage_module(stage_idx)
-
-The returned ``stage_mod`` is a ``nn.Module``, with which you can create an
-optimizer, save or load checkpoints, or apply other parallelisms.
-
-``Pipe`` also allows you to create a distributed stage runtime on a device given
-a ``ProcessGroup``:
-
-.. code-block:: python
-
-  stage = pipe.build_stage(stage_idx, device, group)
-
-Alternatively, if you would like to build the stage runtime later after some
-modification to the ``stage_mod``, you can use a functional version of the
-``build_stage`` API. For example:
-
-.. code-block:: python
-
-  from torch.distributed.pipelining import build_stage
-  from torch.nn.parallel import DistributedDataParallel
-
-  dp_mod = DistributedDataParallel(stage_mod)
-  info = pipe.info()
-  stage = build_stage(dp_mod, stage_idx, info, device, group)
-
-.. note::
-  The ``pipeline`` frontend uses a tracer (``torch.export``) to capture your
-  model into a single graph. If your model is not full-graph'able, you can use
-  our manual frontend below.
-
-
-Hugging Face Examples
-*********************
-
-In the `PiPPy <https://github.com/pytorch/PiPPy>`_ repo where this package was
-original created, we kept examples based on unmodified Hugging Face models.
-See the `examples/huggingface
-<https://github.com/pytorch/PiPPy/tree/main/examples/huggingface>`_ directory.
-
-Examples include:
-
-* `GPT2 <https://github.com/pytorch/PiPPy/tree/main/examples/huggingface/pippy_gpt2.py>`_
-* `Llama <https://github.com/pytorch/PiPPy/tree/main/examples/llama>`_
-
-
-Technical Deep Dive
-*******************
-
-How does the ``pipeline`` API split a model?
-============================================
-
-First, the ``pipeline`` API turns our model into a directed acyclic graph (DAG)
-by tracing the model.  It traces the model using ``torch.export`` -- a PyTorch 2
-full-graph capturing tool.
-
-Then, it groups together the **operations and parameters** needed by a stage
-into a reconstructed submodule: ``submod_0``, ``submod_1``, ...
-
-Different from conventional submodule access methods like ``Module.children()``,
-the ``pipeline`` API does not only cut the module structure of your model, but
-also the **forward** function of your model.
-
-This is necessary because model structure like ``Module.children()`` merely
-captures information during ``Module.__init__()``, and does not capture any
-information about ``Module.forward()``. Said differently, ``Module.children()``
-lacks information about the following aspects key to pipelininig:
-
-* Execution order of child modules in ``forward``
-* Activation flows between child modules
-* Whether there are any functional operators between child modules (for example,
-  ``relu`` or ``add`` operations will not be captured by ``Module.children()``).
-
-The ``pipeline`` API, on the contrary, makes sure that the ``forward`` behavior
-is truly preserved. It also captures the activation flow between the partitions,
-helping the distributed runtime to make correct send/receive calls without human
-intervention.
-
-Another flexibility of the ``pipeline`` API is that split points can be at
-arbitrary levels within your model hierarchy. In the split partitions, the original model
-hierarchy related to that partition will be reconstructed at no cost to you.
-At a result, fully-qualified names (FQNs) pointing to a submodule or parameter
-would be still valid, and services that relies on FQNs (such as FSDP, TP or
-checkpointing) can still run with your partitioned modules with almost zero code
-change.
-
-
-Implementing Your Own Schedule
-******************************
-
-You can implement your own pipeline schedule by extending one of the following two class:
-
-* ``PipelineScheduleSingle``
-* ``PipelineScheduleMulti``
-
-``PipelineScheduleSingle`` is for schedules that assigns *only one* stage per rank.
-``PipelineScheduleMulti`` is for schedules that assigns multiple stages per rank.
-
-For example, ``ScheduleGPipe`` and ``Schedule1F1B`` are subclasses of ``PipelineScheduleSingle``.
-Whereas, ``ScheduleInterleaved1F1B``, ``ScheduleLoopedBFS``, ``ScheduleInterleavedZeroBubble``, and ``ScheduleZBVZeroBubble``
-are subclasses of ``PipelineScheduleMulti``.
-
-
-Logging
-*******
-
-You can turn on additional logging using the `TORCH_LOGS` environment variable from `torch._logging <https://pytorch.org/docs/main/logging.html#module-torch._logging>`_:
-
-* `TORCH_LOGS=+pp` will display `logging.DEBUG` messages and all levels above it.
-* `TORCH_LOGS=pp` will display `logging.INFO` messages and above.
-* `TORCH_LOGS=-pp` will display `logging.WARNING` messages and above.
-
-
-API Reference
-*************
-
-.. automodule:: torch.distributed.pipelining
-
-Model Split APIs
-============================
-
-The following set of APIs transform your model into a pipeline representation.
-
-.. currentmodule:: torch.distributed.pipelining
-
-.. autoclass:: SplitPoint
-
-.. autofunction:: pipeline
-
-.. autoclass:: Pipe
-
-.. autofunction:: pipe_split
-
-Microbatch Utilities
-====================
-
-.. automodule:: torch.distributed.pipelining.microbatch
-
-.. currentmodule:: torch.distributed.pipelining.microbatch
-
-.. autoclass:: TensorChunkSpec
-
-.. autofunction:: split_args_kwargs_into_chunks
-
-.. autofunction:: merge_chunks
-
-Pipeline Stages
-===============
-
-.. automodule:: torch.distributed.pipelining.stage
-
-.. currentmodule:: torch.distributed.pipelining.stage
-
-.. autoclass:: PipelineStage
-
-.. autofunction:: build_stage
-
-Pipeline Schedules
-==================
-
-.. automodule:: torch.distributed.pipelining.schedules
-
-.. currentmodule:: torch.distributed.pipelining.schedules
-
-.. autoclass:: ScheduleGPipe
-
-.. autoclass:: Schedule1F1B
-
-.. autoclass:: ScheduleInterleaved1F1B
-
-.. autoclass:: ScheduleLoopedBFS
-
-.. autoclass:: ScheduleInterleavedZeroBubble
-
-.. autoclass:: ScheduleZBVZeroBubble
-
-.. autoclass:: PipelineScheduleSingle
-  :members:
-
-.. autoclass:: PipelineScheduleMulti
-  :members:

torch/distributed/optim/zero_redundancy_optimizer.py

@@ -284,7 +284,7 @@ class ZeroRedundancyOptimizer(Optimizer, Joinable):
     r"""
     Wrap an arbitrary :class:`optim.Optimizer <torch.optim.Optimizer>` and shards its states across ranks in the group.
 
-    The sharing is done as described by ZeRO_.
+    The sharing is done as described by `ZeRO <https://arxiv.org/abs/1910.02054>`_.
 
     The local optimizer instance in each rank is only
     responsible for updating approximately ``1 / world_size`` parameters and
@@ -365,9 +365,6 @@ class ZeroRedundancyOptimizer(Optimizer, Joinable):
         is to prepend dummy inputs.
 
     .. warning:: ZeroRedundancyOptimizer is experimental and subject to change.
-
-    .. _ZeRO: https://arxiv.org/abs/1910.02054
-
     """
 
     def __init__(