亚马逊搜索如何通过为 Amazon SageMaker 训练作业使用 AWS Batch，将机器学习训练效率提高一倍-青云TOP-AI综合资源站平台|青云聚合API大模型调用平台|全网AI资源导航平台

📢 转载信息

原文链接：https://aws.amazon.com/blogs/machine-learning/how-amazon-search-increased-ml-training-twofold-using-aws-batch-for-amazon-sagemaker-training-jobs/

原文作者：Mona Mona, Bruno Pistone, James Park, and Mayank Jha

在本文中，我们将向您展示亚马逊搜索团队如何通过利用 AWS Batch 来管理 Amazon SageMaker 训练作业，从而优化了对 P5、P4 等 GPU 加速实例家族的利用率。这种托管解决方案使我们能够在这些实例上编排机器学习 (ML) 训练工作负载。我们还将提供一个分步指南来演示该用例的实施过程。

亚马逊搜索中的机器学习

在亚马逊搜索，我们使用数百个 GPU 加速实例来训练和评估 ML 模型，以帮助客户发现他们喜爱的产品。科学家通常会同时训练多个模型，以找到能优化模型性能的最佳特征集、模型架构和超参数设置。我们以前使用先进先出 (FIFO) 队列来协调模型训练和评估作业。然而，我们需要采用更细致的标准来确定作业的运行优先级和顺序。生产模型需要高优先级运行，探索性研究需要中等优先级，而超参数搜索和批量推理则需要低优先级。我们还需要一个能够处理中断的系统。如果作业失败，或者给定实例类型饱和，我们需要作业能够在其他兼容的可用实例类型上运行，同时尊重整体的优先级标准。最后，我们希望有一个托管解决方案，这样我们就可以将更多精力放在模型开发上，而不是管理基础设施。

在评估了多种方案后，我们选择了AWS Batch for Amazon SageMaker Training jobs，因为它最符合我们的要求。该解决方案将 AWS Batch 与 Amazon SageMaker 无缝集成，使我们能够根据优先级标准运行作业。这使得应用科学家无需手动管理资源即可提交多个并发作业。通过利用 AWS Batch 的特性，例如通过公平共享调度实现的先进优先级排序，我们将 GPU 加速实例的峰值利用率从 40% 提高到了 80% 以上。

亚马逊搜索：SageMaker 训练作业的 AWS Batch 实施

我们利用了三项 AWS 技术来设置我们的作业队列。我们使用服务环境 (Service Environments) 来配置 AWS Batch 用于提交和管理 SageMaker 训练作业的 SageMaker AI 参数。我们使用共享标识符 (Share Identifiers) 来确定工作负载的优先级。最后，我们使用Amazon CloudWatch 来监控关键事件或偏离预期行为的情况，并提供警报功能。让我们深入了解这些结构。

服务环境。我们设置了服务环境来表示每种实例家族（如 P5s 和 P4s）可用的总 GPU 容量。每个服务环境都根据我们团队在 AWS Batch 中的预留容量配置了固定的限制。请注意，对于使用SageMaker 训练计划的团队，这些限制可以设置为预留实例的数量，从而使容量规划更加简单。通过定义这些边界，我们确定了服务环境内的总 GPU 实例容量如何在不同生产作业之间分配。每个生产实验都通过共享标识符分配了该容量的一部分。

图 1 提供了一个真实世界的示例，说明我们如何使用 AWS Batch 的公平共享调度在 ShareIDs 之间分配 100 个 GPU 实例。我们将 60 个实例分配给 ProdExp1，40 个分配给 ProdExp2。当 ProdExp2 只使用了 25 个 GPU 实例时，剩余的 15 个可以被 ProdExp1 借用，使其能够扩展到 75 个 GPU 实例。当 ProdExp2 稍后需要全部 40 个 GPU 实例时，调度器会从 ProdExp1 中抢占作业以恢复平衡。此示例使用了 P4 实例家族，但相同的方法可以应用于任何 SageMaker 支持的 EC2 实例家族。这确保了生产工作负载可以保证访问其分配的容量，而探索性或即席实验仍然可以使用任何空闲的 GPU 实例。这种设计保护了关键工作负载，并通过确保没有预留容量被闲置，提高了整体实例利用率。

Figure 1: AWS Batch fair-share scheduling

图 1：AWS Batch 公平共享调度

共享标识符。我们使用共享标识符将服务环境容量的一部分分配给生产实验。共享标识符是在作业提交时应用的字符串标签。AWS Batch 使用这些标签来跟踪使用情况并强制执行公平共享调度。对于需要专用容量的计划，我们在 AWS Batch 中定义了具有配额的预设共享标识符。这为生产流程预留了容量。这些配额充当公平性目标而不是硬性限制。空闲容量仍然可以被借用，但在发生争用时，AWS Batch 通过从过度使用的标识符中抢占资源并将其重新分配给使用不足的标识符来强制执行公平性。

在每个共享标识符内部，从 0 到 99 的作业优先级决定了执行顺序，但仅当 ShareIdentifier 达到其分配的容量限制时，才会触发基于优先级的抢占。图 2 说明了我们如何设置和使用共享标识符。ProdExp1 拥有 60 个 p4d 实例并以各种优先级运行作业。Job A 的优先级为 80，Job B 为 50，Job C 为 30，Job D 的优先级为 10。当所有 60 个实例都被占用，并且提交了一个需要 15 个实例的新高优先级作业（优先级 90）时，系统会抢占优先级最低的正在运行的作业 (Job D) 来腾出空间，同时为该共享标识符保持 60 个实例的总数。这种设计确保了关键工作负载可以保证访问其分配的容量，而探索性或即席实验仍然可以利用任何空闲的 GPU 实例。

Figure 2: Priority scheduling within a Share ID

图 2：共享 ID 内的优先级调度

Amazon CloudWatch。我们使用 Amazon CloudWatch 来检测我们的 SageMaker 训练作业。SageMaker 会自动发布有关作业进度和资源利用率的指标，而 AWS Batch 则提供有关作业调度和执行的详细信息。通过 AWS Batch，我们可以通过AWS Batch API 查询每个作业的状态。这使得跟踪作业在 SUBMITTED、PENDING、RUNNABLE、STARTING、RUNNING、SUCCEEDED 和 FAILED 等状态之间的转换成为可能。我们将这些指标和作业状态发布到 CloudWatch，并配置了仪表板和警报，以便在遇到扩展的等待时间、意外故障或资源利用不足时发出警报。这种内置的集成提供了实时可见性和历史趋势分析，帮助我们的团队在构建定制监控系统的同时保持 GPU 集群的运营效率。

对团队性能的运营影响

通过为 SageMaker 训练作业采用 AWS Batch，我们使得实验可以在无需担心资源可用性或争用的情况下运行。研究人员可以提交作业而无需等待人工调度，这增加了可以并行运行的实验数量。这带来了更短的队列等待时间、更高的 GPU 利用率和更快的训练结果周转时间，直接提高了研究吞吐量和交付时间表。

如何为 SageMaker 训练作业设置 Amazon Batch

要设置类似的解决方案，您可以遵循本教程，其中展示了如何使用多个 GPU 驱动的实例来编排多个 GPU大型语言模型 (LLM) 微调作业。该解决方案也可在 GitHub 上找到。

先决条件

要使用 AWS Batch 对多个 SageMaker 训练作业进行编排，您首先需要完成以下先决条件：

克隆包含此部署资产的 GitHub 存储库。该存储库由引用资产的笔记本组成：

git clone https://github.com/aws/amazon-sagemaker-examples/ cd build_and_train_models/sm-training-queues-pytorch/

创建 AWS Batch 资源

要创建管理 AWS Batch 中 SageMaker 训练作业队列所需的资源，我们在示例中提供了实用程序函数，以自动化创建 Service Environment、Scheduling Policy 和 Job Queue。

服务环境代表了可用于调度的 Amazon SageMaker AI 容量限制，以最大实例数表示。调度策略指示了如何在作业队列中在用户或工作负载之间分配资源计算。作业队列是研究人员与之交互以提交作业和查询作业状态的调度器接口。AWS Batch 提供了我们可以操作的两种不同队列：

FIFO 队列 – 不需要调度策略的队列
公平共享队列 – 需要调度策略Amazon 资源名称 (ARN) 来编排已提交作业的队列

我们建议为每个作业队列创建专用的服务环境，比例为 1:1。FIFO 队列提供基本的消息传递，而公平共享调度 (FSS) 队列提供更复杂的调度，平衡共享标识符、共享权重和作业优先级的内部利用率。对于不需要多个共享但希望能够在作业提交时分配优先级的客户，我们建议创建一个 FSS 队列并在其中使用单个共享来进行所有提交。要创建资源，请执行以下命令：

cd smtj_batch_utils python create_resources.py

您可以浏览以下屏幕截图中显示的 AWS Batch 仪表板，以探索已创建的资源。

Screenshot of AWS Batch Dashboard showing created resources

此自动化脚本创建了两个队列：

ml-c5-xlarge-queue – 用于 CPU 工作负载的 FIFO 队列，优先级为 2
ml-g6-12xlarge-queue – 用于 GPU 工作负载的公平共享队列，优先级为 1

队列 ml-g6-12xlarge-queue 关联的调度策略具有如 HIGHPRI、MIDPRI 和 LOWPRI 等共享属性，以及队列权重。用户可以将作业提交并分配给三个共享之一：HIGHPRI、MIDPRI 或 LOWPRI，并分配如 1（高优先级）、3（中等优先级）和 5（低优先级）等权重。以下屏幕截图显示了调度策略的详细信息：

Screenshot of Scheduling Policy details in AWS Batch

有关如何设置服务环境和作业队列的说明，请参阅“Introducing AWS Batch support for SageMaker Training Jobs”博客中的“入门”部分。

在 SageMaker AI 上运行 LLM 微调作业

我们运行 notebook.ipynb 笔记本以开始使用 AWS Batch 提交 SageMaker 训练作业。该笔记本包含准备工作负载使用的数据、将数据上传到Amazon Simple Storage Service (Amazon S3) 以及定义作业执行所需超参数的代码。

要使用 SageMaker 训练作业运行微调工作负载，此示例使用了 ModelTrainer 类。ModelTrainer 类是一种更新、更直观的模型训练方法，可显著增强用户体验。它支持分布式训练、构建自己的容器 (BYOC) 和配方。

有关 ModelTrainer 的更多信息，您可以参考“加速您的 ML 生命周期：使用全新改进的 Amazon SageMaker Python SDK – 第 1 部分：ModelTrainer”

要设置微调工作负载，请完成以下步骤：

选择实例类型、用于训练作业的容器镜像，并定义模型存储检查点路径：

import sagemaker instance_type = "ml.g6.12xlarge" instance_count = 1 image_uri = sagemaker.image_uris.retrieve(     framework="pytorch",     region=sagemaker_session.boto_session.region_name,     version="2.6",     instance_type=instance_type,     image_scope="training" )

创建 ModelTrainer 函数以封装训练设置。ModelTrainer 类通过封装代码和训练设置来简化体验。在此示例中：

SourceCode – 源代码配置。这用于通过使用本地 Python 脚本来配置运行训练作业的源代码。
Compute – 计算配置。这用于指定训练作业的计算资源。

from sagemaker.modules.configs import Compute, OutputDataConfig, SourceCode, StoppingCondition from sagemaker.modules.distributed import Torchrun from sagemaker.modules.train import ModelTrainer role = sagemaker.get_execution_role() # 定义要运行的脚本 source_code = SourceCode(     source_dir="./scripts",     requirements="requirements.txt",     entry_script="train.py", ) # 定义计算资源 compute_configs = Compute(     instance_type=instance_type,     instance_count=instance_count,     keep_alive_period_in_seconds=0 ) # 定义训练作业名称 job_name = f"train-deepseek-distill-llama-8b-sft-batch" # 定义 OutputDataConfig 路径 output_path = f"s3://{bucket_name}/{job_name}" # 定义 ModelTrainer model_trainer = ModelTrainer(     training_image=image_uri,     source_code=source_code,     base_job_name=job_name,     compute=compute_configs,     distributed=Torchrun(),     stopping_condition=StoppingCondition(max_runtime_in_seconds=7200),     hyperparameters={         "config": "/opt/ml/input/data/config/args.yaml"     },     output_data_config=OutputDataConfig(s3_output_path=output_path),     role=role, )

通过创建 InputData 对象来设置 ModelTrainer 的输入通道，这些对象来自提供的 S3 存储桶路径，用于训练和验证数据集：

from sagemaker.modules.configs import InputData train_input = InputData(     channel_name="train",     data_source=train_dataset_s3_path, ) val_input = InputData(     channel_name="val",     data_source=val_dataset_s3_path, ) config_input = InputData(     channel_name="config",     data_source=train_config_s3_path, ) TRAINING_INPUTS = [train_input, val_input, config_input]

排队 SageMaker 训练作业

本节及后续内容旨在进行交互式操作，以便您可以探索如何使用 Amazon SageMaker Python SDK 将作业提交到您的 Batch 队列。请遵循以下步骤：

选择要使用的队列：

from sagemaker.aws_batch.queue import TrainingQueue SMTJ_BATCH_QUEUE = "ml-g6-12xlarge-queue" queue = TrainingQueue(SMTJ_BATCH_QUEUE)

在下一个单元格中，向队列中提交两个训练作业：
1. LOW PRIORITY (低优先级)
2. MEDIUM PRIORITY (中优先级)

使用 submit API 提交所有作业：

job_name_1 = job_name + "-low-pri" queued_job_1 = queue.submit(     model_trainer, TRAINING_INPUTS, job_name_1, priority=5, share_identifier="LOWPRI" ) job_name_2 = job_name + "-mid-pri" queued_job_2 = queue.submit(     model_trainer, TRAINING_INPUTS, job_name_2, priority=3, share_identifier="MIDPRI" )

显示正在运行和排队中的作业状态

我们可以使用作业队列列表和作业队列快照 API 以编程方式查看队列接下来将运行的作业的快照。对于公平共享队列，此排序是动态的，并且可能会不时需要刷新，因为新作业会提交到队列，或者共享使用情况会随时间变化。

from utils.queue_utils import print_queue_state print_queue_state(queue)

以下屏幕截图显示了以低优先级和中优先级提交的作业处于可运行状态并在队列中。

Screenshot showing jobs in Runnable State and in the queue

您还可以参考 AWS Batch 仪表板（如下所示）来分析作业的状态。

Screenshot of AWS Batch Dashboard analyzing job status

如下所示的屏幕截图所示，第一个执行的 SageMaker 训练作业是“MEDIUM PRIORITY”作业，它遵守了先前定义的调度策略规则。

Screenshot showing the MEDIUM PRIORITY job executing first

您可以在 SageMaker AI 控制台中探索正在运行的训练作业，如下所示的屏幕截图所示。

Screenshot showing the running training job in the SageMaker AI console

提交一个附加作业

您现在可以向队列提交一个带有“HIGH PRIORITY”的附加 SageMaker 训练作业：

job_name_3 = job_name + "-high-pri" queued_job_3 = queue.submit(     model_trainer, TRAINING_INPUTS, job_name_3, priority=1, share_identifier="HIGHPRI" )

您可以从仪表板中查看状态，如下所示的屏幕截图所示。

Screenshot showing the status from the dashboard after submitting the HIGH PRIORITY job

尽管“HIGH PRIORITY”作业是稍后提交到队列的，但它将优先于其他可运行的作业执行，因为它遵守了调度策略规则，如下所示的屏幕截图所示。

Screenshot showing the HIGH PRIORITY job executing before other runnable jobs

如屏幕截图中的调度策略所示，LOWPRI 共享的权重因子（5）高于 MIDPRI 共享（3）。由于较低的权重意味着较高的优先级，因此即使两个作业同时提交，LOWPRI 作业也会在 MIDPRI 作业之后执行。

Screenshot showing scheduling policy weights

清理

为了清理资源以避免产生额外费用，请遵循以下步骤：

验证您的训练作业是否不再运行。为此，请在您的SageMaker 控制台中，选择“训练 (Training)”并检查“训练作业 (Training jobs)”。
通过使用 GitHub 示例中的命令 python create_resources.py --clean 或通过AWS 管理控制台手动删除 AWS Batch 资源。

结论

在本文中，我们演示了亚马逊搜索如何使用 AWS Batch for SageMaker Training Jobs 来优化 GPU 资源利用率和训练作业管理。该解决方案通过实施复杂的队列管理和公平共享调度，将峰值 GPU 利用率从 40% 提高到 80% 以上，从而改变了其训练基础设施。我们建议面临类似 ML 训练基础设施挑战的组织探索 AWS Batch 与 SageMaker 的集成，它提供了内置的队列管理功能和基于优先级的调度。该解决方案消除了手动资源协调，同时通过可配置的调度策略为工作负载提供了适当的优先级。

要开始将 AWS Batch 与 SageMaker 训练作业集成，您可以访问我们在 GitHub 上的 amazon-sagemaker-examples 存储库中的示例代码和实施指南。该示例演示了如何设置AWS 身份和访问管理 (IAM) 权限、创建 AWS Batch 资源，以及使用 ModelTrainer 类编排多个 GPU 驱动的训练作业。

作者感谢 Charles Thompson 和 Kanwaljit Khurmi 在此期间的合作。

关于作者

🚀 想要体验更好更全面的AI调用？

欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。

目录CONTENT

亚马逊搜索如何通过为 Amazon SageMaker 训练作业使用 AWS Batch，将机器学习训练效率提高一倍