Pipeline Parallelism

Model Specification

class deepspeed.runtime.pipe.ProcessTopology(axes, dims)

Manages the mapping of n-dimensional Cartesian coordinates to linear indices. This mapping is used to map the rank of processes to the grid for various forms of parallelism.

Each axis of the tensor is accessed by its name. The provided ordering of the axes defines the layout of the topology. ProcessTopology uses a “row-major” layout of the tensor axes, and so axes=[‘x’, ‘y’] would map coordinates (x,y) and (x,y+1) to adjacent linear indices. If instead axes=[‘y’, ‘x’] was used, coordinates (x,y) and (x+1,y) would be adjacent.

Some methods return ProcessCoord namedtuples.

Create a mapping of n-dimensional tensor coordinates to linear indices.

Parameters

• axes (list) – the names of the tensor axes

• dims (list) – the dimension (length) of each axis of the topology tensor

get_rank(**coord_kwargs)

Return the global rank of a process via its coordinates.

Coordinates are specified as kwargs. For example:

>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])
>>> X.get_rank(x=0, y=1)
1

get_axis_names()

Return a list of the axis names in the ordering of the topology.

get_rank_repr(rank, omit_axes=['data', 'pipe'], inner_sep='_', outer_sep='-')

Return a string representation of a rank.

This method is primarily used for checkpointing model data.

For example:

>>> topo = Topo(axes=['a', 'b'], dims=[2, 2])
>>> topo.get_rank_repr(rank=3)
'a_01-b_01'
>>> topo.get_rank_repr(rank=3, omit_axes=['a'])
'b_01'

Parameters

• rank (int) – A rank in the topology.

• omit_axes (list, optional) – Axes that should not be in the representation. Defaults to [‘data’, ‘pipe’].

• inner_sep (str, optional) – [description]. Defaults to ‘_’.

• outer_sep (str, optional) – [description]. Defaults to ‘-‘.

Returns

A string representation of the coordinate owned by rank.

Return type

str

get_dim(axis)

Return the number of processes along the given axis.

For example:

>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])
>>> X.get_dim('y')
3

get_coord(rank)

Return the coordinate owned by a process rank.

The axes of the returned namedtuple can be directly accessed as members. For .. rubric:: Example

>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])
>>> coord = X.get_coord(rank=1)
>>> coord.x
0
>>> coord.y
1

get_axis_comm_lists(axis)

Construct lists suitable for a communicator group along axis axis.

Example

>>> topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
>>> topo.get_axis_comm_lists('pipe')
[
    [0, 4], # data=0, model=0
    [1, 5], # data=0, model=1
    [2, 6], # data=1, model=0
    [3, 7], # data=1, model=1
]

Returns

A list of lists whose coordinates match in all axes except axis.

filter_match(**filter_kwargs)

Return the list of ranks whose coordinates match the provided criteria.

Example

>>> X = ProcessTopology(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
>>> X.filter_match(pipe=0, data=1)
[2, 3]
>>> [X.get_coord(rank) for rank in X.filter_match(pipe=0, data=1)]
[ProcessCoord(pipe=0, data=1, model=0), ProcessCoord(pipe=0, data=1, model=1)]

Parameters

**filter_kwargs (dict) – criteria used to select coordinates.

Returns

The list of ranks whose coordinates match filter_kwargs.

get_axis_list(axis, idx)

Returns the list of global ranks whose coordinate in an axis is idx.

For example:

>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])
>>> X.get_axis_list(axis='x', idx=0)
[0, 1, 2]
>>> X.get_axis_list(axis='y', idx=0)
[0, 3]

Training

Extending Pipeline Parallelism

class deepspeed.runtime.pipe.schedule.PipeSchedule(micro_batches, stages, stage_id)

Directs the execution of a pipeline engine by generating sequences of PipeInstruction.

Schedules are generators that yield sequences of PipeInstruction to process the micro-batches in one batch. Each yielded step is atomic in the sense that a barrier synchronization can be placed between successive steps without deadlock.

Below is an example schedule that implements data parallelism with gradient accumulation:

class DataParallelSchedule(PipeSchedule):
    def steps(self):
        for step_id in range(self.micro_batches):
            cmds = [
                LoadMicroBatch(buffer_id=0),
                ForwardPass(buffer_id=0),
                BackwardPass(buffer_id=0),
            ]
            if step_id == self.micro_batches - 1:
                cmds.extend([
                    ReduceGrads(),
                    OptimizerStep(),
                ])
            yield cmds

    def num_pipe_buffers(self):
        return 1

Parameters

• micro_batches (int) – The number of micro-batches that comprise a batch.

• stages (int) – The number of pipeline stages.

• stage_id (int) – The pipe stage that will execute the generated schedule.

abstract steps()

Yield a list of PipeInstruction for each step in the schedule.

Note

Schedules must implement steps() to define the schedule.

Returns

Instructions to be executed as one step of the pipeline

num_pipe_buffers()

The number of pipeline buffers that will be used by this stage.

Note

Schedules should specialize num_pipe_buffers() for memory savings at scale.

Returns

The number of buffers for the engine to allocate.

property stage

Stage index used to configure this schedule.

property num_stages

The number of total pipeline stages used to configure this schedule.

property num_micro_batches

The number of total micro_batches used to configure this schedule.

property is_first_stage

True if the configured stage_id is the first stage in the pipeline.

property is_last_stage

True if the configured stage_id is the last stage in the pipeline.

class deepspeed.runtime.pipe.schedule.InferenceSchedule(micro_batches, stages, stage_id)

A schedule for inferencing batches using pipeline parallelism.

num_pipe_buffers()

Only two pipeline buffers are required for inferencing.

Returns

2

class deepspeed.runtime.pipe.schedule.TrainSchedule(micro_batches, stages, stage_id)

A schedule for training a batch using hybrid parallelism.

Pipeline parallelism is extracted through gradient accumulation and thus convergence follows that of a data parallel approach with the same batch size.

num_pipe_buffers()

Return the number of pipeline buffers required for this stage.

This is equivalent to the maximum number of in-flight forward passes, since we need to remember the activations of forward passes in order to run backpropagation. For synchronous 1F1B, this is equivalent to the index difference between this stage and the last stage.

class deepspeed.runtime.pipe.schedule.DataParallelSchedule(micro_batches, stages, stage_id)

An example schedule that trains using traditional data parallelism with gradient accumulation.

num_pipe_buffers()

Only one pipeline buffer needed.

class deepspeed.runtime.pipe.schedule.DualPipeVSchedule(micro_batches, stages, stage_id, forward_only=False)

The DualPipeV schedule of DeepSeek-V3.

GitHub Repo: https://github.com/deepseek-ai/DualPipe

Every rank holds two pipeline stages: stage stage_id (phase 0, activations flow towards the last rank) and stage 2 * stages - 1 - stage_id (phase 1, activations flow back towards rank 0). Micro-batches enter phase 0 on rank 0, turn around on the last rank and finish in phase 1 on rank 0, which is therefore both the first and the loss stage. Interleaving the two phases fills the 1F1B bubble.

Sends and receives are queued and launched as one batch by CommitP2P, so an instruction’s communication is only complete once the next CommitP2P has run. The reference implementation also defers weight gradients (“zero bubble”); this schedule keeps them inside BackwardPass.

Parameters

• micro_batches (int) – Must be at least 2 * stages.

• stages (int) – The number of pipeline ranks (each holds two stages).

• stage_id (int) – The pipeline rank that will execute the generated schedule.

• forward_only (bool) – Skip backward passes and gradient exchanges.

num_pipe_buffers()

One slot per (phase, micro-batch): phase * micro_batches + micro_batch.

steps()

Yield a list of PipeInstruction for each step in the schedule.

Note

Schedules must implement steps() to define the schedule.

Returns

Instructions to be executed as one step of the pipeline

class deepspeed.runtime.pipe.schedule.PipeInstruction(**kwargs)

Base class for all instructions to be executed by the pipeline engine.

All keyword arguments are stored as members similar to a namedtuple. These are then accessible to the PipeEngine during execution.

Parameters

kwargs (optional) – keyword arguments to store as members

class deepspeed.runtime.pipe.schedule.OptimizerStep(**kwargs)

Performs one step with the optimizer and zeros gradients.

Note

Should be issued after ReduceGrads and ReduceTiedGrads.

Note

Can be a synchronization point among data-parallel ranks.

class deepspeed.runtime.pipe.schedule.ReduceGrads(**kwargs)

Reduce the computed gradients among data-parallel processes within the stage.

class deepspeed.runtime.pipe.schedule.ReduceTiedGrads(**kwargs)

Reduce the computed gradients of tied modules within a pipeline-parallel group.

Warning

The stages included in this synchronization point are not known until the model is partitioned among pipeline stages. In the worst case, it includes all pipeline stages. This instruction should be scheduled carefully to avoid deadlocks.

class deepspeed.runtime.pipe.schedule.BufferOpInstruction(buffer_id, **kwargs)

A pipeline instruction that operates on pipeline buffer(s).

Parameters

buffer_id (int) – the index of the pipeline buffer() to modify.

class deepspeed.runtime.pipe.schedule.LoadMicroBatch(buffer_id, **kwargs)

Load a micro-batch into a buffer.

Roughly:

buffers['inputs'][buffer_id] = next(data_iter)

class deepspeed.runtime.pipe.schedule.ForwardPass(buffer_id, **kwargs)

Compute a forward pass.

Roughly:

buffers['outputs'][buffer_id] = forward(buffers['inputs'][buffer_id])

class deepspeed.runtime.pipe.schedule.BackwardPass(buffer_id, **kwargs)

Compute a backward pass and accumulate gradients.

Roughly:

outputs = buffers['outputs'][buffer_id]
gradients = buffers['gradients'][buffer_id]
torch.autograd.backward(tensors=outputs,
                        grad_tensors=gradients)

class deepspeed.runtime.pipe.schedule.SendActivation(buffer_id, **kwargs)

Send activations to the next stage in the pipeline.

Roughly:

send(buffers['outputs'][buffer_id])

Note

The communication is blocking and must be paired with a RecvActivation on the next pipeline stage to avoid deadlock.

class deepspeed.runtime.pipe.schedule.RecvActivation(buffer_id, **kwargs)

Receive activations from the previous stage in the pipeline.

Roughly:

buffers['inputs'][buffer_id] = recv()

Note

The communication is blocking and must be paired with a SendActivation on the previous pipeline stage to avoid deadlock.

class deepspeed.runtime.pipe.schedule.SendGrad(buffer_id, **kwargs)

Send computed gradients to the previous pipeline stage. with respect to the received activations

Note

Only received tensors with requires_grad==True will produce gradients. Missing gradients will be replaced with None on the receiving stage.

Note

The communication is blocking and must be paired with a RecvGrad on the previous pipeline stage to avoid deadlock.

class deepspeed.runtime.pipe.schedule.RecvGrad(buffer_id, **kwargs)

Receive computed gradients the next pipeline stage.

Note

Only activations with requires_grad==True will produce gradients. Missing gradients will be replaced with None.

Note

The communication is blocking and must be paired with a SendGrad on the next pipeline stage to avoid deadlock.

class deepspeed.runtime.pipe.schedule.CommitP2P(**kwargs)

Launch the queued asynchronous sends and receives as one batch and wait for them.

Used by schedules whose communication instructions only queue work, such as DualPipeVSchedule.