TPU Pod 슬라이스에서 TensorFlow 코드 실행

이 문서에서는 TPU Pod에서 TensorFlow를 사용하여 계산을 수행하는 방법을 보여줍니다. 다음 단계를 수행합니다.

  1. TensorFlow 소프트웨어로 TPU Pod 슬라이스 만들기
  2. SSH를 사용하여 TPU VM에 연결
  3. 예시 스크립트 만들기 및 실행

TPU VM은 Cloud TPU API를 호출할 수 있는 권한에 서비스 계정을 사용합니다. 기본적으로 TPU VM에서 필요한 모든 Cloud TPU 권한이 있는 기본 Compute Engine 서비스 계정을 사용합니다. 자체 서비스 계정을 사용하는 경우 서비스 계정에 TPU 뷰어 역할을 추가해야 합니다. Google Cloud 역할에 대한 자세한 내용은 역할 이해를 참조하세요. TPU VM을 만들 때 --service-account 플래그를 사용하여 자체 서비스 계정을 지정할 수 있습니다.

환경 설정하기

  1. Cloud Shell에서 다음 명령어를 실행하여 gcloud의 최신 버전을 실행하세요.

    $ gcloud components update

    gcloud를 설치해야 하는 경우 다음 명령어를 사용합니다.

    $ sudo apt install -y google-cloud-sdk

  2. 다음과 같이 몇 가지 환경 변수를 만듭니다.

    $ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

TensorFlow 런타임으로 v3-32 TPU Pod 슬라이스 만들기

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

명령어 플래그 설명

zone
Cloud TPU를 만들려는 영역입니다.
accelerator-type
가속기 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. 각 TPU 버전에서 지원되는 가속기 유형에 대한 자세한 내용은 TPU 버전을 참조하세요.
version
Cloud TPU 소프트웨어 버전입니다.

SSH를 사용하여 Cloud TPU VM에 연결

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

예시 스크립트 만들기 및 실행

  1. 다음 환경 변수를 설정합니다.

    (vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

  2. 현재 디렉터리에 tpu-test.py이라는 파일을 만들고 이 파일에 다음 스크립트를 복사하여 붙여넣습니다.

    import tensorflow as tf
print("Tensorflow version " + tf.__version__)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

@tf.function
def add_fn(x,y):
  z = x + y
  return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

  3. 다음 명령어를 사용하여 이 스크립트를 실행합니다.

    (vm)$ python3 tpu-test.py

    이 스크립트는 TPU Pod 슬라이스의 각 TensorCore에서 계산을 수행합니다. 출력은 다음과 비슷합니다.

    PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32),
  8: tf.Tensor(2.0, shape=(), dtype=float32),
  9: tf.Tensor(2.0, shape=(), dtype=float32),
  10: tf.Tensor(2.0, shape=(), dtype=float32),
  11: tf.Tensor(2.0, shape=(), dtype=float32),
  12: tf.Tensor(2.0, shape=(), dtype=float32),
  13: tf.Tensor(2.0, shape=(), dtype=float32),
  14: tf.Tensor(2.0, shape=(), dtype=float32),
  15: tf.Tensor(2.0, shape=(), dtype=float32),
  16: tf.Tensor(2.0, shape=(), dtype=float32),
  17: tf.Tensor(2.0, shape=(), dtype=float32),
  18: tf.Tensor(2.0, shape=(), dtype=float32),
  19: tf.Tensor(2.0, shape=(), dtype=float32),
  20: tf.Tensor(2.0, shape=(), dtype=float32),
  21: tf.Tensor(2.0, shape=(), dtype=float32),
  22: tf.Tensor(2.0, shape=(), dtype=float32),
  23: tf.Tensor(2.0, shape=(), dtype=float32),
  24: tf.Tensor(2.0, shape=(), dtype=float32),
  25: tf.Tensor(2.0, shape=(), dtype=float32),
  26: tf.Tensor(2.0, shape=(), dtype=float32),
  27: tf.Tensor(2.0, shape=(), dtype=float32),
  28: tf.Tensor(2.0, shape=(), dtype=float32),
  29: tf.Tensor(2.0, shape=(), dtype=float32),
  30: tf.Tensor(2.0, shape=(), dtype=float32),
  31: tf.Tensor(2.0, shape=(), dtype=float32)
}

삭제

TPU VM 사용이 완료되었으면 다음 단계에 따라 리소스를 삭제하세요.

  1. Compute Engine에서 연결을 해제합니다.

    (vm)$ exit

  2. Cloud TPU를 삭제합니다.

    $ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --zone=${ZONE}

  3. 다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. TPU가 더 이상 나열되지 않았는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다.

    $ gcloud compute tpus tpu-vm list \
  --zone=${ZONE}