設定受管層檢查點 - Amazon SageMaker AI
先決條件步驟 1：為您的叢集啟用受管層檢查點步驟 2：在您的訓練映像中安裝 Python 程式庫步驟 3：建立檢查點組態步驟 4：定義 SageMaker 檔案系統寫入器步驟 5：在您的訓練迴圈中儲存檢查點步驟 6：載入用於復原的檢查點監控和驗證

設定受管層檢查點

本節包含 Amazon SageMaker HyperPod 受管層檢查點的設定程序。您將了解如何在叢集上啟用功能，並在訓練程式碼中實作檢查點。

先決條件

設定受管層檢查點之前，請確定您具有下列先條件：

  • 具有足夠 CPU 記憶體可用於檢查點配置的 Amazon EKS HyperPod 叢集

  • PyTorch 訓練工作負載和 DCP 任務 (兩者都受到支援)

  • 叢集管理的適當 IAM 許可，包括：

    • 訓練 Pod 的 Amazon CloudWatch 和 Amazon S3 寫入許可，用於讀取/寫入檢查點和推送指標

    • 這些許可可以透過 EKS OIDC 設定進行設定

步驟 1：為您的叢集啟用受管層檢查點

重要

您必須選擇加入，才能使用受管層檢查點。

在建立或更新叢集時，透過 HyperPod API 啟用受管層檢查點。當您指定 TieredStorageConfig 參數時，服務會自動安裝記憶體管理系統。對於新叢集，建立：

aws sagemaker update-cluster \
    --cluster-name my-training-cluster \
    --tiered-storage-config {
            "Mode": "Enable"
            "InstanceMemoryAllocationPercentage": percentage
            }

InstanceMemoryAllocationPercentage 參數指定要針對檢查點配置的叢集記憶體 percentage (int)。範圍為 20-100。

步驟 2：在您的訓練映像中安裝 Python 程式庫

在您的訓練映像中安裝 Amazon SageMaker 檢查點程式庫，方法是將其新增至您的 Dockerfile：

# Add this line to your training image Dockerfile
RUN pip install amzn-sagemaker-checkpointing

步驟 3：建立檢查點組態

建立 CheckpointConfig 物件以指定檢查點行為。其中包含：

  • 檢查點位置

  • 檢查點的頻率

  • 命名空間的名稱。

以下範例顯示日誌組態：

from amzn_sagemaker_checkpointing.config.sagemaker_checkpoint_config import SageMakerCheckpointConfig
from amzn_sagemaker_checkpointing.checkpointing.filesystem import SageMakerTieredStorageWriter, SageMakerTieredStorageReader

checkpoint_config = sm_ckpt.CheckpointConfig(
    world_size = 100,
    in_memory_namespace: my-ml-workload, # Logical grouping for checkpoints
    s3_base_path: "s3://bucket-name/checkpointing-path-prefix/",
    s3_every_n_steps: 100,  # Every 100 steps, save to S3
)

步驟 4：定義 SageMaker 檔案系統寫入器

定義您的檢查點檔案系統寫入器。您可以在初始化期間選擇性地指定步驟編號。

基本寫入器 (儲存呼叫中指定的步驟)：

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(checkpoint_config)

具有步驟參數的寫入器 (初始化時指定的步驟)：

smWriter = sagemaker_checkpointing.SageMakerTieredStorageWriter(
    checkpoint_config, 
    step=step_number
)
注意

當您在寫入器初始化期間指定 step 參數時，儲存呼叫中的 checkpoint_id 參數會變成選用參數。步驟參數優先於檢查點目錄格式。

步驟 5：在您的訓練迴圈中儲存檢查點

在您的訓練迴圈中，使用 PyTorch DCP 搭配 FileSystemWriter 儲存檢查點。

使用 PyTorch DCP 搭配 FileSystemWriter

使用 FileSystemWriter 做為輸入來呼叫 dist_cp.save() 方法：

選項 1：使用 checkpoint_id 搭配步驟格式 (在寫入器中未指定步驟時)

# Construct checkpoint directory with step number
checkpoint_dir = f"step_number"

dist_cp.save_state_dict(
    state_dict=state_dict,  # state_dict is a dictionary containing model parameters, optimizer state, etc.
    checkpoint_id=checkpoint_dir,  # Should contain step number
    storage_writer=smWriter
)

選項 2：使用寫入器搭配步驟參數 (checkpoint_id 變成選用)

dist_cp.save_state_dict(
    state_dict=state_dict,
    storage_writer=smWriter  # Step already specified in writer initialization
)
注意

checkpoint_id 值 (或 checkpoint_dir 字串) 的格式必須為 step_number。例如 step_5。在寫入器初始化中使用步驟參數時，checkpoint_id 會變成選用。

步驟 6：載入用於復原的檢查點

當您需要載入檢查點時，請使用 PyTorch DCP 搭配 FileSystemReader。

使用 PyTorch DCP 搭配 FileSystemReader

使用 FileSystemReader 做為輸入來呼叫 DCP 載入方法：

# Define FileSystemReader
smReader = sagemaker_checkpointing.SageMakerTieredStorageReader(
    config=checkpoint_config
)

# Load checkpoint
dist_cp.load_state_dict(
    state_dict=state_dict,
    checkpoint_id=checkpoint_dir,
    storage_reader=smReader
)

監控和驗證

您可以透過指標和日誌來監控和驗證受管層檢查點操作。

自訂記錄 (選用)

您可以透過將自訂記錄器傳遞至程式庫，將檢查點日誌與其他日誌整合。例如，您可以將自訂記錄器新增至訓練程式碼，以便也會在訓練記錄器中收集程式庫中的所有日誌。

增強式服務記錄 (選用)

如需增強偵錯和服務可見性，您可以將檢查點日誌路徑 /var/log/sagemaker_checkpointing 從 Pod 內掛載到主機上的路徑 /var/logs/sagemaker_checkpointing。這可確保僅單獨收集程式庫特定的日誌。這可為服務團隊提供增強的偵錯和支援可見性。