멀티 호스트 TPU 슬라이스 만들기

인스턴스 템플릿 만들기

TPU v6e VM의 인스턴스 템플릿을 만들려면 gcloud compute instance-templates create 명령어를 사용하세요.

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

워크로드 정책 만들기

워크로드 정책은 Compute 인스턴스의 물리적 속성을 정의합니다. TPU 슬라이스에서 가속기 토폴로지는 TPU 칩의 물리적 배열을 정의합니다. 상호 연결된 멀티 호스트 TPU 슬라이스에는 가속기 토폴로지를 지정해야 합니다.

멀티 호스트 TPU 슬라이스의 워크로드 정책을 만들려면 gcloud compute resource-policies create workload-policy 명령어 와 함께 --accelerator-topology 플래그를 사용하세요. 다음 명령어는 2x4 토폴로지가 있는 워크로드 정책을 만듭니다.

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

MIG 만들기

다음 명령어를 실행하여 멀티 호스트 TPU 슬라이스를 구성하는 MIG를 만듭니다.

1. 멀티 호스트 TPU 슬라이스를 구성하는 MIG를 만들려면 gcloud compute instance-groups managed create 명령어를 사용하세요.

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   PROJECT_ID를 Cloud de Confiance by S3NS 프로젝트 ID로 바꿉니다.

2. 선택적으로 다음 명령어를 사용하여 관리형 인스턴스가 실행 중인지 확인합니다.

   • 전반적인 MIG 상태를 보려면 gcloud compute instance-groups managed describe 명령어를 사용하세요.

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • 각 인스턴스의 상태를 보려면 gcloud compute instance-groups managed list-instances 명령어를 사용하세요.

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

JAX 설치

MIG의 모든 TPU VM 인스턴스에 있는 가상 환경에 종속 항목과 JAX 프레임워크를 설치합니다. TPU VM에 3.11 이전 버전의 Python이 설치되어 있는 경우 최신 버전의 JAX를 실행하려면 Python 3.11을 설치해야 합니다.

1. TPU VM에서 실행 중인 Python 버전을 확인합니다.

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   버전이 Python 3.11 이전인 경우 Python 3.11을 설치합니다.

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. 가상 환경을 만듭니다.

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. 가상 환경에 JAX를 설치합니다.

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

슬라이스에서 JAX 코드 실행

TPU 슬라이스에서 JAX 코드를 실행하려면 TPU 슬라이스의 각 호스트에서 코드를 실행해야 합니다. jax.device_count() 함수 호출은 슬라이스의 각 호스트에서 호출될 때까지 응답을 중지합니다. 다음 예에서는 TPU 슬라이스에서 JAX 계산을 실행하는 방법을 보여줍니다.

코드 준비

각 인스턴스에서 example.py라는 파일을 만듭니다.

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

슬라이스에서 코드 실행

슬라이스의 각 TPU VM에서 example.py 프로그램을 실행합니다.

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

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

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]