使用 AWS Trainium 实例驱动的 Amazon ECS 扩展您的机器学习工作负载

使用容器运行机器学习(ML)工作负载正在成为一种常见做法。容器不仅可以完全封装您的训练代码,还包括从硬件库和驱动程序到整个依赖栈的所有内容。您得到的是一个一致且可移植的ML开发环境。使用容器,集群规模扩展变得更加容易。

在2022年底,AWS宣布推出由AWS Trainium加速器驱动的Amazon EC2 Trn1实例的一般可用性,这些实例专为高性能深度学习训练而构建。 Trn1实例相比其他类似的Amazon Elastic Compute Cloud(Amazon EC2)实例可节省高达50%的培训成本。此外,发布了AWS Neuron SDK,以改善此加速,为开发人员提供与此技术交互的工具,例如编译,运行时和配置文件,以实现高性能和成本效益的模型培训。

Amazon Elastic Container Service(Amazon ECS)是一个完全托管的容器编排服务,可简化容器化应用程序的部署,管理和扩展。只需描述您的应用程序和所需资源,Amazon ECS将在灵活的计算选项上启动,监视和扩展您的应用程序,并自动集成其他支持您的应用程序所需的AWS服务。

在本文中,我们将向您展示如何使用Amazon ECS在容器中运行ML培训作业以部署,管理和扩展您的ML工作负载。

解决方案概述

我们将向您介绍以下高级步骤:

  1. 使用AWS CloudFormation提供ECS Trn1实例群集。
  2. 构建带有Neuron SDK的自定义容器映像,并将其推送到Amazon Elastic Container Registry(Amazon ECR)。
  3. 创建任务定义以定义由Amazon ECS运行的ML训练作业。
  4. 在Amazon ECS上运行ML任务。

先决条件

为了跟随本文,需要熟悉Amazon EC2和Amazon ECS等核心AWS服务。

提供ECS Trn1实例群集

要开始,请启动提供的CloudFormation模板,该模板将提供所需的资源,例如VPC,ECS群集和EC2 Trainium实例。

我们使用Neuron SDK在AWS Inferentia和Trainium-based实例上运行深度学习工作负载。它支持您的端到端ML开发生命周期,以创建新模型,对其进行优化,然后将其部署到生产环境中。要使用Trainium训练您的模型,您需要在运行ECS任务的EC2实例上安装Neuron SDK,以映射与硬件相关联的NeuronDevice,以及将推送到Amazon ECR的Docker映像以访问训练模型的命令。

标准版本的Amazon Linux 2或Ubuntu 20不带有安装AWS Neuron驱动程序。因此,我们有两个不同的选项。

第一种选择是使用已安装了Neuron SDK的Deep Learning Amazon Machine Image(DLAMI)。样本可在GitHub repo上找到。您可以根据操作系统选择DLAMI。然后运行以下命令来获取AMI ID:

aws ec2 describe-images --region us-east-1 --owners amazon --filters 'Name=name,Values=Deep Learning AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Name=state,Values=available' --query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' --output text

输出结果如下:

ami-06c40dd4f80434809

此AMI ID可能会随时间而变化,请确保使用命令获取正确的AMI ID。

现在,您可以在CloudFormation脚本中更改此AMI ID,并使用即用型Neuron SDK。要执行此操作,请查找Parameters中的EcsAmiId

"EcsAmiId": { 
    "Type": "String", 
    "Description": "AMI ID", 
    "Default": "ami-09def9404c46ac27c" 
}

第二种选择是在堆栈创建期间填写userdata字段来创建实例。您不需要安装它,因为CloudFormation将设置这一点。有关更多信息,请参阅Neuron设置指南。

本文使用选项2,以防您需要使用自定义图像。完成以下步骤:

  1. 启动提供的 CloudFormation 模板。
  2. 对于 KeyName,输入所需密钥对的名称,它将预加载参数。本文使用的是 trainium-key
  3. 为您的堆栈输入名称。
  4. 如果您在 us-east-1 区域运行,可以将 ALBNameAZIds 的值保留为默认值。

要检查区域中哪个可用区具有 Trn1,请运行以下命令:

aws ec2 describe-instance-type-offerings --region us-east1 --location-type availability-zone --filter Name=instance-type,Values=trn1.2xlarge

  1. 选择下一步并完成堆栈的创建。

堆栈完成后,您可以转到下一步。

准备并推送带有Neuron SDK的ECR镜像

Amazon ECR是一个完全托管的容器注册表,提供高性能的托管,因此您可以在任何地方可靠地部署应用程序镜像和工件。我们使用 Amazon ECR 存储包含我们的脚本和 Neuron 包的自定义 Docker 镜像,这些脚本和 Neuron 包需要使用在 Trn1 实例上运行的 ECS 作业来训练模型。您可以使用 AWS 命令行界面(AWS CLI)或 AWS 管理控制台创建 ECR 存储库。本文使用控制台。完成以下步骤:

  1. 在 Amazon ECR 控制台上,创建一个新的存储库。
  2. 对于可见性设置¸,选择私有
  3. 对于存储库名称,输入名称。
  4. 选择创建存储库

现在您有了一个存储库,让我们构建和推送一个镜像,它可以在本地(在您的笔记本电脑中)或 AWS Cloud9 环境中构建。我们正在训练一个多层感知器(MLP)模型。有关原始代码,请参阅多层感知器训练教程。

  1. 将 train.py 和 model.py 文件复制到项目中。

它已经兼容 Neuron,因此您不需要更改任何代码。

  1. 5. 创建一个 Dockerfile,其中包含安装Neuron SDK和培训脚本的命令:
FROM amazonlinux:2

RUN echo $'[neuron] \n\
name=Neuron YUM Repository \n\
baseurl=https://yum.repos.neuron.amazonaws.com \n\
enabled=1' > /etc/yum.repos.d/neuron.repo

RUN rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB

RUN yum install aws-neuronx-collectives-2.* -y
RUN yum install aws-neuronx-runtime-lib-2.* -y
RUN yum install aws-neuronx-tools-2.* -y
RUN yum install -y tar gzip pip
RUN yum install -y python3 python3-pip
RUN yum install -y python3.7-venv gcc-c++
RUN python3.7 -m venv aws_neuron_venv_pytorch

# Activate Python venv
ENV PATH="/aws_neuron_venv_pytorch/bin:$PATH"
RUN python -m pip install -U pip
RUN python -m pip install wget
RUN python -m pip install awscli

RUN python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
RUN python -m pip install torchvision tqdm torch-neuronx neuronx-cc==2.* pillow
RUN mkdir -p /opt/ml/mnist_mlp
COPY model.py /opt/ml/mnist_mlp/model.py
COPY train.py /opt/ml/mnist_mlp/train.py
RUN chmod +x /opt/ml/mnist_mlp/train.py
CMD ["python3", "/opt/ml/mnist_mlp/train.py"]

如果您想使用Neuron创建自己的Dockerfile,请参阅在AWS ML加速器实例上开发,您可以在该指南中找到其他操作系统和ML框架的指南。

  1. 6. 使用以下代码构建映像,然后将其推送到Amazon ECR(提供您的区域、帐户ID和ECR存储库):
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin {your-account-id}.dkr.ecr.{your-region}.amazonaws.com

docker build -t mlp_trainium .

docker tag mlp_trainium:latest {your-account-id}.dkr.ecr.us-east-1.amazonaws.com/mlp_trainium:latest

docker push {your-account-id}.dkr.ecr.{your-region}.amazonaws.com/{your-ecr-repo-name}:latest

之后,您创建的ECR存储库中应该可以看到您的镜像版本。

作为ECS任务运行ML培训作业

要在Amazon ECS上运行ML培训任务,您首先需要创建一个任务定义。任务定义是在Amazon ECS中运行Docker容器所必需的。

  1. 在Amazon ECS控制台中,选择任务定义
  2. 创建新任务定义菜单中,选择使用JSON创建新任务定义

您可以使用以下任务定义模板作为基线。请注意,在映像字段中,您可以使用上一步生成的映像。确保它包括您的帐户ID和ECR存储库名称。

为了确保安装了Neuron,您可以检查设备块中是否映射了/dev/neuron0卷。这映射到在trn1.2xlarge实例上运行的单个NeuronDevice,该实例有两个核心。

  1. 使用以下模板创建您的任务定义:
{
    "family": "mlp_trainium",
    "containerDefinitions": [
        {
            "name": "mlp_trainium",
            "image": "{your-account-id}.dkr.ecr.us-east-1.amazonaws.com/{your-ecr-repo-name}",
            "cpu": 0,
            "memoryReservation": 1000,
            "portMappings": [],
            "essential": true,
            "environment": [],
            "mountPoints": [],
            "volumesFrom": [],
            "linuxParameters": {
                "capabilities": {
                    "add": [
                        "IPC_LOCK"
                    ]
                },
                "devices": [
                    {
                        "hostPath": "/dev/neuron0",
                        "containerPath": "/dev/neuron0",
                        "permissions": [
                            "read",
                            "write"
                        ]
                    }
                ]
            },
            ,
            "logConfiguration": {
                "logDriver": "awslogs",
                "options": {
                    "awslogs-create-group": "true",
                    "awslogs-group": "/ecs/task-logs",
                    "awslogs-region": "us-east-1",
                    "awslogs-stream-prefix": "ecs"
                }
            }
        }
    ],
    "networkMode": "awsvpc",
    "placementConstraints": [
        {
            "type": "memberOf",
            "expression": "attribute:ecs.os-type == linux"
        },
        {
            "type": "memberOf",
            "expression": "attribute:ecs.instance-type == trn1.2xlarge"
        }
    ],
    "requiresCompatibilities": [
        "EC2"
    ],
    "cpu": "1024",
    "memory": "3072"
}

您还可以使用以下任务定义或以下命令在AWS CLI上完成此步骤:

aws ecs register-task-definition \
--family mlp-trainium \
--container-definitions '[{    
    "name": "my-container-1",    
    "image": "{your-account-id}.dkr.ecr.us-east-1.amazonaws.com/{your-ecr-repo-name}",
    "cpu": 0,
    "memoryReservation": 1000,
    "portMappings": [],
    "essential": true,
    "environment": [],
    "mountPoints": [],
    "volumesFrom": [],
    "logConfiguration": {
        "logDriver": "awslogs",
        "options": {
            "awslogs-create-group": "true",
            "awslogs-group": "/ecs/task-logs",
            "awslogs-region": "us-east-1",
            "awslogs-stream-prefix": "ecs"
        }
    },
    "linuxParameters": {
        "capabilities": {
            "add": [
                "IPC_LOCK"
            ]
        },
        "devices": [{
            "hostPath": "/dev/neuron0",
            "containerPath": "/dev/neuron0",
            "permissions": ["read", "write"]
        }]
    }
}]' \
--requires-compatibilities EC2
--cpu "8192" \
--memory "16384" \
--placement-constraints '[{
    "type": "memberOf",
    "expression": "attribute:ecs.instance-type == trn1.2xlarge"
}, {
    "type": "memberOf",
    "expression": "attribute:ecs.os-type == linux"
}]'

在 Amazon ECS 上运行任务

在创建了 ECS 集群、将镜像推送到 Amazon ECR 并创建了任务定义之后,我们运行任务定义以在 Amazon ECS 上训练模型。

  1. 在 Amazon ECS 控制台中,选择集群
  2. 打开您的集群。
  3. 任务选项卡中,选择运行新任务

  1. 对于启动类型,选择EC2

  1. 对于应用程序类型,选择任务
  2. 对于,选择您创建的任务定义。

  1. 网络部分中,指定由 CloudFormation 栈创建的 VPC、子网和安全组。

  1. 选择创建

您可以在 Amazon ECS 控制台上监视您的任务。

您还可以使用 AWS CLI 运行任务:

aws ecs run-task --cluster <your-cluster-name> --task-definition <your-task-name> --count 1 --network-configuration '{"awsvpcConfiguration": {"subnets": ["<your-subnet-name> "], "securityGroups": ["<your-sg-name> "] }}'

结果将如下所示。

您还可以通过 Amazon CloudWatch 日志组检查培训作业的详细信息。

在训练完模型后,您可以将其存储在 Amazon Simple Storage Service (Amazon S3) 中。

清理

为避免额外支出,您可以将自动缩放组更改为最小容量,并将期望容量更改为零,以关闭 Trainium 实例。 要进行完整的清理,请删除 CloudFormation 堆栈以删除此模板创建的所有资源。

结论

在本篇文章中,我们展示了如何使用 Amazon ECS 部署您的 ML 训练作业。我们创建了一个 CloudFormation 模板来创建 Trn1 实例的 ECS 集群,构建了一个自定义的 Docker 镜像,将其推送到 Amazon ECR,然后在 ECS 集群上使用 Trainium 实例运行了 ML 训练作业。

有关 Neuron 和您可以使用 Trainium 做什么的更多信息,请查看以下资源:

  • 使用 Amazon EC2 Trn1 UltraClusters 扩展大型语言模型 (LLM) 训练
  • 使用 AWS Trainium 和 Amazon EKS 扩展分布式训练
  • Neuron 常见问题解答