Tensor Processing Unit(TPU)는 Google에서 커스텀식으로 개발한 ASIC로, 머신러닝 워크로드를 빠르게 처리하는 데 사용됩니다. Cloud TPU를 사용하여 AI Platform Training에서 학습 작업을 실행할 수 있습니다. AI Platform Training은 TPU를 직접 관리할 필요가 없도록 작업 관리 인터페이스를 제공합니다.
CPU 또는 GPU 학습에 사용하는 방식과 동일하게 AI Platform Training jobs
API를 사용하면 됩니다.
고급 TensorFlow API를 사용하면 Cloud TPU 하드웨어에서 모델을 실행할 수 있습니다.
Google Cloud 환경 설정
시작 가이드의 설정 섹션을 참조하여 Google Cloud 환경을 구성합니다.
프로젝트에 액세스하도록 Cloud TPU 승인
다음 단계를 수행하여 Google Cloud 프로젝트와 연결된 Cloud TPU 서비스 계정 이름을 승인합니다.
projects.getConfig
를 호출하여 Cloud TPU 서비스 계정 이름을 가져옵니다. 예를 들면 다음과 같습니다.PROJECT_ID=PROJECT_ID curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
API에서 반환한
serviceAccountProject
및tpuServiceAccount
필드 값을 저장합니다.Cloud TPU 서비스 계정을 초기화합니다.
curl -H "Authorization: Bearer $(gcloud auth print-access-token)" \ -H "Content-Type: application/json" -d '{}' \ https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
이제 Cloud TPU 서비스 계정을 Cloud ML 서비스 에이전트 역할이 있는 구성원으로 프로젝트에 추가합니다. Google Cloud Console에서 또는 gcloud
명령어를 사용하여 다음 단계를 완료합니다.
Console
- Google Cloud Console에 로그인하고 TPU를 사용 중인 프로젝트를 선택합니다.
- IAM 및 관리자 > IAM을 선택합니다.
- 추가 버튼을 클릭하여 프로젝트에 멤버를 추가합니다.
- 멤버 입력란에 TPU 서비스 계정을 입력합니다.
- 역할 드롭다운 목록을 클릭합니다.
- Cloud ML 서비스 에이전트 역할을 사용 설정합니다(서비스 에이전트 > Cloud ML 서비스 에이전트).
gcloud
프로젝트 ID 및 Cloud TPU 서비스 계정이 포함된 환경 변수를 설정합니다.
PROJECT_ID=PROJECT_ID SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
Cloud TPU 서비스 계정에
ml.serviceAgent
역할을 부여합니다.gcloud projects add-iam-policy-binding $PROJECT_ID \ --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
서비스 계정에 역할을 부여하는 방법에 대한 자세한 내용은 IAM 문서를 참조하세요.
예: 샘플 MNIST 모델 학습
이 섹션에서는 TPU 및 런타임 버전 2.11을 사용하여 샘플 MNIST 모델을 학습시키는 방법을 설명합니다. 예시 작업은 사용 중인 머신 구성에 적합한 사전 정의된 BASIC_TPU
확장 등급을 사용합니다.
이 가이드의 이후 섹션에서는 커스텀 구성을 설정하는 방법을 보여줍니다.
이 예시에서는 gcloud CLI가 설치된 Bash 셸을 사용한다고 가정합니다. 다음 명령어를 실행하여 코드를 가져오고 AI Platform Training에 학습 작업을 제출합니다.
TensorFlow의 참조 모델 코드를 다운로드하고 샘플 코드가 있는 디렉터리로 이동합니다.
git clone https://github.com/tensorflow/models.git \ --branch=v2.11.0 \ --depth=1 cd models
models
디렉터리에setup.py
파일을 만듭니다. 이렇게 하면gcloud ai-platform jobs submit training
명령어가 학습 코드의 tarball을 만들 때models/official
디렉터리에서 필요한 모든 하위 패키지를 포함하고, AI Platform Training이 학습 작업을 실행할 때 종속 항목으로 TensorFlow 데이터 세트를 설치합니다. 이 학습 코드는 TensorFlow 데이터 세트를 사용하여 MNIST 데이터를 로드합니다.setup.py
파일을 만들려면 셸에서 다음 명령어를 실행합니다.cat << END > setup.py from setuptools import find_packages from setuptools import setup setup( name='official', install_requires=[ 'tensorflow-datasets~=3.1', 'tensorflow-model-optimization>=0.4.1' ], packages=find_packages() ) END
gcloud ai-platform jobs submit training
명령어를 사용하여 학습 작업을 제출합니다.gcloud ai-platform jobs submit training tpu_mnist_1 \ --staging-bucket=gs://BUCKET_NAME \ --package-path=official \ --module-name=official.vision.image_classification.mnist_main \ --runtime-version=2.11 \ --python-version=3.7 \ --scale-tier=BASIC_TPU \ --region=us-central1 \ -- \ --distribution_strategy=tpu \ --data_dir=gs://tfds-data/datasets \ --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
BUCKET_NAME을 Google Cloud 프로젝트의 Cloud Storage 버킷 이름으로 바꿉니다. gcloud CLI는 패키징된 학습 코드를 이 버킷에 업로드하고 AI Platform Training은 버킷에 학습 출력을 저장합니다.
학습 작업을 모니터링합니다. 작업이 완료되면
gs://BUCKET_NAME/tpu_mnist_1_output
디렉터리에서 출력을 볼 수 있습니다.
Cloud TPU의 모델 학습에 대한 자세한 내용
이 섹션에서는 Cloud TPU를 사용하여 AI Platform Training에서 작업을 구성하고 모델을 학습시키는 방법에 대해 자세히 설명합니다.
TPU를 제공하는 리전 지정
TPU가 제공되는 리전에서 작업을 실행해야 합니다. 다음 리전은 현재 TPU에 대한 액세스 권한을 제공합니다.
us-central1
europe-west4
모델 학습 및 온라인/일괄 예측을 포함하여 AI Platform Training 서비스에 사용할 수 있는 리전을 충분히 이해하려면 리전 가이드를 읽어보세요.
TensorFlow 및 AI Platform Training 버전 관리
Cloud TPU에서는 모델을 학습시키는 데 AI Platform Training 런타임 버전 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 및 2.11을 사용할 수 있습니다. AI Platform Training 런타임 버전 및 해당 TensorFlow 버전에 대해 자세히 알아보세요.
버전 관리 정책은 Cloud TPU의 경우와 동일합니다. 학습 작업 요청에서 TPU에 사용 가능하고 학습 코드에서 사용되는 TensorFlow 버전과 일치하는 런타임 버전을 지정해야 합니다.
TPU gRPC 서버에 연결
TensorFlow 프로그램에서 TPUClusterResolver
를 사용하여 TPU VM에서 실행되는 TPU gRPC 서버에 연결합니다.
TPU 사용에 대한 TensorFlow 가이드에서는 TPUStrategy
배포 전략에 TPUClusterResolver
를 사용하는 방법을 보여줍니다.
그러나 AI Platform Training에서 실행되는 코드에 TPUClusterResolver
를 사용할 때는 한 가지 중요한 변경 사항을 적용해야 합니다. TPUClusterResolver
인스턴스를 생성할 때는 인수를 제공하면 안 됩니다. tpu
, zone
, project
키워드 인수가 모두 기본값인 None
로 설정되면 AI Platform Training에서 환경 변수를 통해 필요한 연결 세부정보를 클러스터 리졸버에 자동으로 제공합니다.
다음 TensorFlow 2 예시는 AI Platform Training에서 학습하기 위한 클러스터 리졸버 및 배포 전략을 초기화하는 방법을 보여줍니다.
import tensorflow as tf
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)
TensorFlow 코드에서 TPU 사용
머신에서 TPU를 사용하려면 TensorFlow 2의 TPUStrategy
API를 사용합니다. TPU 사용에 대한 TensorFlow 가이드에서 사용하는 방법을 보여줍니다.
TensorFlow 1에서 TPU로 학습시키려면 TPUEstimator
API를 대신 사용할 수 있습니다. TPUEstimator
API에 대한 Cloud TPU 가이드에서 사용하는 방법을 보여줍니다.
Cloud TPU 문서에서는 Cloud TPU에서 사용할 수 있는 하위 수준의 TensorFlow 작업 목록도 제공합니다.
PyTorch 코드에서 TPU 사용
사전 빌드된 PyTorch 컨테이너를 사용할 때 TPU를 사용하려면 torch_xla
패키지를 사용합니다. PyTorch 문서의 학습에서 TPU에 torch_xla
를 사용하는 방법 알아보기 torch_xla
의 사용 예시를 더 확인하려면 PyTorch XLA GitHub 저장소의 가이드를 참조하세요.
AI Platform Training에서 TPU를 사용하여 학습시키는 경우 여러 XLA 기기가 아니라 단일 XLA 기기를 사용하게 됩니다.
이 페이지의 다음 섹션에서 PyTorch 및 TPU에 대한 학습 작업 구성도 참고하세요.
커스텀 TPU 머신 구성
TPU 학습 작업은 두 개의 VM 구성에서 실행됩니다. 한 개의 VM, 즉 마스터는 Python 코드를 실행합니다. 마스터는 TPU 워커에서 실행 중인 TensorFlow 서버를 구동합니다.
AI Platform Training에서 TPU를 사용하려면 다음 세 가지 방법 중 하나를 사용하여 TPU 지원 머신에 액세스할 수 있도록 학습 작업을 구성합니다.
BASIC_TPU
확장 등급을 사용합니다. 이 메서드를 사용해 TPU v2 가속기에 액세스할 수 있습니다.- 마스터 VM에
cloud_tpu
작업자 및 기존 머신 유형을 사용합니다. 이 메서드를 사용해 TPU v2 가속기에 액세스할 수 있습니다. - 마스터 VM에
cloud_tpu
작업자 및 Compute Engine 머신 유형을 사용합니다. 이 메서드를 사용해 TPU v2 또는 TPU v3 가속기에 액세스할 수 있습니다. TPU v3 가속기는 베타 버전으로 제공됩니다.
기본 TPU 지원 머신
이전 예시를 실행할 때와 마찬가지로 확장 등급을 BASIC_TPU
로 설정하여 마스터 VM과 TPU v2 코어가 8개인 TPU 1개가 포함된 TPU VM을 확보합니다.
기존 머신 유형 구성의 TPU 워커
마스터 VM에 더 많은 컴퓨팅 리소스가 필요한 경우에는 커스텀 머신 구성을 설정할 수도 있습니다.
- 확장 등급을
CUSTOM
으로 설정합니다. - 작업 요구사항에 맞는 기존 머신 유형을 사용하도록 마스터 VM을 구성합니다.
workerType
을cloud_tpu
로 설정하여 TPU v2 코어가 8개인 Cloud TPU 한 개가 포함된 TPU VM을 확보합니다.workerCount
를 1로 설정합니다.- Cloud TPU를 사용할 때는 매개변수 서버를 지정해서는 안됩니다.
parameterServerCount
가 0보다 크면 서비스가 작업 요청을 거부합니다.
다음 예시는 이 구성 유형을 사용하는 config.yaml
파일을 보여줍니다.
trainingInput:
scaleTier: CUSTOM
masterType: complex_model_m
workerType: cloud_tpu
workerCount: 1
Compute Engine 머신 유형 구성의 TPU 워커
마스터 VM의 Compute Engine 머신 유형과 TPU VM에 연결된 acceleratorConfig
를 사용하여 커스텀 머신 구성을 설정할 수도 있습니다.
이 구성 유형을 사용하여 TPU v2 코어가 8개인 TPU 워커(acceleratorConfig
가 없는 구성과 유사함) 또는 TPU v3 코어가 8개인 TPU 워커(베타)를 설정할 수 있습니다. TPU v2와 TPU v3 가속기의 차이점에 대해 자세히 알아보세요.
Compute Engine 머신 유형을 사용하면 다음과 같이 보다 유연하게 마스터 VM을 구성할 수 있습니다.
- 확장 등급을
CUSTOM
으로 설정합니다. - 작업 요구사항에 맞는 Compute Engine 머신 유형을 사용하도록 마스터 VM을 구성합니다.
workerType
를cloud_tpu
로 설정합니다.acceleratorConfig
필드가 있는workerConfig
를 추가합니다. 해당acceleratorConfig
안에서type
을TPU_V2
또는TPU_V3
로 설정하거나count
를8
로 설정합니다. 다른 수의 TPU 코어는 연결할 수 없습니다.workerCount
를 1로 설정합니다.- Cloud TPU를 사용할 때는 매개변수 서버를 지정해서는 안됩니다.
parameterServerCount
가 0보다 크면 서비스가 작업 요청을 거부합니다.
다음 예시는 이 구성 유형을 사용하는 config.yaml
파일을 보여줍니다.
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3(베타)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3
count: 8
TPU Pod 사용
TPU Pod는 전용 고속 네트워크 인터페이스로 연결되는 TPU 기기 모음입니다. TPU Pod는 최대 2,048개의 TPU 코어를 포함할 수 있으며, 여러 TPU 간에 처리 부하를 분산할 수 있습니다.
TPU Pod를 사용하려면 먼저 할당량 상향 요청을 제출해야 합니다.
다음 예시 config.yaml
파일에서는 TPU Pod를 사용하는 방법을 보여줍니다.
TPU v2 Pod
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V2_POD
count: 128
TPU v3 Pod
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
workerType: cloud_tpu
workerCount: 1
workerConfig:
acceleratorConfig:
type: TPU_V3_POD
count: 32
각 TPU 유형에 사용할 수 있는 Pod 코어 수에는 제한이 있습니다. 사용 가능한 구성:
TPU Pod 유형 | 사용 가능한 Pod 코어 수 |
---|---|
TPU_V2_POD |
32, 128, 256, 512 |
TPU_V3_POD |
32, 128, 256 |
TPU Pod 코어를 최대한 활용하는 방법에 대한 자세한 내용은 TPU Pod에 대한 Cloud TPU 문서를 참조하세요.
TPU 워커에서 사전 빌드된 PyTorch 컨테이너 사용
TPU를 사용하여 PyTorch 학습을 수행하려면 학습 작업의 trainingInput
에 tpuTfVersion
필드를 지정해야 합니다.
학습에 사용하는 사전 빌드된 PyTorch 컨테이너 버전과 일치하도록 tpuTfVersion
을 설정합니다.
AI Platform Training은 다음과 같은 사전 빌드된 PyTorch 컨테이너에 TPU를 사용한 학습을 지원합니다.
컨테이너 이미지 URI | tpuTfVersion |
---|---|
gcr.io/cloud-ml-public/training/pytorch-xla.1-11 |
pytorch-1.11 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-10 |
pytorch-1.10 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-9 |
pytorch-1.9 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-7 |
pytorch-1.7 |
gcr.io/cloud-ml-public/training/pytorch-xla.1-6 |
pytorch-1.6 |
예를 들어 PyTorch 1.11 사전 빌드된 컨테이너를 사용하여 학습시키려면 다음 config.yaml
파일을 사용하여 학습을 구성할 수 있습니다.
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
tpuTfVersion: pytorch-1.11
acceleratorConfig:
type: TPU_V2
count: 8
이 페이지의 이전 섹션에서 PyTorch 코드에서 TPU 사용도 참고하세요.
TPU 워커에서 커스텀 컨테이너 사용
TPU를 지원하는 AI Platform Training 런타임 버전 중 하나를 사용하는 대신 TPU 워커에서 커스텀 컨테이너를 실행하려면 학습 작업을 제출할 때 추가 구성 필드를 지정해야 합니다. 컨테이너가 사용하는 TensorFlow 버전이 포함된 런타임 버전으로 tpuTfVersion
을 설정합니다. TPU를 사용한 학습에 현재 지원되는 런타임 버전을 지정해야 합니다.
커스텀 컨테이너를 사용하도록 작업을 구성하고 있으므로 AI Platform Training은 학습 작업을 실행할 때 이 런타임 버전의 환경을 사용하지 않습니다. 그러나 AI Platform Training에는 이 필드가 필요하므로 커스텀 컨테이너가 사용하는 TensorFlow 버전에 맞는 TPU 워커를 올바르게 준비할 수 있습니다.
다음 예시는 이전 섹션과 유사한 TPU 구성을 가진 config.yaml
파일을 보여줍니다. 단, 이 경우 마스터 VM과 TPU 워커가 각각 다른 커스텀 컨테이너를 실행합니다.
TPU v2
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V2
count: 8
TPU v3(베타)
trainingInput:
scaleTier: CUSTOM
masterType: n1-highcpu-16
masterConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
workerType: cloud_tpu
workerCount: 1
workerConfig:
imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
tpuTfVersion: 2.11
acceleratorConfig:
type: TPU_V3
count: 8
gcloud beta ai-platform jobs submit training
명령어를 사용하여 학습 작업을 제출하는 경우 tpuTfVersion
API 필드에 config.yaml
파일 대신 --tpu-tf-version
플래그를 지정할 수 있습니다.
TPU가 프로비저닝 후 TPUClusterResolver
사용
커스텀 컨테이너를 사용하는 경우 TPU 프로비저닝 완료를 기다린 후에 TPUClusterResolver
를 호출하여 사용해야 합니다. 다음 샘플 코드는 TPUClusterResolver
로직을 처리하는 방법을 보여줍니다.
def wait_for_tpu_cluster_resolver_ready():
"""Waits for `TPUClusterResolver` to be ready and return it.
Returns:
A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
return None.
Raises:
RuntimeError: if failed to schedule TPU.
"""
tpu_config_env = os.environ.get('TPU_CONFIG')
if not tpu_config_env:
tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
return None
tpu_node = json.loads(tpu_config_env)
tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)
num_retries = 40
for i in range(num_retries):
try:
tpu_cluster_resolver = (
tf.contrib.cluster_resolver.TPUClusterResolver(
tpu=[tpu_node['tpu_node_name']],
zone=tpu_node['zone'],
project=tpu_node['project'],
job_name='worker'))
tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
if 'worker' in tpu_cluster_resolver_dict:
tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
return tpu_cluster_resolver
except Exception as e:
if i < num_retries - 1:
tf.logging.info('Still waiting for provisioning of TPU VM instance.')
else:
# Preserves the traceback.
raise RuntimeError('Failed to schedule TPU: {}'.format(e))
time.sleep(10)
# Raise error when failed to get TPUClusterResolver after retry.
raise RuntimeError('Failed to schedule TPU.')
커스텀 컨테이너를 사용한 분산형 학습에 대해 자세히 알아보세요.
다음 단계
- AI Platform Training의 학습 모델 자세히 알아보기
- AI Platform Training의 초매개변수 조정 및 Cloud TPU를 통한 초매개변수 조정에 대해 자세히 알아보기
- Cloud TPU의 추가 참조 모델 살펴보기
- Cloud TPU 권장사항에 따라 Cloud TPU에 맞게 모델 최적화
- Cloud TPU 문제해결 및 FAQ에서 문제 진단 및 해결에 대한 도움말 참조