TPU를 사용한 모델 학습

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 서비스 계정 이름을 승인합니다.

  1. 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
    
  2. API에서 반환한 serviceAccountProjecttpuServiceAccount 필드 값을 저장합니다.

  3. 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 명령어를 사용하여 다음 단계를 완료합니다.

콘솔

  1. Google Cloud Console에 로그인하고 TPU를 사용 중인 프로젝트를 선택합니다.
  2. IAM 및 관리자 > IAM을 선택합니다.
  3. 추가 버튼을 클릭하여 프로젝트에 멤버를 추가합니다.
  4. 멤버 입력란에 TPU 서비스 계정을 입력합니다.
  5. 역할 드롭다운 목록을 클릭합니다.
  6. Cloud ML 서비스 에이전트 역할을 사용 설정합니다(서비스 에이전트 > Cloud ML 서비스 에이전트).

gcloud

  1. 프로젝트 ID 및 Cloud TPU 서비스 계정이 포함된 환경 변수를 설정합니다.

    PROJECT_ID=PROJECT_ID
    SVC_ACCOUNT=your-tpu-sa-123@your-tpu-sa.google.com.iam.gserviceaccount.com
    
  2. Cloud TPU 서비스 계정에 ml.serviceAgent 역할을 부여합니다.

    gcloud projects add-iam-policy-binding $PROJECT_ID \
        --member serviceAccount:$SVC_ACCOUNT --role roles/ml.serviceAgent
    

서비스 계정에 역할을 부여하는 방법에 대한 자세한 내용은 IAM 문서를 참조하세요.

예: 샘플 MNIST 모델 학습

이 섹션에서는 TPU 및 런타임 버전 2.4을 사용하여 샘플 MNIST 모델을 학습시키는 방법을 보여줍니다. 예시 작업은 사용 중인 머신 구성에 적합한 사전 정의된 BASIC_TPU 확장 등급을 사용합니다. 이 가이드의 이후 섹션에서는 커스텀 구성을 설정하는 방법을 보여줍니다.

이 예시에서는 gcloud 도구가 설치된 Bash 셸을 사용한다고 가정합니다. 다음 명령어를 실행하여 코드를 가져오고 AI Platform Training에 학습 작업을 제출합니다.

  1. TensorFlow의 참조 모델 코드를 다운로드하고 샘플 코드가 있는 디렉터리로 이동합니다.

    git clone https://github.com/tensorflow/models.git \
      --branch=v2.4.0 \
      --depth=1
    
    cd models
    
  2. 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'],
        packages=find_packages()
    )
    END
    
  3. 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.4 \
      --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 도구는 패키징된 학습 코드를 이 버킷에 업로드하고 AI Platform Training은 버킷에 학습 출력을 저장합니다.

  4. 학습 작업을 모니터링합니다. 작업이 완료되면 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를 사용할 수 있습니다. 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을 구성합니다.
  • workerTypecloud_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을 구성합니다.
  • workerTypecloud_tpu로 설정합니다.
  • acceleratorConfig 필드가 있는 workerConfig를 추가합니다. 해당 acceleratorConfig 안에서 typeTPU_V2 또는 TPU_V3로 설정하거나 count8로 설정합니다. 다른 수의 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 워커에서 사전 빌드된 PyTorch 컨테이너 사용

TPU를 사용하여 PyTorch 학습을 수행하려면 학습 작업의 trainingInputtpuTfVersion 필드를 지정해야 합니다. 학습에 사용하는 사전 빌드된 PyTorch 컨테이너 버전과 일치하도록 tpuTfVersion을 설정합니다.

AI Platform Training은 다음과 같은 사전 빌드된 PyTorch 컨테이너에 TPU를 사용한 학습을 지원합니다.

컨테이너 이미지 URI tpuTfVersion
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.7 사전 빌드된 컨테이너를 사용하여 학습시키려면 다음 config.yaml 파일을 사용하여 학습을 구성할 수 있습니다.

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-7
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-7
    tpuTfVersion: pytorch-1.7
    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.4
    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.4
    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.')

커스텀 컨테이너를 사용한 분산형 학습에 대해 자세히 알아보세요.

다음 단계