기계 번역으로 제공되는 번역입니다. 제공된 번역과 원본 영어의 내용이 상충하는 경우에는 영어 버전이 우선합니다.

계층형 체크포인트

HyperPod 관리형 계층형 체크포인트는 클러스터의 CPU 메모리에 체크포인트를 먼저 쓰고 노드 간에 복제합니다. 주기적으로 Amazon S3와 같은 내구성 있는 스토리지에 유지됩니다. 빠른 계층은 객체 스토리지가 아닌 메모리이기 때문에 더 자주 체크포인트를 지정하고 노드가 실패할 때 진행 상황을 줄일 수 있습니다. 작동 방식과 구성 방법은 단원을 참조하십시오HyperPod 관리형 계층형 체크포인트.

설정

먼저 클러스터에서 관리형 계층형 체크포인트를 설정합니다. 클러스터 수준 기능이며 설정은 훈련하는 프레임워크와 동일합니다. 이러한 단계는 관리형 계층형 체크포인트 설정을 참조하세요.

그런 다음 Ray 워크로드를 실행하는 호스트에 체크포인트 패키지를 설치합니다.

pip install amzn-sagemaker-checkpointing

Ray Train과 함께 사용

Ray가 체크포인트를 업로드하도록 하는 SageMakerTieredStorageWriter 대신를 통해 체크포인트를 저장합니다. 훈련 함수 SageMakerCheckpointConfig 내에를 빌드하고 라이터를에 전달합니다async_save. 체크포인트가 Amazon S3에 업로드되는 동안 훈련이 계속되도록 체크포인트를 비동기적으로 보고합니다.

이 예제에서는 Ray 설명서의 비동기 체크포인트 업로드를 사용합니다. 이를 통해 Ray Train은 다음 단계에서 훈련을 계속하는 동안 업로드가 완료될 때까지 백그라운드 스레드를 시작할 수 있습니다. save_checkpoint 함수는 체크포인트를 계층형 스토리지로 스테이징하고를 사용하여 Ray Train에 보고합니다CheckpointUploadMode.ASYNC.는 최상의 체크포인트를 추적하고,를 적용하고num_to_keep, 실패 시 최신 체크포인트에서 복원할 수 있도록 Ray Train을 사용하여 체크포인트를 ray.train.report 기록합니다. Ray Train 체크포인트에 대한 자세한 내용은 Ray 설명서의 체크포인트 저장 및 로드를 참조하세요.

import os
import torch.distributed as dist
from torch.distributed.checkpoint import async_save, load

import ray
import ray.train
from ray.train import (
    Checkpoint, CheckpointConfig, CheckpointUploadMode,
    RunConfig, ScalingConfig, FailureConfig,
)
from ray.train.torch import TorchTrainer

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

S3_PATH = "s3://my-bucket/checkpoints"
EXPERIMENT_NAME = "my-experiment"


def create_checkpoint_config():
    """Create a checkpoint config scoped to this training job."""
    return SageMakerCheckpointConfig(
        namespace=EXPERIMENT_NAME,
        world_size=dist.get_world_size(),
        s3_tier_base_path=S3_PATH,
    )


# Track the previous checkpoint future to avoid async_save deadlock.
# Only one async_save can be in flight at a time because background
# threads perform collectives that require all ranks to participate.
_prev_checkpoint_future = None


def save_checkpoint(model, optimizer, config, step, metrics):
    """Save a checkpoint asynchronously and report it to Ray Train."""
    global _prev_checkpoint_future

    # Wait for the previous checkpoint to finish before starting a new one.
    if _prev_checkpoint_future is not None:
        _prev_checkpoint_future.result()

    config.save_to_s3 = True
    writer = SageMakerTieredStorageWriter(checkpoint_config=config, step=step)
    state_dict = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "step": step,
    }

    future = async_save(state_dict=state_dict, storage_writer=writer)
    _prev_checkpoint_future = future

    def wait_for_upload(checkpoint, name):
        future.result()
        return checkpoint

    ray.train.report(
        metrics=metrics,
        checkpoint=Checkpoint(writer.s3_checkpoint_dir),
        checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
        checkpoint_upload_fn=wait_for_upload,
        delete_local_checkpoint_after_upload=False,
    )


def load_checkpoint(model, optimizer, config):
    """Load the latest checkpoint if one exists. Returns the next step number."""
    reader = SageMakerTieredStorageReader(checkpoint_config=config)
    state_dict = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "step": 0,
    }
    try:
        load(state_dict, storage_reader=reader)
    except FileNotFoundError:
        # No checkpoint found, start from scratch.
        return 0

    model.load_state_dict(state_dict["model"])
    optimizer.load_state_dict(state_dict["optimizer"])
    return state_dict["step"] + 1


def train_func(config):
    device = ray.train.torch.get_device()
    model = build_model().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    ckpt_config = create_checkpoint_config()
    start_step = load_checkpoint(model, optimizer, ckpt_config)

    for step in range(start_step, config["max_steps"]):
        loss = train_step(model, optimizer)

        if step % config["checkpoint_freq"] == 0:
            save_checkpoint(
                model, optimizer, ckpt_config, step,
                metrics={"loss": loss, "step": step},
            )


trainer = TorchTrainer(
    train_func,
    train_loop_config={"max_steps": 1000, "checkpoint_freq": 10},
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
    run_config=RunConfig(
        name=EXPERIMENT_NAME,
        storage_path=S3_PATH,
        failure_config=FailureConfig(max_failures=3),
        checkpoint_config=CheckpointConfig(num_to_keep=3),
    ),
)

중요 사항: