流水线并行

模型规范

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, optional) – 可检查点的层可能不会被检查点。默认为 None，不会进行额外的过滤。

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

forward(forward_input)

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

应由所有子类覆盖。

注意

虽然前向传递的配方需要在此函数中定义，但应随后调用 Module 实例，而不是此函数，因为前者负责运行注册的钩子，而后者则静默地忽略它们。

allreduce_tied_weight_gradients()

在绑定阶段之间对绑定权重的梯度进行全归约

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。

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) – [description]。默认为 ‘_’。

• **outer_sep** (str, optional) – [description]。默认为 ‘-’。

返回

rank 所拥有的坐标的字符串表示形式。

返回类型

str

get_dim(axis)

返回给定轴上进程的数量。

例如

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

get_coord(rank)

返回进程秩所拥有的坐标。

返回的 namedtuple 的轴可以直接作为成员访问。例如 .. rubric:: 例子

>>> 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 (可选) – 作为成员存储的关键字参数

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 配对以避免死锁。