概述

Fold主要面向分布式训练时，需要同时开启流水线并行以及数据并行的情况。Fold的基本思想是提出一个从单独流水线维度视角来看，并非最优的流水线调度，但在多维并行角度（结合数据并行）综合更优的整体调度。Fold面向数据并行通信在多维并行中成为主要瓶颈的场景，能大幅提升多维并行（如，三维并行）时的硬件使用效率。

点击放大

上图展示了Fold模型切分的一个实例，对于一个16层的模型来说，假设将segment的数量设为2，将stage的数量设为4（即，流水线并行维度=4），那么Fold先将一个模型切分为2个segment，然后再将每个segment切分为4个stage。

Fold 3d 数据并行通信隐藏方案：

点击放大

Fold流水线的调度，对于上述切分后的模型，Fold对每一个segment都要进行全部的MicroBatch流水。当Seg1的所有MicroBatch的反向传播都执行完毕后，对于模型Seg1部分的数据并行即可开始，从而Seg1的数据并行Allreduce和Seg2的流水线计算可以在执行序上重叠。注意，此实例2个segment的简化版本，对于s个segment，s-1个segment的数据并行通信可以和计算重叠。

Fold x 流水线并行通信隐藏方案：

点击放大

上图是一种混合调度扩展（Foldx），通过将Fold流水线和原生流水线进行混合调度，将最后一个Segment的反向传播提前，使得该MicroBatch缓存的activation尽快在反向传播中被使用并释放。

功能使用

训练脚本参数配置示例，加粗部分为Fold3D功能配置参数示例：

python -m torch.distributed.launch $DISTRIBUTED_ARGS \
   pretrain_llama.py \
   --DDP-impl local \
   --num-layers 30 \
   --hidden-size 6656 \
   --ffn-hidden-size 17920 \
   --num-attention-heads 52 \
   --micro-batch-size 1 \
   --global-batch-size 32 \
   --seq-length 2048 \
   --max-position-embeddings 4096 \
   --foldx-mode aiao / (--foldx-mode fifo)

参数说明：

模型前向函数(forward_step_func)：用于进行模型的前向输出
数据迭代器(data_iterator)：用于迭代训练数据
模型(model)：指本次训练的具体模型
优化器(optimizer)：指做参数更新使用的优化器
num_microbatches：mini-batch中数据分块的数量

Fold3D与FoldX功能的定义位于“aascendspeed/schedules.py”文件中的“forward_backward_pipelining_with_foldx_aiao()”与“forward_backward_pipelining_with_foldx_fifo()”函数。

开关定义参见“ascendspeed/arguments.py”文件中的“_add_distributed_args”函数。

通过配置--foldx-mode的值来使能，输入值为str型，可选值有“aiao”“fifo”，无默认值。其中aiao表示使用Fold3D功能，fifo表示使用FoldX功能。

代码示例：

def forward_backward_pipelining_with_foldx_aiao(
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,  # unused
    micro_batch_size: int,  # unused
    decoder_seq_length: int = None,  # unused
    forward_only: bool = False,
    collect_non_loss_data: bool = False,):

def forward_backward_pipelining_with_foldx_fifo(
    forward_step_func,
    data_iterator: Union[Iterator, List[Iterator]],
    model: Union[torch.nn.Module, List[torch.nn.Module]],
    num_microbatches: int,
    seq_length: int,  # unused
    micro_batch_size: int,  # unused
    decoder_seq_length: int = None,  # unused
    forward_only: bool = False,
    collect_non_loss_data: bool = False,):

aiao调度算法在流水线并行计算上主要包括两个阶段：all-in阶段和all-out阶段。在all-in阶段，机器上均作前向计算；在all-out阶段，机器做后向求梯度操作，当机器做完模型块的后向求梯度后立刻做数据并行的异步通信操作，并继续做后向求梯度操作，极大的节省了数据并行的通信开销。

fifo调度算法在流水线并行计算上主要包括三个阶段：warmup阶段-1f1b阶段-cooldown阶段。在warmup阶段，机器上均作前向计算；在1f1b阶段，机器做单次前向以及单次后向求梯度操作，能够更好的节省显存开销，Foldx在此过程中做异步通信节省了训练时间；在cooldown阶段，机器做后向求梯度操作（可开启数据并行通信隐藏）。