Cloud TPU v5e 학습

Cloud TPU v5e는 Google Cloud의 최신 세대 AI 가속기입니다. 포드당 칩 공간이 256개로 더 작은 v5e는 변환기, 텍스트-이미지, 컨볼루셔널 신경망(CNN) 학습, 미세 조정, 서빙에 높은 효용성을 제공하는 제품이 되도록 최적화되어 있습니다. 서빙을 위해 Cloud TPU v5e를 사용하는 방법에 대한 자세한 내용은 v5e를 사용한 추론을 참조하세요.

Cloud TPU v5e TPU 하드웨어 및 구성에 대한 자세한 내용은 TPU v5e를 참조하세요.

시작하기

다음 섹션에서는 TPU v5e를 사용하는 방법을 설명합니다.

요청 할당량

학습에 TPU v5e를 사용하려면 할당량이 필요합니다. 주문형 TPU, 예약된 TPU, TPU 스팟 VM에는 다양한 할당량 유형이 있습니다. 추론을 위해 TPU v5e를 사용하는 경우 별도의 할당량이 필요합니다. 할당량에 대한 자세한 내용은 할당량을 참조하세요. TPU v5e 할당량을 요청하려면 Cloud 영업팀에 문의하세요.

Google Cloud 계정 및 프로젝트 만들기

Cloud TPU를 사용하려면 Google Cloud 계정과 프로젝트가 필요합니다. 자세한 내용은 Cloud TPU 환경 설정을 참조하세요.

Cloud TPU 만들기

권장사항은 queued-resource create 명령어를 사용하여 Cloud TPU v5es를 큐에 추가된 리소스로 프로비저닝하는 것입니다. 자세한 내용은 큐에 추가된 리소스 관리를 참조하세요.

Create Node API(gcloud compute tpus tpu-vm create)를 사용하여 Cloud TPU v5es를 프로비저닝할 수도 있습니다. 자세한 내용은 TPU 리소스 관리를 참조하세요.

사용 가능한 학습용 v5e 구성에 대한 자세한 내용은 학습용 Cloud TPU v5e 유형을 참조하세요.

프레임워크 설정

이 섹션에서는 TPU v5e에서 JAX 또는 PyTorch를 사용하는 커스텀 모델 학습을 우한 일반적인 설정 프로세스를 설명합니다. TensorFlow 지원은 tpu-vm-tf-2.17.0-pjrt 및 tpu-vm-tf-2.17.0-pod-pjrt TPU 런타임 버전에서 제공됩니다.

추론 설정 안내는 v5e 추론 소개를 참조하세요.

JAX 설정

칩이 8개를 초과하는 슬라이스 형태이면 슬라이스 하나에 VM이 여러 개 있습니다. 이 경우 SSH를 사용하여 각각 개별적으로 로그인하지 않고 한 번에 모든 TPU VM에서 설치를 실행하려면 --worker=all 플래그를 사용해야 합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

명령어 플래그 설명

변수	설명
TPU_NAME	큐에 추가된 리소스 요청이 할당될 때 생성되는 TPU의 사용자 할당 텍스트 ID입니다.
PROJECT_ID	Google Cloud 프로젝트 이름 Google Cloud 프로젝트 설정에서 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
ZONE	지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참조하세요.
worker	기본 TPU에 액세스할 수 있는 TPU VM입니다.

다음 명령어를 실행하여 기기 수를 확인할 수 있습니다(여기에 표시된 출력은 v5litepod-16 슬라이스로 생성됨). 이 코드는 JAX에서 Cloud TPU TensorCore를 표시하고 기본 작업을 실행할 수 있는지 확인하여 모든 것이 올바르게 설치되었는지 테스트합니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

출력은 다음과 비슷하게 표시됩니다.

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count()에서는 지정된 슬라이스의 칩 총개수를 보여줍니다. jax.local_device_count()에서는 이 슬라이스에서 단일 VM이 액세스할 수 있는 칩 수를 나타냅니다.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

출력은 다음과 비슷하게 표시됩니다.

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

이 문서의 JAX 튜토리얼을 통해 JAX를 사용한 v5e 학습을 시작합니다.

PyTorch 설정

v5e는 PJRT 런타임만 지원합니다. PyTorch 2.1 이상에서는 모든 TPU 버전의 기본 런타임으로 PJRT를 사용합니다.

이 섹션에서는 모든 작업자에 대해 명령어를 사용하여 PyTorch/XLA와 함께 v5e에서 PJRT 사용을 시작하는 방법을 설명합니다.

종속 항목 설치

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

torch, torch_xla, torchvision의 휠을 설치할 때 pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222와 같은 오류가 발생하면 다음 명령어를 사용하여 버전을 다운그레이드합니다.

pip3 install setuptools==62.1.0

PJRT를 사용하여 스크립트 실행

참고: 상당한 크기의 할당이 빈번한 모델의 경우 tcmalloc를 사용하면 기본 malloc 구현에 비해 학습 시간이 크게 향상될 수 있으므로 TPU VM에서 사용되는 기본 malloc는 tcmalloc입니다. 하지만 워크로드에 따라(예: 임베딩 테이블에 대한 대규모 할당이 있는 DLRM) tcmalloc를 사용하면 느려질 수 있습니다. 이 경우 대신 기본 malloc를 사용하여 다음 변수를 설정 해제할 수 있습니다.

unset LD_PRELOAD

다음은 Python 스크립트를 사용하여 v5e VM에서 계산을 수행하는 예시입니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

그러면 다음과 비슷한 출력이 생성됩니다.

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

이 문서의 PyTorch 튜토리얼을 통해 PyTorch를 사용한 v5e 학습을 시작합니다.

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다. 큐에 추가된 리소스를 삭제하려면 슬라이스를 삭제한 후 큐에 추가된 2단계의 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

이 두 단계를 사용하여 FAILED 상태의 큐에 추가된 리소스 요청을 삭제할 수도 있습니다.

JAX/FLAX 예시

다음 섹션에서는 TPU v5e에서 JAX 및 FLAX 모델을 학습시키는 방법의 예시를 설명합니다.

v5e에서 ImageNet 학습

이 튜토리얼에서는 가짜 입력 데이터를 사용하여 v5e에서 ImageNet을 학습시키는 방법을 설명합니다. 실제 데이터를 사용하려면 GitHub의 리드미 파일을 참조하세요.

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   QueuedResource가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

3. JAX 및 jaxlib의 최신 버전을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   참고: jax>=0.4.8을 사용하라는 오류가 표시되면 메시지를 무시해도 됩니다.

4. ImageNet 모델을 클론하고 해당 요구사항을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
   

5. 허위 데이터를 생성하려면 모델에 데이터 세트의 측정기준에 대한 정보가 필요합니다. 이는 ImageNet 데이터 세트의 메타데이터에서 수집할 수 있습니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
   

모델 학습

이전 단계가 모두 완료되면 모델을 훈련할 수 있습니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hugging Face FLAX 모델

FLAX에 구현된 Hugging Face 모델은 Cloud TPU v5e에서 바로 사용할 수 있습니다. 이 섹션에서는 인기 있는 모델을 실행하는 방법에 대한 안내를 제공합니다.

Imagenette에서 ViT 학습

이 튜토리얼에서는 Cloud TPU v5e에서 Fast AI Imagenette 데이터 세트를 사용하여 HuggingFace에서 Vision Transformer(ViT) 모델을 학습시키는 방법을 보여줍니다.

ViT 모델은 컨볼루셔널 네트워크 보다 뛰어난 결과로 ImageNet에서 Transformer 인코더를 성공적으로 학습시킨 첫 번째 모델입니다. 자세한 내용은 다음 리소스를 참조하세요.

• ViT 개요
• ViT 코드 리소스

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

3. JAX 및 라이브러리를 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. 다음과 같이 Hugging Face 저장소를 다운로드하고 요구사항을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.17.0 && pip install -r examples/flax/vision/requirements.txt'
   

5. Imagenette 데이터 세트를 다운로드합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

모델 학습

4GB의 사전 매핑된 버퍼로 모델을 학습시킵니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제하세요.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ViT 벤치마킹 결과

학습 스크립트는 v5litepod-4, v5litepod-16, v5litepod-64에서 실행되었습니다. 다음 표에서는 다양한 가속기 유형의 처리량을 보여줍니다.

가속기 유형	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
전역 배치 크기	32	128	512
처리량(예시/초)	263.40	429.34	470.71

Pokémon에서 확산 학습

이 튜토리얼에서는 Cloud TPU v5e에서 Pokémon 데이터 세트를 사용하여 HuggingFace에서 안정적인 확산 모델을 학습시키는 방법을 보여줍니다.

안정적인 확산 모델은 모든 텍스트 입력에서 실사 이미지를 생성하는 잠재 텍스트 이미지 변환 모델입니다. 자세한 내용은 다음 리소스를 참조하세요.

• 안정적인 확산 개요
• 안정적인 확산 코드 소스

설정

1. 모델 출력을 위한 스토리지 버킷을 설정합니다.

   gcloud storage buckets create gs://your_bucket 
   
   --project=your_project 
   
   --location=us-west1 
   
   export GCS_BUCKET_NAME=your_bucket

2. 환경 변수 만들기

   참고: 변수에 대한 자세한 내용은 명령어 플래그 설명을 참조하세요.

   export GCS_BUCKET_NAME=your_bucket
   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west1-c
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

   명령어 플래그 설명

   변수	설명
   GCS_BUCKET_NAME	Google Cloud 콘솔 -> Cloud Storage -> 버킷에 표시됨
   PROJECT_ID	Google Cloud 프로젝트 이름 Google Cloud 프로젝트 설정에서 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.
   ACCELERATOR_TYPE	TPU 버전은 TPU 버전 페이지를 참조하세요.
   ZONE	지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참조하세요.
   RUNTIME_VERSION	RUNTIME_VERSION에 v2-alpha-tpuv5를 사용합니다.
   SERVICE_ACCOUNT	Google Cloud 콘솔 -> IAM -> 서비스 계정에서 찾을 수 있는 서비스 계정의 주소입니다. 예: [e-mail]
   TPU_NAME	큐에 추가된 리소스 요청이 할당될 때 생성되는 TPU의 사용자 할당 텍스트 ID입니다.
   QUEUED_RESOURCE_ID	큐에 추가된 리소스 요청의 사용자 할당 텍스트 ID입니다. 큐에 추가된 리소스에 대한 자세한 내용은 큐에 추가된 리소스 문서를 참조하세요.
   QUOTA_TYPE	reserved 또는 spot일 수 있습니다. 둘 다 지정되지 않으면 QUOTA_TYPE 기본값은 on-demand입니다. Cloud TPU에서 지원하는 여러 가지 유형의 할당량에 대한 자세한 내용은 할당량을 참조하세요.
   VALID_UNTIL_DURATION	요청의 유효 기간입니다. 서로 다른 유효한 기간에 대한 자세한 내용은 큐에 추가된 리소스를 참조하세요.

3. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved일 수 있으며 예약이 있는 경우 spot일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

4. JAX 및 라이브러리를 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

5. HuggingFace 저장소를 다운로드하고 요구사항을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

모델 학습

4GB의 사전 매핑된 버퍼로 모델을 학습시킵니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

확산 벤치마킹 결과

학습 스크립트는 v5litepod-4, v5litepod-16, v5litepod-64에서 실행되었습니다. 다음 표에서는 처리량을 보여줍니다.

가속기 유형	v5litepod-4	v5litepod-16	v5litepod-64
학습 단계	1500	1500	1500
전역 배치 크기	32	64	128
처리량(예시/초)	36.53	43.71	49.36

OSCAR 데이터 세트에서 GPT2 학습

이 튜토리얼에서는 Cloud TPU v5e에서 OSCAR 데이터 세트를 사용하여 HuggingFace에서 GPT2 모델을 학습시키는 방법을 보여줍니다.

GPT2는 수동 라벨링 없이 원시 텍스트로 사전 학습된 Transformer 모델입니다. 문장의 다음 단어를 예측하도록 학습되었습니다. 자세한 내용은 다음 리소스를 참조하세요.

• GPT2 개요
• GPT2 코드 소스

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

3. JAX 및 라이브러리를 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. HuggingFace 저장소를 다운로드하고 요구사항을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'
   

5. 구성을 다운로드하여 모델을 학습시킵니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'
   

모델 학습

4GB의 사전 매핑된 버퍼로 모델을 학습시킵니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

GPT2 벤치마킹 결과

학습 스크립트는 v5litepod-4, v5litepod-16, v5litepod-64에서 실행되었습니다. 다음 표에서는 처리량을 보여줍니다.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
전역 배치 크기	64	64	64
처리량(예시/초)	74.60	72.97	72.62

PyTorch/XLA

다음 섹션에서는 TPU v5e에서 PyTorch/XLA 모델을 학습시키는 방법의 예시를 설명합니다.

PJRT 런타임을 사용하여 ResNet 학습

PyTorch/XLA는 XRT에서 PjRt로 PyTorch 2.0 이상에서 마이그레이션됩니다. 다음은 PyTorch/XLA 학습 워크로드용 v5e 설정을 위한 업데이트된 안내입니다.

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --{QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   QueuedResource가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

3. Torch/XLA 관련 종속 항목 설치

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

ResNet 모델 학습

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

벤치마크 결과

다음 표는 벤치마크 처리량을 보여줍니다.

가속기 유형	처리량(예시/초)
v5litepod-4	4240ex/s
v5litepod-16	10,810ex/s
v5litepod-64	46,154ex/s

v5e에서 GPT2 학습

이 튜토리얼에서는 wikitext 데이터 세트를 사용하여 PyTorch/XLA에서 HuggingFace 저장소를 사용하여 v5e에서 GPT2를 실행하는 방법을 다룹니다.

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   QueuedResource가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

   state: ACTIVE
   

3. PyTorch/XLA 종속 항목을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get -y update
         sudo apt install -y libopenblas-base
         pip3 install torchvision
         pip3 uninstall -y torch
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. HuggingFace 저장소를 다운로드하고 요구사항을 설치합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         git clone https://github.com/pytorch/xla.git
         pip install --upgrade accelerate
         git clone https://github.com/huggingface/transformers.git
         cd transformers
         git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
         pip install .
         git log -1
         pip install datasets evaluate scikit-learn
         '
   

5. 선행 학습된 모델의 구성을 다운로드합니다.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
         gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
   

모델 학습

16의 배치 크기를 사용하여 2B 모델을 학습시킵니다.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

벤치마크 결과

학습 스크립트는 v5litepod-4, v5litepod-16, v5litepod-64에서 실행되었습니다. 다음 표에서는 다양한 가속기 유형의 벤치마크 처리량을 보여줍니다.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
config	600M	2B	16B
전역 배치 크기	64	128	256
처리량(예시/초)	66	77	31

v5e에서 ViT 학습

이 튜토리얼에서는 cifar10 데이터 세트에서 PyTorch/XLA의 HuggingFace 저장소를 사용하여 v5e에서 VIT를 실행하는 방법을 설명합니다.

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   QueuedResource가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   큐에 추가된 리소스가 ACTIVE 상태이면 출력이 다음과 비슷하게 표시됩니다.

    state: ACTIVE
   

3. PyTorch/XLA 종속 항목 설치

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. HuggingFace 저장소를 다운로드하고 요구사항을 설치합니다.

      gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command="
         git clone https://github.com/suexu1025/transformers.git vittransformers; \
         cd vittransformers; \
         pip3 install .; \
         pip3 install datasets; \
         wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
   

모델 학습

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

TPU 및 큐에 추가된 리소스 삭제

세션이 끝나면 TPU와 큐에 추가된 리소스를 삭제합니다.

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

벤치마크 결과

다음 표에서는 다양한 가속기 유형의 벤치마크 처리량을 보여줍니다.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
전역 배치 크기	32	128	512
처리량(예시/초)	201	657	2,844

TensorFlow 2.x

다음 섹션에서는 TPU v5e에서 TensorFlow 2.x 모델을 학습시키는 방법의 예시를 설명합니다.

단일 호스트 v5e에서 Resnet 학습

이 튜토리얼에서는 가짜 데이터 세트를 사용하여 v5litepod-4 또는 v5litepod-8에서 ImageNet을 학습시키는 방법을 설명합니다. 다른 데이터 세트를 사용하려면 데이터 세트 준비를 참조하세요.

설정

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5litepod-4
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE은 v5litepod-4 또는 v5litepod-8일 수 있습니다.

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다. 큐에 추가된 리소스 상태를 확인하려면 다음 명령어를 사용하세요.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. SSH를 사용하여 TPU에 연결합니다.

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

4. 몇 가지 환경 변수를 설정합니다.

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

5. 모델 저장소 디렉터리로 변경하고 요구사항을 설치합니다.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt
   

모델 학습

학습 스크립트를 실행합니다.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

TPU 및 큐에 추가된 리소스 삭제

1. TPU를 삭제합니다.

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

2. 큐에 추가된 리소스 요청을 삭제합니다.

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

멀티 호스트 v5e에서 Resnet 학습

이 튜토리얼에서는 가짜 데이터 세트를 사용하여 v5litepod-16 이상에서 ImageNet을 학습시키는 방법을 설명합니다. 다른 데이터 세트를 사용하려면 데이터 세트 준비를 참조하세요.

1. 환경 변수를 만듭니다.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE은 v5litepod-16 이상일 수 있습니다.

2. TPU 리소스를 만듭니다.

   gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   참고: QUOTA_TYPE 플래그는 reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 기본값은 on-demand입니다. Cloud TPU에서 지원되는 다양한 할당량 유형에 대한 자세한 내용은 할당량을 참조하세요.

   큐에 추가된 리소스가 ACTIVE 상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다. 큐에 추가된 리소스 상태를 확인하려면 다음 명령어를 사용하세요.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. SSH를 사용하여 TPU(작업자 0)에 연결합니다.

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

4. 몇 가지 환경 변수를 설정합니다.

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0
   export TPU_NAME=your_tpu_name

5. 모델 저장소 디렉터리로 변경하고 요구사항을 설치합니다.

    cd $MODELS_REPO && git checkout r2.15.0
    pip install -r official/requirements.txt
   

모델 학습

학습 스크립트를 실행합니다.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

TPU 및 큐에 추가된 리소스 삭제

1. TPU를 삭제합니다.

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

2. 큐에 추가된 리소스 요청을 삭제합니다.

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet