왜 체크포인팅 주기가 중요할까?

AI 모델 학습, 특히 대규모 언어 모델(LLM)을 학습시킬 때 체크포인트(checkpoint)는 생명줄과 같습니다. 학습 도중 하드웨어 장애나 프리엠션(preemption)이 발생하면 지금까지의 모든 계산이 날아갈 수 있기 때문이죠.

하지만 체크포인트 저장은 단순히 '자주 하면 좋다'가 아닙니다. 저장 주기가 너무 짧으면 I/O 대역폭을 과도하게 소모해 실제 학습 속도가 느려집니다. 반대로 너무 길면 장애 발생 시 손실되는 학습량(예: 100스텝 분량)이 커집니다. 이 전통적인 트레이드오프(trade-off)를 해결하기 위해 Google의 Orbax 팀이 제안한 것이 바로 연속 체크포인팅(Continuous Checkpointing) 입니다.

Google Cloud TPU cluster with continuous checkpointing enabled for AI model training Coding Session Visual

연속 체크포인팅의 핵심 동작 방식

기존 체크포인팅은 일정 스텝(예: 100스텝)마다 동기(synchronous) 방식으로 저장했습니다. 반면 연속 체크포인팅은 이전 저장이 완료되는 즉시 다음 저장을 비동기(asynchronous)로 시작합니다. Orbax가 내부적으로 저장 큐(queue)를 관리하여, 호스트 머신과 I/O 대역폭을 최대한 활용하면서도 성능 저하를 최소화합니다.

MaxText에서 활성화하는 방법

MaxText를 사용한다면 단 한 줄의 설정으로 연속 체크포인팅을 켤 수 있습니다.

# 학습 작업 설정 파일 (예: config.yaml)
enable_checkpointing: True          # 체크포인팅 활성화
async_checkpointing: True           # 비동기 체크포인팅 활성화
enable_continuous_checkpointing: True  # 연속 체크포인팅 활성화 (핵심)
max_num_checkpoints_to_keep: 10     # 최근 10개만 유지 (디스크 절약)

이렇게 설정하면 MaxText는 이전 체크포인트 저장 요청이 백그라운드에서 완료되는 즉시 새로운 저장을 시도합니다. 실제 벤치마크 결과를 보면, 연속 체크포인팅을 활성화했을 때 P50 체크포인트 간격이 현저히 줄어드는 것을 확인할 수 있습니다. 물론 평균 학습 스텝 시간은 소폭 증가하지만, 장애 발생 시 복구 비용을 고려하면 전체적인 Goodput(유효 처리량)은 크게 향상됩니다.

Orbax asynchronous checkpointing mechanism diagram showing storage server and slice communication Programming Illustration

Orbax의 고급 정책: 세밀한 제어

MaxText의 기본 설정 외에도 Orbax는 더 유연한 정책(policy)을 제공합니다. 특히 경량 모델을 학습할 때는 체크포인트가 너무 자주 발생해 불필요한 I/O 오버헤드가 생길 수 있습니다. 이때 minimum_interval_secs를 설정하면 쿨다운 시간을 둘 수 있습니다.

# 최소 30초 간격으로 체크포인트 저장 (과도한 저장 방지)
continuous_checkpointing_policy_with_minimum_interval = \
    save_decision_policy.ContinuousCheckpointingPolicy(minimum_interval_secs=30)

# 180초마다 최소 1개 체크포인트 유지 정책
every_n_seconds_preservation_policy = \
    preservation_policy.EveryNSeconds(180)

또한, 자신만의 보존 정책을 정의할 수도 있습니다.

@dataclasses.dataclass
class CustomizedPreservationPolicy(PreservationPolicy):
    """체크포인트 보존을 위한 사용자 정의 정책"""
    def should_preserve(
        self,
        checkpoints: Sequence[PolicyCheckpointInfo],
        *, context: PreservationContext,
    ) -> Sequence[bool]:
        result = [is_checkpoint_preservable(cp) for cp in checkpoints]
        _log_preservation_decision("사용자 정의 보존 정책", checkpoints, result)
        return result

주의사항 및 한계

연속 체크포인팅이 만능은 아닙니다. 특히 다음 두 가지를 반드시 확인해야 합니다.

  1. 스토리지 버킷과 학습 클러스터의 물리적 위치: 연속 체크포인팅은 네트워크 대역폭에 크게 의존합니다. 스토리지가 학습 클러스터와 동일한 데이터센터(co-located)에 있어야 효과적입니다. 크로스-메트로(cross-metro) 네트워크를 사용하면 체크포인팅 속도가 급격히 저하되어 오히려 안정성을 해칠 수 있습니다.

  2. 멀티-슬라이스 환경에서의 DCN 대역폭: 여러 슬라이스(slice)에 걸쳐 학습할 때, DCN(Data Center Network) 대역폭은 모델 가중치 업데이트와 체크포인팅에 모두 사용됩니다. 다행히 Orbax는 체크포인팅의 주요 부하를 단일 슬라이스(주로 슬라이스 0)와 스토리지 서버 간 통신으로 제한하여, 슬라이스 간 통신을 차단하지 않도록 설계되었습니다. 벤치마크 결과, 멀티-슬라이스 환경에서도 연속 체크포인팅으로 인한 유의미한 속도 저하는 관찰되지 않았습니다.

Benchmark graph comparing P50 checkpoint intervals with and without continuous checkpointing on v5p-128

실무 적용 조언: 체크포인트 전략 수립하기

연속 체크포인팅은 '설정하고 잊는(set-and-forget)' 솔루션에 가깝습니다. 하지만 완벽한 전략을 위해서는 다음을 고려하세요.

  • MTBF(Mean Time Between Failure) 추정: 자신의 클러스터에서 예상되는 평균 장애 간격을 파악하세요. 이 값에 따라 체크포인트 주기를 조정하면 최적의 Goodput을 얻을 수 있습니다.
  • 저장소 비용 vs. 복구 비용: 체크포인트를 너무 많이 유지하면 스토리지 비용이 증가합니다. max_num_checkpoints_to_keep을 적절히 설정하여 비용과 안정성의 균형을 맞추세요.
  • 모니터링: 체크포인트 저장 시간과 간격을 지속적으로 모니터링하여 네트워크나 스토리지 병목이 발생하는지 확인하세요.

이 기술은 특히 대규모 학습(수백~수천 개의 칩)에서 그 진가를 발휘합니다. 작은 클러스터에서는 기존 주기적 체크포인팅과 큰 차이가 없을 수 있지만, 스케일이 커질수록 연속 체크포인팅의 효율은 극대화됩니다.

함께 보면 좋은 글

다음 단계 학습 방향

연속 체크포인팅을 넘어, 더 고급 주제에 관심이 있다면 다음을 살펴보세요.

  1. Orbax의 커스텀 저장/복원 전략: 데이터 병렬, 파이프라인 병렬 등 다양한 병렬화 전략에 따른 체크포인트 최적화.
  2. 다양한 하드웨어(TPU v4, v5p, GPU)에서의 성능 비교: 각 환경에 맞는 최적의 설정 찾기.
  3. Fault-tolerant training의 최신 연구 동향: 체크포인트 없이 장애를 복구하는 기법(예: elastic training).

Google Cloud 공식 블로그 원문에서 더 자세한 벤치마크 데이터를 확인할 수 있습니다.

본 콘텐츠는 신뢰할 수 있는 출처를 바탕으로 AI 도구를 활용하여 초안이 작성되었으며, 편집자의 검토를 거쳐 발행되었습니다. 전문가의 조언을 대체하지 않습니다.