流水线并行

模型规范

class deepspeed.pipe.PipelineModule(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=<function checkpoint>, checkpointable_layers=None, dynamic_shape=False)

使用流水线并行进行并行的模块。

启用流水线并行的关键约束是将前向传播表示为一系列层，并强制在它们之间使用简单的接口。前向传播由模块 layers 隐式定义。关键假设是每层的输出可以直接作为下一层的输入，就像 torch.nn.Sequence 一样。前向传播隐式地

def forward(self, inputs):
    x = inputs
    for layer in self.layers:
        x = layer(x)
    return x

注意

流水线并行与 ZeRO-2 和 ZeRO-3 不兼容。

参数

• layers (Iterable) – 定义流水线结构的层序列。可以是 torch.nn.Sequential 模块。

• num_stages (int, optional) – 流水线并行的程度。如果未指定，则必须提供 topology。

• topology (deepspeed.runtime.pipe.ProcessTopology, optional) – 定义训练并行轴。如果 num_stages 为 None，则必须提供。

• loss_fn (callable, optional) – 损失通过 loss = loss_fn(outputs, label) 计算。

• seed_layers (bool, optional) – 为每个层使用不同的种子。默认为 False。

• seed_fn (type, optional) – 自定义种子生成函数。默认为随机种子生成器。

• base_seed (int, optional) – 起始种子。默认为 1234。

• partition_method (str, optional) – 划分层的方法。默认为 'parameters'。

• activation_checkpoint_interval (int, optional) – 激活检查点的粒度，以层数表示。0 表示禁用激活检查点。

• activation_checkpoint_func (callable, optional) – 用于激活检查点的函数。默认为 deepspeed.checkpointing.checkpoint。

• checkpointable_layers (list[str], optional) – 符合检查点条件的层类名称列表。对于 GPT 模型，ParallelTransformerLayerPipe 总是被检查点，无论此列表如何。如果为 None，则所有带有参数的层都被视为可检查点。默认为 None。

• dynamic_shape – 允许输入具有动态形状。这可能会影响性能。

forward(forward_input)

定义每次调用时执行的计算。

应被所有子类覆盖。

注意

尽管前向传播的实现需要在该函数中定义，但应在之后调用 Module 实例，而不是直接调用此函数，因为前者会处理运行已注册的钩子，而后者会默默忽略它们。

allreduce_tied_weight_gradients()

在绑定阶段之间对绑定权重的梯度进行 All-reduce 操作

topology()

用于查询进程映射的 ProcessTopology 对象。

ckpt_prefix(checkpoints_path, tag)

为该模块写入的所有检查点文件构建前缀。

ckpt_layer_path(ckpt_dir, local_layer_idx)

为特定流水线模块层自定义前缀。

ckpt_layer_path_list(ckpt_dir, local_layer_idx)

获取特定流水线模块层的所有检查点文件列表。

get_additional_losses()

返回模型特定的额外损失以供报告

返回一个字典 {“损失名称”: 损失值}，如果没有额外损失则返回 None。

compile(*args, **kwargs)

使用 torch.compile() 编译此模块的前向传播。

此模块的 __call__ 方法被编译，所有参数都按原样传递给 torch.compile()。

有关此函数参数的详细信息，请参阅 torch.compile()。

class deepspeed.pipe.LayerSpec(typename, *module_args, **module_kwargs)

用于指定流水线并行模块的构建块。

LayerSpec 存储 PipelineModule 中每个阶段的类型信息和参数。例如

nn.Sequence(
    torch.nn.Linear(self.in_dim, self.hidden_dim, bias=False),
    torch.nn.Linear(self.hidden_hidden, self.out_dim)
)

变为

layer_specs = [
    LayerSpec(torch.nn.Linear, self.in_dim, self.hidden_dim, bias=False),
    LayerSpec(torch.nn.Linear, self.hidden_hidden, self.out_dim)]
]

build(log=False)

构建存储的规范。

class deepspeed.pipe.TiedLayerSpec(key, typename, *module_args, forward_fn=None, tied_weight_attr=['weight'], **module_kwargs)

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

管理 n 维笛卡尔坐标到线性索引的映射。此映射用于将进程的秩映射到各种并行形式的网格。

张量的每个轴都通过其名称访问。提供的轴顺序定义了拓扑的布局。ProcessTopology 使用张量轴的“行优先”布局，因此 axes=['x', 'y'] 会将坐标 (x,y) 和 (x,y+1) 映射到相邻的线性索引。如果使用 axes=['y', 'x']，则坐标 (x,y) 和 (x+1,y) 将相邻。

一些方法返回 ProcessCoord 命名元组。

get_rank(**coord_kwargs)

通过坐标返回进程的全局秩。

坐标指定为 kwargs。例如

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

get_axis_names()

按照拓扑顺序返回轴名称列表。

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

返回秩的字符串表示。

此方法主要用于检查点模型数据。

例如

>>> 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'

参数

• rank (int) – 拓扑中的一个秩。

• omit_axes (list, optional) – 不应包含在表示中的轴。默认为 ['data', 'pipe']。

• inner_sep (str, optional) – [描述]。默认为 '_'。

• outer_sep (str, optional) – [描述]。默认为 '-'。

返回

由 rank 拥有的坐标的字符串表示。

返回类型

str

get_dim(axis)

返回给定轴上的进程数。

例如

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

get_coord(rank)

返回进程秩拥有的坐标。

返回的命名元组的轴可以作为成员直接访问。例如：

>>> 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)

构建适合沿 axis 轴的通信器组的列表。

示例

>>> 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
]

返回

一个列表的列表，其中所有轴的坐标都匹配，*除了* axis。

filter_match(**filter_kwargs)

返回坐标与提供标准匹配的秩列表。

示例

>>> 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)]

参数

**filter_kwargs (dict) – 用于选择坐标的标准。

返回

坐标与 filter_kwargs 匹配的秩列表。

get_axis_list(axis, idx)

返回在某个轴上坐标为 idx 的全局秩列表。

例如

>>> 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]

训练

扩展流水线并行

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

通过生成 PipeInstruction 序列来指导流水线引擎的执行。

调度器是生成器，它产生 PipeInstruction 序列以处理一个批次中的微批次。每个生成的步骤都是原子的，即可以在连续步骤之间放置屏障同步而不会发生死锁。

下面是一个实现带有梯度累积的数据并行的调度示例

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

参数

• micro_batches (int) – 构成一个批次的微批次数量。

• stages (int) – 流水线阶段的数量。

• stage_id (int) – 将执行生成调度的流水线阶段。

abstract steps()

为调度中的每个步骤生成一个 PipeInstruction 列表。

注意

调度器必须实现 steps() 来定义调度。

返回

作为流水线一步执行的指令

num_pipe_buffers()

此阶段将使用的流水线缓冲区数量。

注意

调度器应针对大规模内存节省优化 num_pipe_buffers()。

返回

引擎要分配的缓冲区数量。

property stage

用于配置此调度的阶段索引。

property num_stages

用于配置此调度的流水线阶段总数。

property num_micro_batches

用于配置此调度的微批次总数。

property is_first_stage

如果配置的 stage_id 是流水线中的第一个阶段，则为 True。

property is_last_stage

如果配置的 stage_id 是流水线中的最后一个阶段，则为 True。

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

使用流水线并行进行批次推理的调度。

num_pipe_buffers()

推理只需两个流水线缓冲区。

返回

2

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

使用混合并行训练批次的调度。

流水线并行通过梯度累积实现，因此其收敛行为与相同批次大小的数据并行方法一致。

num_pipe_buffers()

返回此阶段所需的流水线缓冲区数量。

这等同于最大在途前向传播次数，因为我们需要记住前向传播的激活值以便进行反向传播。对于同步 1F1B，这等同于此阶段与最后一个阶段之间的索引差。

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

一个使用传统数据并行和梯度累积进行训练的调度示例。

num_pipe_buffers()

只需要一个流水线缓冲区。

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

流水线引擎将执行的所有指令的基类。

所有关键字参数都作为成员存储，类似于 namedtuple。然后，这些参数在执行期间可供 PipeEngine 访问。

参数

kwargs (optional) – 要作为成员存储的关键字参数

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

使用优化器执行一步并清零梯度。

注意

应在 ReduceGrads 和 ReduceTiedGrads 之后发出。

注意

可以是数据并行秩之间的同步点。

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

在阶段内的数据并行进程中规约计算出的梯度。

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

在流水线并行组中规约绑定模块的计算梯度。

警告

此同步点中包含的阶段直到模型在流水线阶段之间划分后才可知。在最坏情况下，它包括所有流水线阶段。应仔细安排此指令以避免死锁。

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

一个操作流水线缓冲区的流水线指令。

参数

buffer_id (int) – 要修改的流水线缓冲区索引。

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

将微批次加载到缓冲区。

大致地

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

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

计算一次前向传播。

大致地

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

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

计算一次反向传播并累积梯度。

大致地

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(buffers['outputs'][buffer_id])

注意

通信是阻塞的，并且必须与下一个流水线阶段上的 RecvActivation 配对，以避免死锁。

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

从流水线中的上一个阶段接收激活。

大致地

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

注意

通信是阻塞的，并且必须与上一个流水线阶段上的 SendActivation 配对，以避免死锁。

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

将计算出的梯度发送到上一个流水线阶段，相对于接收到的激活。

注意

只有 requires_grad==True 的接收张量才会产生梯度。缺失的梯度将在接收阶段被 None 替换。

注意

通信是阻塞的，并且必须与上一个流水线阶段上的 RecvGrad 配对，以避免死锁。

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

从下一个流水线阶段接收计算出的梯度。

注意

只有 requires_grad==True 的激活才会产生梯度。缺失的梯度将被 None 替换。

注意

通信是阻塞的，并且必须与下一个流水线阶段上的 SendGrad 配对，以避免死锁。