MaxText를 사용한 탄력적 훈련으로 TPU 장애를 몇 초 만에 복구

Google Cloud는 MaxText, Pathways, Orbax를 사용하여 Cloud TPU에서 대규모 모델 워크로드를 위한 탄력적 학습 시연을 설명합니다. 이 테스트에서는 Google Kubernetes Engine에서 학습 중 TPU 작업자 하나가 강제로 종료되었고, 전체 워크로드를 다시 시작하지 않고도 작업이 복구되었습니다. 이 설정은 하나의 CPU 컨트롤러 프로세스와 여러 개의 TPU 작업자 포드를 사용했으며, Pathways가 실패를 감지하고 처리 가능한 Python 예외를 발생시켰습니다. 그런 다음 MaxText 재시도 래퍼가 Cloud Storage에서 최신 유효 체크포인트를 복원하고, Kubernetes가 대체 포드를 스케줄한 뒤 학습을 재개했습니다. 이 글은 또한 이 접근 방식이 전체 재시작과 어떻게 다른지, 체크포인팅과 복구가 어떻게 작동하는지, 그리고 더 큰 모델로 확장할 때 무엇을 변경해야 하는지 설명합니다. 또한 Pathways Persistence를 통한 직접 스토리지 액세스를 사용하면 복원 중 컨트롤러 메모리 병목을 피할 수 있다고 언급합니다.