Cloud TPU에서 소프트웨어 버전 전환

개요

TPU에서 실행 중인 프레임워크의 소프트웨어 버전이 로컬 VM에서 실행 중인 버전과 일치해야 합니다. 이제 TPU를 삭제했다가 다시 만들지 않고도 실행 중인 Cloud TPU에서 이 소프트웨어 버전을 전환할 수 있습니다. 또한 특정 나이틀리 버전의 소프트웨어 프레임워크로 Cloud TPU를 구성할 수 있지만 이러한 프레임워크의 지원되는 버전을 선택하는 것이 좋습니다.

사용

버전을 전환하려면 cloud-tpu-client Python 라이브러리를 사용하는 것이 좋습니다.

TensorFlow 사용 예시

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

이렇게 하면 로컬 VM에서 실행되는 TensorFlow 버전과 일치하도록 Cloud TPU가 구성됩니다. 여기에는 공식 출시 버전과 날짜가 지정된 나이틀리 빌드도 포함됩니다.

configure_tpu_version API의 restart_type 매개변수는 버전을 전환할 때 TPU 재시작 동작을 정의합니다. 옵션은 'always'(기본값)와 'ifNeeded'입니다.

  • 'always'를 사용하여 UNHEALTHY_TENSORFLOW 상태가 발생하거나 이전 실행의 리소스 유출로 인해 메모리 부족(OOM) 오류가 반환되는 등의 TPU 문제를 해결할 수 있습니다. 이 옵션을 설정하면 새 프레임워크 버전이 설치되지 않더라도 TPU가 다시 시작됩니다.

  • 'ifneeded'를 사용하면 런타임 버전이 이미 올바른 경우 런타임을 다시 시작하지 않으므로 학습 스크립트에 상당한 시작 시간이 추가되지 않아서 유용할 수 있습니다. 이 옵션을 설정하면 TPU가 올바른 프레임워크 버전이 설치되지 않은 경우에만 다시 시작됩니다.

라이브러리가 Cloud TPU와 직접 통신하므로 이 코드를 동일한 네트워크의 VM에서 실행해야 합니다. 모델의 나머지 부분을 위해 코드 내에서 이 값을 실행하는 것이 좋습니다.

추가 소프트웨어 옵션

TensorFlow에는 올바른 버전을 구성하기 위한 가장 간단한 방법인 tf.__version__ 문자열이 포함됩니다. 기타 소프트웨어 옵션은 다음과 같습니다.

  • PyTorch - pytorch-1.13, pytorch-nightly-dev20220930, pytorch-nightly
  • Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

예를 들어 PyTorch의 최신 나이틀리 빌드로 실행되도록 TPU를 구성합니다.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()