LiBai（李白）模型库覆盖了 Hugging Face、Megatron-LM、DeepSpeed、FairSeq 这些所有主流 Transformer 库的优点，让大模型训练飞入寻常百姓家。
大模型多了去了，告诉我怎么加速？自 2018 年 BERT 诞生，到 GPT-3、ViT 等拥有数以亿计的参数规模的模型不断涌现，AI 模型参数量的爆发式增长已不足为奇，让炼丹师无暇顾及甚至感到麻木。
与此同时，大模型对计算和内存资源提出了巨大的挑战。训练成本急剧上升，比如用一块非常先进的 NVIDIA A100 GPU 训练千亿参数模型的 GPT-3，需要用时 100 多年。
大模型对显存的需求增长远超过 GPU 显存增长的速度，根据 OpenAI 的报告，模型大小的增长速度是每 3.5 月翻一倍，而 GPU 显存则需要 18 个月才能翻倍。受限于 GPU 显存，单个 GPU 无法再容纳大规模模型参数。
因此，业内不得不将计算扩展到多个 GPU 设备上，分布式训练则成为广大开发者的必然选择。
但是分布式训练的门槛太高，即便拥有充足的计算资源，也可能因为搞不定分布式训练而望洋兴叹。分布式并行编程通常需要工程师了解计算机系统和架构的专业知识，再加上相关的实践经验，这进一步增加了探索前沿算法和新模型的难度。凡此种种，使得大模型成为部分科技巨头的特权。如何加速模型训练效率，让更多工程师可以使用、研究大模型成为当务之急。
问题是，市面上那么多支持分布式训练的模型库，选哪个最合适？
近期，由一流科技团队研发的以高效性起家的国产开源深度学习框架 OneFlow 上线了 LiBai（李白）模型库，这个新生代模型库覆盖了 Hugging Face、Megatron-LM、DeepSpeed、FairSeq 这些所有主流 Transformer 库的优点，分布式训练性能一如既往地力压群雄，更重要的是，全局视角编程最大程度降低了分布式使用门槛，让大模型训练飞入寻常百姓家。
LiBai 模型库地址：https://github.com/Oneflow-Inc/libai
那么，它具体是怎么做到的？以下还将从训练性能、易用性等方面对上述分布式训练工具做对比，为你在下一次做分布式训练时的工具选择提供参考指南。
一键自动分布式训练，性能超越 Megatron-LM 和 DeepSpeed
作为一个简单高效的分布式模型训练工具箱，具体而言，LiBai 有如下六大特点：
- 支持单卡代码平滑地扩展到分布式。LiBai 内置的模型与 PyTorch 保持一致风格，大大降低学习和使用成本，只需要简单配置，就可以便捷地扩展至任意规模的并行。这意味着，你可以在单卡上增加新功能，进行模型调试，跑通代码后再丝滑地迁移到分布式上进行训练。如果完全不想配置分布式训练，或是觉得手动配置的分布式训练太慢，那可以试用分布式托管特性，只需安装自动并行的包（https://libai.readthedocs.io/en/latest/tutorials/basics/Auto_Parallel.html），并在 LiBai 里配置一行 graph.auto_parallel=True，就可以专注于模型本身，在完全不用操心分布式的同时获得较快的训练速度。
- 兼容 Hugging Face。OneFlow 和 PyTorch 在 API 层次高度兼容，可以通过简单的代码修改就可以导入 Hugging Face 模型，只须 import oneflow as torch ，基于 LiBai 的数据并行、自动混合精度、Activation Checkpoint、ZeRO 等机制进行一个大规模模型的训练。如果把模型的个别层次替换为 LiBai 内置的 layers ，就可以使用 3D 并行来训练一个大模型。
- 模块化设计。在 LiBai 的实现中，不仅为模型构建提供可复用的基础计算模块，也针对数据加载、训练逻辑、指标计算等做了抽象和模块化处理，方便用户根据自己的需求重写，然后作为插件集成到 LiBai 的训练系统中进行训练。
- 开箱即用。大模型训练通常需要依赖一些技术，LiBai 提供了混合精度训练、梯度重计算、梯度累加、ZeRO 等特性，可以轻松与数据并行、模型并行、流水并行组合使用。
- 快速复现实验。OneFlow 团队参考了 Detectron2 LazyConfig（https://github.com/facebookresearch/detectron2/blob/main/docs/tutorials/lazyconfigs.md） 来构建 LiBai 的配置系统，相比于传统的 argparse 和 yacs-based 配置方式，LiBai 的配置系统更加灵活，使用 Python 语法完成整体构建，所以添加新的参数和模块非常方便，只需要 import 对应的模块即可完成新模块的添加。同时，训练配置还可以序列化成 yaml 文件进行保存，方便直接在文件中进行关键字搜索来查找配置项，如果用户想要复现之前的实验的结果，也直接传入保存的 config.yaml 作为训练配置，保留非常多脚本的文件既不利于查看有效修改，在复现实验的同时也容易弄混实验配置。
- 高效性能。通过和 Megatron-LM 进行严格的 kernel 对齐，实现了多种 kernel fusion 操作，同时得益于 OneFlow 静态图的设计，不管是单卡性能还是各种组合并行的效率，LiBai 都优于英伟达深度优化的 Megatron-LM 和微软的 DeepSpeed。
得益于 OneFlow SBP 原生支持各种并行技术，LiBai 实现了算法描述和并行系统的解耦，仅用 3 万多行代码，它就实现了英伟达 Megatron-LM 和微软 DeepSpeed 两大流行方案共计 10 万行代码才能涵盖的功能 。
数据最具说服力，以下的所有实验数据均在相同的硬件环境、相同的第三方依赖（CUDA、 cuDNN 等）、使用相同的参数和网络结构下，全面对比了 LiBai 和 Megatron-LM 在不同模型下的性能表现（所有性能结果均公开且可复现，https://libai.readthedocs.io/en/latest/tutorials/get_started/Benchmark.html）。未来，OneFlow 团队会公布更大规模集群上 LiBai 的表现。
- Megatron-LM 固定 commit：https://github.com/NVIDIA/Megatron-LM/commit/e156d2fea7fc5c98e645f7742eb86b643956d840
- LiBai commit: https://github.com/Oneflow-Inc/libai/commit/9fc504c457da4fd1e92d854c60b7271c89a55222
- OneFlow commit: https://github.com/Oneflow-Inc/oneflow/commit/55b822e4d3c88757d11077d7546981309125c73f
数据并行
注：以下每组参数的含义： DP 数据并行、MP 模型并行、PP 流水并行、2D 并行、3D 并行 fp16：打开混合精度训练 (amp) nl: num layers (当 Pipeline parallel size = 8 时，为了让每个 stage 有相对数量的 layer 进行计算，我们将 num layers 从 24 调整为 48) ac: enable activation checkpointing mb: micro-batch size per gpu gb: global batch size total dxmxp，其中: d = 数据并行度(data-parallel-size) m = 模型并行度(tensor-model-parallel-size) p = 流水并行度(pipeline-model-parallel-size) 1n1g 表示单机单卡，1n8g 表示单机 8 卡, 2n8g 表示 2 机每机 8 卡共 16 卡， 4n8g 表示 4 机共 32 卡 grad_acc_num_step = global_batch_size / (micro_batch_size * data_parallel_size) 展示的结果为 Throughout
（注：本组 num layers = 24，开启 amp，1n1g micro-batch size = 24, 其余组 micro-batch size = 16）
（注：本组 num layers = 24，开启 amp，1n1g micro-batch size = 6, 其余组 micro-batch size = 4）
模型并行
（注：本组 num layers = 24，开启 amp， 开启 activation checkpointing，
micro-batch size = 128， global batch size = 1024， grad acc step = 8）
（注：本组 num layers = 24，开启 amp）
流水并行
（注：前两组 num layers = 24，grad acc step = 8， 最后一组 num layers = 48, grad acc step = 16，均开启 amp，开启 activation checkpointing）
（注：前两组 num layers = 24，grad acc step = 8， 最后一组 num layers = 48, grad acc step = 16，均开启 amp，开启 activation checkpointing）
2-D 并行
数据 + 模型并行
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing， micro-batch size = 128，grad acc step = 8）
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing， micro-batch size = 32，grad acc step = 8）
数据 + 流水并行
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing， micro-batch size = 128，grad acc step = 8）
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing， micro-batch size = 32，grad acc step = 8）
3-D 并行
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing，grad acc step = 8）
（注：本组均为 num layers = 24，均开启 amp，开启 activation checkpointing，grad acc step = 8）
从上述性能对比数据可以看到，在严格对齐的实验环境下，在 Bert、GPT-2 模型上，LiBai 的训练速度全方位超过 Megatron-LM。
人有我优，人无我有：LiBai vs 其他训练方案
如前所述，为了解决大模型训练的难题，业内已经有 Hugging Face、DeepSpeed 和 Megatron-LM、FairSeq 等流行方案，OneFlow 是否还有必要再研发一个 LiBai 模型库？
接下来，让我们具体比较下上述模型库的优势和不足，也许你就会做出判断。
HuggingFace：提供了全面的 SOTA Transformer 模型 ，便于使用预训练好的模型进行微调，提供强大的社区和生态，方便开发者使用预训练好的模型。但它只支持数据并行的方式，如果模型超过单个 GPU 显存容量就不适用了，而且想用它做从零开始训练速度也比较受限。
FairSeq：主要针对序列模型，在 NLP 和 CV 大一统的趋势下，缺少对 CV 模型的支持。
Megatron-LM：基于 PyTorch 实现了数据并行、模型并行和流水并行，性能高，可以用来真正地训练超大规模模型。
不过它做了大量的定制化，对于不熟悉分布式训练的算法工程师而言学习和使用门槛太高，基本上只能被分布式专家所复用。另外，Megatron-LM 提供的模型也远远少于 Hugging Face，使得想使用 PyTorch 复现大模型的工程师都必须要等待其他分布式高手基于 Megatron-LM 实现了自己想用的模型才行。
DeepSpeed：基于 PyTorch 的模型显存优化相关的深度定制库，提供分布式训练、混合精度训练、ZeRO 等技术，可以有效节约内存的开销，使得在数据并行下也可以有效地训练大模型。但是，DeepSpeed 还不支持模型并行，当模型某些层的参数大到超过了单个 GPU 的显存，或是使用 DeepSpeed 的切分方式导致的通信效率不是最优时，最好还是使用模型并行（Tensor 并行、流水并行)，此时，只能结合 Megetron-LM 加上侵入原代码的改造来满足需求。
除了作为 PyTorch 生态中实现大模型训练之鼻祖的 Megatron-LM 和 DeepSpeed ，来自国内外的多家知名机构也研发和推出了如 FairSeq 等若干大模型训练库，需要特别指出的是，这些库的分布式核心功能无一例外都是基于 Megatron-LM、DeepSpeed 完成的。
LiBai 模型库的与众不同之处在于，它不是以上任一分布式训练工具的简单升级或包装，而是基于 OneFlow 的分布式和图编译器能力构建的大规模预训练模型开发套件。唯有如此，LiBai 不仅在性能上无出其右，在分布式易用性上更是不遑多让：
- 兼容性。可以有效和目前基于 PyTorch 实现的 SOTA 模型兼容，方便用户快速迁移模型。
- 高效性。无论是单卡还是多卡，用户使用 LiBai 都可以提高训练效率。
- 易用性。LiBai 具有优秀的扩展性，可以很方便地根据需求修改模型，增加新功能，更快地完成原型功能的开发。以几乎无感知、零学习成本的方式帮助用户大幅降低分布式深度学习训练的门槛，用户在使用 LiBai 开发新模型和新功能时，只要会单张 GPU 编程就能自动扩展到大规模 GPU 的集群，无须为分布式训练重写代码，从而提高开发的效率。
看完上述对比，相信 LiBai 也会成为 AI 工程师做分布式训练的绝佳选择，你觉得呢？
