TPU elastic training 운영법: 분산 학습 장애를 전체 재시작 없이 복구하는 기준

요약: Google의 MaxText, Pathways, Orbax 기반 elastic training 사례는 분산 AI 학습에서 “한 노드 장애가 전체 작업을 죽이는 구조”를 바꾸려는 시도입니다. 대규모 학습팀이 봐야 할 핵심은 장애 복구를 스케줄러 바깥이 아니라 훈련 루프 안으로 가져오는 방식입니다.

분산 학습 장애가 비싼 이유

대규모 모델 학습은 여러 머신이 함께 움직입니다. 모델 파라미터와 gradient가 shard로 나뉘고, 매 step마다 collective communication이 일어납니다. 이 구조에서 worker 하나가 죽으면 나머지 worker는 필요한 데이터를 기다리다 timeout이 나고, 결국 전체 job이 실패합니다.

전통적인 복구 방식은 단순합니다. Slurm, Kubernetes, Ray 같은 스케줄러가 job 실패를 감지하고 전체를 다시 띄웁니다. 마지막 checkpoint에서 재시작하므로 일부 step은 날아갑니다. 새 pod를 잡고, 컨테이너를 띄우고, Python 프로세스를 시작하고, accelerator에 다시 연결하고, dataloader를 warm-up하는 비용도 다시 냅니다.

Google Developers Blog의 MaxText elastic training 글은 이 문제를 실험적으로 보여줍니다. Cloud TPU와 GKE 위에서 LLM 학습 중 TPU worker를 일부러 죽였고, 전체 job 재시작 없이 같은 프로세스 안에서 복구했습니다. 공개된 사례에서는 kill부터 다음 training step까지 총 downtime이 2분 미만이었고, 대부분은 Kubernetes가 replacement pod를 스케줄링하는 시간이었다고 설명합니다.

MaxText, Pathways, Orbax가 맡는 역할

이 구조를 이해하려면 세 컴포넌트를 나눠봐야 합니다.

MaxText는 JAX 기반 오픈소스 LLM 학습 라이브러리입니다. 모델 설정과 학습 루프를 제공합니다. Orbax는 checkpointing을 맡습니다. TPU host가 각자 model state shard를 Cloud Storage에 병렬로 저장하고, 실패 후 마지막 viable checkpoint로 돌아갈 수 있게 합니다.

가장 중요한 컴포넌트는 Pathways입니다. 일반적인 distributed launcher는 노드마다 Python 프로세스를 띄우고 서로 동등하게 조율합니다. Pathways 구조에서는 단일 Python controller 프로세스가 있고, TPU worker는 compiled program을 받아 실행하는 얇은 실행 계층에 가깝습니다. worker 하나가 죽어도 controller Python 프로세스는 살아 있습니다.

이 차이가 큽니다. 프로세스가 살아 있으면 실패를 Python exception으로 받을 수 있고, exception handler가 checkpoint restore와 재시도를 수행할 수 있습니다. Google 글에서는 pathways-utils의 elastic_retry decorator가 training function을 감싸고, 실패 시 partial state를 정리한 뒤 마지막 checkpoint를 복구해 같은 프로세스에서 다시 호출하는 구조를 설명합니다.

elastic recovery가 실제로 줄이는 비용

여기서 오해하면 안 되는 부분이 있습니다. elastic recovery가 모든 비용을 없애는 건 아닙니다. training function을 다시 들어가면 모델 setup, dataloader 준비, checkpoint restore는 다시 합니다. 죽은 worker를 대체할 pod도 여전히 필요합니다. 컴파일도 마법처럼 사라지지 않습니다. Pathways가 Cloud Storage에 persistent compilation cache를 사용하기 때문에 전체 재시작도 cold compile을 매번 하지 않을 수 있습니다.

그럼에도 elastic recovery가 줄이는 비용은 명확합니다. 전체 workload를 teardown하지 않습니다. controller pod, 살아 있는 worker pod, Python 프로세스, 주변 orchestration을 모두 다시 만드는 대신 죽은 slice만 교체합니다. 공개 사례에서 “수백 초와 몇 분” 차이가 난다고 표현한 이유가 여기에 있습니다.

또 하나의 이점은 replica resize입니다. pause-and-resume은 replacement가 올 때까지 기다렸다가 전체 mesh로 복구합니다. replica resize는 살아 있는 slice만으로 reduced throughput 학습을 계속하다가 replacement가 오면 다시 확장하는 방식입니다. 모든 상황에 맞지는 않지만, 긴 학습 job에서 throughput 손실과 완전 중단 사이의 선택지를 줍니다.

우리 팀이 바로 적용할 수 있는 설계 원칙

대부분의 팀은 내일 바로 Pathways 기반 TPU elastic training을 쓰지는 않을 수 있습니다. 그래도 운영 원칙은 그대로 가져올 수 있습니다.

첫째, checkpoint 주기를 비용 기준으로 정해야 합니다. “30분마다 저장” 같은 고정 규칙보다 한 step 비용, 장애 확률, checkpoint 저장 비용을 같이 봐야 합니다. checkpoint가 너무 드물면 손실 step이 커지고, 너무 잦으면 storage와 I/O가 병목이 됩니다.

둘째, controller와 worker의 책임을 분리해야 합니다. 모든 worker가 동일한 프로세스 생명주기를 공유하면 worker 하나의 장애가 전체 실패로 번집니다. 가능한 구조에서는 control plane을 오래 살아 있게 두고, execution worker를 교체 가능한 단위로 만들어야 합니다.

셋째, 장애를 “job failed” 이벤트로만 보지 말고 “복구 가능한 exception”으로 다뤄야 합니다. 훈련 루프가 실패 유형을 알고 있으면 pause, restore, resize, abort 중 하나를 선택할 수 있습니다.

넷째, 복구 시간을 메트릭으로 봐야 합니다. 학습 loss와 throughput만 보는 팀은 장애 복구 병목을 늦게 발견합니다. mean time to recover, checkpoint restore time, replacement scheduling time, lost steps를 별도 지표로 봐야 합니다.

도입 전에 조심할 점

elastic training은 공짜 안정성이 아닙니다. checkpoint consistency가 약하면 잘못된 상태로 복구할 수 있습니다. 데이터 로더가 deterministic하지 않으면 재시작 후 중복 샘플이나 누락 샘플이 생길 수 있습니다. optimizer state shard가 정확히 복원되지 않으면 학습 품질에 미묘한 영향을 줄 수 있습니다.

Kubernetes 스케줄링도 병목입니다. replacement pod가 빨리 올라오지 못하면 elastic recovery도 기다립니다. TPU slice 확보, node pool 여유, quota, priority class, preemption 정책을 같이 봐야 합니다. Spot이나 preemptible 리소스를 쓰는 경우에는 계획된 중단과 비계획 장애를 구분해야 합니다. Google 글도 suspend-resume과 elastic training은 다른 문제를 푸는 기능이라고 설명합니다.

바로 실행할 체크리스트

• checkpoint 주기를 step 비용과 장애 확률 기준으로 정했는가?
• controller와 worker 생명주기를 분리할 수 있는 구조인가?
• worker 장애를 전체 job 실패가 아니라 복구 가능한 exception으로 다루는가?
• checkpoint restore time과 lost steps를 메트릭으로 수집하는가?
• replacement pod 스케줄링 시간이 복구 병목인지 측정했는가?
• 데이터 로더와 optimizer state가 복구 후 일관성을 유지하는가?
• pause-and-resume과 reduced replica 학습 중 어떤 전략이 서비스 목표에 맞는가?

대규모 학습에서 장애는 예외가 아니라 운영 비용입니다. elastic training의 가치는 장애가 안 나게 하는 데 있지 않습니다. 장애가 나도 전체 시스템을 무너뜨리지 않는 복구 경로를 훈련 루프 안에 넣는 데 있습니다.