관리형 티어 체크포인트 설정

이 섹션에는 Amazon SageMaker HyperPod에 대한 관리형 티어 체크포인트 지정의 설정 프로세스가 포함되어 있습니다. 클러스터에서 기능을 활성화하고 훈련 코드에서 체크포인트 지정을 구현하는 방법을 알아봅니다.

사전 조건

관리형 티어 체크포인트 지정을 설정하기 전에 다음 조건을 충족해야 합니다.

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
)

5단계: 훈련 루프에 체크포인트 저장

훈련 루프에서 FileSystemWriter와 함께 PyTorch DCP를 사용하여 체크포인트를 저장합니다.

FileSystemWriter와 함께 PyTorch DCP 사용

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
)

6단계: 복구를 위해 체크포인트 로드

체크포인트를 로드해야 하는 경우 FileSystemReader와 함께 PyTorch DCP를 사용합니다.

FileSystemReader와 함께 PyTorch DCP 사용

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/logs/sagemaker_checkpointing으로 체크포인트 로그 경로 /var/log/sagemaker_checkpointing을 탑재할 수 있습니다. 이렇게 하면 라이브러리에 한정된 로그만 별도로 수집됩니다. 이를 통해 서비스 팀은 디버깅 및 지원에 대한 향상된 가시성을 확보할 수 있습니다.