通过 AWS Trainium 和 AWS Batch 加速深度学习训练与简化编排 机器学习博客

使用AWS加速深度学习训练与简化编排

重点提取

在训练大型语言模型LLM时，有效的编排和计算资源管理是一个重大挑战。借助AWS的Trainium和AWS Batch，这一挑战可以显著简化，用户可以专注于模型实验、超参数调整和数据分析，而不是繁琐的基础设施管理任务。本文将探讨如何利用AWS Batch和Trainium的强大功能来优化深度学习的训练过程。

在大型语言模型LLM的训练中，如何有效地进行编排和计算资源管理是一项重大挑战。资源的自动化配置、扩展和工作流管理对于优化资源利用率和简化复杂工作流程至关重要，从而实现高效的深度学习训练过程。简化的编排使研究人员和从业人员能够更多地专注于模型实验、超参数调优和数据分析，而不是处理繁琐的基础设施管理任务。这种直接的编排还加速了创新，缩短新模型和应用的上市时间，最终提高了LLM研究和开发工作的整体效率和有效性。

本文将探讨AWS Trainium与AWS Batch的无缝集成，展示如何利用Trainium的强大机器学习ML加速能力，结合AWS Batch提供的高效编排功能。Trainium提供了巨大的可扩展性，使得从小规模模型到大型语言模型的训练任务能够轻松扩展，同时以具有成本效益的方式访问计算能力，降低LLM训练的成本和门槛。AWS Batch是一个管理服务，旨在处理AWS云上的批处理计算工作负载，管理基础设施和任务调度，同时使您能够专注于应用程序开发和结果分析。AWS Batch提供全面的功能，包括托管的批处理计算、容器化工作负载、自定义计算环境和优先级作业队列，并与其他AWS服务无缝集成。

解决方案概述

以下图示展示了解决方案架构。

训练过程如下：

让我们深入了解如何将Trainium与AWS Batch集成。以下示例展示了如何使用AWS Batch和Trainium训练Llama 27B模型。

前提条件

建议不要在本地机器上运行以下脚本。相反，请克隆GitHub仓库，并在x8664基础的实例优选C5xlarge实例类型，使用Linux/Ubuntu操作系统上运行提供的脚本。本文中的示例在Amazon Linux 2023实例上运行。

在开始AWS Batch训练之前，您应该准备好以下资源和工具：

克隆仓库

克隆GitHub仓库并导航至所需目录：

bashgit clone https//githubcom/awsneuron/awsneuronsamplesgit cd awsneuronsamples/torchneuronx/training/awsbatch/llama2

更新配置

首先，更新configtxt文件以指定以下变量的值：

textREGION # 您的AWS地区SUBNET # Trainium实例将要启动的子网SG # 您要与实例关联的安全组ECRREPO # 您的Docker容器镜像将推送到的ECR仓库INSTANCEROLE # 您的IAM实例角色的实例配置ARNDOPRECOMPILATION # 布尔值truefalse，表示是否希望进行Neuron预编译TOKENIZEDDATASETURI # 存储标记数据集的S3 URINEURONCOMPILECACHEURI # 存储Neuron编译缓存的S3 URICHECKPOINTSAVEURI # 存储检查点的S3 URI

提供这些值后，您的configtxt文件应类似于以下代码：

textREGION=useast1SUBNET=subnet012345abcd5689SG=sg012345abcd5689ECRREPO=1010101010dkrecruseast1amazonawscom/yourdockerrepoINSTANCEROLE=arnawsiam1010101010instanceprofile/yourinstanceroleDOPRECOMPILATION=trueTOKENIZEDDATASETURI=s3//your/s3/location/to/store/tokenized/dataset/NEURONCOMPILECACHEURI=s3//your/s3/location/to/store/neuroncompilecache/CHECKPOINTSAVEURI=s3//your/s3/location/to/store/checkpoints/

获取Llama分词器

要对数据集进行标记，您需要从Hugging Face获取分词器。按照说明访问Llama分词器您需要确认并接受许可证条款。获得访问权限后，您可以从Hugging Face下载分词器。成功下载后，将tokenizermodel文件放置在根目录llama2。

设置Llama训练

运行setupsh脚本，该脚本简化了启动AWS Batch训练的先决步骤。该脚本下载Llama 27B模型训练所需的Python文件。此外，它还在提供的模板和脚本中执行环境变量替换，以建立AWS Batch资源。当它运行时，确保您的目录结构符合以下设置：

text build computeenvjson jobdefjson jobqueuejson launchtemplatejson buildandpushdockerimagesh cleanupsh configtxt createresourcessh data getdatasetpy configjson tokenizermodel docker Dockerfile llama2 adamwfp32optimparamspy configjson llamabatchtrainingsh modelingllamanxdpy requirementstxt tpzero1llama27bhfpretrainpy llamabatchtrainingsh downloadandtokenizedatash images awsbatchpng READMEmd scripts buildandpushdockerimagesh cleanupsh createresourcessh downloadandtokenizedatash submitbatchjobsh setupsh submitbatchjobsh templates computeenvjson jobdefjson jobqueuejson launchtemplatejson

标记数据集

接下来，运行downloadandtokenizedatash脚本，完成Llama 27B训练的数据预处理步骤。在此过程中，我们使用从Hugging Face获取的wikicorpus数据集。完成数据集获取后，脚本执行标记并将标记的数据集上传到在configtxt配置文件中指定的预定义S3位置。以下截图展示了预处理结果。

配置资源

接下来，运行createresourcessh脚本，该脚本负责配置训练任务所需的资源。这包括创建放置组、启动模板、计算环境、作业队列和作业定义。以下截图展示了此过程。

构建并推送Docker镜像

现在可以运行buildandpushdockerimagesh脚本，该脚本构建一个针对特定训练任务定制的Docker容器镜像。这个脚本使用Neuron团队发布的深度学习容器镜像，其中内含所需的软件堆栈，并在其上添加运行Llama 27B训练的指令。训练脚本使用neuronxdistributed库进行张量并行，同时采用ZeRO1优化器。随后，生成的Docker容器镜像将上传至您在配置文件configtxt中指定的ECR仓库。

如果您希望修改任何Llama训练的超参数，请在运行buildandpushdockerimagesh之前，在/docker/llamabatchtrainingsh中进行所需的更改。

以下截图展示了构建和推送Docker镜像的过程。

提交训练任务

运行submitbatchjobsh脚本以启动AWS Batch作业并开始Llama2模型训练，如以下截图所示。

提交批处理作业后，将动态创建一个Amazon Elastic Container ServiceAmazon ECS集群。当其运行后，您可以导航到集群以监控所有在Trn132xl实例上主动运行的任务。默认情况下，此示例配置使用4个trn132xl实例。要自定义此设置，可以在submitbatchjobsh脚本中修改numNodes参数。

日志与监控

提交作业后，您可以使用Amazon CloudWatch Logs进行全面的监控、存储和查看由AWS Batch生成的所有日志。请按照以下步骤访问日志：

以下截图展示了一个示例。

检查点

在训练过程中生成的检查点将存储在configtxt文件中指定的预定义S3位置的CHECKPOINTSAVEURI中。默认情况下，训练完成时保存检查点。然而，您可以通过选择在训练循环中的每N步保存检查点来调整此行为。有关此自定义的详细说明，请参阅检查点。

清理

完成后，运行cleanupsh脚本以管理在本文中创建的资源的移除。该脚本负责删除各种组件，如启动模板、放置组、作业定义、作业队列和计算环境。AWS Batch会自动处理