本文為英文版的機器翻譯版本，如內容有任何歧義或不一致之處，概以英文版為準。 # HyperPod 受管分層檢查點 本節說明受管分層檢查點的運作方式，以及它為大規模模型訓練提供的優勢。 Amazon SageMaker HyperPod 受管分層檢查點可協助您更有效率地訓練大規模生成式 AI 模型。它會使用多個儲存層，包括您叢集的 CPU 記憶體。此方法可縮短您的復原時間，並將訓練進度的損失降至最低。它也會在您的訓練基礎設施中使用未充分利用的記憶體資源。 受管分層檢查點可讓您以更高的頻率將檢查點儲存至記憶體。它會定期將其保存至耐久的儲存體。這可在訓練過程中同時維持效能與可靠性。 本指南說明如何在 Amazon EKS HyperPod 叢集上使用 PyTorch 架構來設定、設定和使用受管分層檢查點。 ## 受管分層檢查點的運作方式 受管分層檢查點使用多層儲存方法。CPU 記憶體可做為主要層來存放模型檢查點。次要層包括持久性儲存選項，例如 Amazon S3。 當您儲存檢查點時，系統會跨叢集節點將其存放在配置的記憶體空間中。它會自動跨相鄰運算節點複寫資料，以提高可靠性。此複寫策略可防範單一或多個節點故障，同時提供快速存取以進行復原操作。 系統也會根據您的組態，定期將檢查點儲存至持久性儲存體。這可確保訓練進度的長期耐久性。 重要元件包括： + **記憶體管理系統**：一種記憶體管常駐程式，提供分解的記憶體做為服務進行檢查點儲存 + **HyperPod Python 程式庫**：與分解的儲存 API 互動，並提供跨層儲存、載入和管理檢查點的公用程式 + **檢查點複寫**：自動跨多個節點複寫檢查點以進行容錯 系統透過簡單的 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 許可，包括： + 訓練 Pod 的 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 儲存貯體的寫入存取權，他們可能會產生惡意挑選資料，在受管分層檢查點載入時執行。 **Amazon S3 儲存體的最佳實務** 搭配 Amazon S3 儲存體使用受管分層檢查點時： + **限制 Amazon S3 儲存貯體存取**：確保只有與您訓練叢集相關聯的獲授權使用者和角色才能存取用於檢查點的 Amazon S3 儲存貯體。 + **實作儲存貯體政策**：設定適當的儲存貯體政策，以防止未經授權的存取或修改。 + **驗證存取模式**：實作記錄，以驗證檢查點 Amazon S3 儲存貯體的存取模式。 + **驗證儲存貯體名稱**：請謹慎選擇儲存貯體名稱，以避免發生潛在的儲存貯體劫持。 **網路端點** 受管分層檢查點可在下列連接埠上的每個運算節點上啟用網路端點：9200/TCP、9209/UDP、9210/UDP、9219/UDP、9220/UDP、9229/UDP、9230/UDP、9239/UDP、9240/UDP。這些是檢查點服務運作和維護資料同步所需的連接埠。 根據預設，SageMaker 的網路組態會基於安全目的限制對這些端點的存取。我們建議您維持這些預設限制。 為節點和 VPC 設定網路設定時，請遵循 VPCs、安全群組和 ACLs的 AWS 最佳實務。如需詳細資訊，請參閱下列內容： + [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)