Google, TPU 장애에도 학습 이어가는 MaxText 탄력 학습 사례 공개

Google Developers Blog는 6일(현지시간) JAX 기반 LLM 학습 라이브러리 MaxText와 Pathways를 이용해 TPU 일부가 중단돼도 학습 작업을 전체 재시작 없이 복구하는 탄력 학습 사례를 공개했다. Google은 Cloud TPU와 Google Kubernetes Engine 환경에서 학습 중인 워커를 강제로 종료한 뒤, 단일 컨트롤러와 체크포인트를 활용해 같은 프로세스 안에서 학습을 이어가는 과정을 설명했다.

분산 학습에서는 일반적으로 한 노드가 사라지면 집단 통신이 실패하고 전체 워커가 종료된다. 이후 스케줄러가 작업을 다시 할당하고 마지막 체크포인트부터 재시작하는 방식이 흔하다. Google이 소개한 접근은 TPU 워커 장애를 프로세스 종료가 아니라 Python 예외로 받아 처리한다는 점이 다르다. CPU 쪽 컨트롤러 프로세스가 살아 있기 때문에 장애를 감지하고, 복구 가능한 체크포인트를 고른 뒤, 교체된 TPU 슬라이스와 다시 학습을 시작할 수 있다는 설명이다.

이번 사례의 핵심 구성 요소는 MaxText, Pathways, Orbax다. MaxText는 JAX 기반 LLM 학습 루프를 제공하고, Pathways는 하나의 Python 컨트롤러가 여러 TPU 칩을 다루도록 오케스트레이션한다. Orbax는 체크포인트 저장과 복구를 맡는다. Google은 Pathways가 장애를 감지하면 JAX 런타임 오류가 발생하고, MaxText에 연결된 `elastic_retry`가 이를 잡아 최신 정상 체크포인트를 복원하는 흐름을 제시했다.

실험은 3개의 TPU v5e-16 슬라이스, 총 48개 칩과 CPU 컨트롤러 노드로 구성됐다. 모델은 qwen3-0.6b, 플랫폼은 GKE, 워크로드 구성에는 JobSet이 쓰였다. Google은 워커 하나를 강제 종료한 실험에서 종료 시점부터 다음 학습 단계로 돌아오기까지 약 1분 50초가 걸렸다고 밝혔다. 마지막으로 정상 커밋된 체크포인트는 3300단계였고, 장애 감지 전 학습은 3388단계까지 진행돼 88단계가 되돌아갔다.

Google은 이 결과가 전체 작업을 다시 띄우는 방식과 달리 실패한 슬라이스만 교체한다는 점에서 의미가 있다고 설명했다. 다만 복구가 모든 비용을 없애는 것은 아니다. 교체 워커를 Kubernetes가 다시 배치하는 시간이 여전히 필요하고, 체크포인트 복원과 학습 함수 재진입 비용도 발생한다. 원문은 컴파일 캐시가 있어 전체 재시작과 탄력 복구 모두에서 컴파일 비용 차이는 제한적이며, 차이는 주로 전체 워크로드를 내렸다 다시 올리는 비용을 피하는 데 있다고 봤다.

대형 모델 학습으로 확장할 때의 한계도 함께 제시됐다. Google은 더 큰 모델에서는 체크포인트가 컨트롤러 프록시를 지나면서 메모리 병목이 생길 수 있다며, 각 TPU 호스트가 Cloud Storage의 체크포인트 조각을 직접 읽고 쓰도록 Pathways Persistence를 켜는 방식을 권했다. 앞으로는 Cloud Storage 체크포인트 대신 메모리 스냅샷을 이용해 되감는 단계와 복구 시간을 줄이는 방향도 언급했다.

이번 공개는 AI 모델 학습 비용과 인프라 안정성이 점점 중요한 경쟁 요소가 되는 상황에서 나왔다. 대규모 학습 작업은 실행 시간이 길고 하드웨어 규모가 커 장애 가능성을 완전히 배제하기 어렵다. 탄력 학습은 모델 성능 자체가 아니라 학습 운영의 회복력을 높이는 기술이라는 점에서, 클라우드 AI 인프라와 오픈소스 학습 스택의 실사용 경쟁력을 가르는 요소로 주목된다.

출처: Google Developers Blog https://developers.googleblog.com/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with-maxtext/