Cloud TPU v5p 학습

Cloud TPU v5p는 Google Cloud의 5세대 Cloud TPU이자 v4 TPU의 후속 모델입니다. v5p는 대규모 학습에 최적화되어 있으며 기본 LLM, 확산 모델, 생성형 AI 개발을 위한 선도적인 플랫폼입니다. 개략적으로 v5p는 v4보다 최대 2배 높은 성능을 제공하는 동시에 2배 더 많은 TPU를 포드에 패키징하여(최대 슬라이스 수가 v4에서는 3,000개, v5p에서는 6,000개) 포드 수준에서 최대 4배의 성능을 제공합니다. 또한 더 높은 클록 주파수(1.75Ghz와 1.05Ghz 비교)에서 실행되고, 대규모 임베딩의 경우 SparseCore를 추가하고, 고대역폭 메모리(HBM) 용량을 세 배로 늘립니다.

Cloud TPU v5p 개념

Cloud TPU를 처음 사용하는 경우 TPU 문서 홈을 확인하세요.

모든 Cloud TPU 버전에 대한 Cloud TPU 개념(예: 슬라이스, 호스트, TensorCore)과 Cloud TPU 시스템 아키텍처는 Cloud TPU 시스템 아키텍처 페이지에 설명되어 있습니다.

각 Cloud TPU 버전에는 학습 또는 추론을 위한 특정 가속기 유형이 필요합니다. 이러한 가속기 유형은 v5p 구성에 설명되어 있습니다.

TPU 리소스 관리

TPU VM을 관리하는 데 사용할 수 있는 모든 명령어는 큐에 추가된 리소스 관리에 대한 TPU 관리 또는 큐에 추가된 리소스 사용자 가이드에 설명되어 있습니다.

프레임워크 설정

이 섹션에서는 TPU v5p에서 JAX 또는 PyTorch를 사용하는 모델 학습을 위한 일반적인 설정 프로세스를 설명합니다.

JAX 설정

칩이 4개를 초과하는 슬라이스 형태이면 슬라이스 하나에 VM이 여러 개 있습니다. 이 경우 --worker=all 플래그를 사용하여 단일 명령어로 모든 TPU VM에서 설치를 실행해야 합니다.

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

다음 명령어를 실행하여 기기 수를 확인할 수 있습니다(여기에 표시된 출력은 v5p-32 슬라이스로 생성됨). 이 코드는 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.]

--node=all을 사용하여 모든 멀티슬라이스 작업자에서 명령어를 실행합니다.

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

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

PyTorch 설정

PJRT 런타임은 v5p에 지원되는 유일한 런타임이며 PyTorch 2.1 이상은 모든 TPU 버전의 기본 런타임으로 PJRT를 사용합니다. 이 섹션에서는 모든 작업자를 대상으로 PyTorch/XLA 2.2.0을 사용하여 v5p 포드에서 PJRT 사용을 시작하는 방법을 설명합니다.

종속 항목 설치

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

PJRT와 함께 Python 스크립트를 사용하여 설치 검증을 수행하여 사용 가능한 TPU 기기를 표시합니다(여기서 표시된 출력은 v5p-32 슬라이스에서 생성됨).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

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...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

--node=all을 사용하여 모든 멀티슬라이스 작업자에서 명령어를 실행합니다.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

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

모니터링 및 프로파일링

Cloud TPU v5p는 이전 세대의 Cloud TPU와 동일한 메서드를 사용한 모니터링 및 프로파일링을 지원합니다. Cloud TPU 도구로 모델 프로파일링을 읽고 프로파일링에 대해 알아보고 Cloud TPU VM 모니터링을 읽고 모니터링에 대해 자세히 알아보세요.

학습 튜토리얼

이 섹션에서는 단일 슬라이스의 학습 튜토리얼을 중점적으로 설명합니다. SSH 명령어에 --node=all 플래그를 추가하면 이러한 튜토리얼을 멀티슬라이스 학습에 적용할 수 있습니다. 자세한 내용과 권장사항은 멀티슬라이스 소개를 참조하세요.

JAX 튜토리얼

Diffusion 2.1 학습

이 튜토리얼에서는 Cloud TPU v5p에서 Pokémon 데이터 세트를 사용하여 HuggingFace에서 Stable Diffusion 모델을 학습시키는 방법을 보여줍니다.

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

설정

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

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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

    명령어 플래그 설명

    변수 Description(설명)
    PROJECT_ID Google Cloud 프로젝트 이름
    ACCELERATOR_TYPE TPU 버전은 TPU 버전 페이지를 참조하세요.
    ZONE 지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참조하세요.
    RUNTIME_VERSION v5p의 경우 RUNTIME_VERSION에 v2-alpha-tpuv5를 사용합니다.
    SERVICE_ACCOUNT Google Cloud 콘솔 -> IAM -> 서비스 계정에서 찾을 수 있는 서비스 계정의 주소입니다. 예: tpu-service-account@myprojectID.iam.gserviceaccount.com
    TPU_NAME 큐에 추가된 리소스 요청이 할당될 때 생성되는 TPU의 사용자 할당 텍스트 ID입니다.
    QUEUED_RESOURCE_ID 큐에 추가된 리소스 요청의 사용자 할당 텍스트 ID입니다. 큐에 추가된 리소스에 대한 자세한 내용은 큐에 추가된 리소스 문서를 참조하세요.
    QUOTA_TYPE reserved 또는 best-effort일 수 있습니다. 둘 다 지정되지 않으면 QUOTA_TYPE 기본값은 on-demand입니다. Cloud TPU에서 지원하는 여러 가지 유형의 할당량에 대한 자세한 내용은 할당량을 참조하세요.
    VALID_UNTIL_DURATION 요청의 유효 기간입니다. 서로 다른 유효한 기간에 대한 자세한 내용은 큐에 추가된 리소스를 참조하세요.

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

    gcloud alpha 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 상태가 되면 SSH를 통해 TPU VM에 연결할 수 있습니다. 다음 명령어를 실행하여 큐에 추가된 리소스의 상태를 확인합니다.

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

    큐에 추가된 리소스가 ACTIVE 상태가 되면 출력은 다음과 비슷합니다.

    state: ACTIVE

  3. JAX 및 종속 항목을 설치합니다.

    # compatible with v5p: only jax version 0.4.19 and later \
# jax 0.4.19 requires py 3.10 \

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
--command='pip install "jax[tpu]==0.4.20" -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/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

  5. 모델 학습

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

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

삭제

세션이 끝나면 TPU 및 큐에 추가된 리소스 요청을 삭제하거나 'FAILED' 상태의 큐에 추가된 리소스 요청을 삭제합니다. 큐에 추가된 리소스를 삭제하려면 슬라이스를 삭제한 후 큐에 추가된 리소스 요청을 다음과 같이 두 단계로 삭제합니다.

   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

또는 --force를 사용하여 한 번에 슬라이스 및 큐에 추가된 리소스 요청을 삭제합니다.

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

벤치마크 결과

Stable Diffusion 학습 스크립트는 v5p-8, v5p-32, v5p-128에서 실행되었습니다. 다음 표에서는 처리량을 보여줍니다.

v5p-8

v5p-32

v5p-128

학습 단계

150

150

150

전역 배치 크기

32

64

64

처리량(예시/초)

12.10

18.08

19.10

MaxText

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

MaxText는 Cloud TPU를 대상으로 순수 Python/JAX로 작성되어 임의로 확장 가능하고 테스트를 철저하게 거친 고성능 오픈소스 LLM입니다. MaxText는 연구원과 개발자에게 자연어 처리(NLP) 연구 및 개발의 발전을 위해 이용하기 쉽고 조정 가능한 도구를 제공합니다.

이 튜토리얼을 실행하기 전에 Cloud TPU 환경을 설정해야 합니다.

  1. 환경 변수 설정

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

    명령어 플래그 설명

    변수 Description(설명)
    PROJECT_ID Google Cloud 프로젝트 이름
    TPU_NAME TPU에 대한 사용자 정의 이름입니다.
    ACCELERATOR_TYPE TPU 버전은 TPU 버전 페이지를 참조하세요.
    ZONE 지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참조하세요.
    RUNTIME_VERSION v5p의 경우 런타임 버전에 v2-alpha-tpuv5를 사용합니다.
    RUN_NAME 사용자가 제공한 실험 실행 이름입니다.

    멀티슬라이스에 권장되는 선택적 설정:

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    멀티슬라이스 워크로드를 실행 중이고 최적의 네트워크 성능을 원하는 경우 최대 전송 단위(MTU) 8,896바이트로 전용 네트워크를 만들고 적절한 방화벽 규칙을 구성하는 것이 좋습니다. 이 단계는 선택사항이지만 특히 데이터 센터 네트워크(DCN)에서 슬라이스 수를 확장하는 경우 성능을 크게 향상시킬 수 있습니다. 네트워크를 만들려면 프로젝트에서 compute.networks.create 권한이 필요합니다. 다음 예시에서는 전용 네트워크 및 방화벽 규칙을 만드는 방법을 보여줍니다.

    전용 네트워크를 만듭니다.

    gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT_ID} \
--subnet-mode=auto \
--bgp-routing-mode=regional

    방화벽 규칙 만들기

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

  2. MaxText 저장소 클론

    git clone https://github.com/google/maxtext.git

  3. 모델 학습

    다음 섹션에서는 MaxText를 학습시키는 두 가지 옵션을 설명합니다.

    옵션 1

    Cloud TPU 프로비저닝 및 종속 항목 설치부터 모델 실행 및 리소스 해제에 이르기까지 스크립트에서 전체 워크플로를 관리하려면 multihost_job.py를 사용하면 됩니다.

    cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

    스크립트를 시작한 후에는 로그에 다음과 비슷한 메시지가 표시됩니다. 출력 메시지에서 로그 위치를 참조합니다. TPU 프로비저닝이 완료되면 첫 번째 링크를 클릭하여 모든 작업자 로그에 액세스합니다.

    ------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID
옵션 2

프로비저닝된 Cloud TPU에서 학습 스크립트를 여러 번 실행하려면 multihost_runner.py 스크립트를 사용하여 리소스를 사용합니다.

  1. TPU를 만들기 위한 변수를 설정합니다.

    export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type

    --node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

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

    gcloud alpha 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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    QueuedResource가 ACTIVE 상태가 되면 SSH를 사용하여 TPU VM에 연결할 수 있습니다.

    describe 명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.

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

    큐에 추가된 리소스가 ACTIVE 상태이면 출력은 다음과 비슷합니다.

     state: ACTIVE

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

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

  4. 종속 항목 설치

    export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path

    cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'

  5. 32b.sh, 64b.sh와 같은 다양한 구성 스크립트를 사용하여 모델을 실행합니다. TPU VM에서 스크립트를 실행하는 경우 --INTERNAL_IP=true 플래그를 추가해야 합니다.

    python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

삭제

TPU 및 큐에 추가된 리소스를 삭제합니다.

벤치마크 결과

MaxText 학습 스크립트는 bf16 정밀도로 32B에서 1160B까지 실행되었습니다. 이러한 실행 결과는 다음 표에 나와 있습니다.

매개변수 수

가속기 유형

TFLOP/칩/초

모델 FLOPS 사용률

(MFU)

32B

v5p-128

3.28E+02

71.47%

64B

v5p-128

3.23E+02

70.31%

128B

v5p-256

3.15E+02

68.68%

128B

v5p-512

3.15E+02

68.53%

256B

v5p-1024

3.16E+02

68.82%

512B

v5p-1024

2.94E+02

63.99%

1024B

v5p-2048

2.49E+02

64.05%

1024B

v5p-4096

2.97E+02

64.80%

1160B

v5p-7680

2.95E+02

64.27%

1160B

v5p-12288

3.04E+02

66.23%

256B 매개변수 모델은 bf16 및 int8 가중치를 모두 사용하여 v5p-512 및 v5p-1024에서 테스트되었습니다. 다음 표는 이러한 테스트의 결과를 보여줍니다.

v5p-512

v5p-512

v5p-1024

v5p-1024

전역 배치 크기

(토큰)

5.24E+05

5.24E+05

1.05E+06

1.05E+06

정밀도

bf16

int8

bf16

int8

TFLOP/칩/초

307

408

308

414

모델 FLOPS 사용률

(MFU)

66.98%

88.85%

67.09%

90.23%

TensorFlow 튜토리얼

단일 호스트 v5p에서 ResNet 학습

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

설정

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

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

    이 튜토리얼에서는 v5p-8을 ACCELERATOR_TYPE으로 사용합니다.

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

    gcloud alpha 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 상태가 되면 SSH를 사용하여 TPU VM에 연결할 수 있습니다. 큐에 추가된 리소스 상태를 확인하려면 다음 명령어를 사용하세요.

    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

모델 학습

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

    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 및 큐에 추가된 리소스를 삭제합니다.

멀티 호스트 v5p에서 Resnet 학습

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

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

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

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

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

    gcloud alpha 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 상태가 되면 SSH를 사용하여 TPU VM에 연결할 수 있습니다.

    describe 명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.

    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 TPU_NAME=your_tpu_name
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

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

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

모델 학습

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

    python3 official/vision/train.py \
  --tpu=${TPU_NAME} \
  --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 및 큐에 추가된 리소스를 삭제합니다.

PyTorch/XLA

Llama 2

이 튜토리얼에서는 General and Scalable Parallelization for ML Computation Graphs(GSPMD)를 통해 PyTorch/XLA의 HugdingFace 저장소 포크를 사용하여 v5p에서 Llama 2 7B 모델을 학습시키는 방법을 설명합니다.

설정

  1. 프로젝트 ID, 가속기 유형, 영역, 런타임 버전, TPU 이름에 대한 변수를 만듭니다.

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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_DURATION=1d

  2. TPU 리소스 만들기

    gcloud alpha 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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    QueuedResource가 ACTIVE 상태가 되면 SSH를 사용하여 TPU VM에 연결할 수 있습니다.

    describe 명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.

    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
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -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 -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

  5. 7B 모델 구성을 다운로드합니다.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

  6. 모델 학습

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

멀티슬라이스 환경에서 실행하는 경우 --spmd_dcn_parallelism 플래그를 슬라이스 수로 설정해야 합니다.

SPMD_USER_GUIDE는 HF 스크립트의 모든 다양한 환경 변수와 전환을 설명하는 심층적인 사용자 가이드를 제공합니다. LIBTPU_INIT_ARGS는 PyTorch/XLA에 통합되고 향후 출시 버전에서는 기본적으로 사용 설정될 예정입니다.

삭제

TPU 및 큐에 추가된 리소스를 삭제합니다.

벤치마크 결과

다음 표에는 3가지 Llama 2 모델 크기 처리량이 모두 포함되어 있습니다.

v5p-8

v5p-128

v5p-128

모델 크기

7B

13B

70B

전역 배치 크기

96

1024

128

샤딩 메시 형태

(4, 1)

(64, 1)

(16, 4)

모델 FLOPS 사용률

(MFU)

56.67%

55.80%

51.85%

지원 및 의견

여러분의 모든 의견을 환영합니다. 의견을 공유하거나 지원을 요청하려면 Cloud TPU 지원 또는 의견 제출 양식을 작성하세요.