启动 LLM 大规模预训练

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

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

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

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

本示例将使用 DeepSpeedJob 以简便迅速的方式启动 Megatron-LM GPT-3 系列（125M、1.3B、13B 和 175B）模型的预训练任务。

准备

在你的项目中创建一个名为 megatron、大小 250 GiB 以上的 PVC，然后创建一个同样名为 megatron 的 Notebook 挂载该 PVC（镜像类型和模板不限）。

进入 Notebook 或远程连接到 Notebook，启动一个终端，执行以下命令以克隆必要的仓库：

cd ~
git clone https://github.com/microsoft/Megatron-DeepSpeed.git
git clone https://github.com/NVIDIA/Megatron-LM.git
git clone https://github.com/t9k/examples.git

选用 enwiki（英文维基百科）作为数据集，使用脚本下载数据集并进行预处理：

pip install wikiextractor
cd examples/deepspeed/megatron-gpt/dataset
python download_wiki.py en
cd ../dataset
./preprocess_wiki.sh

然后使用这一数据集重新训练类 GPT-2 的 tokenzier：

cd ../tokenizer
python train_tokenizer.py ../dataset/wiki-en/all wiki-en-tokenizer

启动训练

使用如下 YAML 配置文件创建 DeepSpeedJob 以启动 125M 模型的训练：

# 数据并行训练 125M 参数的 GPT 模型
kubectl create -f \
  examples/deepspeed/megatron-gpt/training/gpt-125m-4xdp.yaml

通过以下命令查看训练过程中打印的日志：

kubectl logs gpt-125m-worker-0 -f

训练过程中产生的 TensorBoard 日志文件保存在 output/gpt-1-3b/tensorboard 路径下，同样可以在 Notebook 中创建一个 TensorBoard 实例以查看：

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

使用其他训练配置

examples/deepspeed/megatron-gpt/training 路径下有多个 YAML 配置文件，分别对应不同模型参数量的不同并行训练策略。例如 gpt-13b-4xdp-4xpp-4xtp.yaml 对应 13B 模型、4 度数据并行、4 度流水线并行以及 4 度张量并行。用户可以选择其中任意一个 YAML 配置文件创建 DeepSpeedJob，只要确保有足够的计算资源可用。

下表总结了训练 125M 和 1.3B 模型的主要参数上的差异：

*预计时间为参考值，与具体的硬件环境有关，这里以 A100 SXM4 80GB GPU 节点 + IB 网络连接为例。