Cloud TPU Pod로 유지보수 이벤트 관리

개요

TPU VM은 TPU 하드웨어가 연결된 Compute Engine VM의 인스턴스입니다. Compute Engine VM은 Compute Engine VM 유지보수 이벤트를 따릅니다. 각 TPU가 Compute Engine VM에 연결되므로 TPU Pod에서와 같이 TPU를 더 많이 사용하면 VM 중 하나에서 유지보수 이벤트가 발생할 가능성이 증가합니다.

이 문서에서는 Cloud TPU에서 장기 실행 학습 작업의 유지보수 이벤트를 처리하는 다양한 방법을 설명합니다.

체크포인트를 사용하여 유지보수 이벤트에서 빠르게 복구

체크포인트는 유지보수 이벤트를 짧게 복구하기 위한 핵심으로, 자주 저장해야 합니다. 약 1시간마다 체크포인트를 저장하는 것이 좋습니다. 체크포인트가 없는 경우 유지보수 이벤트나 기타 학습 중단으로 인해 학습 과정이 상당수 손실될 수 있습니다.

체크포인트는 일반적으로 학습에 사용된 모든 저장된 매개변수(예: 모델 가중치)를 나타냅니다. 체크포인트를 저장하는 데 걸리는 시간은 초 단위부터 분 단위까지 다양합니다.

대부분의 유지보수 이벤트는 자동으로 복구되고 학습 작업은 직접 개입하지 않고도 계속되지만 작업이 다시 시작하지 않고 자동으로 계속되는 극단적인 경우가 있을 수 있습니다. 이 경우 TPU 리소스를 삭제하여 다시 만들고 저장된 체크포인트에서 학습 작업을 다시 시작해야 합니다. 자동 복구 실패를 감지하고 복구하는 방법에 대한 자세한 내용은 TPU 실패 감지 및 복구를 참조하세요.

체크포인트 저장 및 로드에 사용되는 메커니즘은 각 ML 프레임워크마다 다릅니다. 지원되는 Cloud TPU 모델은 일반적으로 체크포인트가 내장되어 있습니다. 체크포인트에 대한 자세한 내용은 TensorFlow 2.x, PyTorch, JAX/flax를 참조하세요.

유지보수 이벤트 감지

다음 gcloud describe 명령어를 사용해서 TPU에서 유지보수 이벤트가 발생하는지 여부 및 시간을 감지할 수 있습니다.

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

이 명령어의 출력에는 TPU의 현재 상태와 최근 유지보수 이벤트에 대한 설명이 표시됩니다. 출력은 다음과 비슷하게 표시됩니다.

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

유지보수 이벤트 로그

TPU의 유지보수 이벤트 로그 기록은 시스템 이벤트 감사 로그에서 확인할 수 있습니다.

Google Cloud 콘솔 탐색 메뉴에서 Compute Engine > VM 인스턴스를 클릭하고 다음 예시를 검색합니다.

"tpu.nodes.terminate" OR "tpu.nodes.restart"

검색 기간 내에 TPU 워커의 중단 및 수리가 표시됩니다. 로그에는 이벤트 날짜 및 시간, 이벤트 유형, '종료' 이벤트의 경우, protoPayload.metadata.terminateReason에 대한 종료 이유가 표시됩니다.

유지보수 이벤트 처리

유지보수 이벤트 중단을 해결하는 방법은 몇 가지가 있습니다.

  1. 정기적으로 체크포인트 저장

    이상적인 시나리오에서는 '중단 이벤트'가 발생하면 최신 체크포인트에서 학습이 다시 시작합니다.

  2. 학습 스크립트 재시도

    학습 스크립트는 '중단 이벤트'의 결과로 중지될 수 있습니다. bash 스크립트를 사용하여 학습이 완료될 때까지 학습 스크립트를 계속 재시도할 수 있습니다. 각 재시도는 최근 체크포인트부터 계속되므로, 재시도 스크립트는 항상 체크포인트와 함께 사용되어야 합니다.

    프로덕션에 즉시 사용 가능한 학습 파이프라인은 Google Kubernetes Engine(GKE)과 같은 리소스 관리 시스템을 사용해야 합니다. TPU VM과 함께 Google Kubernetes Engine을 사용하는 방법에 대한 자세한 내용은 TPU 워크로드 배포를 참조하세요. bash 스크립트를 사용하여 완료될 때까지 학습 스크립트를 계속 재시도할 수 있습니다. 예를 들면 다음과 같습니다.

    while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 [training command]"; do sleep 1; done

  3. TPU 실패 감지 및 복구

    TPU가 유지보수 이벤트에서 복구되지 않는 경우 복구 스크립트를 사용하여 TPU 상태를 감지하고 TPU를 삭제하고 다시 만들 수 있습니다. 이 스크립트 예시는 여기에서 확인할 수 있습니다. 수동으로 TPU 삭제하고 다시 만드는 방법에 대한 자세한 내용은 TPU 관리를 참조하세요.

    TPU VM을 만들거나 다시 만들 경우 --metadata startup-script 매개변수로 시작 스크립트를 지정할 수 있습니다. 시작 스크립트는 TPU VM이 생성될 때마다 실행됩니다. 자세한 내용은 표준 설치 스크립트 실행을 참조하세요.