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
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.
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.
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
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 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.
To see how we can split a model into a pipeline, let's first take an example trivial neural network:
import torch
class MyNetworkBlock(torch.nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
self.lin = torch.nn.Linear(in_dim, out_dim)
def forward(self, x):
x = self.lin(x)
x = torch.relu(x)
return x
class MyNetwork(torch.nn.Module):
def __init__(self, in_dim, layer_dims):
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)
in_dim = 512
layer_dims = [512, 1024, 256]
mn = MyNetwork(in_dim, layer_dims).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:
from pippy import pipeline, annotate_split_points, Pipe, SplitPoint
annotate_split_points(mn, {'layer0': SplitPoint.END,
'layer1': SplitPoint.END})
batch_size = 32
example_input = torch.randn(batch_size, in_dim, device=device)
chunks = 4
pipe = pipeline(mn, chunks, example_args=(example_input,))
print(pipe)
"""
************************************* pipe *************************************
GraphModule(
(submod_0): PipeStageModule(
(L__self___layer0_mod_lin): Linear(in_features=512, out_features=512, bias=True)
)
(submod_1): PipeStageModule(
(L__self___layer1_mod_lin): Linear(in_features=512, out_features=1024, bias=True)
)
(submod_2): PipeStageModule(
(L__self___layer2_lin): Linear(in_features=1024, out_features=256, bias=True)
(L__self___output_proj): Linear(in_features=256, out_features=10, bias=True)
)
)
def forward(self, arg0):
submod_0 = self.submod_0(arg0); arg0 = None
submod_1 = self.submod_1(submod_0); submod_0 = None
submod_2 = self.submod_2(submod_1); submod_1 = 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.
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:
# 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"])
world_size = int(os.environ["WORLD_SIZE"])
# Initialize distributed environment
import torch.distributed as dist
dist.init_process_group(rank=rank, world_size=world_size)
# 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:
# Input data
x = torch.randn(batch_size, in_dim, 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
PiPPy is 3-clause BSD licensed, as found in the LICENSE file.
If you use PiPPy in your publication, please cite it by using the following BibTeX entry.
@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}
}