LiBai 支持所有常见并行训练策略
分布式训练大模型是个复杂问题，涉及到数据并行（data parallel），模型并行（tensor/model parallel），流水并行（pipeline parallel）等多种并行策略，LiBai 模型库支持这三种常见的并行策略以及这些并行策略的任意组合（并行策略的基本概念：https://docs.oneflow.org/master/parallelism/01_introduction.html）。
自行实现这些并行策略让人十分头疼，比如以前为了使用自动混合精度训练，需要学习配置 Apex；为了支持数据加载流水线，需要学习配置 DALI；为了使用 ZeRO 减少显存占用，需要学习配置 DeepSpeed …… 但用 LiBai 就完全不用担心这类问题，它内置了多种并行策略且具备良好的可扩展性。
以下是 LiBai 中各类并行方法的实例。
万能并行的实现方式
借助 OneFlow 的 SBP 接口，用户可以很方便地根据自身的需求，依照 GPU 的分组排布情况对网络中的输入或者权重进行切分，以实现数据或张量并行。
在 LiBai 的 layers 模块（libai.layers）下，已内置一系列可自适应不同并行策略的网络层，包括常用的 Linear、MLP、Transformer 模块等，使用 LiBai 的 layers 搭建的神经网络, 只需调整配置文件中关于分布式配置的超参，就可以轻松实现纯数据并行、纯张量并行以及数据 & 张量混合并行的训练策略。
关于分布式配置的格式如下：
# configs/common/train.py# Distributed argumentsdist=dict( data_parallel_size=1, tensor_parallel_size=1, pipeline_parallel_size=1,)
通过 data_parallel_size 与 tensor_parallel_size 来控制输入数据与模型权重在不同 GPU 组上的切分方式，当用户使用 LiBai 的内置 layers 模块搭建好神经网络后，可以在自己的训练配置文件中修改分布式超参, 以实现不同的并行训练策略，上图所有值都取为 1 表示在单卡上运行。假设用户拥有一台 8 卡机器，下面介绍一下如何通过修改此配置文件实现数据并行、张量并行以及流水并行训练。
具体操作可参考 LiBai 分布式配置文档：https://libai.readthedocs.io/en/latest/tutorials/basics/Distributed_Configuration.html
纯数据并行 & 纯模型并行
当用户要在 8 卡上进行纯数据（或模型）并行训练, 只需要在训练配置文件中对分布式超参进行覆写即可:
- 纯数据并行
# your config.pyfrom libai.config import get_configtrain = get_config("common/train.py").train train.dist.data_parallel_size = 8
训练时，在不同的 rank 上会复制一份相同的模型，每个 rank 会分别处理一部分的输入数据, 以实现数据并行训练。
- 纯模型并行
# your config.pyfrom libai.config import get_configtrain = get_config("common/train.py").train train.dist.tensor_parallel_size = 8
在这种情况下, 模型会自动在 8 个 GPU 上进行切分, 每个 GPU 仅包含整体模型结构的一部分, 以实现模型并行训练。
数据 & 模型混合并行训练
当用户要在 8 卡上进行数据与模型混合并行训练, 只需要在训练配置文件中对分布式超参进行以下简单改动：
# your config.pyfrom libai.config import get_configtrain = get_config("common/train.py").train train.dist.data_parallel_size = 2train.dist.tensor_parallel_size = 4
这种情况下, LiBai 会自动对 GPU 进行分组, 我们以 [0, 1, 2, 3, 4, 5, 6, 7] 对 8 个 GPU 进行编号，当设置了 data_parallel_size=2 以及 tensor_parallel_size=4 后，在执行时，会自动将 8 个 GPU 进行分组，可以表示为 [[0, 1, 2, 3], [4, 5, 6, 7]], 其中[0, 1, 2, 3] 为一组，[4, 5, 6, 7]为一组，执行时，会在组之间进行数据并行训练，在组内进行模型并行训练。
流水并行的配置
流水并行的核心概念可以简单总结为：将网络分为多个阶段（stage）， 不同的 stage 被分发到不同的 GPU 上, 每个 stage 的计算结果传递给下一个 stage 进行计算，最终按接力的方式完成训练。关于流水并行的具体内容可参考：https://docs.oneflow.org/master/parallelism/01_introduction.html#_6。
朴素流水并行配置
在 LiBai 下可以通过设置 placement 参数，将网络的不同层分配到不同的 GPU 上，placement 参数的值可以通过 libai.utils.distributed 下的 get_layer_placement()接口轻松配置，LiBai 会自动根据配置文件（config）中的分布式配置，来做 stage 的切分，将不同的 placement 自动分配到不同的 stage 上，所以只需要为网络的每一层配置好 placement，再结合分布式配置，便可以轻松实现流水并行配置。
在大部分网络中，往往用一层 Linear 层作为网络的头部（head）, 产生网络的最终结果用作分类或者其他任务, 所以以 Linear 层为例, 简要介绍 LiBai 中最简单的流水并行配置方法：
from libai.layers import Linear self.head = Linear(hidden_size, num_classes)
配置网络模块（module）的 placement
在 LiBai 中可以通过两种方式将一层网络分配到对应的 placement 上：
1、通过 to_global 接口结合 get_layer_placement()来手动指定 placement， 这里通过设置 get_layer_placement(-1)来将 head 层配置到最后一组接力的 placement 上。
from libai.layers import Linearimport libai.utils.distributed as dist self.head = Linear(hidden_size, num_classes).to_global(placement=dist.get_layer_placement(-1))
2、（Recommended） 在 libai.layers 中实现的 module 自带 layer_idx 参数, 可以直接设置 layer_idx 参数来指定这一层的 placement
from libai.layers import Linear self.head = Linear(hidden_size, num_classes, layer_idx=-1)
配置输入数据的 placement
在配置好了网络中模块的 placement 后, 还需要指定输入数据的 placement, 因为只有当输入和网络在同一个 stage 的时候才可以进行计算, 最直观的方式就是为输入和网络配置相同的 placement, 可以结合 to_global 与 get_layer_placement()实现：
class MyModule(nn.Module): def __init__(self, ... *, layer_idx): ... self.layer_idx = layer_idx ... def forward(self, input_data): input_data = input_data.to_global(placement=dist.get_layer_placement(self.layer_idx)) ...
结合配置文件轻松实现朴素流水并行
在配置好网络中不同层的 placement 以及输入的 placement 后，在执行流水并行前，用户只需要调整配置文件（config）即可，需要提前知道网络中的层数，并且调整配置文件中的 pipeline_num_layers：
# set the number of pipeline stages to be 2train.dist.pipeline_parallel_size = 2 # set model layers for pipelinetrain.dist.pipeline_num_layers = hidden_layers
1F1B 是在 PipeDream（https://arxiv.org/pdf/1806.03377.pdf）中介绍的一种新的流水并行训练方式，可以更好地节省显存与利用资源。LiBai 也可以比较容易地支持这种 1F1B 的策略（https://github.com/Oneflow-Inc/libai/blob/main/docs/source/tutorials/advanced_tutorials/customize_dataloader.md）
3D 并行的实现
掌握了数据 & 模型混合并行，以及流水并行以后，配置数据 + 模型 + 流水并行也只是综合一下上述各种并行的改动即可。
# your config.pyfrom libai.config import get_configtrain = get_config("common/train.py").train train.dist.data_parallel_size = 2train.dist.tensor_parallel_size = 2train.dist.pipeline_parallel_size = 2 hidden_layers = 8 #网络的层数train.dist.pipeline_num_layers = hidden_layers
还是以 8 卡作为例子，在设置 data_parallel_size，tensor_parallel_size， pipeline_parallel_size 都为 2 以后，在执行时，模型将根据用户设置的 pinepine_num_layers 在 GPU 上自动进行划分。
以上述配置为例，模型将在 [0, 1, 2, 3] 和[4, 5, 6, 7]号 GPU 上拆分为 2 个 stage。其中，stage0 会在 [0, 2] 和[1, 3]号 GPU 上数据并行；在 [0, 1] 和[2, 3]号 GPU 上模型并行；stage1 会在 [4, 6] 和[5, 7]号 GPU 上数据并行；在 [4, 5] 和[6, 7]号 GPU 上模型并行。
自定义并行训练
根据上文的介绍，LiBai 在 libai/layers / 下提供了封装好的模块供用户调用。通过这些模块的组合，用户可以拼凑出自己的并行网络。
当 LiBai 中的模块无法满足用户需求时，用户也可以非常方便地自定义并行策略。不同于 PyTorch 下需要手工插入 scatter -> forward -> reduce 等一系列复杂的通信操作，在 LiBai 中，用户只需在初始化 tensor 时定义 sbp 和 placement，便可像写单机运行的代码一样跑起来自己的并行代码。（sbp 和 placement 的详情可参考：https://docs.oneflow.org/master/parallelism/04_2d-sbp.html）。
举例来说，在用户进行 4 卡训练时，网络的中间结果有一个 shape 为 (16, 8) 的 2D Parallel 的 tensor 在 GPU 上的划分方式为如下图, 在 LiBai 中。该 tensor 的 placement 分布为 ranks=[[0, 1],[2, 3]]，SBP 为 (S[0], S[1]) 或(S[1], S[0])。
[ | X00 gpu0 | X01 gpu1-------------------------- X10 gpu2 | X11 gpu3 | ]
其中， Xij 的 shape 都为 (8, 4) 均匀的分布在每张卡上, 如果你想对这个 tensor 加入一些随机噪声，那么在 LiBai 中可以非常方便地加上如下代码:
LiBai 中封装 dist.get_nd_sbp()是为了兼容 1D parallel 的需求，同时 dist.get_layer_placement()是为了方便配置 pipeline parallel。大多数情况下，用户可以直接参照以下代码：
# test.pyimport oneflow as flowfrom omegaconf import DictConfigfrom oneflow import nn from libai.utils import distributed as dist cfg = DictConfig( dict(data_parallel_size=2, tensor_parallel_size=2, pipeline_parallel_size=1))dist.setup_dist_util(cfg) class Noise(nn.Module): def __init__(self): super().__init__() self.noise_tensor = flow.randn( 16, 8, sbp=dist.get_nd_sbp([flow.sbp.split(0), flow.sbp.split(1)]), placement=dist.get_layer_placement(layer_idx=0) ) # 也可以换成以下的写法 # self.noise_tensor = flow.randn( # 16, 8, # sbp=(flow.sbp.split(0), flow.sbp.split(1)), # placement=flow.placement("cuda", ranks=[[0, 1],[2, 3]]) # ) def forward(self, x): return x + self.noise_tensor Noise = Noise() x = flow.zeros( 16, 8, sbp=(flow.sbp.split(0), flow.sbp.split(1)), placement=flow.placement("cuda", ranks=[[0, 1],[2, 3]]))y = Noise(x) print(f"rank: {flow.env.get_rank()}, global tensor: shape {y.shape} sbp {y.sbp} placement {y.placement}, local tensor shape: {y.to_local().shape}")
运行指令：
python3 -m oneflow.distributed.launch --nproc_per_node 4 test.py
以下显示输出，根据 shape 可以看到每个 rank 下 tensor 的分布，以及在 global 视角下该 tensor 的信息。
rank: 2, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])rank: 3, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])rank: 1, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])rank: 0, global tensor: shape oneflow.Size([16, 8]) sbp (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=1)) placement oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]]), local tensor shape: oneflow.Size([8, 4])
未来计划
LiBai 目前已支持 BERT、GPT、ViT、Swin-Transformer、T5 等常见模型，以及 MoCoV3、MAE 等最新研究，开箱即用，并且可以很方便地在下游任务上进行微调。
此外，OneFlow 也会更好地兼容 Hugging Face 的模型，接入其生态，同时再利用 OneFlow 自动并行功能，帮助用户享受只写单卡代码即自动扩展到分布式系统的一劳永逸的激爽体验。
未来，在支持更多模型训练的基础上，OneFlow 也会持续完善推理和 Serving 相关的功能，从而打通训练和部署的全流程，让 OneFlow 成为用户的一站式开发平台。
- LiBai 模型库地址：https://github.com/Oneflow-Inc/libai
- LiBai 文档地址：https://libai.readthedocs.io/en/latest
- OneFlow 项目地址：https://github.com/Oneflow-Inc/oneflow
