本文属于机器翻译版本。若本译文内容与英语原文存在差异，则一律以英文原文为准。 # HyperPod 托管分层检查点 本节介绍托管分层检查点的工作原理及其为大规模模型训练带来的好处。 Amazon SageMaker HyperPod 托管分层检查点可帮助您更高效地训练大规模生成式 AI 模型。它使用多个存储层，包括集群的 CPU 内存。此方法可以缩短恢复时间，并最大限度地减少训练进度损失；还能有效利用训练基础设施中未充分利用的内存资源。 托管分层检查点功能允许以更高的频率将检查点保存到内存中。此功能会定期将检查点保存到持久性存储中。这有助于在训练过程中同时保障性能与可靠性。 本指南介绍如何在 Amazon EKS HyperPod 集群上通过 PyTorch 框架设置、配置和使用托管分层检查点功能。 ## 托管分层检查点的工作原理 托管分层检查点使用多层存储方法。CPU 内存用作存储模型检查点的主层。辅助层包括 Amazon S3 等持久性存储选项。 保存检查点时，系统会将其存储在跨集群节点分配的内存空间中。它会在相邻的计算节点间自动复制数据，以提高可靠性。此复制策略可防范单个或多个节点故障，并提供快速访问权限以进行恢复操作。 系统还会根据您的配置定期将检查点保存到持久性存储中。这可确保训练进度的长期持久性。 关键组件包括： + **内存管理系统**：一种内存管理进程守护程序，可提供解耦式内存即服务以用于检查点存储 + **HyperPod Python 库**：与分解存储接口， APIs 并提供用于跨层保存、加载和管理检查点的实用工具 + **检查点复制**：自动在多个节点间复制检查点以实现容错能力 该系统通过简单的 API 调用与 PyTorch 训练循环无缝集成。这只需对现有代码进行极少的更改。 ## 优势 托管分层检查点为大规模模型训练提供了多种优势： + **提高了可用性**：管理检查点的保存、复制、持久化和恢复操作 + **更快的检查点操作**：与基于磁盘的检查点相比，基于内存的存储缩短了保存和加载时间，从而加快了恢复速度 + **容错能力**：跨节点自动复制检查点，可防范硬件节点故障 + **极少的代码更改**：简单的 API 集成只需对现有训练脚本进行极少的修改 + **提高训练吞吐量**：减少检查点开销，这意味着更多时间可用于实际训练 **Topics** + [托管分层检查点的工作原理](#managed-tier-checkpointing-works) + [优势](#managed-tier-checkpointing-benefits) + [设置托管分层检查点](managed-tier-checkpointing-setup.md) + [移除托管的分层检查点](managed-tier-checkpointing-remove.md) + [托管分层检查点的安全注意事项](managed-tier-security-considerations.md) # 设置托管分层检查点 本节包含 Amazon SageMaker HyperPod 托管分层检查点的设置流程。您将了解如何在集群上启用该功能并在训练代码中实现检查点。 **Topics** + [先决条件](#managed-tier-checkpointing-setup-prerequisites) + [步骤 1：为集群启用托管分层检查点](#managed-tier-checkpointing-setup-step-enable-for-cluster) + [步骤 2：在训练映像中安装 Python 库](#managed-tier-checkpointing-setup-step-install-library) + [第 3 步：在训练循环中保存检查点](#managed-tier-checkpointing-setup-step-save-checkpoint-in-loop) + [步骤 4：加载恢复检查点](#managed-tier-checkpointing-setup-step-load-checkpoint) + [验证您的托管分层检查点操作](#managed-tier-checkpointing-setup-validation) ## 先决条件 在设置托管分层检查点操作之前，请确保您具有： + 具有足够的 CPU 内存可用于分配检查点的 Amazon EKS HyperPod 集群 + PyTorch 训练工作负载和 DCP 作业（两者均受支持） + 用于管理集群的相应 IAM 权限，包括： + Amazon CloudWatch 和 Amazon S3 为训练容器写入权限，以读取/写入检查点和推送指标 + 这些权限可通过 [EKS OIDC 设置](https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html)进行配置 ## 步骤 1：为集群启用托管分层检查点 **重要** 您必须选择使用托管分层检查点功能。 在创建或更新集群 HyperPod APIs 时，通过启用托管分层检查点功能。当您指定 `TieredStorageConfig` 参数时，该服务会自动安装内存管理系统。 对于新集群，您可以使用[https://docs.aws.amazon.com/cli/latest/reference/sagemaker/create-cluster.html](https://docs.aws.amazon.com/cli/latest/reference/sagemaker/create-cluster.html) AWS CLI。 ``` aws sagemaker create-cluster \ --cluster-name cluster-name \ --orchestrator "Eks={ClusterArn=eks-cluster-arn}" \ --instance-groups '{ "InstanceGroupName": "instance-group-name", "InstanceType": "instance-type", "InstanceCount": instance-count, "LifeCycleConfig": { "SourceS3Uri": "s3-path-to-lifecycle-scripts", "OnCreate": "lifecycle-script-name" }, "ExecutionRole": "instance-group-iam-role", "ThreadsPerCore": threads-per-core, "InstanceStorageConfigs": [ { "EbsVolumeConfig": {"VolumeSizeInGB": volume-size} } ] }' \ --vpc-config '{ "SecurityGroupIds": ["security-group-ids"], "Subnets": ["subnets"] }' \ --tiered-storage-config '{ "Mode": "Enable" }' ``` `InstanceMemoryAllocationPercentage` 参数指定要为检查点分配的集群内存的 `percentage`（int）。范围是 20-100。 ## 步骤 2：在训练映像中安装 Python 库 将 [Amazon SageMaker 检查点库](https://pypi.org/project/amzn-sagemaker-checkpointing/)及其依赖项添加到您的 Dockerfile 中，将其安装到您的训练镜像中： ``` # Add this line to your training image Dockerfile RUN pip install amzn-sagemaker-checkpointing s3torchconnector tenacity torch boto3 s3torchconnector ``` ## 第 3 步：在训练循环中保存检查点 在训练循环中，您可以使用 PyTorch DCP 异步保存检查点。以下是如何执行此操作的示例。 ``` import torch import torch.distributed as dist from torch.distributed.checkpoint import async_save, load from amzn_sagemaker_checkpointing.checkpointing.filesystem.filesystem import ( SageMakerTieredStorageWriter, SageMakerTieredStorageReader ) # Initialize distributed training dist.init_process_group(backend="nccl") # Configure checkpointing checkpoint_config = SageMakerCheckpointConfig( # Unique ID for your training job # Allowed characters in ID include: alphanumeric, hyphens, and underscores namespace=os.environ.get('TRAINING_JOB_NAME', f'job-{int(time.time())}'), # Number of distributed processes/available GPUs world_size=dist.get_world_size(), # S3 storage location, required for SageMakerTieredStorageReader for read fallbacks # Required for SageMakerTieredStorageWriter when save_to_s3 is True s3_tier_base_path="s3://my-bucket/checkpoints" ) # Your model and optimizer model = MyModel() optimizer = torch.optim.AdamW(model.parameters()) # Training loop future = None in_memory_ckpt_freq = 10 s3_ckpt_freq = 50 for training_step in range(1000): # ... training code ... # Save checkpoint if (training_step % in_memory_ckpt_freq == 0 or training_step % s3_ckpt_freq == 0): # Create state dictionary state_dict = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "step": training_step, "epoch": epoch } # Create storage writer for current step checkpoint_config.save_to_s3 = training_step % s3_ckpt_freq == 0 storage_writer = SageMakerTieredStorageWriter( checkpoint_config=checkpoint_config, step=training_step ) # wait for previous checkpoint to get completed if future is not None: exc = future.exception() if exc: print(f"Failure in saving previous checkpoint:{str(exc)}") # Handle failures as required else: result = future.result() # Process results from save, if required # Async save checkpoint using PyTorch DCP future = async_save(state_dict=state_dict, storage_writer=storage_writer) # Continue training while checkpoint saves in background ``` ## 步骤 4：加载恢复检查点 以下是加载检查点的示例。 ``` # Create state dictionary template state_dict = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "step": 0, "epoch": 0 } # Load latest checkpoint storage_reader = SageMakerTieredStorageReader(checkpoint_config=checkpoint_config) load(state_dict, storage_reader=storage_reader) # Load specific checkpoint step storage_reader = SageMakerTieredStorageReader( checkpoint_config=checkpoint_config, step=500 # Or don't pass step if you have to load the latest available step. ) try: load(state_dict, storage_reader=storage_reader) except BaseException as e: print(f"Checkpoint load failed: {str(e)}") # Add additional exception handling ``` ## 验证您的托管分层检查点操作 您可以使用日志验证托管分层检查点操作。 **自定义日志记录（可选）** 您可以通过将自定义记录器传递给库来将检查点日志与其他日志集成。例如，您可以将自定义记录器添加到训练代码，这样库中的所有日志也会被收集到训练记录器中。 **增强型服务日志记录（可选）** 要增强调试和服务可见性，可以将检查点日志路径 `/var/log/sagemaker_checkpointing` 从容器组（pod）中挂载到主机上的路径 `/var/logs/sagemaker_checkpointing`。这可确保仅单独收集库特定的日志，并为服务团队提供更高的调试和支持可见性。 # 移除托管的分层检查点 本节介绍如何在不再需要托管分层检查点功能时将其禁用。 要禁用托管分层检查点，[https://docs.aws.amazon.com/cli/latest/reference/sagemaker/update-cluster.html](https://docs.aws.amazon.com/cli/latest/reference/sagemaker/update-cluster.html) AWS CLI 请使用更新您的集群配置： ``` aws sagemaker update-cluster \ --cluster-name cluster-name \ --tiered-storage-config '{ "Mode": "Disable" }' ``` 这会从集群中移除内存管理进程守护程序。该守护程序是作为标准的 Kubernetes 实现的， DaemonSet 并遵循标准的 Kubernetes 生命周期管理。 # 托管分层检查点的安全注意事项 本节介绍使用托管分层检查点操作时的重要安全注意事项。它包括 Python pickle 的使用、Amazon S3 加密和网络端点安全。 **Python pickle 的使用** 托管分层检查点使用 Python 的 pickle 模块来反序列化存储在 Amazon S3 中的检查点数据。此实现会产生重大安全影响： + **扩展信任边界**：在 Amazon S3 中使用托管分层检查点功能时，Amazon S3 存储桶将成为集群信任边界的一部分。 + **代码执行风险**：Python 的 pickle 模块可以在反序列化期间执行任意代码。如果未经授权的用户获得了对您的检查点 Amazon S3 存储桶的写入权限，则他们可能会制作恶意 pickle 数据，这些数据在通过托管分层检查点加载时执行。 **Amazon S3 存储的最佳实践** 在 Amazon S3 存储中使用托管分层检查点功能时： + **限制 Amazon S3 存储桶访问权限**：确保仅与训练集群关联的授权用户和角色有权访问用于检查点的 Amazon S3 存储桶。 + **实施存储桶策略**：配置适当的存储桶策略以防止未经授权的访问或修改。 + **验证访问模式**：实施日志记录以验证检查点 Amazon S3 存储桶的访问模式。 + **验证存储桶名称**：选择存储桶名称时需保持谨慎，以避免潜在的存储桶劫持风险。 **网络端点** 托管分层检查点在每个计算节点上启用以下端口上的网络终端节点：9200/TCP、9209/UDP、9210/UDP、9219/UDP、9220/UDP、9229/UDP、9230/UDP、9230/UDP、9239/UDP、9240/UDP。这些端口是检查点服务正常运行并保持数据同步所必需的。 默认情况下，出于安全考虑， SageMaker的网络配置限制对这些端点的访问。建议您保留这些默认限制。 为节点和 VPC 配置网络设置时，请遵循安全组和 AWS 的最佳 VPCs实践 ACLs。有关更多信息，请参阅下列内容： + [Amazon SageMaker HyperPod 先决条件](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-prerequisites.html#sagemaker-hyperpod-prerequisites-optional-vpcCluster) + [VPC 安全最佳实践](https://docs.aws.amazon.com/vpc/latest/userguide/vpc-security-best-practices.html)