PiPPy: Pipeline Parallelism for PyTorch

[!NOTE] PiPPy has been migrated into PyTorch as a subpackage: torch.distributed.pipelining. You can find the detailed documentation here. The current repo mainly serves as a land of examples. The PiPPy library code will be removed. Please use the APIs in torch.distributed.pipelining instead. Thank you!

Why PiPPy? | Install guide | Examples | PiPPy Explained

Why PiPPy?

One of the most important techniques for advancing the state of the art in deep learning is scaling. Common techniques for scaling neural networks include data parallelism, tensor/operation parallelism, and pipeline parallelism. In many cases, pipeline parallelism in particular can be an effective technique for scaling, however it is often difficult to implement, requiring intrusive code changes to model code and difficult-to-implement runtime orchestration code. PiPPy aims to provide a toolkit that does said things automatically to allow high-productivity scaling of models.

What is PiPPy?

The PiPPy project consists of a compiler and runtime stack for automated parallelism and scaling of PyTorch models. Currently, PiPPy focuses on pipeline parallelism, a technique in which the code of the model is partitioned and multiple micro-batches execute different parts of the model code concurrently. To learn more about pipeline parallelism, see this article.

Figure: Pipeline parallel. "F", "B" and "U" denote forward, backward and weight update, respectively. Different colors represent different micro-batches.

PiPPy provides the following features that make pipeline parallelism easier:

• Automatic splitting of model code by tracing the model. The goal is for the user to provide model code as-is to the system for parallelization, without having to make heavyweight modifications to make parallelism work.
• Related to the last point, PiPPy supports non-trivial topologies, including skip connections and tied weights/layers. PiPPy provides configurable behavior for tied weights, allowing for transmission across pipeline stages or replication and gradient synchronization.
• First-class support for cross-host pipeline parallelism, as this is where PP is typically used (over slower interconnects). This is currently missing from the torchgpipe-based torch.distributed.pipeline.sync.Pipe.
• Composability with other parallelism schemes such as data parallelism or tensor splitting model parallelism (overall, known as "3d parallelism"). Currently, pipelining and data parallelism can be composed. Other compositions will be available in the future.
• Support for pipeline scheduling paradigms, including schedules like fill-drain (GPipe), 1F1B and interleaved 1F1B. More schedules will be added too.

For in-depth technical architecture, see ARCHITECTURE.md.

Install

PiPPy requires PyTorch version newer than 2.2.0.dev to work. To quickly install, for example, PyTorch nightly, run the following command from the same directory as this README:

pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html

You can also select the CUDA build of PyTorch if your system has NVIDIA GPUs, for example:

pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cu118/torch_nightly.html

To install PiPPy from source, run the following command in the same directory as this README:

python setup.py install

To expose PiPPy for development such that changes to this repo are reflected in the imported package, run:

python setup.py develop

Examples

In this repo, we provide rich examples based on realistic models. In particular, we show how to apply PiPPy without any code change to the model. Please refer to the HuggingFace examples directory. Examples include: BERT, GPT2, T5, LLaMA, etc.

PiPPy Explained

PiPPy consists of two parts: a compiler and a runtime. The compiler takes your model code, splits it up, and transforms it into a Pipe, which is a wrapper that describes the model at each pipeline stage and their data-flow relationship. The runtime executes the PipelineStages in parallel, handling things like micro-batch splitting, scheduling, communication, and gradient propagation, etc. We will cover the APIs for these concepts in this section.

Splitting a Model with Pipe

To see how we can split a model into a pipeline, let's first take an example trivial neural network:

```python import torch

class MyNetworkBlock(torch.nn.Module): def init(self, indim, outdim): super().init() self.lin = torch.nn.Linear(indim, outdim)

def forward(self, x):
    x = self.lin(x)
    x = torch.relu(x)
    return x

class MyNetwork(torch.nn.Module): def init(self, indim, layerdims): super().init()

    prev_dim = in_dim
    for i, dim in enumerate(layer_dims):
        setattr(self, f'layer{i}', MyNetworkBlock(prev_dim, dim))
        prev_dim = dim

    self.num_layers = len(layer_dims)
    # 10 output classes
    self.output_proj = torch.nn.Linear(layer_dims[-1], 10)

def forward(self, x):
    for i in range(self.num_layers):
        x = getattr(self, f'layer{i}')(x)

    return self.output_proj(x)

indim = 512 layerdims = [512, 1024, 256] mn = MyNetwork(indim, layerdims).to(device) ```

This network is written as free-form Python code; it has not been modified for any specific parallelism technique.

Let us see our first usage of the pippy.Pipe interface:

```python from pippy import pipeline, annotatesplitpoints, Pipe, SplitPoint

annotatesplitpoints(mn, {'layer0': SplitPoint.END, 'layer1': SplitPoint.END})

batchsize = 32 exampleinput = torch.randn(batchsize, indim, device=device) chunks = 4

pipe = pipeline(mn, chunks, exampleargs=(exampleinput,)) print(pipe)

""" ************************************* pipe ************************************* GraphModule( (submod0): PipeStageModule( (Lselflayer0modlin): Linear(infeatures=512, outfeatures=512, bias=True) ) (submod_1): PipeStageModule( (Lselflayer1modlin): Linear(infeatures=512, outfeatures=1024, bias=True) ) (submod2): PipeStageModule( (Lselflayer2lin): Linear(infeatures=1024, outfeatures=256, bias=True) (Lself__outputproj): Linear(infeatures=256, outfeatures=10, bias=True) ) )

def forward(self, arg0): submod0 = self.submod0(arg0); arg0 = None submod1 = self.submod1(submod0); submod0 = None submod2 = self.submod2(submod1); submod1 = None return [submod_2] """ ```

So what's going on here? First, pipeline turns our model into a directed acyclic graph (DAG) by tracing the model. Then, it groups together the operations and parameters into pipeline stages. Stages are represented as submod_N submodules, where N is a natural number.

We used annotate_split_points to specify that the code should be split and the end of layer0 and layer1. Our code has thus been split into three pipeline stages. PiPPy also provides SplitPoint.BEGINNING if a user wants to split before certain annotation point.

While the annotate_split_points API gives users a way to specify the split points without modifying the model, PiPPy also provides an API for in-model annotation: pipe_split(). For details, you can read this example.

This covers the basic usage of the Pipe API. For more information, see the documentation.

Using PipelineStage for Pipelined Execution

Given the above Pipe object, we can use one of the PipelineStage classes to execute our model in a pipelined fashion. First off, let us instantiate a PipelineStage instance:

```python

We are using torchrun to run this example with multiple processes.

torchrun defines two environment variables: RANK and WORLD_SIZE.

rank = int(os.environ["RANK"]) worldsize = int(os.environ["WORLDSIZE"])

Initialize distributed environment

import torch.distributed as dist dist.initprocessgroup(rank=rank, worldsize=worldsize)

Pipeline stage is our main pipeline runtime. It takes in the pipe object,

the rank of this process, and the device.

from pippy.PipelineStage import PipelineStage stage = PipelineStage(pipe, rank, device) ```

We can now run the pipeline by passing input to the first PipelineStage:

```python

Input data

x = torch.randn(batchsize, indim, device=device)

Run the pipeline with input x. Divide the batch into 4 micro-batches

and run them in parallel on the pipeline

if rank == 0: stage(x) elif rank == world_size - 1: output = stage() else: stage() ```

Note that since we split our model into three stages, we must run this script with three workers. For this example, we will use torchrun to run multiple processes within a single machine for demonstration purposes. We can collect up all of the code blocks above into a file named example.py and then run it with torchrun like so:

torchrun --nproc_per_node=3 example.py

License

PiPPy is 3-clause BSD licensed, as found in the LICENSE file.

Citing PiPPy

If you use PiPPy in your publication, please cite it by using the following BibTeX entry.

bibtex @Misc{pippy2022, author = {James Reed, Pavel Belevich, Ke Wen, Howard Huang, Will Constable}, title = {PiPPy: Pipeline Parallelism for PyTorch}, howpublished = {\url{https://github.com/pytorch/PiPPy}}, year = {2022} }