Pengantar Trillium (v6e)

v6e digunakan untuk merujuk ke Trillium dalam dokumentasi, TPU API, dan log ini. v6e mewakili TPU generasi ke-6 Google.

Dengan 256 chip per Pod, arsitektur v6e memiliki banyak kesamaan dengan v5e. Sistem ini dioptimalkan untuk pelatihan, penyesuaian, dan penayangan transformer, text-to-image, dan jaringan saraf konvolusi (CNN).

Lihat dokumen v6e untuk mengetahui informasi tentang konfigurasi dan arsitektur sistem v6e.

Dokumen pengantar ini berfokus pada proses pelatihan dan inferensi model menggunakan framework JAX, PyTorch, atau TensorFlow. Dengan setiap framework, Anda dapat menyediakan TPU menggunakan resource yang diantrekan atau Google Kubernetes Engine (GKE). Penyiapan GKE dapat dilakukan menggunakan perintah XPK atau GKE.

Prosedur umum untuk melatih atau menayangkan model menggunakan v6e

1. Menyiapkan Google Cloud project
2. Kapasitas aman
3. Menyiapkan lingkungan TPU
4. Menyediakan lingkungan Cloud TPU
5. Menjalankan workload pelatihan atau inferensi model
6. Pembersihan

Menyiapkan project Google Cloud

1. Login ke Akun Google Anda. Jika Anda belum melakukannya, daftar untuk membuat akun baru.
2. Di konsol Google Cloud, pilih atau buat project Cloud dari halaman pemilih project.
3. Aktifkan penagihan untuk project Google Cloud Anda. Penagihan diperlukan untuk semua penggunaan Google Cloud.
4. Instal komponen gcloud alpha.

5. Jalankan perintah berikut untuk menginstal versi terbaru komponen gcloud.

   gcloud components update
   

6. Aktifkan TPU API melalui perintah gcloud berikut di Cloud Shell. Anda juga dapat mengaktifkannya dari Konsol Google Cloud.

   gcloud services enable tpu.googleapis.com
   

7. Mengaktifkan izin dengan akun layanan TPU untuk Compute Engine API

   Akun layanan memungkinkan layanan Cloud TPU mengakses layanan Google Cloud lainnya. Akun layanan yang dikelola pengguna adalah praktik Google Cloud yang direkomendasikan. Ikuti panduan ini untuk membuat dan memberikan peran. Peran berikut diperlukan:

   • TPU Admin
   • Storage Admin
   • Penulis Log
   • Penulis Metrik Pemantauan

   a. Siapkan izin XPK dengan akun pengguna Anda untuk GKE: XPK.

8. Lakukan autentikasi dengan Akun Google Anda dan tetapkan ID project dan zona default.
   auth login memberikan otorisasi ke gcloud untuk mengakses Google Cloud dengan kredensial pengguna Google.
   PROJECT_ID adalah Google Cloud nama project.
   ZONE adalah zona tempat Anda ingin membuat TPU.

    gcloud auth login
    gcloud config set project ${PROJECT_ID}
    gcloud config set compute/zone ${ZONE}
   

9. Buat identitas layanan untuk VM TPU.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}
   

Kapasitas aman

Hubungi tim penjualan/akun dukungan Cloud TPU Anda untuk meminta kuota TPU dan menjawab pertanyaan apa pun tentang kapasitas.

Menyediakan lingkungan Cloud TPU

TPU v6e dapat disediakan dan dikelola dengan GKE, dengan GKE dan XPK (alat CLI wrapper melalui GKE), atau sebagai resource dalam antrean.

Prasyarat

• Pastikan project Anda memiliki kuota TPUS_PER_TPU_FAMILY yang cukup, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud .
• v6e telah diuji dengan konfigurasi berikut:
  • python 3.10 atau yang lebih baru
  • Versi software harian:
    • JAX harian 0.4.32.dev20240912
    • LibTPU harian 0.1.dev20240912+nightly
  • Versi software stabil:
    • JAX + JAX Lib v0.4.37

• Pastikan project Anda memiliki cukup kuota TPU untuk:

  • Kuota VM TPU
  • Kuota Alamat IP

  • Kuota hyperdisk-balance

    Catatan: v6e mendukung Hyperdisk Balance dan Hyperdisk ML. Tidak mendukung saldo persistent disk (PD).

• Izin project pengguna

  • Jika Anda menggunakan GKE dengan XPK, lihat Izin Konsol Cloud di akun pengguna atau layanan untuk mengetahui izin yang diperlukan untuk menjalankan XPK.

Variabel lingkungan

Di Cloud Shell, buat variabel lingkungan berikut:

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-16
export ZONE=us-east1-d
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for provisioning Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the default network becoming overloaded.

export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Deskripsi flag perintah

Variabel	Deskripsi
NODE_ID	ID TPU yang ditetapkan pengguna yang dibuat saat permintaan resource yang diantrekan dialokasikan.
PROJECT_ID	Google Cloud Nama Project. Gunakan project yang sudah ada atau buat project baru di
ZONA	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
ACCELERATOR_TYPE	Lihat Jenis Akselerator.
RUNTIME_VERSION	v2-alpha-tpuv6e
SERVICE_ACCOUNT	Ini adalah alamat email untuk akun layanan yang dapat Anda temukan di Konsol Google Cloud -> IAM -> Akun Layanan Misalnya: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES	Jumlah slice yang akan dibuat (hanya diperlukan untuk Multislice).
QUEUED_RESOURCE_ID	ID teks yang ditetapkan pengguna untuk permintaan resource yang diantrekan.
VALID_DURATION	Durasi validitas permintaan resource yang diantrekan.
NETWORK_NAME	Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME	Nama firewall jaringan sekunder yang akan digunakan.

Pengoptimalan performa jaringan

Untuk performa terbaik,gunakan jaringan dengan MTU 8.896 (unit transmisi maksimum).

Secara default, Virtual Private Cloud (VPC) hanya menyediakan MTU sebesar 1.460 byte yang akan memberikan performa jaringan yang kurang optimal. Anda dapat menetapkan MTU jaringan VPC ke nilai apa pun antara 1.300 byte dan 8.896 byte (inklusif). Ukuran MTU kustom umum adalah 1.500 byte (Ethernet standar) atau 8.896 byte (maksimum yang memungkinkan). Untuk mengetahui informasi selengkapnya, lihat Ukuran MTU jaringan VPC yang valid.

Untuk informasi selengkapnya tentang cara mengubah setelan MTU untuk jaringan yang ada atau default, lihat Mengubah setelan MTU jaringan VPC.

Contoh berikut membuat jaringan dengan 8.896 MTU.

Catatan: RESOURCE_NAME adalah nama dasar jaringan dan firewall. Nama ini dapat berupa nama apa pun yang Anda pilih.

export RESOURCE_NAME=RESOURCE_NAME
export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork
export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall
export PROJECT=X
gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \
 --subnet-mode=auto --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network ${NETWORK_NAME}
 --allow tcp,icmp,udp --project=${PROJECT}

Menggunakan multi-NIC (Opsi untuk Multislice)

Variabel lingkungan berikut diperlukan untuk subnet sekunder saat Anda menggunakan lingkungan Multislice.

export NETWORK_NAME_2=${RESOURCE_NAME}
export SUBNET_NAME_2=${RESOURCE_NAME}
export FIREWALL_RULE_NAME=${RESOURCE_NAME}
export ROUTER_NAME=${RESOURCE_NAME}-network-2
export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2
export REGION=us-central2

Gunakan perintah berikut untuk membuat perutean IP kustom untuk jaringan dan subnet.

Catatan: Kedua subnetwork yang digunakan harus memiliki rentang IP yang saling eksklusif (termasuk rentang CIDR IP utama dan sekunder di subnetwork).

gcloud compute networks create "${NETWORK_NAME_2}" --mtu=8896
   --bgp-routing-mode=regional --subnet-mode=custom --project=${PROJECT_ID}
gcloud compute networks subnets create "${SUBNET_NAME_2}" \
   --network="${NETWORK_NAME_2}" \
   --range=10.10.0.0/18 --region="${REGION}" \
   --project=$PROJECT

gcloud compute firewall-rules create "${FIREWALL_RULE_NAME}" \
   --network "${NETWORK_NAME_2}" --allow tcp,icmp,udp \
   --source-ranges 10.10.0.0/18 --project=${PROJECT_ID}

gcloud compute routers create "${ROUTER_NAME}" \
  --project="${PROJECT_ID}" \
  --network="${NETWORK_NAME_2}" \
  --region="${REGION}"

gcloud compute routers nats create "${NAT_CONFIG}" \
  --router="${ROUTER_NAME}" \
  --region="${REGION}" \
  --auto-allocate-nat-external-ips \
  --nat-all-subnet-ip-ranges \
  --project="${PROJECT_ID}" \
  --enable-logging

Setelah slice multi-jaringan dibuat, Anda dapat memvalidasi bahwa kedua NIC sedang digunakan dengan menyiapkan cluster XPK dan menjalankan --command ifconfig sebagai bagian dari beban kerja XPK.

Gunakan perintah xpk workload berikut untuk menampilkan output perintah ifconfig di log konsol Cloud dan pastikan eth0 dan eth1 memiliki mtu=8896.

python3 xpk.py workload create \
   --cluster CLUSTER_NAME \
   (--base-docker-image maxtext_base_image|--docker-image CLOUD_IMAGE_NAME \
   --workload ${USER}-xpk-$ACCELERATOR_TYPE-$NUM_SLICES \
   --tpu-type=${ACCELERATOR_TYPE} \
   --num-slices=${NUM_SLICES}  \
   --on-demand \
   --zone $ZONE \
   --project $PROJECT_ID \
   [--enable-debug-logs] \
   [--use-vertex-tensorboard] \
   --command "ifconfig"

Pastikan eth0 dan eth1 memiliki mtu=8.896. cara untuk memverifikasi bahwa Anda memiliki multi-nic yang berjalan adalah dengan menjalankan perintah --command "ifconfig" sebagai bagian dari beban kerja XPK. Kemudian, lihat output yang dicetak dari beban kerja xpk tersebut di log konsol cloud dan pastikan eth0 dan eth1 memiliki mtu=8896.

Peningkatan setelan TCP

Untuk TPU yang dibuat menggunakan antarmuka resource yang diantrekan, Anda dapat menjalankan perintah berikut untuk meningkatkan performa jaringan dengan meningkatkan batas buffering penerimaan TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
  --project "$PROJECT" \
  --zone "$ZONE" \
  --node=all \
  --command='sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"' \
  --worker=all

Penyediaan dengan resource dalam antrean

Kapasitas yang dialokasikan dapat disediakan menggunakan perintah create resource dalam antrean.

1. Buat permintaan resource yang diantrekan TPU.

   Flag --reserved hanya diperlukan untuk resource yang dicadangkan, bukan untuk resource on demand.

   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} \
     [--reserved]
   
     # The following flags are only needed if you are using Multislice.
     --node-count node-count  # Number of slices in a Multislice \
     --node-prefix node-prefix # An optional user-defined node prefix;
      the default is QUEUED_RESOURCE_ID.

   Jika permintaan resource dalam antrean berhasil dibuat, status dalam kolom "response" akan berupa "WAITING_FOR_RESOURCES" atau "FAILED". Jika permintaan resource dalam antrean berada dalam status "WAITING_FOR_RESOURCES", resource tersebut telah ditambahkan ke antrean dan akan disediakan jika ada cukup kapasitas TPU yang dialokasikan. Jika permintaan resource yang diantrekan berada dalam status "FAILED", alasan kegagalan akan ada dalam output. Masa berlaku permintaan resource dalam antrean akan berakhir jika v6e tidak disediakan dalam durasi yang ditentukan, dan statusnya menjadi "GAGAL". Lihat dokumentasi publik Resource dalam Antrean untuk mengetahui informasi selengkapnya.

   Saat permintaan resource yang diantrekan berada dalam status "AKTIF", Anda dapat terhubung ke VM TPU menggunakan SSH. Gunakan perintah list atau describe untuk membuat kueri status resource yang diantrekan.

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

   Jika resource yang diantrekan berada dalam status "ACTIVE", output-nya akan mirip dengan berikut ini:

     state:
      state: ACTIVE
   

2. Mengelola VM TPU Anda. Untuk mengetahui opsi guna mengelola VM TPU, lihat mengelola VM TPU.

3. Menghubungkan ke VM TPU menggunakan SSH

   Anda dapat menginstal biner di setiap VM TPU dalam slice TPU dan menjalankan kode. Lihat bagian Jenis VM untuk menentukan jumlah VM yang akan dimiliki slice Anda.

   Untuk menginstal biner atau menjalankan kode, Anda dapat menggunakan SSH untuk terhubung ke VM menggunakan perintah tpu-vm ssh.

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
      --node=all # add this flag if you are using Multislice
   

   Untuk menggunakan SSH guna terhubung ke VM tertentu, gunakan flag --worker yang mengikuti indeks berbasis 0:

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --worker=1
   

   Jika memiliki bentuk slice lebih dari 8 chip, Anda akan memiliki beberapa VM dalam satu slice. Dalam hal ini, gunakan parameter --worker=all dan --command dalam perintah gcloud alpha compute tpus tpu-vm ssh untuk menjalankan perintah di semua VM secara bersamaan. Contoh:

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
     --zone  ${ZONE} --worker=all \
     --command='pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
     -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Menghapus resource yang diantrekan

   Menghapus resource yang diantrekan di akhir sesi atau menghapus permintaan resource yang diantrekan yang berada dalam status "GAGAL". Untuk menghapus resource dalam antrean, hapus slice, lalu permintaan resource dalam antrean dalam 2 langkah:

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

   Catatan: Anda juga dapat menghapus resource yang diantrekan dalam satu langkah menggunakan flag --force dengan perintah delete resource yang diantrekan:

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

Menyediakan TPU v6e dengan GKE atau XPK

Jika menggunakan perintah GKE dengan v6e, Anda dapat menggunakan perintah Kubernetes atau XPK untuk menyediakan TPU dan melatih atau menayangkan model. Lihat Merencanakan TPU di GKE untuk mempelajari cara merencanakan konfigurasi TPU di cluster GKE. Bagian berikut memberikan perintah untuk membuat cluster XPK dengan dukungan NIC tunggal dan dukungan multi NIC.

Perintah untuk membuat cluster XPK dengan dukungan NIC tunggal

export CLUSTER_NAME xpk-cluster-name
export ZONE=us-central2-b
export PROJECT=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME=${CLUSTER_NAME}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Catatan: Gunakan jaringan kustom untuk performa yang lebih baik serta menghindari jaringan default yang kelebihan beban.

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

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

export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

   python3 xpk.py cluster create --cluster $CLUSTER_NAME \
   --cluster-cpu-machine-type=n1-standard-8 \
   --num-slices=$NUM_SLICES \
   --tpu-type=$TPU_TYPE \
   --zone=$ZONE  \
   --project=$PROJECT \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"  \
   --create-vertex-tensorboard

Deskripsi flag perintah

Variabel	Deskripsi
CLUSTER_NAME	Nama yang ditetapkan pengguna untuk cluster XPK.
PROJECT_ID	Google Cloud Nama Project. Gunakan project yang sudah ada atau buat project baru di
ZONA	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
TPU_TYPE	Lihat Jenis Akselerator.
NUM_SLICES	Jumlah slice yang ingin Anda buat
CLUSTER_ARGUMENTS	Jaringan dan subnetwork yang akan digunakan. Misalnya: "--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
NUM_SLICES	Jumlah slice yang akan dibuat.
NETWORK_NAME	Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME	Nama firewall jaringan sekunder yang akan digunakan.

Perintah untuk membuat cluster XPK dengan dukungan multi-NIC

export CLUSTER_NAME xpk-cluster-name
export ZONE=us-central2-b
export PROJECT=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export exportSUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}

Catatan: Gunakan jaringan kustom untuk performa yang lebih baik serta menghindari jaringan default yang kelebihan beban.

   gcloud compute networks create "${NETWORK_NAME_1}" \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=$PROJECT

   gcloud compute networks subnets create "${SUBNET_NAME_1}" \
   --network="${NETWORK_NAME_1}" \
   --range=10.11.0.0/18 \
   --region="${REGION}" \
   --project=$PROJECT

   gcloud compute firewall-rules create "${FIREWALL_RULE_NAME}" \
   --network "${NETWORK_NAME_1}" \
   --allow tcp,icmp,udp \
   --project="${PROJECT}"

  gcloud compute routers create "${ROUTER_NAME}" \
    --project="${PROJECT}" \
    --network="${NETWORK_NAME_1}" \
    --region="${REGION}"

  gcloud compute routers nats create "${NAT_CONFIG}" \
     --router="${ROUTER_NAME}" \
     --region="${REGION}" \
     --auto-allocate-nat-external-ips \
     --nat-all-subnet-ip-ranges \
     --project="${PROJECT}" \
     --enable-logging

Secondary subnet for multi-nic experience. Need custom ip routing to be different from the first network's subnet.

export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}

   gcloud compute networks create "${NETWORK_NAME_2}" \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=$PROJECT

   gcloud compute networks subnets create "${SUBNET_NAME_2}" \
   --network="${NETWORK_NAME_2}" \
   --range=10.10.0.0/18 \
   --region="${REGION}" \
   --project=$PROJECT

   gcloud compute firewall-rules create "${FIREWALL_RULE_NAME}" \
   --network "${NETWORK_NAME_2}" \
   --allow tcp,icmp,udp \
   --project="${PROJECT}"

   gcloud compute routers create "${ROUTER_NAME}" \
     --project="${PROJECT}" \
     --network="${NETWORK_NAME_2}" \
     --region="${REGION}"

   gcloud compute routers nats create "${NAT_CONFIG}" \
     --router="${ROUTER_NAME}" \
     --region="${REGION}" \
     --auto-allocate-nat-external-ips \
     --nat-all-subnet-ip-ranges \
     --project="${PROJECT}" \
     --enable-logging

export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking
--network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"

export NODE_POOL_ARGUMENTS="--additional-node-network
network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

python3 ~/xpk/xpk.py cluster create \
--cluster $CLUSTER_NAME \
--num-slices=$NUM_SLICES \
--tpu-type=$TPU_TYPE \
--zone=$ZONE  \
--project=$PROJECT \
--on-demand \
--custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
--custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
--create-vertex-tensorboard

Deskripsi flag perintah

Variabel	Deskripsi
CLUSTER_NAME	Nama yang ditetapkan pengguna untuk cluster XPK.
PROJECT_ID	Google Cloud Nama Project. Gunakan project yang ada atau buat project baru di
ZONA	Lihat dokumen Region dan zona TPU untuk zona yang didukung.
TPU_TYPE	Lihat Jenis Akselerator.
NUM_SLICES	Jumlah slice yang ingin Anda buat
CLUSTER_ARGUMENTS	Jaringan dan subnetwork yang akan digunakan. Misalnya: "--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
NODE_POOL_ARGUMENTS	Jaringan node tambahan yang akan digunakan. Misalnya: "--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"
NUM_SLICES	Jumlah slice yang akan dibuat (hanya diperlukan untuk Multislice).
NETWORK_NAME	Nama jaringan sekunder yang akan digunakan.
NETWORK_FW_NAME	Nama firewall jaringan sekunder yang akan digunakan.

Penyiapan framework

Bagian ini menjelaskan proses penyiapan umum untuk pelatihan model ML menggunakan framework JAX, PyTorch, atau TensorFlow. Anda dapat menyediakan TPU menggunakan resource dalam antrean atau GKE. Penyiapan GKE dapat dilakukan menggunakan perintah XPK atau Kubernetes.

Penyiapan untuk JAX

Bagian ini memberikan contoh untuk menjalankan workload JAX di GKE, dengan atau tanpa XPK, serta menggunakan resource yang diantrekan.

Menyiapkan JAX menggunakan GKE

Contoh berikut menyiapkan satu host 2X2 menggunakan file YAML Kubernetes.

Satu slice di satu host

apiVersion: v1
kind: Pod
metadata:
  name: tpu-pod-jax-v6e-a
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
  containers:
  - name: tpu-job
    image: python:3.10
    securityContext:
      privileged: true
    command:
    - bash
    - -c
    - |
      pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python3 -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 4
      limits:
        google.com/tpu: 4

Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:

Total TPU chips: 4

Satu slice di multi-host

Contoh berikut menyiapkan node pool multi-host 4X4 menggunakan file YAML Kubernetes.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python -c 'import jax; print("Total TPU chips:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:

Total TPU chips: 16

Multi-slice di multi-host

Contoh berikut menyiapkan dua node pool multi-host 4X4 menggunakan file YAML Kubernetes.

Sebagai prasyarat, Anda harus menginstal JobSet v0.2.3 atau yang lebih baru.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          parallelism: 4
          completions: 4
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              containers:
              - name: jax-tpu
                image: python:3.10
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true python -c 'import jax; print("Total TPU chips:", jax.device_count())'
                resources:
                  limits:
                   google.com/tpu: 4
                  requests:
                   google.com/tpu: 4

Setelah berhasil diselesaikan, Anda akan melihat pesan berikut di log GKE:

Total TPU chips: 32

Untuk mengetahui informasi selengkapnya, lihat Menjalankan workload Multislice dalam dokumentasi GKE.

Untuk performa yang lebih baik, Aktifkan hostNetwork.

Multi-NIC

Untuk memanfaatkan multi-NIC di GKE, manifes Pod Kubernetes harus memiliki anotasi tambahan. Berikut adalah contoh manifes workload multi-NIC non-TPU.

apiVersion: v1
kind: Pod
metadata:
  name: sample-netdevice-pod-1
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth1","network":"netdevice-network"}
      ]
spec:
  containers:
  - name: sample-netdevice-pod
    image: busybox
    command: ["sleep", "infinity"]
    ports:
    - containerPort: 80
  restartPolicy: Always
  tolerations:
  - key: "google.com/tpu"
    operator: "Exists"
    effect: "NoSchedule"

Jika exec ke dalam Pod Kubernetes, Anda akan melihat NIC tambahan menggunakan kode berikut.

$ k exec --stdin --tty sample-netdevice-pod-1 -- /bin/sh
/ # ip a
1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue qlen 1000
    link/loopback 00:00:00:00:00:00 brd 00:00:00:00:00:00
    inet 127.0.0.1/8 scope host lo
       valid_lft forever preferred_lft forever
2: eth0@if11: <BROADCAST,MULTICAST,UP,LOWER_UP,M-DOWN> mtu 1460 qdisc noqueue
    link/ether da:be:12:67:d2:25 brd ff:ff:ff:ff:ff:ff
    inet 10.124.2.6/24 brd 10.124.2.255 scope global eth0
       valid_lft forever preferred_lft forever
3: eth1: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1460 qdisc mq qlen 1000
    link/ether 42:01:ac:18:00:04 brd ff:ff:ff:ff:ff:ff
    inet 172.24.0.4/32 scope global eth1
       valid_lft forever preferred_lft forever

Menyiapkan JAX menggunakan GKE dengan XPK

Lihat contoh di README xpk.

Untuk menyiapkan dan menjalankan XPK dengan MaxText, lihat: Cara menjalankan MaxText.

Menyiapkan JAX menggunakan resource yang diantrekan

Instal JAX di semua VM TPU dalam satu atau beberapa slice secara bersamaan menggunakan gcloud alpha compute tpus tpu-vm ssh. Untuk Multislice, tambahkan --node=all.

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
 --zone ${ZONE} --worker=all \
 --command='pip install -U --pre jax jaxlib libtpu-nightly requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html</code>'

Anda dapat menjalankan kode Python berikut untuk memeriksa jumlah core TPU yang tersedia di slice dan menguji apakah semuanya diinstal dengan benar (output yang ditampilkan di sini dihasilkan dengan slice v6e-16):

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

Outputnya mirip dengan hal berikut ini:

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() menunjukkan jumlah total chip dalam slice yang diberikan. jax.local_device_count() menunjukkan jumlah chip yang dapat diakses oleh satu VM dalam slice ini.

gcloud alpha compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} --zone=${ZONE} --worker=all  \
   --command='git clone -b mlperf4.1 https://github.com/google/maxdiffusion.git &&
   cd maxdiffusion && git checkout 975fdb7dbddaa9a53ad72a421cdb487dcdc491a3 &&
   && pip install -r requirements.txt  && pip install . '

Memecahkan masalah penyiapan JAX

Tips umum adalah mengaktifkan logging panjang di manifes beban kerja GKE Anda. Kemudian, berikan log ke dukungan GKE.

TPU_MIN_LOG_LEVEL=0 TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0

Pesan error

no endpoints available for service 'jobset-webhook-service'

Error ini berarti set tugas tidak diinstal dengan benar. Periksa untuk melihat apakah Pod Kubernetes deployment jobset-controller-manager sedang berjalan. Untuk mengetahui informasi selengkapnya, lihat dokumentasi pemecahan masalah JobSet untuk mengetahui detailnya.

TPU initialization failed: Failed to connect

Pastikan versi node GKE Anda adalah 1.30.4-gke.1348000 atau yang lebih baru (GKE 1.31 tidak didukung).

Penyiapan untuk PyTorch

Catatan: v6e hanya mendukung runtime PJRT dan PyTorch 2.1+ menggunakan PJRT sebagai runtime default untuk semua versi TPU.

Bagian ini menjelaskan cara mulai menggunakan PJRT di v6e dengan PyTorch/XLA. Python 3.10 adalah versi Python yang direkomendasikan.

Menyiapkan PyTorch menggunakan GKE dengan XPK

Anda dapat menggunakan penampung Docker berikut dengan XPK yang telah menginstal dependensi PyTorch:

us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_20241028

Untuk membuat beban kerja XPK, gunakan perintah berikut:

python3 xpk.py workload create \
--cluster ${CLUSTER_NAME} \
[--docker-image | --base-docker-image] us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_20241028 \
--workload ${USER} -xpk-${ACCELERATOR_TYPE} -${NUM_SLICES} \
--tpu-type=${ACCELERATOR_TYPE} \
--num-slices=${NUM_SLICES}  \
--on-demand \
--zone ${ZONE} \
--project ${PROJECT_ID} \
--enable-debug-logs \
--command 'python3 -c "import torch; import torch_xla; import torch_xla.runtime as xr; print(xr.global_runtime_device_count())"'

Menggunakan --base-docker-image akan membuat image Docker baru dengan direktori kerja saat ini yang di-build ke dalam Docker baru.

Menyiapkan PyTorch menggunakan resource yang diantrekan

Ikuti langkah-langkah berikut untuk menginstal PyTorch menggunakan resource yang diantrekan dan menjalankan skrip kecil di v6e.

Menginstal dependensi menggunakan SSH untuk mengakses VM

Untuk Multislice, tambahkan --node=all:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --worker=all \
    --command='sudo apt install -y libopenblas-base pip3 \
    install --pre torch==2.6.0.dev20241028+cpu torchvision==0.20.0.dev20241028+cpu \
    --index-url https://download.pytorch.org/whl/nightly/cpu
    pip install "torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20241028-cp310-cp310-linux_x86_64.whl" -f https://storage.googleapis.com/libtpu-releases/index.html
    pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'

Meningkatkan performa model dengan alokasi yang cukup besar dan sering

Untuk model yang memiliki alokasi yang cukup besar dan sering, kami telah mengamati bahwa penggunaan tcmalloc meningkatkan performa secara signifikan dibandingkan dengan implementasi malloc default, sehingga malloc default yang digunakan di VM TPU adalah tcmalloc. Namun, bergantung pada workload Anda (misalnya, dengan DLRM yang memiliki alokasi yang sangat besar untuk tabel penyematan), tcmalloc dapat menyebabkan pelambatan. Dalam hal ini, Anda dapat mencoba menghapus setelan variabel berikut menggunakan malloc default:

unset LD_PRELOAD

Gunakan skrip Python untuk melakukan penghitungan di VM v6e:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}
   --project ${PROJECT_ID} \
   --zone ${ZONE} --worker all --command='
   unset LD_PRELOAD
   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)"
'

Tindakan ini akan menghasilkan output yang mirip dengan berikut ini:

SSH: Attempting to connect to worker 0...
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
xla:0
tensor([[ 0.3355, -1.4628, -3.2610],
        [-1.4656,  0.3196, -2.8766],
        [ 0.8668, -1.5060,  0.7125]], device='xla:0')

Penyiapan untuk TensorFlow

Untuk Pratinjau Publik v6e, hanya versi runtime tf-nightly yang didukung.

Anda dapat mereset tpu-runtime dengan versi TensorFlow yang kompatibel dengan v6e dengan menjalankan perintah berikut:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} \
    --zone  ${ZONE} --worker=all --command="sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images\/grpc_tpu_worker:v6e\"/' /etc/systemd/system/tpu-runtime.service"
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
    --zone ${ZONE} --worker=all --command='sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime'

Gunakan SSH untuk mengakses worker-0:

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

Instal TensorFlow di pekerja-0:

sudo apt install -y libopenblas-base
pip install cloud-tpu-client
pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310
pip install cloud-tpu-client

pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
-f https://storage.googleapis.com/libtpu-tf-releases/index.html --force

Ekspor variabel lingkungan TPU_NAME:

export TPU_NAME=v6e-16

Anda dapat menjalankan skrip Python berikut untuk memeriksa jumlah core TPU yang tersedia di slice Anda dan untuk menguji apakah semuanya diinstal dengan benar (output yang ditampilkan dibuat dengan slice v6e-16):

import TensorFlow as tf
print("TensorFlow version " + tf.__version__)

@tf.function
  def add_fn(x,y):
  z = x + y
  return z

  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
  tf.config.experimental_connect_to_cluster(cluster_resolver)
  tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
  strategy = tf.distribute.TPUStrategy(cluster_resolver)

  x = tf.constant(1.)
  y = tf.constant(1.)
  z = strategy.run(add_fn, args=(x,y))
  print(z)

Outputnya mirip dengan hal berikut ini:

PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32)
}

v6e dengan SkyPilot

Anda dapat menggunakan TPU v6e dengan SkyPilot. Gunakan langkah-langkah berikut untuk menambahkan informasi lokasi/harga terkait v6e ke SkyPilot.

1. Tambahkan kode berikut ke akhir ~/.sky/catalogs/v5/gcp/vms.csv :

   ,,,tpu-v6e-1,1,tpu-v6e-1,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-1,1,tpu-v6e-1,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-1,1,tpu-v6e-1,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-4,1,tpu-v6e-4,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-4,1,tpu-v6e-4,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-4,1,tpu-v6e-4,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-8,1,tpu-v6e-8,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-8,1,tpu-v6e-8,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-8,1,tpu-v6e-8,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-16,1,tpu-v6e-16,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-16,1,tpu-v6e-16,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-16,1,tpu-v6e-16,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-32,1,tpu-v6e-32,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-32,1,tpu-v6e-32,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-32,1,tpu-v6e-32,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-64,1,tpu-v6e-64,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-64,1,tpu-v6e-64,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-64,1,tpu-v6e-64,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-128,1,tpu-v6e-128,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-128,1,tpu-v6e-128,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-128,1,tpu-v6e-128,us-east5,us-east5-b,0,0
   ,,,tpu-v6e-256,1,tpu-v6e-256,us-south1,us-south1-a,0,0
   ,,,tpu-v6e-256,1,tpu-v6e-256,europe-west4,europe-west4-a,0,0
   ,,,tpu-v6e-256,1,tpu-v6e-256,us-east5,us-east5-b,0,0
   

2. Tentukan resource berikut dalam file YAML:

   # tpu_v6.yaml
   resources:
     accelerators: tpu-v6e-16 # Fill in the accelerator type you want to use
     accelerator_args:
       runtime_version: v2-alpha-tpuv6e # Official suggested runtime
   

3. Luncurkan cluster dengan TPU v6e:

      sky launch tpu_v6.yaml -c tpu_v6
   

4. Hubungkan ke TPU v6e menggunakan SSH: ssh tpu_v6

Tutorial inferensi

Tutorial berikut menunjukkan cara menjalankan inferensi di TPU v6e:

• Inferensi MaxText JetStream pada v6e

• Inferensi PyTorch JetStream di v6e

• Inferensi MaxDiffusion di v6e

• Inferensi vLLM di v6e

Contoh pelatihan

Bagian berikut memberikan contoh untuk melatih model MaxText, MaxDiffusion, dan PyTorch di TPU v6e.

Pelatihan MaxText dan MaxDiffusion di VM Cloud TPU v6e

Bagian berikut membahas siklus proses pelatihan model MaxText dan MaxDiffusion.

Secara umum, langkah-langkah tingkat tingginya adalah:

1. Build image dasar workload.
2. Jalankan workload Anda menggunakan XPK.
   1. Buat perintah pelatihan untuk beban kerja.
   2. Deploy workload.
3. Ikuti beban kerja dan lihat metrik.
4. Hapus workload XPK jika tidak diperlukan.
5. Hapus cluster XPK jika tidak diperlukan lagi.

Membangun image dasar

Instal MaxText atau MaxDiffusion dan build image Docker:

1. Clone repositori yang ingin Anda gunakan dan ubah ke direktori untuk repositori tersebut:

   MaxText:

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

   MaxDiffusion:

   git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
   

2. Konfigurasikan Docker untuk menggunakan Google Cloud CLI:

   gcloud auth configure-docker
   

3. Build image Docker menggunakan perintah berikut atau menggunakan JAX Stable Stack. Untuk mengetahui informasi selengkapnya tentang JAX Stable Stack, lihat Mem-build image Docker dengan JAX Stable Stack.

   bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.37
   

4. Jika Anda meluncurkan beban kerja dari mesin yang tidak memiliki image yang di-build secara lokal, upload image:

   bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${USER}_runner
   

Mem-build image Docker dengan JAX Stable Stack

Penting: JAX Stable Stack masih dalam fase eksperimental.

Anda dapat mem-build image Docker MaxText dan MaxDiffusion menggunakan image dasar JAX Stable Stack.

JAX Stable Stack menyediakan lingkungan yang konsisten untuk MaxText dan MaxDiffusion dengan menggabungkan JAX dengan paket inti seperti orbax, flax, dan optax, beserta libtpu.so yang memenuhi syarat yang mendorong utilitas program TPU dan alat penting lainnya. Library ini diuji untuk memastikan kompatibilitas dan memberikan fondasi yang stabil untuk mem-build dan menjalankan MaxText dan MaxDiffusion. Hal ini akan menghilangkan potensi konflik karena versi paket yang tidak kompatibel.

JAX Stable Stack menyertakan libtpu.so yang dirilis sepenuhnya dan memenuhi syarat, library inti yang mendorong kompilasi, eksekusi, dan konfigurasi jaringan ICI program TPU. Rilis libtpu menggantikan build harian yang sebelumnya digunakan oleh JAX, dan memastikan fungsi komputasi XLA yang konsisten di TPU dengan pengujian kualifikasi tingkat PJRT di IR HLO/StableHLO.

Untuk mem-build image Docker MaxText dan MaxDiffusion dengan JAX Stable Stack, saat Anda menjalankan skrip docker_build_dependency_image.sh, tetapkan variabel MODE ke stable_stack dan tetapkan variabel BASEIMAGE ke image dasar yang ingin Anda gunakan.

Contoh berikut menentukan us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.37-rev1 sebagai image dasar:

bash docker_build_dependency_image.sh MODE=stable_stack
BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.37-rev1

Untuk mengetahui daftar image dasar JAX Stable Stack yang tersedia, lihat Image JAX Stable Stack di Artifact Registry.

Menjalankan workload menggunakan XPK

1. Tetapkan variabel lingkungan berikut jika Anda tidak menggunakan nilai default yang ditetapkan oleh MaxText atau MaxDiffusion:

   export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
   export PER_DEVICE_BATCH_SIZE=2
   export NUM_STEPS=30
   export MAX_TARGET_LENGTH=8192

2. Build skrip model Anda. Skrip ini akan disalin sebagai perintah pelatihan di langkah berikutnya.

   Jangan jalankan skrip model terlebih dahulu.

   Catatan: Pemeriksaan titik otomatis tidak diaktifkan. Jika Anda melihat error seperti, Could not find a known environment for initializing distributed JAX, gunakan JAX_FORCE_TPU_INIT=true.

   MaxText

   MaxText adalah LLM open source berperforma tinggi dan sangat skalabel yang ditulis dalam Python dan JAX murni serta menargetkan Google Cloud TPU dan GPU untuk pelatihan dan inferensi.

   JAX_PLATFORMS=tpu,cpu \
   ENABLE_PJRT_COMPATIBILITY=true \
   TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
   TPU_SLICE_BUILDER_DUMP_ICI=true && \
   python /deps/MaxText/train.py /deps/MaxText/configs/base.yml \
           base_output_directory=${BASE_OUTPUT_DIR} \
           dataset_type=synthetic \
           per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
           enable_checkpointing=false \
           gcs_metrics=true \
           profiler=xplane \
           skip_first_n_steps_for_profiler=5 \
           steps=${NUM_STEPS}  # attention='dot_product'"
   

   Gemma2

   Gemma adalah serangkaian LLM bobot terbuka yang dikembangkan oleh Google DeepMind, berdasarkan riset dan teknologi Gemini.

   Catatan: Gemma2 memerlukan TPU v6e-256 atau yang lebih besar, jadi Sediakan TPU v6e dengan 256 chip:

   python3 MaxText/train.py MaxText/configs/base.yml \
       model_name=gemma2-27b \
       run_name=gemma2-27b-run \
       base_output_directory=${BASE_OUTPUT_DIR} \
       max_target_length=${MAX_TARGET_LENGTH} \
       per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
       steps=${NUM_STEPS} \
       enable_checkpointing=false \
       use_iota_embed=true \
       gcs_metrics=true \
       dataset_type=synthetic \
       profiler=xplane \
       attention=flash
   

   Mixtral 8x7b

   Mixtral adalah model AI canggih yang dikembangkan oleh Mistral AI, yang memanfaatkan arsitektur campuran pakar (MoE) yang jarang.

   Catatan: Mixtral memerlukan TPU v6e-256 atau yang lebih besar, jadi Sediakan TPU v6e dengan 256 chip:

   python3 MaxText/train.py MaxText/configs/base.yml \
       base_output_directory=${BASE_OUTPUT_DIR} \
       per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
       model_name=mixtral-8x7b \
       steps=${NUM_STEPS} \
       max_target_length=${MAX_TARGET_LENGTH} \
       tokenizer_path=assets/tokenizer.mistral-v1 \
       attention=flash \
       dtype=bfloat16 \
       dataset_type=synthetic \
       profiler=xplane
   

   Llama3-8b

   Llama adalah serangkaian LLM bobot terbuka yang dikembangkan oleh Meta.

   Catatan: Llama3-8b memerlukan TPU v6e-256 atau yang lebih besar, jadi Sediakan TPU v6e dengan 256 chip:

   python3 MaxText/train.py MaxText/configs/base.yml \
       model_name=llama3-8b \
       base_output_directory=${BASE_OUTPUT_DIR} \
       dataset_type=synthetic \
       tokenizer_path=assets/tokenizer_llama3.tiktoken \
       per_device_batch_size=${PER_DEVICE_BATCH_SIZE} # set to 4 \
       gcs_metrics=true \
       profiler=xplane \
       skip_first_n_steps_for_profiler=5 \
       steps=${NUM_STEPS} \
       max_target_length=${MAX_TARGET_LENGTH} \
       attention=flash"
   

   MaxDiffusion

   MaxDiffusion adalah kumpulan implementasi referensi dari berbagai model difusi laten yang ditulis dalam Python murni dan JAX yang berjalan di perangkat XLA, termasuk Cloud TPU dan GPU. Stable Diffusion adalah model teks ke gambar laten yang menghasilkan gambar fotorealistik dari input teks apa pun.

   Anda perlu menginstal cabang Git tertentu untuk menjalankan MaxDiffusion seperti yang ditunjukkan dalam perintah git checkout berikut.

   git clone https://github.com/google/maxdiffusion.git
   && cd maxdiffusion
   && git checkout e712c9fc4cca764b0930067b6e33daae2433abf0
   && pip install -r requirements.txt
   && pip install .
   

   Skrip pelatihan:

       cd maxdiffusion && OUT_DIR=${BASE_OUTPUT_DIR} \
       python src/maxdiffusion/train_sdxl.py \
       src/maxdiffusion/configs/base_xl.yml \
       revision=refs/pr/95 \
       activations_dtype=bfloat16 \
       weights_dtype=bfloat16 \
       resolution=1024 \
       per_device_batch_size=1 \
       output_dir=${OUT_DIR}  \
       jax_cache_dir=${OUT_DIR}/cache_dir/ \
       max_train_steps=200 \
       attention=flash run_name=sdxl-ddp-v6e
   
       

   Catatan: Secara opsional, tetapkan dcn_data_parallelism untuk Pemrosesan multislice.

3. Jalankan model menggunakan skrip yang Anda buat pada langkah sebelumnya. Anda harus menentukan flag --base-docker-image untuk menggunakan gambar dasar MaxText atau menentukan flag --docker-image dan gambar yang ingin Anda gunakan.

   Opsional: Anda dapat mengaktifkan logging debug dengan menyertakan flag --enable-debug-logs. Untuk informasi selengkapnya, lihat Men-debug JAX di MaxText.

   Opsional: Anda dapat membuat Eksperimen Vertex AI untuk mengupload data ke Vertex AI TensorBoard dengan menyertakan tanda --use-vertex-tensorboard. Untuk mengetahui informasi selengkapnya, lihat Memantau JAX di MaxText menggunakan Vertex AI.

   python3 xpk.py workload create \
       --cluster CLUSTER_NAME \
       {--base-docker-image maxtext_base_image|--docker-image ${CLOUD_IMAGE_NAME}} \
       --workload ${USER}-xpk-$ACCELERATOR_TYPE-$NUM_SLICES \
       --tpu-type=$ACCELERATOR_TYPE \
       --num-slices=$NUM_SLICES  \
       --on-demand \
       --zone $ZONE \
       --project $PROJECT_ID \
       [--enable-debug-logs] \
       [--use-vertex-tensorboard] \
       --command $YOUR-MODEL-SCRIPT

   Ekspor variabel berikut:

   export ClUSTER_NAME=CLUSTER_NAME: Nama cluster XPK Anda. export ACCELERATOR_TYPEACCELERATOR_TYPE: Versi dan ukuran TPU Anda. Misalnya, v6e-256. export NUM_SLICES=NUM_SLICES: Jumlah slice TPU. export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT: Skrip model yang akan dieksekusi sebagai perintah pelatihan.

   Output-nya menyertakan link untuk mengikuti beban kerja Anda, mirip dengan berikut ini:

   [XPK] Follow your workload here: https://console.cloud.google.com/kubernetes/service/zone/project_id/default/workload_name/details?project=project_id
   

   Buka link dan klik tab Logs untuk melacak beban kerja Anda secara real time.

Men-debug JAX di MaxText

Gunakan perintah XPK tambahan untuk mendiagnosis alasan cluster atau beban kerja tidak berjalan.

• Daftar workload XPK
• XPK inspector
• Aktifkan logging panjang dalam log beban kerja menggunakan flag --enable-debug-logs saat Anda membuat beban kerja XPK.

Memantau JAX di MaxText menggunakan Vertex AI

Lihat data skalar dan profil melalui TensorBoard terkelola Vertex AI.

1. Tingkatkan permintaan pengelolaan resource (CRUD) untuk zona yang Anda gunakan dari 600 menjadi 5.000. Hal ini mungkin bukan masalah untuk beban kerja kecil yang menggunakan kurang dari 16 VM.

2. Instal dependensi seperti cloud-accelerator-diagnostics untuk Vertex AI:

   # xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
   cd ~/xpk
   pip install .

3. Buat cluster XPK menggunakan flag --create-vertex-tensorboard, seperti yang didokumentasikan dalam Membuat Vertex AI TensorBoard. Anda juga dapat menjalankan perintah ini di cluster yang ada.

4. Buat eksperimen Vertex AI saat menjalankan beban kerja XPK menggunakan tanda --use-vertex-tensorboard dan tanda --experiment-name opsional. Untuk mengetahui daftar langkah lengkap, lihat Membuat Vertex AI Experiment untuk mengupload data ke Vertex AI TensorBoard.

Catatan: Jika Anda menjalankan beban kerja non-MaxText, gunakan contoh MaxText ini untuk membuat perubahan upload data yang sesuai.

Log menyertakan link ke Vertex AI TensorBoard, mirip dengan hal berikut:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Anda juga dapat menemukan link Vertex AI TensorBoard di Konsol Google Cloud. Buka Vertex AI Experiments di konsol Google Cloud. Pilih wilayah yang sesuai dari drop-down.

Direktori TensorBoard juga ditulis ke bucket Cloud Storage yang Anda tentukan dengan ${BASE_OUTPUT_DIR}.

Menghapus workload XPK

Gunakan perintah xpk workload delete untuk menghapus satu atau beberapa workload berdasarkan awalan tugas atau status tugas. Perintah ini mungkin berguna jika Anda mengirim beban kerja XPK yang tidak perlu lagi dijalankan, atau jika Anda memiliki tugas yang macet di antrean.

Menghapus cluster XPK

Gunakan perintah xpk cluster delete untuk menghapus cluster:

python3 xpk.py cluster delete --cluster ${CLUSTER_NAME} \
--zone $ZONE --project $PROJECT_ID

Pelatihan Llama dan PyTorch/XLA di VM Cloud TPU v6e

Tutorial ini menjelaskan cara melatih model Llama menggunakan PyTorch/XLA di TPU v6e menggunakan set data WikiText.

Dapatkan akses ke Hugging Face dan model Llama 3

Anda memerlukan token akses pengguna Hugging Face untuk menjalankan tutorial ini. Untuk informasi tentang cara membuat dan menggunakan token akses pengguna, lihat dokumentasi Hugging Face tentang token akses pengguna.

Anda juga memerlukan izin untuk mengakses model Llama 3 8B di Hugging Face. Untuk mendapatkan akses, buka model Meta-Llama-3-8B di Hugging Face dan minta akses.

Membuat VM TPU

Buat TPU v6e dengan 8 chip untuk menjalankan tutorial.

1. Siapkan variabel lingkungan:

   export ACCELERATOR_TYPE=v6e-8
   export VERSION=v2-alpha-tpuv6e
   export TPU_NAME=$USER-$ACCELERATOR_TYPE
   export PROJECT=YOUR_PROJECT
   export ZONE=YOUR_ZONE

2. Buat VM TPU:

   gcloud alpha compute tpus tpu-vm create $TPU_NAME --version=$VERSION \
       --accelerator-type=$ACCELERATOR_TYPE --zone=$ZONE --project=$PROJECT

Penginstalan

Instal pytorch-tpu/transformers fork dari Hugging Face Transformers dan dependensinya. Tutorial ini diuji dengan versi dependensi berikut yang digunakan dalam contoh ini:

• torch: kompatibel dengan 2.5.0
• torch_xla[tpu]: kompatibel dengan 2.5.0
• jax: 0.4.33
• jaxlib: 0.4.33

gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT --zone $ZONE \
    --worker=all --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
    cd transformers
    sudo pip3 install -e .
    pip3 install datasets
    pip3 install evaluate
    pip3 install scikit-learn
    pip3 install accelerate
    pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
    pip install jax==0.4.33 jaxlib==0.4.33 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'

Menyiapkan konfigurasi model

Perintah pelatihan di bagian berikutnya, Jalankan model, menggunakan dua file konfigurasi JSON untuk menentukan parameter model dan konfigurasi FSDP (Fully Sharded Data Parallel). Sharding FSDP digunakan untuk bobot model agar sesuai dengan ukuran batch yang lebih besar selama pelatihan. Saat melatih dengan model yang lebih kecil, Anda mungkin cukup menggunakan paralelisme data dan mereplikasi bobot di setiap perangkat. Untuk mengetahui informasi selengkapnya tentang cara membuat shard tensor di seluruh perangkat di PyTorch/XLA, lihat Panduan Pengguna SPMD PyTorch/XLA.

1. Buat file konfigurasi parameter model. Berikut adalah konfigurasi parameter model untuk Llama3-8B. Untuk model lainnya, temukan konfigurasi di Hugging Face. Misalnya, lihat konfigurasi Llama2-7B.

   Penting: Tetapkan use_cache ke False untuk pelatihan.

   cat > llama-config.json <
   {
       "architectures": [
           "LlamaForCausalLM"
       ],
       "attention_bias": false,
       "attention_dropout": 0.0,
       "bos_token_id": 128000,
       "eos_token_id": 128001,
       "hidden_act": "silu",
       "hidden_size": 4096,
       "initializer_range": 0.02,
       "intermediate_size": 14336,
       "max_position_embeddings": 8192,
       "model_type": "llama",
       "num_attention_heads": 32,
       "num_hidden_layers": 32,
       "num_key_value_heads": 8,
       "pretraining_tp": 1,
       "rms_norm_eps": 1e-05,
       "rope_scaling": null,
       "rope_theta": 500000.0,
       "tie_word_embeddings": false,
       "torch_dtype": "bfloat16",
       "transformers_version": "4.40.0.dev0",
       "use_cache": false,
       "vocab_size": 128256
   }
   EOF

2. Buat file konfigurasi FSDP:

   cat > fsdp-config.json <
   {
       "fsdp_transformer_layer_cls_to_wrap": [
           "LlamaDecoderLayer"
       ],
       "xla": true,
       "xla_fsdp_v2": true,
       "xla_fsdp_grad_ckpt": true
   }
   EOF

   Untuk informasi selengkapnya tentang FSDP, lihat FSDPv2.

3. Upload file konfigurasi ke VM TPU menggunakan perintah berikut:

   gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json $TPU_NAME:. \
       --worker=all \
       --project=$PROJECT \
       --zone $ZONE

Menjalankan model

Dengan menggunakan file konfigurasi yang Anda buat di bagian sebelumnya, jalankan skrip run_clm.py untuk melatih model Llama 3 8B pada set data WikiText. Skrip pelatihan memerlukan waktu sekitar 10 menit untuk dijalankan di TPU v6e-8.

1. Login ke Hugging Face di TPU Anda menggunakan perintah berikut:

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT \
       --zone $ZONE \
       --worker=all \
       --command='
       pip3 install "huggingface_hub[cli]"
       huggingface-cli login --token HUGGING_FACE_TOKEN'

2. Jalankan pelatihan model:

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT \
       --zone $ZONE \
       --worker=all \
       --command='
       export PJRT_DEVICE=TPU
       export XLA_USE_SPMD=1
       export ENABLE_PJRT_COMPATIBILITY=true
           # Optional variables for debugging:
       export XLA_IR_DEBUG=1
       export XLA_HLO_DEBUG=1
       export PROFILE_EPOCH=0
       export PROFILE_STEP=3
       export PROFILE_DURATION_MS=100000
           # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
       export PROFILE_LOGDIR=PROFILE_PATH
       python3 transformers/examples/pytorch/language-modeling/run_clm.py \
       --dataset_name wikitext \
       --dataset_config_name wikitext-2-raw-v1 \
       --per_device_train_batch_size 16 \
       --do_train \
       --output_dir /home/$USER/tmp/test-clm \
       --overwrite_output_dir \
       --config_name /home/$USER/llama-config.json \
       --cache_dir /home/$USER/cache \
       --tokenizer_name meta-llama/Meta-Llama-3-8B \
       --block_size 8192 \
       --optim adafactor \
       --save_strategy no \
       --logging_strategy no \
       --fsdp "full_shard" \
       --fsdp_config /home/$USER/fsdp-config.json \
       --torch_dtype bfloat16 \
       --dataloader_drop_last yes \
       --flash_attention \
       --max_steps 20'

Memecahkan masalah PyTorch/XLA

Jika Anda menetapkan variabel opsional untuk proses debug di bagian sebelumnya, profil untuk model akan disimpan di lokasi yang ditentukan oleh variabel PROFILE_LOGDIR. Anda dapat mengekstrak file xplane.pb yang disimpan di lokasi ini dan menggunakan tensorboard untuk melihat profil di browser menggunakan petunjuk TensorBoard Jika PyTorch/XLA tidak berperforma seperti yang diharapkan, lihat panduan pemecahan masalah, yang berisi saran untuk men-debug, membuat profil, dan mengoptimalkan model Anda.

Pelatihan DLRM DCN v2 di v6e

Tutorial ini menunjukkan cara melatih model DLRM DCN v2 di TPU v6e. Anda perlu Menyediakan TPU v6e dengan 64, 128, atau 256 chip.

Jika Anda menjalankan di multi-host, reset tpu-runtime dengan versi TensorFlow yang sesuai dengan menjalankan perintah berikut. Jika menjalankan di satu host, Anda tidak perlu menjalankan dua perintah berikut.

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID}
--zone  ${ZONE} --worker=all \
--command="sudo sed -i 's/TF_DOCKER_URL=.*/TF_DOCKER_URL=gcr.io\/cloud-tpu-v2-images\/grpc_tpu_worker:v6e\"/' /etc/systemd/system/tpu-runtime.service"

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
 --zone  ${ZONE}   \
 --worker=all \
 --command='sudo systemctl daemon-reload && sudo systemctl restart tpu-runtime'

SSH ke worker-0

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

Menetapkan nama TPU

export TPU_NAME=${TPU_NAME}

Menjalankan DLRM v2

Catatan: Salin cuplikan kode berikut dalam file bernama script.sh

pip install --user setuptools==65.5.0

pip install cloud-tpu-client

pip install gin-config && pip install tensorflow-datasets && pip install tf-keras-nightly --no-deps

pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -f https://storage.googleapis.com/libtpu-tf-releases/index.html --force

git clone https://github.com/tensorflow/recommenders.git
git clone https://github.com/tensorflow/models.git

export PYTHONPATH=~/recommenders/:~/models/
export TF_XLA_FLAGS='--tf_mlir_enable_mlir_bridge=true --tf_xla_sparse_core_disable_table_stacking=true --tf_mlir_enable_convert_control_to_data_outputs_pass=true --tf_mlir_enable_merge_control_flow_pass=true'

TF_USE_LEGACY_KERAS=1 TPU_LOAD_LIBRARY=0 python3 ./models/official/recommendation/ranking/train.py  --mode=train     --model_dir=gs://ptxla-debug/tf/sc/dlrm/runs/2/ --params_override="
runtime:
  distribution_strategy: tpu
  mixed_precision_dtype: 'mixed_bfloat16'
task:
  use_synthetic_data: false
  use_tf_record_reader: true
  train_data:
    input_path: 'gs://trillium-datasets/criteo/train/day_*/*'
    global_batch_size: 16384
    use_cached_data: true
  validation_data:
    input_path: 'gs://trillium-datasets/criteo/eval/day_*/*'
    global_batch_size: 16384
    use_cached_data: true
  model:
    num_dense_features: 13
    bottom_mlp: [512, 256, 128]
    embedding_dim: 128
    interaction: 'multi_layer_dcn'
    dcn_num_layers: 3
    dcn_low_rank_dim: 512
    size_threshold: 8000
    top_mlp: [1024, 1024, 512, 256, 1]
    use_multi_hot: true
    concat_dense: false
    dcn_use_bias: true
    vocab_sizes: [40000000,39060,17295,7424,20265,3,7122,1543,63,40000000,3067956,405282,10,2209,11938,155,4,976,14,40000000,40000000,40000000,590152,12973,108,36]
    multi_hot_sizes: [3,2,1,2,6,1,1,1,1,7,3,8,1,6,9,5,1,1,1,12,100,27,10,3,1,1]
    max_ids_per_chip_per_sample: 128
    max_ids_per_table: [280, 128, 64, 272, 432, 624, 64, 104, 368, 352, 288, 328, 304, 576, 336, 368, 312, 392, 408, 552, 2880, 1248, 720, 112, 320, 256]
    max_unique_ids_per_table: [104, 56, 40, 32, 72, 32, 40, 32, 32, 144, 64, 192, 32, 40, 136, 32, 32, 32, 32, 240, 1352, 432, 120, 80, 32, 32]
    use_partial_tpu_embedding: false
    size_threshold: 0
    initialize_tables_on_host: true
trainer:
  train_steps: 10000
  validation_interval: 1000
  validation_steps: 660
  summary_interval: 1000
  steps_per_loop: 1000
  checkpoint_interval: 0
  optimizer_config:
    embedding_optimizer: 'Adagrad'
    dense_optimizer: 'Adagrad'
    lr_config:
      decay_exp: 2
      decay_start_steps: 70000
      decay_steps: 30000
      learning_rate: 0.025
      warmup_steps: 0
    dense_sgd_config:
      decay_exp: 2
      decay_start_steps: 70000
      decay_steps: 30000
      learning_rate: 0.00025
      warmup_steps: 8000
  train_tf_function: true
  train_tf_while_loop: true
  eval_tf_while_loop: true
  use_orbit: true
  pipeline_sparse_and_dense_execution: true"

Jalankan script.sh:

chmod +x script.sh
./script.sh

Catatan: Jika Anda menjalankan Tensorflow di penyiapan GKE, instal TF-TPU whl + libtpu melalui perintah berikut:

pip install https://storage.googleapis.com/tensorflow-public-build-artifacts/prod/tensorflow/official/release/nightly/linux_x86_tpu/wheel_py310/749/20240915-062017/github/tensorflow/build_output/tf_nightly_tpu-2.18.0.dev20240915-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl \
-f https://storage.googleapis.com/libtpu-tf-releases/index.html --force

Flag berikut diperlukan untuk menjalankan workload rekomendasi (DLRM DCN):

ENV TF_XLA_FLAGS='--tf_mlir_enable_mlir_bridge=true \
--tf_mlir_enable_tpu_variable_runtime_reformatting_pass=false \
--tf_mlir_enable_convert_control_to_data_outputs_pass=true \
--tf_mlir_enable_merge_control_flow_pass=true --tf_xla_disable_full_embedding_pipelining=true' \
ENV LIBTPU_INIT_ARGS="--xla_sc_splitting_along_feature_dimension=auto \
--copy_with_dynamic_shape_op_output_pjrt_buffer=true"

Hasil benchmark

Bagian berikut berisi hasil benchmark untuk DLRM DCN v2 dan MaxDiffusion di v6e.

DLRM DCN v2

Skrip pelatihan DLRM DCN v2 dijalankan pada skala yang berbeda. Lihat throughput dalam tabel berikut.

	v6e-64	v6e-128	v6e-256
Langkah-langkah pelatihan	7.000	7.000	7.000
Ukuran batch global	131072	262144	524288
Throughput (contoh/dtk)	2975334	5111808	10066329

MaxDiffusion

Kami menjalankan skrip pelatihan untuk MaxDiffusion di v6e-4, v6e-16, dan 2xv6e-16. Lihat throughput dalam tabel berikut.

	v6e-4	v6e-16	Dua v6e-16
Langkah-langkah pelatihan	0,069	0,073	0,13
Ukuran batch global	8	32	64
Throughput (contoh/dtk)	115,9	438,4	492,3

Penjadwalan pengumpulan

Trillium (v6e) menyertakan fitur baru yang disebut "penjadwalan pengumpulan". Fitur ini menawarkan cara untuk mengelola beberapa slice TPU yang menjalankan beban kerja inferensi satu host di GKE dan Cloud TPU API. Mengelompokkan slice ini ke dalam koleksi akan memudahkan penyesuaian jumlah replika agar sesuai dengan permintaan. Update software dikontrol dengan cermat untuk memastikan bahwa sebagian slice dalam koleksi selalu tersedia untuk menangani traffic yang masuk.

Lihat dokumentasi GKE untuk mengetahui informasi selengkapnya tentang penggunaan penjadwalan pengumpulan dengan GKE.

Fitur penjadwalan pengumpulan hanya berlaku untuk v6e.

Catatan: Untuk GKE, koleksi harus ditentukan pada saat node GKE disediakan. Anda tidak dapat menambahkan node yang sudah disediakan ke dalam koleksi. Koleksi dapat diskalakan (lebih banyak node dapat ditambahkan ke dalamnya) atau diskalakan ke bawah (node yang ada dapat dihapus darinya).

Menggunakan penjadwalan pengumpulan dari Cloud TPU API

Koleksi satu host di Cloud TPU API adalah resource yang diantrekan yang menerapkan tanda khusus (--workload-type = availability-optimized) untuk menunjukkan ke infrastruktur yang mendasarinya bahwa resource tersebut dimaksudkan untuk digunakan untuk menayangkan workload.

Perintah berikut menyediakan koleksi host tunggal menggunakan Cloud TPU API:

gcloud alpha compute tpus queued-resources create my-collection \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --accelerator-type $ACCELERATOR_TYPE \
   --node-count ${NODE_COUNT} \
   --workload-type=availability-optimized

Catatan: Jenis akselerator dapat berupa v6e-1, v6e-4, atau v6e-8 seperti yang dijelaskan dalam konfigurasi yang didukung v6e.

Memantau dan membuat profil

Cloud TPU v6e mendukung pemantauan dan pembuatan profil menggunakan metode yang sama dengan Cloud TPU generasi sebelumnya. Untuk mengetahui informasi selengkapnya tentang pemantauan, lihat Memantau VM TPU.