利用 Metaflow 和 AWS Trainium 经济高效地开发和训练大型模型

关键要点

本篇文章探讨了如何利用 AWS Trainium 加速器和开源工具 Metaflow，经济高效地训练和微调大型机器学习模型。我们将介绍 Metaflow 的基本概念，Trainium 与 Metaflow 如何集成，以及如何设置训练基础设施。

当前，构建高质量的人工智能AI系统面临多项技术挑战，尤其是如何在数据库、计算资源与基础设施流畅运行中保持效率。

引言

在当前的 AI 领域，构建一种具有生产级别的系统例如进行多语言情感分析面临着诸多技术挑战。过去，处理自然语言处理NLP可能是主要的研发支出，但到了 2024 年，组织们更倾向于使用大型语言模型LLMs，这使得重点从建模转向支持 LLM 工作流所需的基础设施。对于 AWS 和 Outerbounds 的客户来说，目标是构建差异化的机器学习和人工智能ML/AI系统，并在时间上持续改进。这通常意味着，出于安全、控制和规模的考虑，单纯使用第三方 LLM API 并不现实。因此，拥有控制基础设施和管理 AI 系统工作流的能力显得至关重要。

回到最初的问题，我们可能会面临三项 MLOps 挑战：

训练和微调模型需要高质量的数据。实验、训练、跟踪和编排生产系统需要多样的云基础设施。大规模计算也是推动系统运行的重要因素。

在本文中，我们强调了 Outerbounds 和 AWS 之间的合作，旨在解决上述后两项挑战。首先，AWS Trainium 加速器为训练和微调大型模型提供了高性能、经济高效且易于获得的解决方案。其次，开源的 Metaflow 提供了必要的软件基础设施，使开发者能够以更友好的方式构建生产级别的 ML/AI 系统。Metaflow 提供了一种易于访问且稳健的 Python API，涵盖了从数据和计算到工作流和可观察性的整个基础设施栈。

接下来的篇幅中，我们将首先介绍 Metaflow 和 Trainium 的集成。然后将指导您如何设置所需的基础设施栈，以利用您的数据资产在 Trainium 硬件上进行 Llama2 模型的预训练或微调。

Metaflow 概述

Metaflow 最初是在 Netflix 开发的，旨在帮助数据科学家和 ML 工程师快速构建 ML/AI 系统，并将它们部署在生产级基础设施上。Netflix 在 2019 年将该框架开源，并与 AWS 服务如 AWS Batch、AWS Step Functions 等整合，使您能够在 AWS 帐户中构建类似 Netflix 规模的 ML/AI 环境。

Metaflow 的核心动机是满足所有 ML/AI 项目的典型需求，通过简单易用的 APIs 从原型到生产实施再回到原型。以下是其工作流程示意图。

Metaflow 的一致性 API 简化了团队在构建真实世界的 ML/AI 系统时的工作。它帮助科学家和工程师有效访问、移动和处理数据；跟踪和版本控制实验和模型；协调和集成与周边系统的工作流；并轻松扩展到云计算。此外，它支持团队使用命名空间和版本生产分支来部署工作流。

如今，借助于本次发布，您拥有了一个简单的计算选项来处理需要训练或微调复杂深度学习模型的工作流程：在 Trainium 上运行它们。

Metaflow 如何与 Trainium 集成

从 Metaflow 开发者的角度来看，使用 Trainium 与使用其他加速器相似。在配置好 Metaflow 部署以通过所使用的计算平台访问 Trainium 芯片后，ML 工程师和数据科学家们可以在深度学习代码的领域内自主操作。科学家们可以编写 PyTorch、Hugging Face 代码，并使用 AWS Neuron SDK 和 NeuronX 分布式 SDK，使这些框架专门针对 Trainium 设备进行优化，而 Metaflow 则与底层的 AWS 服务整合，解放用户对如何高效运行代码的担忧。

如以下图所示，您可以用几行 Python 代码声明以下内容：

启动的节点数量每个节点使用的 Trainium 设备数量节点之间的互联方式Elastic Fabric Adapter资源利用率检查频率torchrun 过程应在每个节点上运行的训练脚本

您可以在 start 步骤中初始化训练过程，指示下一个 train 步骤在两个并行实例上运行numparallel=2。train 步骤的装饰器配置了所需的训练设置：

@torchrun 设置跨两个实例的 PyTorch 分布式@batch 配置由 AWS Batch 管理的 Trainium 节点@neuronmonitor 激活监控用户界面，以便监控 Trainium 核心的利用率

Metaflow 允许您用几行代码配置所有这些功能，主要好处在于能将基于 Trainium 的训练代码嵌入到更大的生产系统中，并利用 Metaflow 提供的框架。

使用 Trainium 和 Metaflow 的好处

Trainium 和 Metaflow 的结合解决了本文前面讨论的问题。Trainium 设备和 Neuron 软件堆栈使团队能够轻松访问和有效使用所需的高性能硬件，以支持前沿 AI。

Trainium 在构建真实世界的 AI 系统时提供了几项关键优势：

优势说明成本优势Trainium 实例的生成 AI 模型训练和微调成本可比 AWS 上类似实例下降多达 50。可用性它在许多 AWS 区域均有提供，通常比基于 GPU 的实例类型更常见，同时在全球最流行的区域内都能扩展。软件及硬件成熟性硬件和软件都由 AWS 进行积极开发。

如果您一直在面临 GPU 资源的可用性和成本问题，您肯定会对此好处感到欣赏。有效使用 Trainium 可能需要一些基础设施构建和知识，而这也是本次集成的关键动机。通过 Metaflow 和在本文中提供的部署脚本，您应该能轻松开始使用 Trainium。

除了轻松访问外，使用 Trainium 和 Metaflow 还带来了一些额外的好处：

基础设施可及性

Metaflow 以其开发者友好的 API 而闻名，让 ML/AI 开发者可以专注于构建模型和应用，而无需担心基础设施。Metaflow 帮助工程师管理基础设施，确保其与现有系统和政策的无缝整合。

数据、模型和配置管理

Metaflow 提供了内置的无缝工件持久性、跟踪和版本控制，涵盖工作流的完整状态，从而确保您遵循 MLOps 的最佳实践。得益于 Metaflow 的 高吞吐量 S3 客户端，您可以极快地加载和保存数据集及模型检查点，而无需担心额外的基础设施，如共享文件系统。您可以使用工件来管理配置，这样从超参数到集群规模的所有内容都可以在单个文件中管理，并与结果一起追踪。

轻蜂加速器下载官网

可观察性

Metaflow 提供了一种方便的用户界面，您可以 自定义以观察您用例中的重要指标和数据。在 Trainium 的场景中，我们提供了一个定制的可视化选项，使您能够监控 Trainium 实例内 NeuronCores 的利用率，确保资源得到高效使用。下图展示了核心上方和内存下方利用率的可视化示例。

多节点计算

最后，Metaflow 的一个巨大好处是它可以用来 管理先进的多实例训练集群，否则将需要大量复杂的工程工作。例如，您可以使用 Metaflow 的 @torchrun 和 @batch 装饰器，训练一个在 Trainium 实例上划分的大型 PyTorch 模型。

在后台，装饰器使用 AWS Batch 多节点 设置一个训练集群，并配置指定数量的 Trainium 实例以跨实例训练 PyTorch 模型。利用本篇文章提供的启动模板，您可以享受到通过 Elastic Fabric Adapter (EFA) 网络接口实现低延迟、高吞吐量网络。

解决方案概述

作为一个实际示例，让我们设置在 Trainium 上预训练 Llama2 所需的完整基础架构。相同的配方适用于库中的微调示例。

部署和配置 Metaflow

如果您已经在使用 Metaflow 部署，可以跳过此步骤，直接部署 Trainium 计算环境。

部署

要使用 AWS CloudFormation 部署 Metaflow 栈，请完成以下步骤：

下载 CloudFormation 模板。在 CloudFormation 控制台中，选择导航窗格中的 Stacks。选择 Create new stack。在 Prepare template 中选择 Template is ready。在 Template source 中选择 Upload a template file。上传模板。选择 Next。

如果您是 Metaflow 的新用户，或者作为概念验证尝试此配方，建议您将 APIBasicAuth 参数更改为 false，并保留所有其他默认参数设置。完成栈创建过程。

在您创建 CloudFormation 栈并配置 Metaflow 使用该栈资源后，无需其他设置。如需获取有关 AWS CloudFormation 部署的 Metaflow 组件的更多信息，请查看 AWS Managed with CloudFormation。

配置

要使用您刚刚从笔记本电脑或云工作站部署的栈，请完成以下步骤：

准备一个 Python 环境并在其中安装 Metaflow：

bashpip install metaflow

在终端中运行 metaflow configure aws。

bashmetaflow configure aws

在 CloudFormation 栈部署完成后，栈详细信息页面的 Outputs 部分将包含资源名称和其值的列表，您可以在 Metaflow AWS 配置提示中使用。

部署 Trainium 计算环境

先前步骤中默认的 Metaflow 部署包含 AWS Batch 计算环境，但无法安排作业在 Amazon Elastic Compute Cloud (Amazon EC2) 实例上运行 Trainium 设备。要部署用于 Trainium 加速器的 AWS Batch 计算环境，可以使用以下 CloudFormation 模板。完整以下步骤：

下载 CloudFormation 模板。在 CloudFormation 控制台中，选择导航窗格中的 Stacks。选择 Create new stack。在 Prepare template 中选择 Template is ready。在 Template source 中选择 Upload a template file。上传模板。选择 Next。完成栈创建过程。

请记住您创建的 AWS Batch 作业队列的名称，以便在后续步骤中使用。

准备运行 Metaflow 任务的基础 Docker 映像

当使用 AWS Batch 作为计算后端时，Metaflow 任务会在 Docker 容器内运行。为了运行 Trainium 作业，开发者需要构建自定义镜像并在 Metaflow 开发者使用的 @batch 装饰器中指定该镜像作为任务资源：

python@batch(trainium=16 efa=8 image=YOURIMAGEINECR)@stepdef trainllama2(self) # neuron 分布式训练代码

要创建镜像，请完成以下步骤：

创建 Amazon Elastic Container Registry (Amazon ECR) 注册表以存储您的镜像。创建并登录到具有足够内存的 EC2 实例。本文中，我们使用的是 Ubuntu x86 OS 的 C54xlarge 实例。安装 Docker。将以下 Dockerfile 拷贝到您的实例中。使用上游基础镜像提供者进行身份验证：

bashaws ecr getloginpassword region REGION docker login username AWS passwordstdin 763104351884dkrecrREGIONamazonawscom

构建镜像：

bashdocker build t YOURIMAGENAMEYOURIMAGETAG

在 Amazon ECR 控制台中，导航到您创建的 ECR 注册表，您将找到从 EC2 实例进行身份验证并推送镜像所需的命令。

在工作站上克隆存储库

现在您准备验证基础设施是否正确工作，之后可以运行复杂的分布式训练代码，如 Llama2 训练。要开始，请将示例 存储库 克隆到配置了 Metaflow 和 AWS 的工作站上：

bashgit clone https//githubcom/outerbounds/metaflowtrainium

使用 allreduce 示例验证基础设施

要验证您的基础设施配置，请完成以下步骤：

转到 allreduce 示例：

bashcd allreducetrn

打开 flowpy 文件，并确保将作业队列和镜像设置为您通过 AWS CloudFormation 部署的队列名称和您推送到 Amazon ECR 的镜像。运行 allreduce 代码，请执行以下 Metaflow 命令：

bashpython flowpy packagesuffixes=sh run

您可以在 Metaflow UI 中找到日志为可读性起见进行了截断：

plaintextTask is starting (status SUBMITTED)Task is starting (status RUNNABLE) (parallel node status [SUBMITTED3])Task is starting (status STARTING) (parallel node status [SUBMITTED3])Task is starting (status RUNNING) (parallel node status [SUBMITTED3])Setting up task environmentDownloading code packageCode package downloadedTask is startingCompiler status PASSresult OK step 0 tensor([[64 64 64][64 64 64]] device