大型语言模型（或称LLMs）已成为日常对话的话题。从“Facebook花费4.5年时间”到“ChatGPT仅用2个月时间”，LLMs的广泛应用可见一斑。生成式预训练转换器（GPT）使用因果自回归更新进行预测。这些模型架构在语音识别、文本生成和问答等多种任务上表现出惊人的性能。近期的几个模型，如NeoX、Falcon、Llama等都使用了GPT架构作为骨架。LLMs的训练需要大量的计算时间和数以百万计的资金。在本文中，我们将总结在NeoX上使用AWS Trainium进行训练的过程，AWS Trainium是一种专为深度学习训练进行优化的机器学习（ML）加速器。我们将概述如何在不损失任何模型质量的前提下，以成本效益高（每320万个标记/$）的方式使用AWS Trainium进行这类模型的训练。
解决方案概述
GPT NeoX和Pythia模型
GPT NeoX和Pythia是Eleuther-AI开源的因果语言模型，NeoX模型大约有200亿个参数，Pythia模型有69亿个参数。两者都是遵循类似Chat GPT3的解码器模型的设计。然而，它们还有一些补充，这些补充也被广泛应用于最近的模型，如Llama。特别地，它们具有部分在头维度上旋转的旋转位置嵌入（ROPE）。原始模型（NeoX和Pythia 6.9B）是使用可公开获取的Pile数据集进行训练，并使用Megatron和Deepspeed后端进行去重。
我们使用基于AWS Trainium的Trn1实例，使用Neuron NeMo库来演示这些模型的预训练和微调。为了建立概念验证和快速复现，我们将使用较小规模的维基百科数据集子集，使用GPT2字节对编码（BPE）分词器进行标记化。
步骤
按以下方式下载预标记的维基百科数据集：
export DATA_DIR=~/examples_datasets/gpt2mkdir -p ${DATA_DIR} && cd ${DATA_DIR}wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.jsonwget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txtaws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-requestaws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-requestaws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-requestNeoX 20B和Pythia 6.9B都使用部分旋转的ROPE，例如，旋转25%的头维度，保持其余部分不旋转。为了在AWS Trainium加速器上高效实现部分旋转，我们不是将旋转和非旋转维度进行串联，而是对非旋转维度附加零频率，然后旋转完整的头维度集合。这个简单的技巧帮助我们提高了在AWS Trainium上的吞吐量（每秒处理的序列数）。
训练步骤
为了运行训练，我们使用由SLURM管理的多节点Amazon Elastic Compute Cloud (Amazon EC2) Trn1集群，每个节点包含一个trn1.32xl实例。每个trn1.32xl有16个加速器，每个加速器有两个工作人员。在下载最新的Neuron NeMo软件包之后，使用提供的neox和pythia预训练和微调脚本，并执行以下操作进行四节点训练。
-
编译：使用三个训练迭代预编译模型以生成和保存图表：
sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh -
运行：加载从第一步缓存的图表执行训练
sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh -
监控结果
tensorboard --logdir=nemo_experiments/megatron_neox
使用相同的步骤运行Pythia 6.9B模型，将neox_20B_slurm.sh替换为pythia_6.9B_slurm.sh。
预训练和微调实验
我们使用Neuron NeMo库在AWS Trainium上进行GPT-NeoX和Pythia模型的预训练，进行了10k次迭代，并展示了对这些模型进行1k步微调。对于预训练，我们使用了NeMo内部的GPT2 BPE tokenizer，并遵循与原始模型相同的配置文件。在AWS Trainium上进行微调需要更改一些参数（例如词汇大小的分割因子），这些参数在微调脚本中提供，以适应Megatron与NeMo的差异和GPU与AWS Trainium的更改。表-1展示了在不同节点数量下的多节点分布式训练吞吐量。
| 模型 | 张量并行 | 管道并行 | 实例数量 | 成本（每小时$） | 序列长度 | 全局批量大小 | 吞吐量（seq/sec） | 成本-吞吐量比（Tokens/$） |
| Pythia 6.9B | 8 | 1 | 1 | 7.59 | 2048 | 256 | 10.4 | 10,102,387 |
| 8 | 1 | 4 | 30.36 | 2048 | 256 | 35.8 | 8,693,881 | |
| NeoX 20B | 8 | 4 | 4 | 30.36 | 2048 | 16384 | 13.60 | 3,302,704 |
| 8 | 4 | 8 | 60.72 | 2048 | 16384 | 26.80 | 3,254,134 | |
| 8 | 4 | 16 | 121.44 | 2048 | 16384 | 54.30 | 3,296,632 | |
| 8 | 4 | 32 | 242.88 | 2048 | 16384 | 107.50 | 3,263,241 | |
| 8 | 4 | 64 | 485.76 | 2048 | 16384 | 212.00 | 3,217,708 |
表格1。 将GPT NeoX和Pythia模型的平均吞吐量与训练步数增加的节点数量进行比较。 trn1.32xl的定价基于3年预留的每小时有效费率。
接下来，我们还评估了在AWS Trainium上进行模型训练的损失轨迹，并将其与在P4d（Nvidia A100 GPU核心）集群上进行的相应运行进行比较。除了训练损失外，我们还比较了一些有用的指标，如梯度范数，它是每次训练迭代时计算的模型梯度的2范数，用于监控训练进度。训练结果显示在图1、2和NeoX 20B的微调结果显示在图3中。
图1。 训练损失在所有工作节点上的平均值（左）和梯度范数（右）在每次训练步骤中。NeoX 20B在4个节点上使用小型维基数据集在GPU和Trainium上进行训练，使用相同的训练超参数（全局批量大小=256）。GPU使用BF16和默认的混合精度，而AWS Trainium使用完整的BF16和随机舍入。GPU和AWS Trainium的损失和梯度范数轨迹匹配。
图2。 训练损失在所有工作节点上的平均值（左）和梯度范数（右）在每次训练步骤中。与图1中的GPT NeoX类似，Pythia 6.9B在4个节点上使用小型维基数据集在GPU和Trainium上进行训练，使用相同的训练超参数（全局批量大小=256）。GPU和Trainium的损失和梯度范数轨迹匹配。
图3。 在GPU和AWS Trainium上对GPT NeoX 20B模型进行微调，训练损失在所有工作节点上的平均值（左）和梯度范数（右）。使用小型维基数据集进行微调演示。GPU和AWS Trainium的损失和梯度范数轨迹匹配。
在本文中，我们展示了在AWS深度学习硬件上进行成本高效的LLM训练。我们使用Neuron NeMo库在AWS Trn1上训练了GPT NeoX 20B和Pythia 6.9B模型。使用AWS Trainium进行的200亿模型的成本归一化吞吐量约为每美元花费3.2M个标记。除了在AWS Trainium上进行成本高效的训练外，我们还获得了类似的模型准确性，这从训练步骤的损失和梯度范数轨迹可以看出。我们还对AWS Trainium上的NeoX 20B模型的可用检查点进行了微调。有关使用NeMo Megatron在AWS Trainium上进行分布式训练的更多信息，请参阅AWS Trainium Neuron参考。您可以在这里找到开始对Llama模型进行微调的良好资源，Llama2微调。要开始使用托管的AWS Trainium在Amazon SageMaker上部署ML模型，请参阅使用AWS Trainium和Amazon SageMaker训练您的ML模型。