대규모 AI 모델 학습을 위한 SageMaker HyperPod 체크포인트 전략

[블로그 본문]

대규모 인공지능 모델 학습을 위한 효율적인 체크포인팅 전략
SageMaker HyperPod의 Managed Tiered Checkpointing 활용 가이드

도입부
초거대 AI 모델의 시대가 도래한 지금, 수천억 개 파라미터를 가진 모델 학습은 막대한 인프라 자원을 요구합니다. 모델 학습 중 장애가 발생했을 때 재시작을 방지하고 이전 상태를 복구할 수 있는 "체크포인팅(Checkpointing)" 전략은 필수가 되었습니다. 그러나 너무 자주 체크포인트를 저장하면 스토리지 비용과 시스템 부담이 커지고, 너무 드물게 저장하면 장애 시 며칠 간의 학습 데이터가 손실될 수 있습니다.

AWS는 이러한 딜레마를 해결하기 위해 Amazon SageMaker HyperPod에서 “Managed Tiered Checkpointing(계층형 체크포인트 관리)” 기능을 공개했습니다. 본 글에서는 이 기능이 어떻게 작동하는지, 실제 어떻게 활용하고 구성하는지를 실용적인 가이드와 함께 설명합니다.

본론
체크포인팅은 모델 학습 중간에 현재 상태(모델 파라미터, 옵티마이저 상태, 학습 메타데이터 등)를 저장하는 방식입니다. 장애 발생 시 저장된 체크포인트에서 학습을 재개함으로써 시간을 절약하고 비용을 절감할 수 있습니다.

하지만 Llama 3(70B)와 같은 대규모 모델은 체크포인트의 크기가 130GB에 달하며, 옵티마이저 상태까지 포함하면 약 521GB 수준입니다. DeepSeek-R1(671B)은 옵티마이저를 포함할 경우 무려 5TB에 이르기 때문에, 기존 스토리지 기반의 체크포인팅은 병목 현상을 초래하거나 높은 비용을 야기합니다.

이를 해결하기 위해 SageMaker HyperPod는 다음과 같은 레이어 구조의 계층형 체크포인트 시스템을 제공합니다:

1. 가장 빠른 RAM(CPU 메모리)에 우선 저장
2. RDMA를 통해 인접 노드로 비동기 복제
3. 주기적으로 Amazon S3와 같은 영구 스토리지에 저장

체크포인팅은 PyTorch DCP(Distributed Checkpointing) 라이브러리와 통합되어 자동화와 병렬 처리 효율을 극대화합니다. 단 몇 줄의 코드만 추가하면 체크포인팅 전략을 구현할 수 있습니다.

예시 코드:

pip install amzn-sagemaker-checkpointing torch boto3 s3torchconnector

import torch
from torch.distributed.checkpoint import async_save, load
from amzn_sagemaker_checkpointing.checkpointing.filesystem.filesystem import SageMakerTieredStorageWriter

state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
checkpoint_config = SageMakerCheckpointConfig(namespace="train-run", world_size=dist.get_world_size())
storage_writer = SageMakerTieredStorageWriter(checkpoint_config=checkpoint_config, step=100)
async_save(state_dict=state_dict, storage_writer=storage_writer)

이처럼 학습 중 GPU에서 CPU로 체크포인트 데이터를 복사하고, RAM에 저장 후 Amazon S3로 백업하는 자동화 절차를 통해 복구는 빠르게, 스토리지 비용은 최소화됩니다.

또한 IAM 권한 정책을 설정하면 SageMaker HyperPod 클러스터가 S3 버킷에 자유롭게 접근할 수 있으며, AWS CLI를 통해 간편하게 클러스터를 생성하거나 업데이트할 수 있습니다:

aws sagemaker create-cluster –cluster-name "ml-cluster" –tiered-storage-config '{ "Mode": "Enable" }'

운영 후기
AWS 내부 테스트 결과, 수백 개에서 1만 5천 개 이상의 GPU가 연결된 분산 클러스터에서도 수 초 내에 체크포인트가 저장되며, 성능 저하 없이 안정적이고 빠르게 학습을 복원할 수 있음이 확인되었습니다.

또한 로그 파일(/var/log/sagemaker_checkpointing/)이나 콘솔 출력으로 체크포인트 저장 상태, 스텝 번호, 네트워크 대역폭 사용량 등을 손쉽게 추적하고 문제를 해결할 수 있습니다.

결론
Amazon SageMaker HyperPod의 Managed Tiered Checkpointing 기능은 고성능 모델 학습 환경에서의 자동화, 활용성, 비용 절감이라는 측면에서 강력한 도구입니다. CPU RAM, RDMA, S3 등 다양한 계층형 스토리지 전략을 통해 더 나은 복원 전략을 구현하고, 병목없이 학습 속도를 유지하며 개발 생산성을 높일 수 있습니다.

모델 크기나 환경에 따라 체크포인팅 주기를 유연하게 설정하고 코드를 간편하게 통합함으로써, 여러분의 인공지능 프로젝트에 최적화된 자동화된 체크포인트 저장 전략을 적용할 수 있습니다.

더 많은 기술 세부 사항은 아래 AWS 공식 문서를 확인해 보세요.

[1] https://aws.amazon.com/blogs/machine-learning/accelerate-your-model-training-with-managed-tiered-checkpointing-on-amazon-sagemaker-hyperpod/