来源｜机器之心
大模型多了去了，告诉我怎么加速？自 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-In...
那么，它具体是怎么做到的？以下还将从训练性能、易用性等方面对上述分布式训练工具做对比，为你在下一次做分布式训练时的工具选择提供参考指南。
1
一键自动分布式训练，性能超越 Megatron-LM 和 DeepSpeed
作为一个简单高效的分布式模型训练工具箱，具体而言，LiBai 有如下六大特点：
- 支持单卡代码平滑地扩展到分布式。LiBai 内置的模型与 PyTorch 保持一致风格，大大降低学习和使用成本，只需要简单配置，就可以便捷地扩展至任意规模的并行。这意味着，你可以在单卡上增加新功能，进行模型调试，跑通代码后再丝滑地迁移到分布式上进行训练。如果完全不想配置分布式训练，或是觉得手动配置的分布式训练太慢，那可以试用分布式托管特性，只需安装自动并行的包（https://libai.readthedocs.io/...），并在 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/facebookre...） 来构建 LiBai 的配置系统，相比于传统的 argparse 和 yacs-based 配置方式，LiBai 的配置系统更加灵活，使用 Python 语法完成整体构建，所以添加新的参数和模块非常方便，只需要 import 对应的模块即可完成新模块的添加。同时，训练配置还可以序列化成 yaml 文件进行保存，方便直接在文件中进行关键字搜索来查找配置项，如果用户想要复现之前的实验的结果，也直接传入保存的 config.yaml 作为训练配置，保留非常多脚本的文件既不利于查看有效修改，在复现实验的同时也容易弄混实验配置。
- 高效性能。通过和 Megatron-LM 进行严格的 kernel 对齐，实现了多种 kernel fusion 操作，同时得益于 OneFlow 静态图的设计，不管是单卡性能还是各种组合并行的效率，LiBai 都优于英伟达深度优化的 Megatron-LM 和微软的 DeepSpeed。