节俭与准确性：使用 AWS Trainium 对 GPT NeoX 和 Pythia 模型进行经济高效的训练 | 亚马逊网络服务

大型语言模型（或 LLM）已成为日常对话的话题。 它们的快速采用从达到 100 亿用户所需的时间可见一斑，该时间已从“Facebook 的 4.5 年”缩短到“ChatGPT 的 2 个月”的历史最低水平。 生成式预训练 Transformer (GPT) 使用因果自回归更新来进行预测。 这些模型架构证明了语音识别、文本生成和问答等各种任务具有惊人的性能。 最近的几个型号，例如 NeoX, 鹘, 骆驼 使用GPT架构作为骨干。 培训法学硕士需要大量的计算时间，这需要花费数百万美元。 在这篇文章中，我们将总结 GPT 的训练过程 NeoX on AWS 培训，一款专为深度学习训练而优化的机器学习 (ML) 加速器。 我们将概述如何使用 AWS Trainium 经济高效地（3.2 万个代币/美元）训练此类模型，而不损失任何模型质量。

解决方案概述

GPT NeoX 和 Pythia 模型

GPT NeoX 和 皮提亚 是 Eleuther-AI 的开源因果语言模型，NeoX 中有大约 20 亿个参数，Pythia 中有 6.9 亿个参数。 两者都是遵循与 Chat GPT3 类似的架构设计的解码器模型。 然而，它们还添加了一些附加功能，这些附加功能也广泛应用于 Llama 等最新型号中。 特别是，它们具有旋转位置嵌入 (ROPE)，可在头部尺寸上部分旋转。 原始模型（NeoX 和 Pythia 6.9B）是在公开可用的模型上进行训练的 桩数据集 具有重复数据删除功能并使用 Megatron 和 Deepspeed 后端。

我们使用以下方法在基于 AWS Trainium 的 Trn1 实例上演示了这些模型的预训练和微调 神经元尼莫 图书馆。 为了建立概念验证和快速复制，我们将使用一个较小的维基百科数据集子集，并使用 GPT2 字节对编码 (BPE) 标记器进行标记。

演练

下载预标记化的维基百科数据集，如下所示：

export DATA_DIR=~/examples_datasets/gpt2 mkdir -p ${DATA_DIR} && cd ${DATA_DIR} wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-request

NeoX 20B 和 Pythia 6.9B 都使用带有部分旋转的 ROPE，例如，旋转头部尺寸的 25% 并保持其余部分不旋转。 为了在 AWS Trainium 加速器上高效实现部分旋转，我们不是串联旋转维度和非旋转维度，而是为非旋转维度附加零频率，然后旋转完整的头部维度集。 这个简单的技巧帮助我们提高了 AWS Trainium 上的吞吐量（每秒处理的序列）。

训练步骤

为了运行训练，我们使用 SLURM 托管的多节点 Amazon Elastic Compute Cloud (Amazon EC2) Trn1 集群，每个节点包含一个 trn1.32xl 实例。 每个 trn1.32xl 有 16 个加速器，每个加速器有 XNUMX 名工作人员。 下载最新的后 神经元尼莫 包，使用提供的 新氧 和 PYTHIA 使用优化的超参数预训练和微调脚本，并执行以下四节点训练。

1. 编译：通过三个训练迭代预编译模型以生成并保存图形：

   sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh

2. 运行：通过加载第一步中的缓存图表来执行训练

   sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh

3. 监控结果

   tensorboard --logdir=nemo_experiments/megatron_neox

运行 Pythia 6.9B 模型需要遵循相同的步骤，并替换 neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

预训练和微调实验

我们使用 AWS Trainium 演示了 GPT-NeoX 和 Pythia 模型的预训练 神经元尼莫 10k 次迭代的库，并且还展示了这些模型 1k 步的微调。 对于预训练，我们使用 NeMo 内部的 GPT2 BPE 分词器并遵循相同的方法 配置 与原始模型中使用的一样。 在 AWS Trainium 上进行微调需要更改几个参数（例如 词汇大小划分因子），在微调脚本中提供，以适应 Megatron 与 NeMo 的差异以及 GPU 与 AWS Trainium 的变化。 不同节点数的多节点分布式训练吞吐量如表1所示。

表1。 比较 GPT NeoX 和 Pythia 模型在节点数量变化的情况下训练最多 500 个步骤的平均吞吐量。 这 trn1.32xl 的定价 基于 3 年预留有效每小时费率。

接下来，我们还评估了 AWS Trainium 上模型训练的损失轨迹，并将其与 P4d（Nvidia A100 GPU 核心）集群上的相应运行进行比较。 除了训练损失之外，我们还比较有用的指标，例如梯度范数，它是每次训练迭代时计算的模型梯度的 2-范数，用于监控训练进度。 NeoX 1B 的训练结果如图 2、图 20 所示，微调结果如图 3 所示。

图1。 所有工人的训练损失平均值（左）和训练每一步的梯度范数（右）。 NeoX 20B 在 GPU 和 Trainium 上使用小型 wiki 数据集在 4 个节点上进行训练，具有相同的训练超参数（全局批量大小 = 256）。 GPU 使用 BF16 和默认混合精度，而 AWS Trainium 使用带有随机舍入的完整 BF16。 GPU 和 AWS Trainium 的损失和梯度范数轨迹匹配。

图2。 所有工人的训练损失平均值（左）和训练每一步的梯度范数（右）。 与图 1 中的 GPT NeoX 类似，Pythia 6.9B 在 GPU 和 Trainium 上具有小型 wiki 数据集的 4 个节点上进行训练，具有相同的训练超参数（全局批量大小 = 256）。 GPU 和 Trainium 的损失和梯度范数轨迹匹配。

图3。 在 GPU 和 AWS Trainium 上微调 GPT NeoX 20B 模型，使用所有工作人员的平均训练损失（左）和梯度范数（右）。 一个小型 wiki 数据集用于微调演示。 GPU 和 AWS Trainium 的损失和梯度范数轨迹匹配。

在这篇文章中，我们展示了在 AWS 深度学习硬件上对法学硕士进行经济高效的培训。 我们使用 Neuron NeMo 库在 AWS Trn20 上训练了 GPT NeoX 6.9B 和 Pythia 1B 模型。 使用 AWS Trainium 的 20 亿个模型的标准化吞吐量成本约为 3.2 万个代币/花费的美元。 除了在 AWS Trainium 上进行经济高效的训练之外，我们还获得了类似的模型准确性，这从训练步骤损失和梯度范数轨迹中可以明显看出。 我们还微调了 AWS Trainium 上 NeoX 20B 模型的可用检查点。 有关 AWS Trainium 上的 NeMo Megatron 分布式训练的更多信息，请参阅 NeMo Megatron 的 AWS Neuron 参考。 可以在这里找到开始微调 Llama 模型的好资源， Llama2 微调。 开始使用托管 AWS Trainium 亚马逊SageMaker，请参阅 使用 AWS Trainium 和 Amazon SageMaker 训练您的 ML 模型.

作者简介

高拉夫·古普塔（Gaurav Gupta） 目前是 Amazon Web Services (AWS) AI 实验室的应用科学家。 Gupta 博士在南加州大学维特比分校获得了博士学位。 他的研究兴趣涵盖序列数据建模、学习偏微分方程、机器学习信息论、分数动力学模型和复杂网络等领域。 他目前正在研究法学硕士训练行为、偏微分方程视觉模型、信息论多模态模型的应用和数学问题。 Gupta 博士在 Neurips、ICLR、ICML、Nature、IEEE Con​​trol Society、ACM cyber-physical Society 等顶级期刊/会议上发表过论文。

本·斯奈德 是一位使用 AWS 深度学习的应用科学家。 他的研究兴趣包括基础模型、强化学习和异步优化。 工作之余，他喜欢骑自行车和野外露营。

阿米斯 (R) 马米达拉 是 AWS Annapurna Labs 的高级机器学习应用工程师。 Mamidala 博士在俄亥俄州立大学获得了高性能计算和通信博士学位。 在 IBM 研究部门任职期间，Mamidala 博士为 BlueGene 类计算机做出了贡献，该类计算机经常在最强大、最节能的超级计算机 Top500 排名中名列前茅。 该项目荣获2009年度国家科技创新奖章。 在金融对冲基金担任短暂的人工智能工程师后，Mamidala 博士加入了安纳普尔纳实验室，专注于大语言模型训练。

俊焕 是 AWS AI 实验室的首席科学家。 欢博士致力于人工智能和数据科学。 他在领先的会议和期刊上发表了 180 多篇经过同行评审的论文。 他于 2009 年获得 NSF 教师早期职业发展奖。在加入 AWS 之前，他曾在百度研究部门工作，担任杰出科学家和百度大数据实验室负责人。 他创立了人工智能初创公司StylingAI Inc.，并于2019-2021年担任首席执行官兼首席科学家。 在加入行业之前，他是堪萨斯大学 EECS 系的 Charles E. 和 Mary Jane Spahr 教授。

斯鲁蒂·科帕卡 是 AWS 的高级产品营销经理。 她帮助客户探索、评估和采用 Amazon EC2 加速计算基础设施来满足他们的机器学习需求。

• SEO 支持的内容和 PR 分发。 今天得到放大。
• PlatoData.Network 垂直生成人工智能。 赋予自己力量。 访问这里。
• 柏拉图爱流。 Web3 智能。 知识放大。 访问这里。
• 柏拉图ESG。 碳， 清洁科技, 能源， 环境， 太阳能， 废物管理。 访问这里。
• 柏拉图健康。 生物技术和临床试验情报。 访问这里。
• Sumber: https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/