为了更好地发挥 GPU 集群的计算能力，训练具有惊人功能的、强大的万亿参数模型，一个高效简洁的大模型训练工具十分必要。幻方-深度求索最近研发了一款深度学习训练工具，名为 HAI-LLM，其实现了四种并行训练方式：ZeRO 支持的数据并行、流水线并行、张量切片模型并行和序列并行。这种并行能力适应了不同工作负载的需求，可以支持数万亿规模的超大模型并扩展到数千个 GPU。基于萤火集群的特性而自研的 haiscale 高性能算子，可以帮助 HAI-LLM 极大优化大模型训练的显存效率和计算效率。
本文将为大家介绍 HAI-LLM 工具的背后技术和使用方式。
显存与计算效率
首先，我们先了解训练巨大模型的显存和计算效率的挑战。
显存效率
训练万亿参数模型所需的显存远远超出了单张 GPU 的显存大小。在使用 Adam 优化器进行混合精度训练时，存储模型状态量（参数、梯度和优化器状态量）就需要约 12TB 的显存。作为比较，使用有 40 GB 的显存的英伟达 A100 GPU 训练大模型，仅仅为了存储模型状态，就需要 400 张这样的 GPU。
模型前向传播过程中会保存中间激活值，直到反向传播计算完 Loss 后才释放，这种额外消耗的显存会随 batch 大小而增加。假设 batch 设置为 1 的情况下，训练万亿参数模型就会产生超过 400 GiB 的激活显存。即使可以用 Checkpoint 方式（计算换显存）提前处理掉激活显存，但是对于训练而言，显存需求仍然非常之高。
必须在多个 GPU 设备之间有效地划分模型状态量和激活显存，才能让这种大模型在不耗尽显存的情况下开始训练。
计算效率
经估算端到端训练一个万亿参数的模型大约需要 5000 Zflops（即 5 后面带有 24 个零；这个估算结果基于 OpenAI 的研究《law of scaling》。这意味着训练这样一个模型需要 4000 张 A100 以 50％ 的计算效率运行大约 100 天。
尽管大型超级计算 GPU 集群可以拥有超过 4000 个 GPU，但是由于 batch 大小的限制，要在这种规模上实现高计算效率仍然是一项挑战。计算效率随着计算时间对通信时间的比例的增加而增加。该比例与 batch 大小成正比。但是，训练模型的 batch 大小有一个上限：超过这个上限收敛情况会明显变差。
实际上最大的模型之一，GPT-3 的训练 batch 大小约 1500。如果使用大约 4000 张 GPU，即使我们可以自由设置 batch 大小为 4000，每张卡上的 batch 大小也只有 1，这将影响扩展性。
3D并行
设计高效的并行方式对训练大模型至关重要。HAI-LLM 结合萤火集群的特性，可以实现三种并行训练方式的灵活组合：数据并行、流水并行、张量并行，并使用序列并行方式对 Transformer 进一步并行优化，极大地提升显存利用和计算效果。
数据并行
为了减少显存的占用，HAI-LLM 采用 ZeRO 数据并行的方式，把训练优化器的状态均分到 1/N 个 GPU 上 （N 是数据并行的数量），如下图所示：
在做前向和反向传播时，ZeRO 数据并行会先做 allgather 获得完整的参数，然后在前向和反向传播结束后释放掉，只保留 1/N 的参数和梯度。
使用比较简单，HAI-LLM 提供了hai_llm.optim.ZeroOptimizer 工具，该工具可以自动切分优化器状态并执行梯度同步。代码示例如下：
from hai_llm.optim import AdamW, ZeroOptimizer
opt_cls = partial(AdamW, lr=1e-3, weight_decay=0.1)
optimizer = ZeroOptimizer(
model.parameters(),
opt_cls,
group=dp_group,
max_grad_norm=cfg.optim.max_grad_norm,
grad_scaler=scaler
)训练的时候只需要执行 optimizer.step() 就可以完成数据并行的梯度同步和参数更新。
流水并行
流水并行把模型切分成 M 个阶段，分别放在 M 个 GPU 上，能够大幅节省显存占用和提升扩展性能。训练的时候每个阶段接收上一阶段的输出作为输入，然后传给下一个阶段。如下图所示：
可以看到，这种并行方式会造成大量的 GPU 空闲（Idle）。为了减少浪费，Gpipe 和 PipeDream 的概念被提出：如果同时进行多个迭代，每个节点在同一时刻负责不同的迭代的计算，就可以避免数据依赖，不用在原地干等了：
在 HAI-LLM 中，模型的切分是在 Builder 里由用户切分的（比如 GPTBuilder），用户可以继承 hai_llm.builders.base_builder.BaseBuilder 并实现自己的 build_model 方法，在该方法中创建每个 GPU 对应的模型实例并传入 _blocks2model 中。Builder 创建出来的模型包含一个 forward_backward 方法，用户可以传进输入、损失函数、标签，该方法会执行前向传播和反向传播并返回 Loss：
loss, _ = gpt.forward_backward(x, criterion=criterion, labels=(y,))上面等同于以下过程：
output = gpt(x)
loss = criterion(output, y)
loss.backward()张量并行
张量并行把全连接层的参数和计算切分到 K 个 GPU 上，能够节省显存占用和提升扩展性能。假设 K=2，张量并行把矩阵乘法 X @ A @ B 变成 X @ A1 @ B1 + X @ A2 @ B2，其中 A 沿着纵轴切成 A1 和 A2，B 沿着横轴切成 B1 和 B2；GPU 0 会存放参数 A1 和 B1 并负责 X @ A1 @ B1 的计算，GPU 1 则存放参数 A2 和 B2 并负责 X @ A2 @ B2 的计算。
HAI-LLM 中的模型均支持张量并行的方式，如果用户需要自定义模型结构，我们在 haiscale.tensor 中提供了以下四个用来构建张量并行模型的工具：
ColumnParallelLinear：沿着纵轴切分的全连接层算子；RowParallelLinear：沿着横轴切分的全连接层算子；VocabParallelEmbedding：沿着字典维度切分的 Embedding；vocab_parallel_cross_entropy：接收沿着字典维度切分的输入，进行交叉熵损失函数计算。
用户使用 haiscale.tensor 构建模型可以不需要关心具体的张量并行通讯的实现。具体的使用方法可以参考 GPT 和 LLaMA 的模型代码。
序列并行
Megatron-LM 提出的序列并行方案旨在通过一些方式分摊张量并行中无法分摊的显存。
如上图所示，在张量并行的基础上，将 Transformer 核的 LayerNorm 以及 Dropout 层的输入按序列长度维度进行了切分，使得每个 GPU 上面只需要做一部分的 Dropout 和 LayerNorm 。这样做进一步降低了计算资源和显存的开销。
HAI-LLM 实现了 allreduce_sequence_parallel 方法，用户可以通过一行代码实现模型中 Transformer 算子的序列并行。如下代码所示：
cfg.gpt.sequence_parallel = True
for step in range(resume_step, cfg.train.steps):
...
allreduce_sequence_parallel(gpt.module)
...训练配置
HAI-LLM 通过 Python 文件进行模型图纸、超参数配置，通过 load_config 函数加载后会把字典转换成 AttrDict 对象，能够直接通过属性访问。使用方式如下例所示：
from hai_llm.utils import load_config
cfg = load_config(os.environ.get('CONFIG', 'toy/toy_config.py'))
model = cfg.model_builder.cls(cfg).build_model()配置文件中至少需要描述清楚如下几个部分：3D 并行设置 parallel，模型图纸 model_builder，超参数 model_param 和中间结果保存 checkpoint。HAI-LLM 目前提供了 GPT-3，LLaMA，PaLM 三种模型的复现和配置案例。
效果评估
HAI-LLM 提供了 /eval_benchmark/eval.py 脚本在验证集上测试模型的 Loss，其需要通过 hfai 提交到集群中运行。输入模型的配置文件、Checkpoint 位置、验证集数据等信息，指定对应的评估方法即可。如下所示：
hfai python eval_benchmark/eval.py \
--cfg <config-file> \
--ckpt_path <checkpoint-dir> \
--load_step <load_step> \
--val_ckpt_path <val-checkpoint-dir> \
--data_path <val-data-path> \
-- -n xx -p 40 -f -i ubuntu2004-cu113同时，/eval_benchmark/ 下也提供了几种常用的模型效果评估方法，方便用户对训练好的模型进行评测。评估方法包括：
- hellaswag
- winogrande
- piqa
- triviaqa
- mmlu
- race