提升在 Amazon EKS 上的生成 AI 分布式培训工作负载效率

关键要点

在当今快速发展的人工智能AI领域，训练大规模语言模型LLMs面临着重大挑战。本文介绍如何使用 NVIDIA NeMo 框架在 Amazon Elastic Kubernetes Service (EKS) 集群上运行分布式培训工作负载，提供详细步骤和建议。

引言

在人工智能的快速发展中，训练大型语言模型LLMs面临着显著的挑战。这些模型通常需要大量的计算资源和复杂的基础设施来处理海量数据和复杂算法。如果没有一个结构化的框架，整个过程可能会变得非常耗时、昂贵且复杂。企业在管理分布式培训工作负载、有效利用资源以及优化模型准确性和性能方面常常面临困难。此时，NVIDIA NeMo 框架可以发挥重要作用。本文将提供如何在 Amazon EKS 集群上运行分布式培训工作负载的逐步指导。

NVIDIA NeMo 框架

NVIDIA NeMo 是一个以云为中心的端到端框架，用于训练和部署具有数十亿和数万亿参数的生成 AI 模型。该框架提供了一整套工具、脚本和配方，支持 LLM 工作的每个阶段，从数据准备到训练和部署。NVIDIA NeMo 为语言和图像应用的模型推理提供了多种自定义技术，并优化为大规模推理。通过提供端到端的工作流、高级并行技术、内存节省策略和分布式检查点，NVIDIA NeMo 确保 AI 模型培训更加高效、可扩展且高性能。

虎跃加速器链接

使用 NVIDIA NeMo 的优势

优势描述端到端管道包括数据准备、培训等不同阶段，支持定制化数据的即插即用方法并行技术包含数据、张量、流水线、序列、专家及上下文并行等内存节省技术包括选择性激活再计算、CPU卸载、多种类型的注意力及分布式优化器等数据加载工具支持不同架构分布式检查点

解决方案概览

可以通过 Slurm 或 Kubernetes 管理 NVIDIA NeMo。Amazon EKS 是一个托管的 Kubernetes 服务，使得在 AWS 上运行 Kubernetes 集群变得简单。它管理 Kubernetes 控制平面的可用性和可伸缩性，并支持计算节点的自动伸缩和生命周期管理，帮助您运行高可用的容器应用程序。

Amazon EKS 是运行分布式培训工作负载的理想平台，因为它与 AWS 服务具有强大集成和性能特征。它与 Amazon FSx for Lustre 无缝集成，启用快速的数据访问和管理。Amazon EKS 还与 Amazon CloudWatch 集成，提供全面的日志记录和监控，帮助您洞悉集群性能和资源利用情况。此外，它支持 Amazon Simple Storage Service (Amazon S3) 进行可扩展和持久的数据存储和管理。通过增强网络性能与 Elastic Fabric Adapter (EFA)，实现了低延迟、高吞吐量的节点间连接。这些特性使得 Amazon EKS 成为优化 AI 和机器学习 (ML) 培训工作流的强大和高效选择。

下图展示了解决方案架构：

本文将介绍在 EKS 集群上运行分布式培训工作负载的步骤。高层步骤如下：

设置一个启用 EFA 的 2 节点 24xlarge 集群。设置 FSx for Lustre 文件系统，以便共享数据存储训练数据集和模型检查点。设置 NVIDIA NeMo 的环境。修改 NVIDIA NeMo Kubernetes 清单，以准备数据集并训练模型。

先决条件

需要能发起一个基于 CPU 的 Amazon Elastic Compute Cloud (Amazon EC2) 实例，以便用于创建 EKS 集群。启动并运行实例后，SSH 登录到 EC2 实例并安装以下 CLI 工具：

最新版本的 AWS Command Line Interface (AWS CLI)kubectleksctlhelm

如在非 Linux 平台，安装步骤可能有所不同。请参考上述文档，按照对应平台的说明安装 CLI。同样需要确保您已预留 p4de24xlarge 实例的容量，且拥有 capacityReservationID。

启动 EKS 集群

ECR p4de24xlarge 实例配备 NVIDIA A100 80GB GPU，这在分布式培训生成 AI 工作负载中极为流行。更多信息，请参阅 Amazon EC2 实例类型。本节将展示如何为 p4de24xlarge 实例创建 EKS 集群，并进行 按需容量保留。

提供集群创建配置文件：p4declusterconfigyaml，参考以下代码：

bash git clone https//githubcom/awssamples/awsomedistributedtraininggit cd awsomedistributedtraining/3testcases/2nemolauncher/EKS

eksctl create cluster f p4declusterconfigyaml

在创建集群时，有以下几点需要注意：

确保 kubectl 版本及指定区域正确。更新 capacityReservationID 字段，并确保在 managedNodeGroups 部分指定 availabilityZones，这应与您容量所在的可用区 ID 相同。此配置将创建两个托管节点组：一个用于系统节点，使用 c52xlarge 实例，另一个用于在 p4de24xlarge 实例上运行分布式培训。托管节点组将使用 Amazon EKS 优化的 AMI。如果您希望提供自定义 AMI，可以创建自托管节点组并指定自定义 AMI。有关查找 AMI ID 的更多详情，请参阅 检索 Amazon EKS 优化的 Amazon Linux AMI ID。有关 Amazon EKS 优化 AMI 的更多详细信息，请查看 GitHub 仓库。确保 efaEnabled 设置为 true。您可以使用相同的配置为其他节点组创建集群。有关支持 EFA 的实例类型列表，请参见 支持的实例类型。

另一个在生成 AI 分布式培训工作负载中较为流行的实例是 p548xlarge 实例，配备 NVIDIA H100 80GB GPU。有关如何为现有 EKS 集群添加 P5 节点组的信息，请参见 AWS CLI 脚本用于 EKS 管理。

创建集群后，可以通过向 kubectl 配置文件添加新上下文来启用 kubectl 与集群进行通信：

bash aws eks updatekubeconfig region regioncode name mycluster

通过运行以下命令确认与集群的通信：

bash kubectl get svc

接下来，您可以安装 AWS EFA Kubernetes 设备插件。EFA 是一种网络接口，增强 EC2 实例间通信性能，这对涉及 GPU 的分布式培训工作负载至关重要。此插件允许 Kubernetes 识别和利用 EFA 设备，从而促进了高吞吐量和低延迟的网络连接，以实现有效的分布式培训和深度学习应用程序。

使用以下代码安装插件：

bash helm repo add eks https//awsgithubio/ekscharts

helm install efa eks/awsefak8sdeviceplugin n kubesystem

NVIDIA 设备插件 for Kubernetes 通过 kubelet 将 GPU 支持引入到您的 EKS 集群，向 Kubernetes API 服务器公开 GPU，广告可用 GPU 资源，允许 Kubernetes 调度和管理 GPU 加速的工作负载。

使用以下代码安装插件：

bash wget https//rawgithubusercontentcom/NVIDIA/k8sdeviceplugin/v0143/nvidiadevicepluginyml

kubectl apply f nvidiadevicepluginyml

运行以下命令以验证所有 Pods：

bash kubectl get pods allnamespaces

您可以运行 kubectl get nodes 验证节点。

另外，您可以使用 EKS node viewer 工具查看节点、成本和状态。安装后，输入 eksnodeviewer 以获得以下视图。

节点视图显示了两个 p4de24xlarge 计算节点的 IP 地址。

我们可以选择其中一个私有 IP DNS 名称进一步检查和描述节点，如下所示：

bashkubectl describe node ip19216816537uswest2computeinternal

前述命令将描述节点的详细信息。确保 EFA 正确安装，您应当看到如下截图中的详细信息。

对于 p4 节点，您会看到 vpcamazonawscom/efa4，而对于 p548xlarge 节点，您应该会看到 vpcamazonawscom/efa32。

如果节点组中启用了 EFA，请确保一个安全组附加到节点，允许来自同一安全组的所有流量规则。这是 EFA 正常工作的前提。有关说明，请参见 开始使用 EFA 和 MPI。此安全组仅用于测试目的。对于生产环境，建议创建仅允许来自您连接的 IP 地址例如您的计算机 IP 地址或本地网络中的 IP 地址范围的入站 SSH 规则。

创建 FSx for Lustre 文件系统

对于分布式培训应用程序，通常使用数百个 GPU 实例，每个节点中包含多个 GPU，确保所有节点能够高效访问共享文件系统以共同训练同一数据集至关重要。为此，需要高性能的文件系统，具备高吞吐量和低延迟。我们建议使用 FSx for Lustre 文件系统用于大规模分布式培训，因为其满足这些要求，并为参与培训过程的所有节点提供无缝的数据访问。

要在您的 EKS 集群上挂载 FSx for Lustre 文件系统，请完成以下步骤：

使用以下脚本创建一个 AWS 身份与访问管理 (IAM) 角色，并附加 FSx 策略：

bash export FSXPOLICYNAME=fsxcsi

wget https//githubcom/awssamples/awsdoeks/blob/main/ContainerRoot/eks/deployment/csi/fsx/fsxpolicyjson export FSXPOLICYDOC=file//fsxpolicyjson

# 从 EC2 自动扩展组 export EKSINSTANCEPROFILENAME=(eks1ec6fc6b1a19d65d66ac293ff0a20eb9)

POLICYARN=(aws iam createpolicy policyname {FSXPOLICYNAME} policydocument FSXPOLICYDOC query PolicyArn output text)

INSTANCEPROFILE=(aws iam listinstanceprofiles query InstanceProfiles[InstanceProfileName=={EKSINSTANCEPROFILENAME}]{InstanceProfileNameInstanceProfileName} output text)

ROLENAME=(aws iam getinstanceprofile instanceprofilename {INSTANCEPROFILE} query InstanceProfileRoles[0]RoleName output text)

# 将 FSx 策略附加到角色 {ROLENAME} aws iam attachrolepolicy policyarn {POLICYARN} rolename {ROLENAME}

使用以下脚本创建一个安全组，允许 EKS 节点访问文件系统：

bash # 从 EC2 控制台 export MYREGION=uswest2 # FSXSUBNETID 应与计算节点的相同 ID。您可以从 EKS 控制台获得此信息 export FSXSUBNETID=subnet0edecd850cff2cfad # 从 EC2 自动扩展组 export FSXSECURITYGROUPNAME=eksfsxsg

# 从 EKS 控制台获取 VPCID export VPCID=vpc04411d49af198a6ea

# 创建安全组 export SECURITYGROUPID=(aws ec2 createsecuritygroup vpcid {VPCID} region {MYREGION} groupname {FSXSECURITYGROUPNAME} description FSx for Lustre Security Group query GroupId output text)

export SUBNETCIDR=(aws ec2 describesubnets region {MYREGION} query Subnets[SubnetId=={FSXSUBNETID}]{CIDRCidrBlock} output text)

# 入站规则 aws ec2 authorizesecuritygroupingress region {MYREGION} groupid {SECURITYGROUPID} protocol tcp port 988 cidr {SUBNETCIDR}

从 FSx for Lustre 控制台创建 12 TB 的 Persistent2 FSx for Lustre 文件系统，确保在与计算实例 (FSXSUBNETID) 相同的可用区，Amazon EKS 的 VPC (VPCID) 和您创建的安全组 (SECURITYGROUPID)。

文件系统创建后，请注意文件系统 ID、DNS 名称和挂载名。

在挂载文件系统之前，您需要安装 FSx CSI 驱动程序，该驱动程序允许 EKS 集群管理 FSx for Lustre 文件系统的生命周期。

安装 FSx CSI 驱动程序：

bash echo Installing FSx CSI driver kubectl apply k githubcom/kubernetessigs/awsfsxcsidriver/deploy/kubernetes/overlays/stable/ref=master

echo FSx pods in kubesystem namespace kubectl n kubesystem get pods grep fsx

接下来，为了挂载文件系统，提供以下脚本在 fsxstorageclassyaml、fsxpvyaml 和 fsxpvcyaml 文件中：

bash # 存储类 kubectl apply f fsxstorageclassyaml kubectl get sc

# 持久卷 kubectl apply f fsxpvyaml

# 持久卷声明 kubectl apply f fsxpvcyaml

确保卷