使用Torchtune在Amazon SageMaker上微调Meta Llama 31模型

作者 Kanwaljit Khurmi Trevor Harvey Less Wright Matthias Reso Roy Allela时间 2024年9月19日分类 高级 (300) Amazon SageMaker永久链接评论

关键要点

在不断发展的人工智能领域，企业寻求定制大型语言模型LLMs以满足特定需求。本篇文章将探讨如何利用Amazon SageMaker和Torchtune库微调Meta Llama 31模型，实现模型的推理、量化及评估，为用户提供更灵活的操作流程。

在当今快速发展的AI领域，企业不断寻求使用先进的大型语言模型LLMs来满足自身特定需求。虽然基础模型FMs提供了令人印象深刻的即用功能，但真正的竞争优势通常在于通过微调实现深度模型定制。然而，对于复杂任务，微调LLMs通常需要高级的AI专业知识以有效对齐和优化。意识到这一挑战，PyTorch团队开发了Torchtune，这是一个原生的PyTorch库，旨在简化LLMs的创建、微调和实验过程，使更广泛的用户和应用能够获得便利。

本文中，AWS与Meta的PyTorch团队合作，展示如何使用PyTorch的Torchtune库在Amazon SageMaker的完全托管环境中微调Meta Llama类似架构。我们通过一步一步的实现过程展示模型的微调、推理、量化及评估。我们在单个p4d24xlarge工作节点提供8个Nvidia A100 GPU上使用LoRA微调策略对Meta Llama 31 8B模型进行操作。

在深入步骤指南之前，我们首先探索了技术栈的性能，通过在不同配置和实例类型中微调Meta Llama 31 8B模型。

以下图表显示，单个p4d24xlarge实例的性能比两个g548xlarge实例每个配备8个NVIDIA A10 GPU高出70，而成本减少近47。因此，本文中的例子已经针对p4d24xlarge配置进行了优化。当然，您也可以根据需要更改SageMaker 估算器的参数，使用相同的代码在不同的实例配置上进行单节点或多节点训练。您还可以通过使用SageMaker托管热池和访问预下载的模型使用Amazon弹性文件系统进一步优化训练时间。

微调LLMs的挑战

生成性AI模型提供了许多有前景的商业用例。然而，为了保持这些LLMs在特定业务领域的事实准确性和相关性，微调是必不可少的。随着模型参数数量和现代LLMs的上下文长度不断增加，这一过程变得内存密集。为了应对这些挑战，微调策略如LoRA低秩适配和QLoRA量化低秩适配通过在变换器层中添加低秩并行结构来限制可训练参数的数量。这使得即使在内存较低的系统上如普通GPU也能训练LLMs。然而，这增加了复杂性，因为必须处理新依赖，并对新的技术调整训练配方和超参数。

今天的企业需要的是用户友好的训练配方，为这些受欢迎的微调技术提供抽象，解决常见的陷阱，以最有利的方式开展整个调优过程。

Torchtune的帮助

Torchtune是一个原生的PyTorch库，旨在使LLMs的微调过程民主化和简化。这样，它使研究人员、开发人员和组织能够轻松地将这些强大的LLMs适应他们特定的需求和约束。它为多种微调技术提供训练配方，可以通过YAML文件进行配置。这些配方实现了常见的微调方法全权重、LoRA、QLoRA以及其他常见任务，如推理和评估。它们自动应用一系列重要特性FSDP、激活检查点、梯度累积、混合精度，并且特定于给定的模型系列如Meta Llama 3/31或Mistral以及计算环境单节点与多节点。

此外，Torchtune与主要库和框架集成，例如Hugging Face数据集、EleutherAI的Eval Harness和Weights amp Biases。这有助于满足生成性AI微调生命周期的需求，从数据摄取和多节点微调到推理和评估。以下图示展示了本文描述的各个步骤的可视化。

有关Torchtune及其概念的更多信息，请参阅安装说明和PyTorch文档。

解决方案概述

本文展示了如何利用SageMaker Training通过任务特定的训练作业在独立的计算集群上运行Torchtune配方。SageMaker Training是一个全面、完全托管的ML服务，能够进行可扩展的模型训练。它提供灵活的计算资源选择、对自定义库的支持、按需付费定价模型和自我修复能力。通过管理工作负载调度、健康检查和基础设施，SageMaker帮助降低训练时间和总拥有成本。

解决方案架构整合了以下关键组件，以增强微调工作流的安全性和效率：

安全增强 训练作业在您的虚拟私有云VPC的私有子网内运行，显著提高了机器学习ML工作流的安全性。高效存储解决方案 使用Amazon EFS在ML工作流的各个阶段加速模型存储和访问。可定制环境 我们在训练作业中使用自定义容器。SageMaker对自定义容器的支持允许您将所有必要的依赖项、专用框架和库打包到一个单一工件中，提供对ML环境的完全控制。

以下图示展示了解决方案架构。用户可通过API或命令行界面CLI或使用SageMaker SDK为每个单独步骤发起流程。随后，SageMaker将启动训练作业，使用请求的计算实例数和类型来运行特定任务。图示中的每个步骤都会从Amazon简单存储服务Amazon S3存储桶访问Torchtune配方，并在工作流的不同阶段使用Amazon EFS保存和访问模型工件。

通过将每个Torchtune步骤解耦，我们在灵活性和集成之间取得平衡，允许步骤独立执行，并为通过无缝管道集成功能自动化此过程提供可能。

在这个用例中，我们使用LoRA微调Meta Llama 31 8B模型。随后，我们运行模型推理，并根据需要使用Torchtune和SageMaker Training对模型进行量化和评估。

配方、配置、数据集和提示模板都是完全可配置的，可使您能够根据需要调整Torchtune。为了演示这一点，我们在此用例中使用了自定义提示模板，并结合Hugging Face库的开源数据集Samsung/samsum。

我们使用Torchtune的多设备LoRA配方lorafinetunedistributed和SageMaker定制版本的Meta Llama 31 8B默认配置llama31/8Blora进行微调。

前提条件

在运行SageMaker Jupyter笔记本之前，您需要完成以下前提条件：

创建一个Hugging Face 访问令牌，以便访问Hugging Face上被保护的存储库metallama/MetaLlama318B。创建一个Weights amp Biases API密钥，以便访问Weights amp Biases仪表盘进行日志记录和监控。请求一个SageMaker服务配额，获取1个mlp4d24xlarge和1个mlg52xlarge实例。创建一个AWS身份与访问管理IAM角色，并附加托管策略AmazonSageMakerFullAccess、AmazonEC2FullAccess、AmazonElasticFileSystemFullAccess和AWSCloudFormationFullAccess，以便提供所需访问权限给SageMaker运行示例。这只是演示用途，您应根据特定的安全需求进行调整。创建一个Amazon SageMaker Studio域详见快速设置Amazon SageMaker，以便以先前角色访问Jupyter笔记本。参考说明来设置Docker构建权限。

登录到笔记本控制台并克隆GitHub存储库：

bash git clone https//githubcom/awssamples/sagemakerdistributedtrainingworkshopgit cd sagemakerdistributedtrainingworkshop/13torchtune

运行笔记本ipynb来使用AWS CloudFormation堆栈设置VPC和Amazon EFS。

审查Torchtune配置

以下图示展示了我们的工作流步骤。

您可以通过直接使用tune CLI查找符合您用例的Torchtune配置。针对此文，我们提供与SageMaker目录路径结构对齐的修改后的配置文件：

bashsh42 cd config/sh42 ls ltrrwrwr 1 ec2user ec2user 1151 Aug 26 1834 configl318bgenorigyamlrwrwr 1 ec2user ec2user 1172 Aug 26 1834 configl318bgentrainedyamlrwrwr 1 ec2user ec2user 644 Aug 26 1849 configl318bquantyamlrwrwr 1 ec2user ec2user 2223 Aug 28 1453 configl318blorayamlrwrwr 1 ec2user ec2user 1223 Sep 4 1428 configl318bevaltrainedyamlrwrwr 1 ec2user ec2user 1213 Sep 4 1429 configl318bevaloriginalyaml

Torchtune使用这些配置文件在执行配方时选择并配置组件例如模型和分词器。

构建容器

根据我们的示例，我们创建了一个自定义容器，以提供自定义库，如Torch Nightlies和Torchtune。请完成以下步骤：

bashsh42 cat Dockerfile

1小时试用加速器

设置REGION构建参数的默认值

ARG REGION=uswest2

SageMaker PyTorch训练镜像

FROM {ACCOUNTID}dkrecr{REGION}amazonawscom/pytorchtraining230gpupy311cu121ubuntu2004sagemaker

卸载现有的PyTorch库

RUN pip uninstall torch torchvision transformerengine y

安装最新版本的PyTorch和torchvision

RUN pip install forcereinstall torch==241 torchao==040 torchvision==0191

运行1buildcontaineripynb笔记本，直到以下命令将此文件推送到ECR存储库：

bash!smdocker build repository acceleratelatest

smdocker是一个用于通过AWS CodeBuild在SageMaker Studio中构建Docker镜像的CLI工具。我们在笔记本中作为安装库的一部分。

接下来，我们将运行2torchtunellama31ipynb笔记本以处理所有微调工作流任务。

对于每个任务，我们将审查三个工件：

Torchtune配置文件包含计算和Torchtune配方详细信息的SageMaker任务配置SageMaker任务输出

运行微调任务

在本部分中，我们将逐步讲解如何运行和监控微调任务。

运行微调作业

以下代码展示了一个简化的Torchtune配方配置，突出了微调作业的几个关键组件：

包含LoRA秩配置的模型组件用于数据分词的Meta Llama 3分词器用于读取和写入检查点的检查点功能加载数据集的组件

bashsh42 cat configl318blorayaml

模型参数

model component torchtunemodelsllama31lorallama318b loraattnmodules [qproj vproj] lorarank 8 loraalpha 16

分词器

tokenizer component torchtunemodelsllama3llama3tokenizer path /opt/ml/input/data/model/hfmodel/original/tokenizermodel

checkpointer component torchtuneutilsFullModelMetaCheckpointer checkpointfiles [ consolidated00pth ]

数据集和采样器

dataset component torchtunedatasetssamsumdataset trainoninput Truebatchsize 13

训练

epochs 1gradientaccumulationsteps 2

更多

我们使用Weights amp Biases来记录和监控训练作业，这有助于我们跟踪模型的性能：

bashmetriclogger component torchtuneutilsmetricloggingWandB