大模型训练难于上青天？效率超群、易用的“李白”模型库来了

来源｜机器之心

大模型多了去了，告诉我怎么加速？自 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

那么，它具体是怎么做到的？以下还将从训练性能、易用性等方面对上述分布式训练工具做对比，为你在下一次做分布式训练时的工具选择提供参考指南。

1

一键自动分布式训练，性能超越 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&