Implementation:Deepspeedai DeepSpeed PipelineModule Init

Overview

Concrete tool for constructing pipeline-parallel model partitions provided by the DeepSpeed library. PipelineModule takes a sequence of layers, partitions them across pipeline stages, builds only the layers assigned to the local stage, and establishes the communication grid for point-to-point data transfer between stages.

Description

PipelineModule takes a sequence of layers (nn.Module, LayerSpec, or callable), partitions them across pipeline stages using the specified method, builds only the layers assigned to the local stage, and establishes the communication grid for point-to-point data transfer between stages.

The construction process:

Creates or accepts a PipeDataParallelTopology defining the pipeline and data parallel dimensions.
Constructs a PipelineParallelGrid for inter-stage and intra-stage communication.
Partitions layers using the chosen method ('parameters', 'uniform', 'type:regex').
Builds only local-stage layers by calling LayerSpec.build() or registering existing nn.Module instances.
Indexes tied modules and creates communication groups for weight synchronization across stages.
Moves constructed layers to the local GPU device.

Code Reference

Repository: https://github.com/deepspeedai/DeepSpeed
File: deepspeed/runtime/pipe/module.py
Lines: L86-220

Signature:

class PipelineModule(nn.Module):
    def __init__(self,
                 layers,
                 num_stages=None,
                 topology=None,
                 loss_fn=None,
                 seed_layers=False,
                 seed_fn=None,
                 base_seed=1234,
                 partition_method='parameters',
                 activation_checkpoint_interval=0,
                 activation_checkpoint_func=checkpointing.checkpoint,
                 checkpointable_layers=None,
                 dynamic_shape=False)

Import:

from deepspeed.pipe import PipelineModule

I/O Contract

Inputs

Parameter	Type	Required	Default	Description
layers	Iterable	Yes	—	Sequence of `LayerSpec`, `nn.Module`, or callable objects defining the pipeline
num_stages	int	No	None	Pipeline parallelism degree; must divide world size. Either this or `topology` required.
topology	ProcessTopology	No	None	Custom topology for multi-dimensional parallelism. Either this or `num_stages` required.
loss_fn	callable	No	None	Loss function for the last stage; signature `loss_fn(outputs, labels)`
seed_layers	bool	No	False	Use a different random seed for each layer during construction
seed_fn	callable	No	None	Custom seed-setting function
base_seed	int	No	1234	Starting seed value
partition_method	str	No	'parameters'	One of `'parameters'`, `'uniform'`, `'type:regex'`, or `'profile'`
activation_checkpoint_interval	int	No	0	Checkpoint every N layers; 0 disables
activation_checkpoint_func	callable	No	checkpointing.checkpoint	Function used for activation checkpointing
checkpointable_layers	list[str]	No	None	Layer class names eligible for checkpointing
dynamic_shape	bool	No	False	Allow dynamic input shapes (may impact performance)

Outputs

Output	Type	Description
PipelineModule	PipelineModule (nn.Module)	Module with layers partitioned across stages, communication grid established, tied weights synchronized

Usage Example

from deepspeed.pipe import PipelineModule, LayerSpec
import torch.nn as nn

# Define a 24-layer model
layers = [LayerSpec(nn.Linear, 1024, 1024) for _ in range(24)]

# Create a 4-stage pipeline with parameter-balanced partitioning
model = PipelineModule(
    layers=layers,
    num_stages=4,
    loss_fn=nn.CrossEntropyLoss(),
    partition_method='parameters',
    activation_checkpoint_interval=1
)

# Using type-based partitioning for transformer models
transformer_layers = [
    LayerSpec(nn.Embedding, vocab_size, hidden_size),
    *[LayerSpec(TransformerLayer, hidden_size) for _ in range(12)],
    LayerSpec(nn.Linear, hidden_size, vocab_size),
]
model = PipelineModule(
    layers=transformer_layers,
    num_stages=4,
    partition_method='type:TransformerLayer',
    loss_fn=nn.CrossEntropyLoss(),
)

Related Pages

Knowledge Sources

Last updated: 2026-02-09 00:00 GMT

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment