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

개요

TPU 노드 및 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를 참조하세요.

유지보수 이벤트 처리

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

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

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

  2. 학습 스크립트 재시도

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

    프로덕션에 즉시 사용 가능한 학습 파이프라인은 GKE 또는 기타 리소스 관리 시스템을 사용해야 합니다. 그렇지 않으면 bash 스크립트를 사용하여 완료될 때까지 학습 스크립트를 계속 재시도할 수 있습니다. 예를 들면 다음과 같습니다.

    TPU 노드 사용:

    (VM에서) while ! python3 [training command]; do sleep 1; done

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

    (TPU VM에서가 아니라 GCP 또는 터미널에서 TPU VM 명령어를 실행해야 합니다.)

  3. TPU 실패 감지 및 복구

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