大模型训练难于上青天?效率超群、易用的“李白”模型库来了
来源|机器之心
大模型多了去了,告诉我怎么加速?自 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,其中:
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,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)
(注:本组均为 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,grad acc step = 8)
兼容性。可以有效和目前基于 PyTorch 实现的 SOTA 模型兼容,方便用户快速迁移模型。 高效性。无论是单卡还是多卡,用户使用 LiBai 都可以提高训练效率。 易用性。LiBai 具有优秀的扩展性,可以很方便地根据需求修改模型,增加新功能,更快地完成原型功能的开发。以几乎无感知、零学习成本的方式帮助用户大幅降低分布式深度学习训练的门槛,用户在使用 LiBai 开发新模型和新功能时,只要会单张 GPU 编程就能自动扩展到大规模 GPU 的集群,无须为分布式训练重写代码,从而提高开发的效率。
# configs/common/train.py
# Distributed arguments
dist=dict(
data_parallel_size=1,
tensor_parallel_size=1,
pipeline_parallel_size=1,
)
纯数据并行
# your config.py
from libai.config import get_config
train = get_config("common/train.py").train
train.dist.data_parallel_size = 8
纯模型并行
# your config.py
from libai.config import get_config
train = get_config("common/train.py").train
train.dist.tensor_parallel_size = 8
# your config.py
from libai.config import get_config
train = get_config("common/train.py").train
train.dist.data_parallel_size = 2
train.dist.tensor_parallel_size = 4
from libai.layers import Linear
self.head = Linear(hidden_size, num_classes)
from libai.layers import Linear
import libai.utils.distributed as dist
self.head = Linear(hidden_size, num_classes).to_global(placement=dist.get_layer_placement(-1))
from libai.layers import Linear
self.head = Linear(hidden_size, num_classes, layer_idx=-1)
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))
...
# set the number of pipeline stages to be 2
train.dist.pipeline_parallel_size = 2
# set model layers for pipeline
train.dist.pipeline_num_layers = hidden_layers
# your config.py
from libai.config import get_config
train = get_config("common/train.py").train
train.dist.data_parallel_size = 2
train.dist.tensor_parallel_size = 2
train.dist.pipeline_parallel_size = 2
hidden_layers = 8 #网络的层数
train.dist.pipeline_num_layers = hidden_layers
[ |
X00 gpu0 | X01 gpu1
--------------------------
X10 gpu2 | X11 gpu3
| ]
LiBai 中封装 dist.get_nd_sbp()是为了兼容 1D parallel 的需求,同时 dist.get_layer_placement()是为了方便配置 pipeline parallel。大多数情况下,用户可以直接参照以下代码:
# test.py
import oneflow as flow
from omegaconf import DictConfig
from 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
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 模型库地址:https://github.com/Oneflow-Inc/libai LiBai 文档地址:https://libai.readthedocs.io/en/latest OneFlow 项目地址:https://github.com/Oneflow-Inc/oneflow