Halaman ini diterjemahkan oleh Cloud Translation API.

Men-deploy TPU Multislice di GKE

Autopilot Standard

Halaman ini menunjukkan cara men-deploy workload di Google Kubernetes Engine (GKE) menggunakan konfigurasi Multislice Cloud TPU untuk pelatihan skala besar yang hemat biaya.

Tutorial ini ditujukan untuk engineer Machine Learning (ML) serta Admin dan operator platform yang ingin menggunakan orkestrasi container Kubernetes untuk mengelola beban kerja pelatihan, penyesuaian, dan inferensi model berskala besar menggunakan TPU. Untuk mempelajari lebih lanjut peran umum dan contoh tugas yang dirujuk dalam Google Cloud konten, lihat Peran dan tugas pengguna GKE umum.

Sebelum mengonfigurasi Multislice di GKE, pastikan Anda memahami konsep berikut:

Apa itu TPU Multislice

Multislice TPU adalah organisasi arsitektur VM dalam slice TPU yang memungkinkan dua atau beberapa slice Cloud TPU berkomunikasi melalui Jaringan Pusat Data (DCN). Multislice memungkinkan pelatihan skala besar, hemat biaya, dan full-stack dengan penskalaan near-linear hingga puluhan ribu chip TPU. Dalam konfigurasi Multislice, GKE men-deploy workload Multislice di beberapa slice TPU. Komunikasi antara chip TPU dalam slice terjadi melalui interkoneksi antar-chip (ICI). Komunikasi antar-slice terjadi melalui DCN.

Sebaiknya gunakan Multislice jika Tugas Anda terlalu besar untuk dimuat dalam satu slice TPU.

Ketersediaan multislice di GKE

Standard mendukung Multislice di versi 1.27.4-gke.900 dan yang lebih baru.
Autopilot mendukung Multislice di versi 1.29.2-gke.1521000 dan yang lebih baru.
Multislice mendukung framework JAX dan PyTorch. Versi JAX minimum yang didukung adalah 2.1.
Multislice hanya mendukung node pool slice TPU multi-host. Misalnya, Anda tidak dapat menggunakan Multislice dengan ct4p-hightpu-4t dengan topologi 2x2x1 atau ct5lp-hightpu-4t dengan topologi 2x2, karena ini adalah node pool slice TPU host tunggal.
Multislice hanya mendukung pelatihan multi-pengontrol sinkron.
Workload multislice hanya dapat berjalan di seluruh slice TPU yang memiliki jenis, ukuran, dan topologi TPU yang sama.
Multislice tidak mendukung TPU v3.

Sebelum memulai

Sebelum memulai, pastikan Anda telah melakukan tugas berikut:

Aktifkan Google Kubernetes Engine API.

Aktifkan Google Kubernetes Engine API

Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan perintah gcloud components update. gcloud CLI versi sebelumnya mungkin tidak mendukung menjalankan perintah dalam dokumen ini.
Catatan: Untuk penginstalan gcloud CLI yang ada, pastikan untuk menyetel properti compute/region. Jika Anda terutama menggunakan cluster zona, tetapkan compute/zone. Dengan menyetel lokasi default, Anda dapat menghindari error di gcloud CLI yang seperti ini: One of [--zone, --region] must be supplied: Please specify location. Anda mungkin perlu menentukan lokasi dalam perintah tertentu jika lokasi cluster Anda berbeda dengan lokasi default yang Anda tetapkan.

Buat cluster Standard atau cluster Autopilot yang menjalankan versi yang mendukung Multislice. Untuk versi yang didukung, lihat Ketersediaan multi-slice di GKE.
Pastikan project Anda memiliki kuota yang cukup untuk Cloud TPU di GKE.
Instal JobSet v0.2.3 atau yang lebih baru.

Menjalankan workload di Multislice

Bagian ini menunjukkan cara menjalankan workload di Multislice. Jika Anda menggunakan mode GKE Autopilot, lanjutkan ke bagian Menjalankan workload Multislice. Cluster Autopilot yang menjalankan versi 1.29.2-gke.1521000 atau yang lebih baru mengaktifkan TPU secara default.

Menyiapkan node pool mode Standar

Bagian ini membahas langkah-langkah berikut:

Buat tiga node pool slice TPU multi-host
Memverifikasi status node pool

Buat node pool slice TPU

Anda dapat membuat lebih dari satu node pool slice TPU multi-host. Untuk tujuan panduan ini, buat tiga node pool slice TPU multi-host untuk menjalankan workload Multislice. Langkah-langkah untuk membuat node pool slice TPU multi-host berbeda-beda, bergantung pada apakah Anda menggunakan Ironwood (TPU7x) atau versi TPU yang lebih lama.

Ironwood (TPU7x)

Anda dapat membuat node pool slice TPU multi-host di versi Ironwood (TPU7x) menggunakan Google Cloud CLI atau Terraform:

gcloud

Untuk membuat node pool slice TPU multi-host dengan Ironwood (TPU7x), Anda harus membuat kebijakan workload terlebih dahulu.

Buat kebijakan beban kerja:
```
gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION
```
Ganti kode berikut:
- WORKLOAD_POLICY_NAME: nama untuk kebijakan workload Anda.
- TPU_TOPOLOGY: topologi TPU Ironwood (TPU7x). Contoh, 2x2x2. Untuk melihat semua topologi Ironwood (TPU7x) yang didukung, lihat bagian topologi.
- PROJECT_ID: Google Cloud Project ID Anda.
- REGION: region untuk kebijakan workload. Kebijakan workload adalah resource regional dan dapat digunakan kembali di seluruh node pool yang memiliki topologi yang sama.
Buat node pool dengan kebijakan workload:
```
gcloud container node-pools create NODE_POOL_NAME \
  --cluster=CLUSTER_NAME \
  --machine-type=tpu7x-standard-4t \
  --placement-policy=WORKLOAD_POLICY_NAME \
  --location=CONTROL_PLANE_LOCATION \
  --node-locations=NODE_ZONE \
  --project=PROJECT_ID \
  --reservation=RESERVATION_NAME \
  --reservation-affinity=specific
```
Ganti kode berikut:
- NODE_POOL_NAME: nama untuk node pool baru.
- CLUSTER_NAME: nama cluster GKE Anda.
- WORKLOAD_POLICY_NAME: nama kebijakan workload yang Anda buat.
- CONTROL_PLANE_LOCATION: lokasi Compute Engine bidang kontrol cluster Anda. Berikan region untuk cluster regional, atau zona untuk cluster zona.
- NODE_ZONE: nama zona berdasarkan versi TPU yang ingin Anda gunakan. Untuk mengidentifikasi lokasi yang tersedia, lihat Ketersediaan TPU di GKE.
- PROJECT_ID: Google Cloud Project ID Anda.
- RESERVATION_NAME: nama reservasi yang akan digunakan.
Dalam perintah ini, flag --tpu-topology telah diganti dengan flag --placement-policy.

Terraform

Pastikan Anda menggunakan versi 4.84.0 atau yang lebih baru dari penyedia google.
Buat kebijakan beban kerja:
```
resource "google_compute_resource_policy" {
  name   = "WORKLOAD_POLICY_NAME"
  region = CLUSTER_LOCATION
  workload_policy {
    type = "HIGH_THROUGHPUT"
    accelerator_topology = "TPU_TOPOLOGY"
  }
}
```
Ganti kode berikut:
- WORKLOAD_POLICY_NAME: nama untuk kebijakan workload Anda.
- CLUSTER_LOCATION: Lokasi komputasi untuk cluster. Sebaiknya Anda memiliki cluster regional untuk keandalan panel kontrol Kubernetes yang lebih tinggi. Anda juga dapat menggunakan cluster zonal. Untuk mengetahui informasi selengkapnya, lihat Memilih topologi dan versi TPU.
- TPU_TOPOLOGY: topologi TPU Ironwood (TPU7x). Contoh, 2x2x2. Untuk melihat semua topologi Ironwood (TPU7x) yang didukung, lihat Merencanakan TPU.
Untuk mengetahui informasi selengkapnya tentang referensi google_compute_resource_policy, lihat Penyedia Terraform.
Dalam konfigurasi Terraform, tambahkan blok berikut:
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    flex_start = false
    spot = true
  }

  placement_policy {
    policy_name = WORKLOAD_POLICY_NAME
  }
}
```
Ganti kode berikut:
- NODE_POOL_RESOURCE_NAME: nama resource node pool di template Terraform.
- PROJECT_ID: Project ID Anda.
- CLUSTER_NAME: nama cluster yang ada tempat node pool akan ditambahkan.
- POOL_NAME: nama node pool yang akan dibuat.
- NODE_ZONES: daftar yang dipisahkan koma untuk satu atau beberapa zona tempat GKE membuat node pool.
- NUM_NODES: jumlah node dalam node pool. Nilainya harus nol atau hasil kali jumlah TPU chip yang dibagi empat, karena dalam slice TPU multi-host, setiap node slice TPU memiliki empat chip. Misalnya, jika TPU_TOPOLOGY adalah 4x8, maka ada 32 chip, yang berarti NUM_NODES harus 8. Untuk mempelajari topologi TPU lebih lanjut, gunakan tabel di Memilih versi TPU.
- TPU_TOPOLOGY: ini menunjukkan topologi fisik yang dipilih untuk slice TPU. Format topologi bergantung pada versi TPU yang Anda gunakan. Untuk mempelajari lebih lanjut topologi TPU, gunakan tabel di Memilih topologi.
Secara opsional, Anda juga dapat menggunakan variabel berikut:
- RESERVATION_NAME: jika Anda menggunakan pemesanan TPU, berikan daftar label resource pemesanan yang akan digunakan saat membuat node pool. Untuk mempelajari cara mengisi RESERVATION_LABEL_VALUES di kolom reservation_affinity lebih lanjut, silakan melihat Penyedia Terraform.
- autoscaling: membuat node pool dengan penskalaan otomatis yang diaktifkan. Saat GKE menskalakan node pool slice TPU multi-host, GKE secara atomik akan meningkatkan skala node pool dari nol hingga ukuran maksimum.
  - MAX_NODES: ukuran maksimum node pool. Nilai harus sama dengan hasil perkalian nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM. Misalnya, jika TPU_TOPOLOGY adalah 2x2x2, produknya adalah 8. Karena setiap VM di tpu7x-standard-4t memiliki 4 chip, jumlah node adalah 2.
- spot: node pool yang akan menggunakan Spot VM untuk node slice TPU. Setelan ini tidak dapat diubah setelah node pool dibuat. Untuk mengetahui informasi selengkapnya, silakan melihat Spot VM.
- flex_start: node pool yang akan menggunakan opsi konsumsi flex-start. Setelan ini tidak dapat disetel ke true jika spot diaktifkan.

Versi TPU lainnya

Anda dapat membuat node pool slice TPU multi-host di versi v3, v4, v5p, v5e, dan Trillium (v6e) menggunakan Google Cloud CLI, Terraform, atau konsol Google Cloud .

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      [--num-nodes=NUM_NODES] \
      [--spot \]
      [--flex-start \]
      [--enable-autoscaling \
        --max-nodes MAX_NODES]
      [--reservation-affinity=specific \
      --reservation=RESERVATION_NAME] \
      [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
      [--placement-type=COMPACT]

Ganti kode berikut:

POOL_NAME: nama node pool baru.
CONTROL_PLANE_LOCATION: lokasi Compute Engine untuk bidang kontrol cluster Anda. Berikan region untuk cluster regional, atau zona untuk cluster zona.
CLUSTER_NAME: nama cluster.
NODE_ZONES: nama zona berdasarkan versi TPU yang ingin Anda gunakan. Untuk mengidentifikasi lokasi yang tersedia, lihat Ketersediaan TPU di GKE.
MACHINE_TYPE: jenis mesin yang akan digunakan untuk node. Untuk mempelajari jenis mesin yang tersedia lebih lanjut, lihat Memilih versi TPU.
TPU_TOPOLOGY: topologi fisik untuk slice TPU. Format topologi bergantung pada versi TPU. Untuk mengetahui informasi selengkapnya tentang topologi TPU, gunakan tabel di Memilih topologi.

Untuk mengetahui informasi selengkapnya, lihat Topologi.

Secara opsional, Anda juga dapat menggunakan flag berikut:
NUM_NODES: jumlah node dalam node pool. Nilainya harus nol atau hasil kali nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM. Untuk TPU v4 dan TPU v5e multi-host, jumlah chip di setiap VM adalah empat. Oleh karena itu, jika TPU_TOPOLOGY Anda adalah 2x4x4 (TPU v4 dengan empat chip di setiap VM), maka NUM_NODES adalah 32/4 yang sama dengan 8. Jika Anda menghapus tanda ini, jumlah node akan dihitung dan ditetapkan secara default berdasarkan topologi dan jenis mesin.
RESERVATION_NAME: nama reservasi yang digunakan GKE saat membuat node pool. Jika Anda menghapus flag ini, GKE akan menggunakan node pool slice TPU yang tersedia. Untuk mengetahui informasi selengkapnya tentang pemesanan TPU, lihat Pemesanan TPU.
--spot: menetapkan node pool agar menggunakan Spot VM untuk node slice TPU. Nama ini tidak dapat diubah setelah pembuatan node pool. Untuk mengetahui informasi selengkapnya, silakan membaca Spot VM.
--flex-start: menetapkan node pool agar menggunakan VM mulai fleksibel. VM mulai fleksibel dibuat menggunakan opsi konsumsi flex-start, yang didukung di GKE versi 1.33.0-gke.1712000 atau yang lebih baru.
--enable-autoscaling: Membuat node pool dengan penskalaan otomatis yang diaktifkan. Saat GKE menskalakan node pool slice TPU multi-host, GKE secara atomik akan meningkatkan skala node pool dari nol hingga ukuran maksimum.
- MAX_NODES: ukuran maksimum node pool. Flag --max-nodes diperlukan jika --enable-autoscaling diberikan dan harus sama dengan hasil perkalian nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM.
--node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: Memberi tahu GKE bahwa node pool slice TPU multi-host adalah kumpulan. Gunakan tanda ini jika kondisi berikut berlaku:
- Node pool menjalankan workload inferensi.
- Node pool menggunakan TPU Trillium.
- Spot VM tidak mendukung penjadwalan pengumpulan.
Untuk mengetahui informasi selengkapnya tentang pengelolaan penjadwalan pengumpulan, lihat Mengelola penjadwalan pengumpulan di slice TPU multi-host.
--placement-type=COMPACT: Membuat node pool dengan penempatan rapat yang diaktifkan. Opsi ini harus digunakan dengan tanda --tpu-topology. Untuk mengetahui informasi selengkapnya, lihat Membuat kebijakan penempatan yang ringkas dan Topologi TPU.

Terraform

Pastikan Anda menggunakan versi 4.84.0 atau yang lebih baru dari penyedia google.
Tambahkan blok berikut ke konfigurasi Terraform Anda:
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    flex_start = false
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
Ganti kode berikut:
- NODE_POOL_RESOURCE_NAME: nama resource node pool di template Terraform.
- PROJECT_ID: Project ID Anda.
- CLUSTER_NAME: nama cluster yang ada tempat node pool akan ditambahkan.
- POOL_NAME: nama node pool yang akan dibuat.
- CLUSTER_LOCATION: lokasi komputasi untuk cluster. Sebaiknya Anda memiliki cluster regional untuk keandalan panel kontrol Kubernetes yang lebih tinggi. Anda juga dapat menggunakan cluster zonal. Untuk mempelajari lebih lanjut, silakan membaca Memilih topologi dan versi TPU.
- NODE_ZONES: daftar yang dipisahkan koma untuk satu atau beberapa zona tempat GKE membuat node pool.
- NUM_NODES: jumlah node dalam node pool. Nilainya harus nol atau hasil kali jumlah TPU chip yang dibagi empat, karena dalam slice TPU multi-host, setiap node slice TPU memiliki 4 chip. Misalnya, jika TPU_TOPOLOGY adalah 4x8, maka ada 32 chip, yang berarti harus ada 8 NUM_NODES. Untuk mempelajari topologi TPU lebih lanjut, gunakan tabel di Memilih versi TPU.
- TPU_TOPOLOGY: ini menunjukkan topologi fisik untuk slice TPU. Format topologi bergantung pada versi TPU yang Anda gunakan. Untuk mempelajari lebih lanjut topologi TPU, gunakan tabel di Memilih topologi.
Secara opsional, Anda juga dapat menggunakan variabel berikut:
- RESERVATION_NAME: jika Anda menggunakan pemesanan TPU, ini adalah daftar label resource pemesanan yang akan digunakan saat membuat node pool. Untuk mempelajari cara mengisi RESERVATION_LABEL_VALUES di kolom reservation_affinity lebih lanjut, silakan melihat Penyedia Terraform.
- autoscaling: Membuat node pool dengan penskalaan otomatis yang diaktifkan. Saat GKE menskalakan node pool slice TPU multi-host, GKE secara atomik akan meningkatkan skala node pool dari nol hingga ukuran maksimum.
  - MAX_NODES: ini adalah ukuran maksimum node pool. Nilainya harus sama dengan hasil kali nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM).
- spot: memungkinkan node pool menggunakan Spot VM untuk node slice TPU. Nama ini tidak dapat diubah setelah pembuatan node pool. Untuk mengetahui informasi selengkapnya, silakan melihat Spot VM.
- flex_start: Menetapkan node pool agar menggunakan opsi konsumsi flex-start. Tidak dapat disetel ke true jika spot diaktifkan.

Konsol

Untuk membuat node pool dengan TPU:

Buka halaman Google Kubernetes Engine di konsol Google Cloud .

Buka Google Kubernetes Engine
Di daftar cluster, klik nama cluster yang ingin diubah.
Klik Add node pool.
Di bagian Node pool details, centang kotak Specify node locations.
Pilih nama zona berdasarkan versi TPU yang ingin Anda gunakan. Untuk mengidentifikasi lokasi yang tersedia, lihat Ketersediaan TPU di GKE.
Dari panel navigasi, klik Node.
Di bagian Konfigurasi Mesin, pilih TPU.
Di menu drop-down Seri, pilih salah satu opsi berikut:
- CT3: Perangkat TPU v3, host tunggal
- CT3P: Slice pod multi-host TPU v3
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
- CT6E: TPU Trillium (v6e)
Di menu drop-down Jenis mesin, pilih nama mesin yang akan digunakan untuk node. Gunakan tabel Pilih versi TPU untuk mempelajari cara menentukan jenis mesin dan topologi TPU yang membuat TPU slice node pool multi-host.
Di menu drop-down Topologi TPU, pilih topologi fisik untuk slice TPU.
Dalam dialog Perubahan yang diperlukan, klik Buat perubahan.
Pastikan Jenis boot disk adalah Persistent disk standar atau Persistent disk SSD.
Secara opsional, pilih kotak centang Aktifkan node di Spot VM agar dapat menggunakan Spot VM untuk node di node pool.
Klik Buat.

Memverifikasi status node pool

Dapatkan kredensial agar Anda dapat menggunakan kubectl untuk mengakses cluster:
```
gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=CONTROL_PLANE_LOCATION
```
Ganti kode berikut:
- CLUSTER_NAME: Nama cluster.
- PROJECT_ID: Project ID Anda.
- CONTROL_PLANE_LOCATION: lokasi Compute Engine bidang kontrol cluster Anda. Berikan region untuk cluster regional, atau zona untuk cluster zona.

Gunakan kubectl, di Cloud Shell, untuk melihat node slice TPU Anda:

kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY

Ganti kode berikut:

TPU_ACCELERATOR: Jenis akselerator TPU yang Anda gunakan saat membuat node pool. Misalnya, tpu-v4-podslice, atau tpu-v5-lite-podslice.
TPU_TOPOLOGY: Topologi fisik untuk slice TPU.

Outputnya mirip dengan hal berikut ini:

 NAME                                    STATUS   ROLES    AGE    VERSION
 gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000

Menjalankan workload Multislice

Di bagian ini, Anda akan menjalankan workload JAX yang menampilkan jumlah chip TPU global dalam slice TPU, lalu keluar.

Untuk menjalankan workload JAX, lakukan hal berikut:

Buat manifes tpu-multislice.yaml berikut:

Autopilot

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: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                 limits:
                    google.com/tpu: NUM_CHIPS

Standar

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: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                   google.com/tpu: NUM_CHIPS

Ganti kode berikut:

NUM_SLICES: Jumlah node pool slice TPU. Dalam hal ini, NUM_SLICES sama dengan 3.
ACCELERATOR_TYPE: Jenis akselerator TPU yang Anda gunakan saat membuat node pool. Misalnya, tpu-v4-podslice atau tpu-v5-lite-podslice.
TPU_TOPOLOGY: Topologi fisik untuk slice TPU. Misalnya, 4x4x4 atau 2x2, bergantung pada versi TPU.
NUM_NODES: Jumlah node di node pool. Nilainya harus nol atau hasil kali nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah TPU chip di setiap VM. Untuk TPU v4 multi-host, jumlah chip TPU di setiap VM adalah empat. Untuk TPU v5e multi-host, jumlah chip TPU di setiap VM adalah satu, empat, atau delapan. Oleh karena itu, jika TPU_TOPOLOGY Anda adalah 2x4x4 (TPU v4 dengan empat chip TPU di setiap VM), maka NUM_NODES adalah 32/4 yang sama dengan 8.
NUM_CHIPS: Untuk TPU v4 multi-host, jumlah chip TPU di setiap VM adalah empat. Untuk TPU v5e multi-host, jumlah chip TPU di setiap VM adalah satu, empat, atau delapan. Untuk mempelajari lebih lanjut, lihat Chip TPU di VM dalam slice TPU.

Dalam manifes ini:

JobSet adalah Layanan Headless dengan nama yang sama dengan nama JobSet, dalam hal ini adalah multislice-job.
Anotasi alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool mengonfigurasi afinitas pod untuk memastikan semua pod dijadwalkan pada slice yang sama.
Kolom image ditetapkan ke image AI JAX terbaru. Untuk menyetel versi yang berbeda, lihat daftar image AI JAX saat ini.
maxRestarts: 4 menunjukkan jumlah maksimum GKE memulai ulang JobSet saat Job turunan gagal. Jika JobSet dimulai ulang mencapai maksimum yang ditentukan, JobSet akan ditandai sebagai gagal.
Kolom parallelism dan completions sama dengan jumlah node di setiap node pool.
backoff adalah 0 karena Multislice hanya mendukung pelatihan multi-pengontrol sinkron. Harus ditetapkan ke 0. Gagal menjalankan tugas jika ada pod yang gagal.
Nilai di bagian afinitas memastikan bahwa hanya ada satu workload Multislice TPU yang berjalan dalam grup Multislices.
containerPort: 8080 adalah port untuk koordinator MXLA
containerPort: 8431 adalah port untuk mengekspor metrik penggunaan TPU
securityContext: privileged: true menunjukkan bahwa node telah mengaktifkan mode hak istimewa untuk mengakses TPU. Node di GKE versi 1.28 atau yang lebih baru tidak perlu mengaktifkan mode hak istimewa untuk mengakses TPU. Untuk mempelajari lebih lanjut, lihat bagian Menjalankan penampung tanpa mode istimewa.

Terapkan manifes:
```
kubectl apply -f tpu-multislice.yaml
```

Pastikan beban kerja diterima:

kubectl get jobsets

Outputnya mirip dengan hal berikut ini:

NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

Pantau status Pod yang disediakan:

kubectl get pods

Outputnya mirip dengan hal berikut ini:

 NAME                                READY   STATUS      RESTARTS   AGE
 multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
 multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
 multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
 multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
 multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
 multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s

JobSet multislice-job menjadwalkan, membuat, lalu menjalankan Pod hingga selesai. Nama Pod menggunakan format <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. Awalan jobsetName menentukan JobSet tempat Pod berada.

Opsional: Hapus workload JAX:
```
kubectl delete -f tpu-multislice.yaml
```

Mengonfigurasi setelan tambahan

Bagian berikut menjelaskan konfigurasi tambahan yang dapat Anda terapkan ke Multislice.

Meningkatkan performa jaringan dengan hostNetwork

Untuk meningkatkan performa jaringan di antara slice TPU, sebaiknya aktifkan hostNetworking. Gunakan hostNetwork: true dalam spesifikasi Pod untuk melewati semua stack jaringan Kubernetes dan memungkinkan Pod Kubernetes Anda menggunakan jaringan host secara langsung untuk komunikasi VM-ke-VM.

Untuk mengaktifkan hostNetworking, tambahkan dua baris berikut ke spesifikasi Pod Anda:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Untuk terus menggunakan podHostnames untuk penemuan node pekerja dengan hostNetwork, tetapkan dnsPolicy: ClusterFirstWithHostNet. Hal ini penting saat Anda menjalankan tugas pelatihan yang dilanjutkan otomatis dan Anda harus memiliki nama yang sama untuk memuat ulang checkpoint yang sama.

Jika Anda menggunakan TPU Trillium (v6e) atau Ironwood (TPU7x) (Pratinjau) dan Pod Anda menggunakan hostNetworking, instal DaemonSet berikut untuk menyetel /proc/sys/net/ipv4/tcp_rmem di node.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Meningkatkan performa jaringan tanpa hostNetwork di TPU Trillium atau Ironwood (TPU7x)

Jika Anda menggunakan TPU Trillium atau Ironwood (TPU7x) (Pratinjau) dan Pod Anda tidak dapat menggunakan hostNetworking, aktifkan multi-jaringan dengan mode netdevice untuk performa jaringan terbaik. Dukungan NIC mode netdevice dengan multi-jaringan meneruskan NIC VM langsung ke Pod, mengabaikan Kubernetes dan GKE Dataplane V2.

Jenis mesin ct6e-standard-4t didukung oleh dua NIC fisik. Kubernetes memerlukan satu vNIC yang tidak dapat diteruskan ke Pod. Oleh karena itu, setiap node harus memiliki tiga vNIC agar Pod dapat memiliki akses langsung ke dua vNIC untuk mencapai performa NIC fisik yang terbaik.

Untuk mengaktifkan mode netdevice untuk ct6e-standard-4t, selesaikan langkah-langkah berikut:

Buat dua VPC tambahan yang mendukung mode netdevice
Membuat cluster GKE dengan kemampuan multi-jaringan

Konfigurasi dua jaringan netdevice. Misalnya, Anda dapat menggunakan objek GKENetworkParamSet dan Network berikut (SECOND_VPC dan THIRD_VPC adalah VPC yang dibuat pada langkah sebelumnya):

apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-second
spec:
  vpc: SECOND_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-third
spec:
  vpc: THIRD_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-second
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-second
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-third
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-third

Hubungkan Pod Anda dengan tiga jaringan. Misalnya, Anda dapat menggunakan anotasi berikut dalam spesifikasi Pod:

metadata:
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth1","network":"tpu-second"},
        {"interfaceName":"eth2","network":"tpu-third"},
      ]

Terapkan sysctl jaringan di dalam Pod, di container init atau container aplikasi. Misalnya, Anda dapat menambahkan container init berikut ke spesifikasi Pod:

initContainers:
- name: "network-optimization-sysctls"
  image: "busybox"
  securityContext:
    privileged: true
  command:
  - bash
  - -c
  - |
    echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
    echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
    echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
    echo 131072 > /proc/sys/net/core/optmem_max
    echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem

Praktik terbaik:

Gunakan antarmuka eth1 dan eth2 untuk performa jaringan yang lebih baik, bukan antarmuka eth0. Hal ini dapat dilakukan dengan menambahkan export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" ke spesifikasi workload.

Aktifkan logging

Log yang dikeluarkan oleh container yang berjalan di node GKE, termasuk node slice TPU, dapat dilihat di Logs Explorer, jika Anda telah mengaktifkan logging sistem GKE di cluster Anda.

Anda dapat melihat log dari GKE menggunakan Logs Explorer dengan filter berikut untuk melihat log container untuk beban kerja Anda:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

Gunakan filter berikut untuk slice dan pekerja TPU:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

Untuk mempelajari lebih lanjut, lihat Melihat log TPU GKE.

Mengaktifkan metrik tambahan

Selain metrik TPU umum, ada 4 metrik runtime TPU khusus multislice tambahan. Metrik ini tersedia di GKE versi 1.29.1-gke.1016000 atau yang lebih baru. Workload TPU harus menggunakan JAX versi 0.4.24

Berikut adalah metrik multiris yang tersedia:

Latensi transfer DCN (Data Center Network): Distribusi latensi transfer jaringan untuk traffic multislice.
Latensi kolektif: Distribusi latensi kolektif end-to-end untuk traffic multislice.
Latensi transfer Host-ke-Perangkat: Distribusi latensi transfer host ke perangkat untuk setiap bagian data untuk traffic multislice.
Latensi transfer Perangkat ke Host: Distribusi latensi transfer perangkat ke host untuk setiap bagian data untuk traffic multislice.

Metrik ini berada dalam skema container Kubernetes (k8s_container):

kubernetes.io/container/multislice/network/dcn_transfer_latencies
kubernetes.io/container/multislice/network/collective_end_to_end_latencies
kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

Slice TPU versus Multislice

Tabel berikut membedakan organisasi arsitektur slice TPU dan Multislice:

	Slice TPU	Multi-slice
Interkonektivitas	Workload berjalan di satu slice TPU. Semua chip TPU dalam slice terhubung dengan ICI.	Workload berjalan di beberapa slice TPU. Komunikasi dalam slice terjadi melalui ICI. Komunikasi antar-slice terjadi melalui DCN.
Node pool yang didukung	Slice TPU host tunggal dan slice TPU multi-host	Grup slice TPU multi-host
Jenis beban kerja yang direkomendasikan	IndexedJob atau JobSet	JobSet

Membersihkan resource

Cara termudah untuk menghilangkan penagihan adalah dengan menghapus Google Cloud project yang Anda buat untuk tutorial ini. Atau, Anda dapat menghapus setiap resource.

Menghapus project

Perhatian: Penghapusan project memiliki efek berikut:

Semua hal dalam project akan dihapus. Jika Anda menggunakan project yang ada untuk mengerjakan tugas di dokumen ini, saat Anda menghapusnya, pekerjaan lain dalam project tersebut juga akan dihapus.
Project ID kustom akan hilang. Saat membuat project ini, Anda mungkin juga membuat project ID kustom yang masih ingin digunakan pada masa mendatang. Agar tidak kehilangan URL yang menggunakan project ID tersebut, seperti URL appspot.com, hapus resource yang dipilih di dalam project, alih-alih menghapus seluruh project.

Jika berencana mempelajari beberapa arsitektur, tutorial, atau panduan memulai, Anda dapat menggunakan kembali project untuk membantu Anda agar tidak melampaui batas kuota project.

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

Menghapus resource satu per satu

Hapus cluster GKE:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

Langkah berikutnya

Pelajari cara Mengorkestrasi workload Multislice dengan slice TPU

Men-deploy TPU Multislice di GKE Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.

Apa itu TPU Multislice

Ketersediaan multislice di GKE

Sebelum memulai

Menjalankan workload di Multislice

Menyiapkan node pool mode Standar

Buat node pool slice TPU

Ironwood (TPU7x)

gcloud

Terraform

Versi TPU lainnya

gcloud

Terraform

Konsol

Memverifikasi status node pool

Menjalankan workload Multislice

Autopilot

Standar

Mengonfigurasi setelan tambahan

Meningkatkan performa jaringan dengan hostNetwork

Meningkatkan performa jaringan tanpa hostNetwork di TPU Trillium atau Ironwood (TPU7x)

Aktifkan logging

Mengaktifkan metrik tambahan

Slice TPU versus Multislice

Membersihkan resource

Menghapus project

Menghapus resource satu per satu

Langkah berikutnya

Men-deploy TPU Multislice di GKE