节约与精准的结合：使用 AWS Trainium 对 GPT NeoX 和 Pythia 模型进行成
成本效益与准确性：利用 AWS Trainium 高效训练 GPT NeoX 和 Pythia 模型
作者：Gaurav Gupta Amith Mamidala Ben Snyder Luke Huan Shruti Koparkar发布日期：2023年12月12日分类：人工智能、生成式 AI、中级(200)文章链接评论区
ip加速器安卓关键要点
大型语言模型LLMs如 GPT NeoX 和 Pythia 迅速流行，训练成本高昂。使用 AWS Trainium，可以以成本低至每千个令牌 32 美元的价格训练这些模型。文章总结了在 AWS Trainium 上高效训练 GPT NeoX 和 Pythia 的步骤，保留了模型的高质量。使用 AWS 提供的 Neuron NeMo 库，实现了高效的模型训练和微调。大型语言模型LLMs已成为日常讨论的话题。它们的迅速普及可通过达到一亿用户所需的时间来看，Facebook 曾花费了45年，而 ChatGPT 则仅仅用了两个月。生成式预训练变换器GPT通过因果自回归更新来进行预测，能够在语音识别、文本生成和问题回答等多项任务上展现出卓越的表现。最近的多个模型，如 NeoX、Falcon 和 Llama，均以 GPT 架构为基础。LLMs 的训练需要大量的计算时间，成本高达数百万美元。在本文中，我们将总结在 AWS Trainium一种针对深度学习训练优化的机器学习加速器上训练 GPT NeoX 的过程，并展现如何以低成本32M tokens/进行训练而不损失模型质量。
解决方案概述
GPT NeoX 和 Pythia 模型
GPT NeoX 和 Pythia 是 Eleuther AI 发布的开源因果语言模型，参数分别约为 200 亿和 69 亿。两者均为解码器模型，遵循与 ChatGPT3 类似的架构设计，且包含许多最近模型广泛采用的附加特性，如部分旋转的旋转位置嵌入ROPE。原始模型NeoX 和 Pythia 69B在公开提供的 Pile 数据集 上训练，已去重，并使用 Megatron 和 Deepspeed 后端。
我们将在基于 AWS Trainium 的 Trn1 实例上演示这些模型的预训练和微调，使用 Neuron NeMo 库。为了建立概念证明和快速重现，我们将使用经过 GPT2 字节对编码BPE标记化的小型 Wikipedia 数据集子集。
实施步骤
下载经过预处理的 Wikipedia 数据集，如下所示：
bashexport DATADIR=/examplesdatasets/gpt2
mkdir p {DATADIR} ampamp cd {DATADIR}
wget https//s3amazonawscom/modelshuggingfaceco/bert/gpt2vocabjsonwget https//s3amazonawscom/modelshuggingfaceco/bert/gpt2mergestxtaws s3 cp s3//neurons3/trainingdatasets/gpt/wikipedia/mygpt2textdocumentbin nosignrequestaws s3 cp s3//neurons3/trainingdatasets/gpt/wikipedia/mygpt2textdocumentidx nosignrequestaws s3 cp s3//neurons3/trainingdatasets/gpt/wikipedia/licensetxt nosignrequest
NeoX 20B 和 Pythia 69B 使用 ROPE 进行部分旋转，例如，旋转 25 的头部维度，保持其余部分不旋转。为了在 AWS Trainium 加速器上有效地实现部分旋转，我们采用了在非旋转维度上添加零频率，在旋转整个头部维度集合的简单技巧，这帮助我们提高了 AWS Trainium 上的吞吐量每秒处理的序列数。
训练步骤
为运行训练，我们使用由 SLURM 管理的多节点 Amazon Elastic Compute Cloud (Amazon EC2) Trn1 集群，每个节点包含一个 trn132xl 实例。每个 trn132xl 拥有 16 个加速器及每个加速器两个工作线程。下载最新的 Neuron NeMo 包后，使用提供的 neox 和 pythia 预训练和微调脚本，并使用经过优化的超参数执行以下操作以进行四节点训练。
编译：预编译模型，通过三次训练迭代生成和保存图形：bash sbatch nodes 4 compileslurm /neoX20Bslurmsh
运行：通过加载第一步缓存的图形执行训练：bash sbatch nodes 4 runslurm /neoX20Bslurmsh
监控结果：bash tensorboard logdir=nemoexperiments/megatronneox
运行 Pythia 69B 模型时，同样的步骤适用，只需将 neox20Bslurmsh 替换为 pythia69Bslurmsh。
预训练和微调实验
我们在 AWS Trainium 上使用 Neuron NeMo 库进行了 GPTNeoX 和 Pythia 模型的预训练，迭代达 10000 次，并进行了 1000 步的微调。预训练使用 NeMo 内的 GPT2 BPE 标记器，遵循与原始模型相同的 配置。在 AWS Trainium 上微调模型需要改变一些参数如词汇表大小除数，这些在微调脚本中提供，以适应 Megatron 与 NeMo 之间的差异以及 GPU 与 AWS Trainium 的变化。表格 1 显示了多节点分布式训练的吞吐量与节点数量的关系。
模型张量并行管道并行实例数量成本/小时序列长度全局批量大小吞吐量seq/sec成本吞吐量比tokens/Pythia 69B811759204825610410102387814303620482563588693881NeoX 20B84430362048163841360330270484860722048163842680325413484161214420481638454303296632843224288204816384107503263241846448576204816384212003217708表 1 比较不同节点数量下 GPT NeoX 和 Pythia 模型的平均吞吐量。 trn132xl 的定价 是基于三年预留的有效每小时费率。
接下来，我们还评估了 AWS Trainium 上模型训练的损失轨迹，并与在 P4dNvidia A100 GPU 核心集群的相应运行进行比较。除了训练损失外，我们还比较有用指标，例如梯度范数，即在每次训练迭代中计算的模型梯度的 2范数，以监控训练进度。训练结果显示在图 1 和 2 中，NeoX 20B 的微调在图 3 中。
图 1 训练损失在所有工作线程平均左和每一步训练的梯度范数右。NeoX 20B 在 4 个节点上，使用小型 wiki 数据集在 GPU 和 Trainium 上训练，采用相同的训练超参数全局批量大小=256。GPU 使用 BF16 和默认混合精度，而 AWS Trainium 则使用完整的 BF16 和随机舍入。损失和梯度范数的轨迹在 GPU 和 AWS Trainium 上一致。
图 2 训练损失在所有工作线程平均左和每一步训练的梯度范数右。与图 1 中的 GPT NeoX 类似，Pythia 69B 在 4 个节点上，使用小型 wiki 数据集在 GPU 和 Trainium 上训练，采用相同的训练超参数全局批量大小=256。损失和梯度范数的轨迹在 GPU 和 Trainium 上一致。
图 3 在 GPU 和 AWS Trainium 上微调 GPT NeoX 20B 模型，训练损失在所有工作线程之间的平均值左和梯度范数右。微调演示使用小型 wiki 数据集。损失和梯度范数的轨迹在 GPU 和 AWS Trainium 上一致。
在本文中，我们展示了如何在 AWS 深度学习硬件上以成本高效的方式训练 LLM。我们利用 AWS Trainium 和 Neuron NeMo 库训练了 GPT NeoX 20B 和 Pythia 69B 模型。使用 AWS Trainium，20 亿参数模型的成本归一化吞吐量约为每美元 32M 令牌。通过在 AWS Trainium 上进行成本效益高的训练，我们获得了相似的模型精度，训练步骤损失和梯度范数轨迹的相似性验证了这一点。我们还在 AWS Trainium 上微调了 NeoX 20B 模型的可用检查点。有关在 AWS Trainium 上使用 NeMo Megatron 进行分布式训练的更多信息，请参阅 AWS Neuron Reference for NeMo Megatron。有关 Llama 模型微调的资源，可以参考 Llama2 微调。要了解如何在 Amazon SageMaker 上使用管理的 AWS Trainium，可以查看 使用 AWS Trainium 和 Amazon SageMaker 训练 ML 模型。
作者简介
Gaurav Gupta 目前是亚马逊网络服务AWSAI 实验室的应用科学家。Gupta 博士在南加州大学维特比学院获得博士学位。他的研究兴趣包括顺序数据建模、学习偏微分方程、机器学习中的信息论、分数动态模型和复杂网络。当前，他致力于大型语言模型训练行为的应用与数学问题、与 PDEs 的视觉模型以及信息论多模态模型。Gupta 博士在 Neurips、ICLR、ICML、Nature、IEEE 控制学会及 ACM 网络物理社会等顶级期刊/会议上发表了多篇文章。
Ben Snyder 是 AWS 深度学习领域的一名应用科学家。他的研究兴趣包括基础模型、强化学习和异步优化。工作之余，他喜欢骑自行车和进行野外露营。
Amith (R) Mamidala 是 AWS Annapurna Labs 的高级机器学习应用工程师。Mamidala 博士在俄亥俄州立大学获得高性能计算和通信的博士学位。在 IBM 研究期间，Mamidala 博士参与开发了 BlueGene 类计算机，该计算机曾在 Top500 超级计算机中多次排名第一，并以高能效闻名。该项目于 2009 年获得国家技术与创新奖。经过短暂的在金融对冲基金担任 AI 工程师的工作后，Mamidala 博士加入 Annapurna Labs，专注于大型语言模型的训练。
Jun (Luke) Huan 是 AWS AI 实验室的首席科学家。Huan 博士从事 AI 和数据科学方面的研究，已在各大领先会议和期刊上发表超过 180 篇同行评审的论文。2009 年，他获得了国家科学基金会早期职业发展奖。在加入 AWS 之前，他曾在百度研究担任杰出科学家及百度大数据实验室负责人，20192021 年期间，他创办了 AI 初创公司 StylingAI Inc，并担任 CEO 及首席科学家。加入工业界之前，他是堪萨斯大学电
