进行 LLM 大规模预训练

LLM(大型语言模型)是当前 AI 领域备受瞩目的研究和应用领域。预训练是构建 LLM 的初始步骤,通常需要在 GPU 集群上进行大规模(10^3~10^4 卡)的并行训练,从而使模型获得对语言结构、语义直至语言所描述的世界和各种规律的全面理解。

在集群上运行大规模模型训练面临着一些关键挑战:

  • 有效地分配和管理计算资源以满足模型训练的计算需求是一项复杂任务,需要避免资源浪费和不必要的等待时间。
  • 数据的并行加载和分发以及模型参数的同步和更新,需要高效的通信和协调机制,以确保训练的一致性和速度。
  • 故障容错和中断恢复必不可少,因为在大规模集群上运行时硬件故障和网络问题可能导致训练中断。

DeepSpeed 是目前最受欢迎的大规模分布式训练框架,而平台提供了 DeepSpeedJob,这是专为使用 DeepSpeed 框架的分布式训练而设计的 Job 类型。

本教程演示使用 Megatron-DeepSpeed 和 DeepSpeedJob 以简单快速的方式进行 Megatron-LM GPT-3 系列(125M、1.3B、13B 和 175B)模型的预训练。本教程对应示例 examples/deepspeed/megatron-gpt

本教程的 Apps 架构如下图所示:

app-arch

运行示例

创建一个名为 megatron、大小 250 GiB 的 PVC,然后安装一个同样名为 megatron 的 JupyterLab App 挂载该 PVC,镜像选用 PyTorch 2 类型,计算资源申请 16 个 CPU(核心)、32 GiB 内存。

进入 JupyterLab,启动一个终端,前往本教程对应的示例,参照其 README 文档进行操作。

查看训练信息

训练开始后,安装一个 Job Manager App(如有 Job Manager 则直接复用),进入 Job Manager,可以看到名为 gpt-125m 的 DeepSpeedJob 正在运行,点击其名称进入详情页面:

running

可以看到刚才创建的 DeepSpeedJob 的基本信息、状况信息和事件信息:

details

点击上方的副本标签页,查看 DeepSpeedJob 的 Pod 信息;点击副本右侧的  > 日志以查看训练脚本执行过程中的日志输出:

replicas
view-log

点击上方的资源监测标签页,查看 DeepSpeedJob 运行过程中使用集群计算资源、网络资源和存储资源的情况:

resources

一段时间之后,DeepSpeedJob 的状态变为 Succeeded,表示训练成功完成。

done

训练完成之后,模型文件将保存到 output/gpt-125m/model 路径下,后续用于文本生成或进一步的微调。

若 DeepSpeedJob 在运行过程中出错,其状态会变为 Error,并在事件信息和 Pod 信息部分显示错误信息,此时需要根据给出的错误信息进行问题排查。

查看训练指标

训练产生的 tfevents 文件保存在 output/gpt-125m/tensorboard 路径下,可以在 JupyterLab 中创建一个 TensorBoard 实例或安装一个 TensorBoard App以查看:

tensorboard

使用其他训练配置

examples/deepspeed/megatron-gpt/training 路径下有多个 YAML 配置文件,分别对应不同模型参数量的不同并行训练策略,细节如下表所示:

参数量训练 token 量数据集配置文件并行策略GPU 使用(参考)预计时间*
125M2.5Benwikigpt-125m.yaml-1x A100 40GB~8h
gpt-125m-2xdp.yaml2x数据并行2x A100 40GB~4h
gpt-125m-4xdp.yaml4x数据并行4x A100 40GB~2h
1.3B26Benwikigpt-1-3b-4xdp.yaml4x数据并行4x A100 40GB~8d
gpt-1-3b-8xdp.yaml8x数据并行8x A100 40GB~4d
gpt-1-3b-4xdp-4xpp.yaml4x数据并行 + 4x流水线并行2x 8x A100 40GB
13B260Benwikigpt-13b-4xdp-4xpp.yaml4x数据并行 + 4x流水线并行2x 8x A100 80GB
gpt-13b-4xdp-8xpp.yaml4x数据并行 + 8x流水线并行4x 8x A100 80GB
gpt-13b-4xdp-4xpp-4xtp.yaml4x数据并行 + 4x流水线并行 + 4x张量并行8x 8x A100 80GB
175B3.5T(WIP)

*预计时间为参考值,与具体的硬件环境有关,这里以 A100 PCIE 40GB GPU 节点为例

参考