Checkpointless Training

디스크 I/O 병목을 완전히 제거하고, 메모리 기반 P2P 상태 복제로 패러다임을 전환한 AWS의 혁신적인 Fault Tolerance 기술을 심층 분석합니다.

1. 개요: 패러다임의 전환

1.1 전통적 Checkpoint-Restart

기존 방식은 주기적으로 모델 상태를 디스크(S3, FSx)에 저장하고, 장애 발생 시 전체 클러스터를 종료한 후 디스크에서 상태를 복구합니다.

1.2 Checkpointless: Always-On In-Memory Recovery

Checkpointless Training은 디스크 저장을 제거하고, 클러스터 메모리 내에서 P2P로 상태를 복제합니다.

2. Component 1: Rootless NCCL

2.1 기존 TCPStore의 문제점

표준 NCCL/Gloo 초기화는 중앙의 Rank 0 (Master) 노드가 TCPStore를 관리하여 프로세스를 동기화합니다.

# 기존 방식: Rank 0가 TCPStore 호스팅
import torch.distributed as dist

dist.init_process_group(
    backend="nccl",
    init_method="tcp://MASTER_ADDR:MASTER_PORT",  # SPOF!
    rank=rank,
    world_size=world_size
)

2.2 Rootless 솔루션: Decentralized Handshake

중앙 집중식 TCPStore를 제거하고, 각 Rank가 독립적으로 피어(Peer) 연결 정보를 계산하는 Symmetric Address Pattern을 도입했습니다.

# Rootless NCCL 활성화
export HPCT_USE_ROOTLESS=1

# TCPStore 우회
export TORCH_SKIP_TCPSTORE=1

# 포트 범위 확장 (다중 프로세스 통신용)
sysctl -w net.ipv4.ip_local_port_range="20000 65535"

3. Component 2: MMAP Data Loading

3.1 Linux mmap 원리

mmap()은 파일이나 메모리 영역을 프로세스의 가상 주소 공간에 매핑합니다. 여러 프로세스가 동일한 물리 메모리를 공유할 수 있어 메모리 효율이 극대화됩니다.

3.2 /dev/shm 공유 메모리

/dev/shm은 Linux의 tmpfs 기반 공유 메모리 영역입니다. RAM에 직접 데이터를 저장하므로 디스크 I/O 없이 초고속 접근이 가능합니다.

3.3 MMAPDataModule API

from hyperpod_checkpointless_training.dataloader.mmap_data_module import MMAPDataModule
from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig

# MMAP 기반 데이터 로더 설정
data_module = MMAPDataModule(
    data_module=MY_DATA_MODULE(...),
    mmap_config=CacheResumeMMAPConfig(
        cache_dir="/dev/shm/pdl_cache",    # 공유 메모리 경로
        prefetch_length=10,                # 미리 가져올 배치 수
        val_prefetch_length=10,            # 검증용 prefetch
        lookback_length=2,                 # 최근 소비된 배치 보관 수
        checkpoint_frequency=None,         # 캐시 체크포인트 주기
        enable_batch_encryption=False     # 배치 암호화 여부
    ),
)

3.4 CacheResumeMMAPConfig 파라미터

4. Component 3: In-Process Recovery (IPR)

4.1 기존 방식: 전체 Kill

기존에는 하나의 GPU에서 CUDA Error가 발생하면 전체 클러스터의 모든 프로세스를 강제로 종료하고 처음부터 다시 시작했습니다 (Job-level Restart).

4.2 IPR: 프로세스 격리

장애를 프로세스 레벨(Process-level)로 격리합니다. 정상 노드들은 프로세스를 종료하지 않고 기존의 CUDA Context와 컴파일러 캐시를 유지한 채 대기 상태(Wait state)에 진입합니다.

4.3 Recovery 유형

4.4 @HPWrapper Decorator

Re-executable Code Block (RCB)을 지정하는 데코레이터입니다. 장애 발생 시 해당 블록만 재실행됩니다.

from hyperpod_checkpointless_training.inprocess.wrap import HPWrapper
from hyperpod_checkpointless_training.inprocess.health_check import CudaHealthCheck
from hyperpod_checkpointless_training.inprocess.train_utils import HPAgentK8sAPIFactory
from hyperpod_checkpointless_training.inprocess.abort import (
    CheckpointlessAbortManager,
    CheckpointlessFinalizeCleanup
)

# In-Process Recovery Wrapper
@HPWrapper(
    health_check=CudaHealthCheck(),                    # CUDA 상태 검사
    hp_api_factory=HPAgentK8sAPIFactory(),            # K8s API 연동
    abort_timeout=60.0,                              # Abort 타임아웃
    abort=CheckpointlessAbortManager.get_default_checkpointless_abort(),
    finalize=CheckpointlessFinalizeCleanup(),
)
def run_main(cfg, caller=None):
    # 이 블록 전체가 Re-executable Code Block (RCB)
    trainer = Trainer(...)
    trainer.wrapper = caller  # wrapper 연결 필수
    trainer.fit(model, data_module)

4.5 CudaHealthCheck

CUDA 디바이스의 상태를 주기적으로 검사합니다. 기본 타임아웃은 30초입니다.

4.6 wait_rank() API

HyperPod Training Operator로부터 Rank 정보를 수신할 때까지 대기합니다.

from hyperpod_checkpointless_training.inprocess.train_utils import wait_rank

# Training Operator가 RANK, WORLD_SIZE 환경변수를 설정할 때까지 대기
wait_rank()

# 이후 분산 학습 초기화
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

5. Component 4: P2P State Replication

5.1 N+1 Redundancy / M-copy Scheme

모델의 Weights와 Optimizer State를 여러 노드 그룹에 중복 복제합니다. num_distributed_optimizer_instances >= 2 설정으로 활성화됩니다.

5.2 데이터 전송 경로

5.3 CheckpointlessMegatronStrategy

from hyperpod_checkpointless_training.nemo_plugins.strategy import CheckpointlessMegatronStrategy

trainer = Trainer(
    strategy=CheckpointlessMegatronStrategy(
        # N+1 Redundancy: 최소 2개 인스턴스 필요
        num_distributed_optimizer_instances=2,

        # 기타 Megatron 설정
        tensor_model_parallel_size=8,
        pipeline_model_parallel_size=4,
        context_parallel_size=1,
    ),
)

5.4 CheckpointManager API

from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import CheckpointManager

checkpoint_manager = CheckpointManager(
    enable_checksum=False,   # 모델 상태 체크섬 검증
    enable_offload=False,    # CPU 메모리로 오프로드
)

# 주요 메서드
checkpoint_manager.save_checkpoint(trainer)              # 메모리에 저장
checkpoint_manager.try_checkpointless_load(trainer)      # P2P 복구 시도
checkpoint_manager.checkpointless_recovery_feasible(trainer)  # 복구 가능성 검증
checkpoint_manager.store_rng_states()                    # RNG 상태 저장
checkpoint_manager.load_rng_states()                     # RNG 상태 복원

5.5 PEFTCheckpointManager

PEFT(Parameter-Efficient Fine-Tuning) 워크플로우용 확장 매니저입니다. Base 모델 가중치와 Adapter 파라미터를 분리하여 관리합니다.

from hyperpod_checkpointless_training.nemo_plugins.checkpoint_manager import PEFTCheckpointManager

# PEFT (LoRA 등) 학습용
peft_checkpoint_manager = PEFTCheckpointManager(
    enable_checksum=False,
    enable_offload=False,
)

6. Component 5: Training Operator

6.1 개요

HyperPod Training Operator는 모든 복구 과정을 조율하는 Kubernetes-native Operator입니다.

6.2 K8s CRD (Custom Resource Definition)

HyperPodPytorchJob CRD를 통해 학습 Job을 정의합니다.

apiVersion: sagemaker.amazonaws.com/v1
kind: HyperPodPytorchJob
metadata:
  name: checkpointless-training-job
  namespace: default
spec:
  nprocPerNode: "8"
  runPolicy:
    cleanPodPolicy: "None"
  replicaSpecs:
    - name: worker
      replicas: 16
      template:
        spec:
          containers:
            - name: training
              image: <hyperpod-checkpointless-container-image>
              env:
                - name: HPCT_USE_ROOTLESS
                  value: "1"
              command:
                - hyperpodrun
                - --nproc_per_node=8
                - --inprocess-restart   # IPR 활성화
                - training_script.py

6.3 Health Monitoring & Hot Spare Pool

Training Operator는 클러스터 상태를 지속적으로 모니터링하고, 미리 준비된 Hot Spare 노드를 즉시 투입합니다.

6.4 장애 복구 흐름

6.5 hyperpodrun 명령

# entrypoint.sh 예시
#!/bin/bash
set -ex

# Rootless NCCL 설정
export HPCT_USE_ROOTLESS=1
sysctl -w net.ipv4.ip_local_port_range="20000 65535"

# HyperPod Elastic Agent로 학습 실행
hyperpodrun \
    --nproc_per_node=8 \
    --inprocess-restart \
    --inprocess-timeout=120 \
    --shutdown-signal=SIGKILL \
    --shutdown-timeout=15 \
    training_script.py

7. Recovery Flow 전체

8. Managed Tiered Checkpointing

Checkpointless Training과 별개로, AWS는 Tiered Checkpointing 기능도 제공합니다. RDMA over EFA와 S3 통합을 통해 기존 Checkpoint I/O를 최적화합니다.

8.1 amzn-sagemaker-checkpointing 라이브러리

from amzn_sagemaker_checkpointing import (
    SageMakerTieredStorageWriter,
    SageMakerTieredStorageReader
)
import torch.distributed.checkpoint as dcp

# Tiered Storage Writer 설정
writer = SageMakerTieredStorageWriter(
    path="s3://my-bucket/checkpoints/",
    storage_path="/local/nvme/staging/",  # 로컬 NVMe 스테이징
    thread_count=32,                       # 병렬 업로드 스레드
)

# 분산 Checkpoint 저장
dcp.save(
    state_dict=state_dict,
    storage_writer=writer,
)

# Tiered Storage Reader 설정
reader = SageMakerTieredStorageReader(
    path="s3://my-bucket/checkpoints/step_1000/",
    storage_path="/local/nvme/staging/",
    thread_count=32,
)

# 분산 Checkpoint 로드
dcp.load(
    state_dict=state_dict,
    storage_reader=reader,
)

8.2 Tiered Storage 아키텍처