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 "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
다음 명령어를 실행하여 기기 수를 확인할 수 있습니다(여기에 표시된 출력은 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를 사용하여 스크립트 실행
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의 리드미 파일을 참조하세요.
설정
환경 변수를 만듭니다.
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
-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}QueuedResource가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE JAX 및 jaxlib의 최신 버전을 설치합니다.
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'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'허위 데이터를 생성하려면 모델에 데이터 세트의 측정기준에 대한 정보가 필요합니다. 이는 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 인코더를 성공적으로 학습시킨 첫 번째 모델입니다. 자세한 내용은 다음 리소스를 참조하세요.
설정
환경 변수를 만듭니다.
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
-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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'다음과 같이 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'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에서 안정적인 확산 모델을 학습시키는 방법을 보여줍니다.
안정적인 확산 모델은 모든 텍스트 입력에서 실사 이미지를 생성하는 잠재 텍스트 이미지 변환 모델입니다. 자세한 내용은 다음 리소스를 참조하세요.
설정
환경 변수를 만듭니다.
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
-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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'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} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'
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 모델입니다. 문장의 다음 단어를 예측하도록 학습되었습니다. 자세한 내용은 다음 리소스를 참조하세요.
설정
환경 변수를 만듭니다.
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
-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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'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'구성을 다운로드하여 모델을 학습시킵니다.
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 설정을 위한 업데이트된 안내입니다.
설정
환경 변수를 만듭니다.
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
-
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}QueuedResource가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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를 실행하는 방법을 다룹니다.
설정
환경 변수를 만듭니다.
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
-
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}QueuedResource가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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'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 '선행 학습된 모델의 구성을 다운로드합니다.
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를 실행하는 방법을 설명합니다.
설정
환경 변수를 만듭니다.
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
-
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}QueuedResource가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}큐에 추가된 리소스가
ACTIVE상태이면 출력이 다음과 비슷하게 표시됩니다.state: ACTIVE 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'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을 학습시키는 방법을 설명합니다. 다른 데이터 세트를 사용하려면 데이터 세트 준비를 참조하세요.
설정
환경 변수를 만듭니다.
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일 수 있습니다.-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다. 큐에 추가된 리소스 상태를 확인하려면 다음 명령어를 사용하세요.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} SSH를 사용하여 TPU에 연결합니다.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}몇 가지 환경 변수를 설정합니다.
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모델 저장소 디렉터리로 변경하고 요구사항을 설치합니다.
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 및 큐에 추가된 리소스 삭제
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
멀티 호스트 v5e에서 Resnet 학습
이 튜토리얼에서는 가짜 데이터 세트를 사용하여 v5litepod-16 이상에서 ImageNet을 학습시키는 방법을 설명합니다. 다른 데이터 세트를 사용하려면 데이터 세트 준비를 참조하세요.
환경 변수를 만듭니다.
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이상일 수 있습니다.-
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}큐에 추가된 리소스가
ACTIVE상태가 되면 TPU VM에 SSH를 통해 연결할 수 있습니다. 큐에 추가된 리소스 상태를 확인하려면 다음 명령어를 사용하세요.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} SSH를 사용하여 TPU(작업자 0)에 연결합니다.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}몇 가지 환경 변수를 설정합니다.
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모델 저장소 디렉터리로 변경하고 요구사항을 설치합니다.
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 및 큐에 추가된 리소스 삭제
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