流水线并行
模型规范
- 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)[source]
使用流水线并行进行并行的模块。
启用流水线并行的关键约束是将前向传播表示为一系列层,并强制在它们之间使用简单的接口。前向传播由模块
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 – 允许输入具有动态形状。这可能会影响性能。
- class deepspeed.pipe.LayerSpec(typename, *module_args, **module_kwargs)[source]
用于指定流水线并行模块的构建块。
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)] ]
- class deepspeed.pipe.TiedLayerSpec(key, typename, *module_args, forward_fn=None, tied_weight_attr=['weight'], **module_kwargs)[source]
- class deepspeed.runtime.pipe.ProcessTopology(axes, dims)[source]
管理 n 维笛卡尔坐标到线性索引的映射。此映射用于将进程的秩映射到各种并行形式的网格。
张量的每个轴都通过其名称访问。提供的轴顺序定义了拓扑的布局。ProcessTopology 使用张量轴的“行优先”布局,因此 axes=['x', 'y'] 会将坐标 (x,y) 和 (x,y+1) 映射到相邻的线性索引。如果使用 axes=['y', 'x'],则坐标 (x,y) 和 (x+1,y) 将相邻。
一些方法返回 ProcessCoord 命名元组。
- get_rank(**coord_kwargs)[source]
通过坐标返回进程的全局秩。
坐标指定为 kwargs。例如
>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) >>> X.get_rank(x=0, y=1) 1
- get_rank_repr(rank, omit_axes=['data', 'pipe'], inner_sep='_', outer_sep='-')[source]
返回秩的字符串表示。
此方法主要用于检查点模型数据。
- 例如
>>> 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)[source]
返回给定轴上的进程数。
- 例如
>>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3]) >>> X.get_dim('y') 3
- get_coord(rank)[source]
返回进程秩拥有的坐标。
返回的命名元组的轴可以作为成员直接访问。例如:
>>> 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)[source]
构建适合沿
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)[source]
返回坐标与提供标准匹配的秩列表。
示例
>>> 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 匹配的秩列表。
训练
扩展流水线并行
- class deepspeed.runtime.pipe.schedule.PipeSchedule(micro_batches, stages, stage_id)[source]
通过生成
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()[source]
为调度中的每个步骤生成一个
PipeInstruction
列表。注意
调度器必须实现
steps()
来定义调度。- 返回
作为流水线一步执行的指令
- 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)[source]
使用流水线并行进行批次推理的调度。
- class deepspeed.runtime.pipe.schedule.TrainSchedule(micro_batches, stages, stage_id)[source]
使用混合并行训练批次的调度。
流水线并行通过梯度累积实现,因此其收敛行为与相同批次大小的数据并行方法一致。
- class deepspeed.runtime.pipe.schedule.DataParallelSchedule(micro_batches, stages, stage_id)[source]
一个使用传统数据并行和梯度累积进行训练的调度示例。
- class deepspeed.runtime.pipe.schedule.PipeInstruction(**kwargs)[source]
流水线引擎将执行的所有指令的基类。
所有关键字参数都作为成员存储,类似于
namedtuple
。然后,这些参数在执行期间可供PipeEngine
访问。- 参数
kwargs (optional) – 要作为成员存储的关键字参数
- class deepspeed.runtime.pipe.schedule.OptimizerStep(**kwargs)[source]
使用优化器执行一步并清零梯度。
注意
应在
ReduceGrads
和ReduceTiedGrads
之后发出。注意
可以是数据并行秩之间的同步点。
- class deepspeed.runtime.pipe.schedule.ReduceTiedGrads(**kwargs)[source]
在流水线并行组中规约绑定模块的计算梯度。
警告
此同步点中包含的阶段直到模型在流水线阶段之间划分后才可知。在最坏情况下,它包括所有流水线阶段。应仔细安排此指令以避免死锁。
- class deepspeed.runtime.pipe.schedule.BufferOpInstruction(buffer_id, **kwargs)[source]
一个操作流水线缓冲区的流水线指令。
- 参数
buffer_id (int) – 要修改的流水线缓冲区索引。
- class deepspeed.runtime.pipe.schedule.LoadMicroBatch(buffer_id, **kwargs)[source]
将微批次加载到缓冲区。
大致地
buffers['inputs'][buffer_id] = next(data_iter)
- class deepspeed.runtime.pipe.schedule.ForwardPass(buffer_id, **kwargs)[source]
计算一次前向传播。
大致地
buffers['outputs'][buffer_id] = forward(buffers['inputs'][buffer_id])
- class deepspeed.runtime.pipe.schedule.BackwardPass(buffer_id, **kwargs)[source]
计算一次反向传播并累积梯度。
大致地
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)[source]
将激活发送到流水线中的下一个阶段。
大致地
send(buffers['outputs'][buffer_id])
注意
通信是阻塞的,并且必须与下一个流水线阶段上的
RecvActivation
配对,以避免死锁。
- class deepspeed.runtime.pipe.schedule.RecvActivation(buffer_id, **kwargs)[source]
从流水线中的上一个阶段接收激活。
大致地
buffers['inputs'][buffer_id] = recv()
注意
通信是阻塞的,并且必须与上一个流水线阶段上的
SendActivation
配对,以避免死锁。