Men-deploy TPU Multislice di GKE

Autopilot Standard

Halaman ini memperkenalkan konfigurasi TPU Multislice di Google Kubernetes Engine (GKE). Sebelum mengonfigurasi Multislice di GKE, Anda harus memahami konsep berikut:

Apa itu TPU Multislice

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

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

Ketersediaan multislice di GKE

Standard mendukung Multislice dalam versi 1.27.4-gke.900 dan yang lebih baru.
Autopilot mendukung Multislice dalam 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 kumpulan node 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 multipengontrol sinkron.
Workload multislice hanya dapat berjalan di seluruh slice TPU yang memiliki jenis, ukuran, dan topologi TPU yang sama.

Sebelum memulai

Sebelum memulai, pastikan Anda telah menjalankan tugas berikut:

Aktifkan Google Kubernetes Engine API.

Aktifkan Google Kubernetes Engine API

Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu initialize gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.
Catatan: Untuk penginstalan gcloud CLI yang ada, pastikan untuk menyetel properti compute/region dan 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.

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

Menjalankan beban kerja pada Multislice

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

Menyiapkan kumpulan node mode Standar

Bagian ini membahas langkah-langkah berikut:

Membuat tiga kumpulan node slice TPU multi-host
Memverifikasi status kumpulan node

Membuat kumpulan node TPU

Anda dapat membuat lebih dari satu kumpulan node TPU multi-host. Untuk tujuan panduan ini, buat tiga kumpulan node TPU multi-host untuk menjalankan beban kerja Multislice. Anda dapat membuat kumpulan node slice TPU multi-host menggunakan Google Cloud CLI, Terraform, atau Konsol Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

Ganti kode berikut:

POOL_NAME: Nama node pool baru.
LOCATION: Nama zona berdasarkan versi TPU yang ingin Anda gunakan:
- Untuk TPU v4, gunakan us-central2-b.
- Jenis mesin TPU v5e yang dimulai dengan ct5l- tidak pernah berfungsi multi-host.
- Untuk jenis mesin TPU v5e yang dimulai dengan ct5lp-, gunakan us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, atau europe-west4-a.
- Untuk jenis mesin TPU v5p yang dimulai dengan ct5p-, gunakan us-east1-d, us-east5-a, atau us-east5-c.
Untuk mempelajari lebih lanjut, lihat ketersediaan TPU di GKE.
CLUSTER_NAME: Nama cluster.
NODE_ZONE: Daftar yang dipisahkan koma untuk satu atau beberapa zona tempat GKE membuat node pool.
MACHINE_TYPE: Jenis mesin yang akan digunakan untuk node. Untuk mempelajari jenis mesin yang tersedia lebih lanjut, silakan membaca Pemetaan konfigurasi TPU.
TPU_TOPOLOGY: Topologi fisik untuk slice TPU. Format topologi bergantung pada versi TPU sebagai berikut:
- TPU v4 atau v5p: Menentukan topologi dalam 3 tuple ({A}x{B}x{C}), misalnya 4x4x4.
- TPU v5e: Menentukan topologi dalam 2 tuple ({A}x{B}), misalnya 2x2.
Untuk mempelajari lebih lanjut, lihat Topologi.
NUM_NODES: Jumlah node di node pool. Nilai ini harus nol atau merupakan hasil dari 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), NUM_NODES adalah 32/4 yang sama dengan 8.

Secara opsional, Anda juga dapat menggunakan flag berikut:

RESERVATION_NAME: Nama pemesanan yang digunakan GKE saat membuat node pool. Jika flag ini dihapus, GKE akan menggunakan TPU node pool yang tersedia. Untuk mempelajari pemesanan TPU lebih lanjut, silakan melihat Pemesanan TPU.
--spot: Menetapkan node pool agar menggunakan Spot VM untuk TPU node. Nama ini tidak dapat diubah setelah pembuatan node pool. Untuk mengetahui informasi selengkapnya, silakan membaca Spot VM.
--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. Tanda --max-nodes diperlukan jika --enable-autoscaling diberikan dan harus sama dengan hasil dari nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM.

Terraform

Pastikan Anda menggunakan penyedia google versi 4.84.0 atau yang lebih baru.
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]
    }
    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: Menghitung lokasi untuk cluster. Sebaiknya Anda memiliki cluster regional untuk keandalan panel kontrol Kubernetes yang lebih tinggi. Anda juga dapat menggunakan cluster zona. 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 TPU node memiliki 4 chip. Misalnya, jika TPU_TOPOLOGY adalah 4x8, maka ada 32 chip, yang berarti harus ada 8 NUM_NODES. Untuk mempelajari lebih lanjut topologi TPU, gunakan tabel di bagian Pemetaan konfigurasi TPU.
- TPU_TOPOLOGY: Ini menunjukkan topologi fisik yang diinginkan untuk slice TPU. Format topologi bergantung pada versi TPU yang Anda gunakan:
  - Untuk TPU v4: Tentukan topologi dalam 3 tuple ({A}x{B}x{C}), misalnya 4x4x4.
  - Untuk TPU v5e: Tentukan topologi dalam 2 tuple ({A}x{B}), misalnya 2x2.
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. Ini harus sama dengan hasil dari 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 TPU node. Nama ini tidak dapat diubah setelah pembuatan node pool. Untuk mengetahui informasi selengkapnya, silakan melihat Spot VM.

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 zona berdasarkan versi TPU yang ingin Anda gunakan:
- Untuk TPU v4, gunakan us-central2-b.
- Jenis mesin TPU v5e yang dimulai dengan ct5l- tidak pernah berfungsi multi-host.
- Untuk jenis mesin TPU v5e yang dimulai dengan ct5lp-, gunakan us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, atau europe-west4-a.
- Untuk jenis mesin TPU v5p yang dimulai dengan ct5p-, gunakan us-east1-d, us-east5-a, atau us-east5-c.
Dari panel navigasi, klik Node.
Di bagian Konfigurasi Mesin, pilih TPU.
Di menu drop-down Seri, pilih salah satu opsi berikut:
- CT4P: Untuk TPU v4.
- CT5LP: Untuk TPU v5e.
Di menu drop-down Jenis mesin, pilih nama mesin yang akan digunakan untuk node. Gunakan tabel Pemetaan konfigurasi TPU untuk mempelajari cara menentukan jenis mesin dan topologi TPU yang membuat TPU 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 Create.

Memverifikasi status kumpulan node

Dapatkan kredensial agar Anda dapat menggunakan kubectl untuk mengakses cluster:
```
gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID
```
Ganti kode berikut:
- CLUSTER_NAME: Nama cluster.
- PROJECT_ID: Project ID Anda.
Untuk melihat node TPU Anda, gunakan kubectl di Cloud Shell:
```
kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
   -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, tpu-v5-lite-device, 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 beban kerja JAX yang menunjukkan jumlah chip TPU global di slice TPU, lalu keluar.

Untuk menjalankan beban kerja 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: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  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: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  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 kumpulan node TPU. Dalam hal ini, NUM_SLICES sama dengan 3.
ACCELERATOR_TYPE: Jenis akselerator TPU yang Anda gunakan saat membuat kumpulan node. Misalnya, tpu-v4-podslice, tpu-v5-lite-device, 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. Nilai ini harus nol atau merupakan hasil dari nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM. Untuk TPU v4 multi-host, jumlah chip di setiap VM adalah empat. Untuk TPU v5e multi-host, jumlah chip di setiap VM adalah satu, empat, atau delapan. Oleh karena itu, jika TPU_TOPOLOGY Anda adalah 2x4x4 (TPU v4 dengan empat chip di setiap VM), NUM_NODES adalah 32/4 yang sama dengan 8.
NUM_CHIPS: Untuk TPU v4 multi-host, jumlah chip di setiap VM adalah empat. Untuk TPU v5e multi-host, jumlah chip di setiap VM adalah satu, empat, atau delapan. Untuk mempelajari lebih lanjut, baca chip TPU di VM TPU.

Dalam manifes ini:

JobSet adalah Layanan Headless dengan nama yang sama dengan nama JobSet, dalam hal ini adalah multislice-job.
maxRestarts: 4 menunjukkan frekuensi maksimum GKE memulai ulang JobSet saat Job turunan gagal. Jika mulai ulang JobSet mencapai batas maksimum yang ditetapkan, JobSet akan ditandai sebagai gagal.
Kolom parallelism dan completions sama dengan jumlah node di setiap kumpulan node.
backoff adalah 0 karena Multislice hanya mendukung pelatihan multi-pengontrol sinkron. Harus ditetapkan ke 0. Menggagalkan tugas saat ada pod yang gagal.
Nilai di bagian afinitas memastikan bahwa hanya ada satu workload TPU Multislice yang berjalan di grup Multislice.
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 istimewa untuk mengakses TPU. Untuk mempelajari lebih lanjut, lihat Menjalankan container tanpa mode hak 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

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

Konfigurasi tambahan

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

Mengaktifkan hostNetwork di Pod Standar GKE

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

Untuk mengaktifkan hostNetworking, hapus dua baris berikut dari 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 secara otomatis dan Anda harus memiliki nama yang sama untuk memuat ulang checkpoint yang sama.

Logging

Log yang dikeluarkan oleh container yang berjalan di node GKE, termasuk VM TPU, akan terlihat di Logs Explorer, jika Anda telah mengaktifkan logging sistem GKE di cluster.

Anda dapat melihat log dari GKE menggunakan Logs Explorer dengan filter berikut guna melihat log container untuk workload 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 pekerja dan slice 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>

Kemampuan observasi dan metrik

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 multislice yang tersedia:

Latensi transfer DCN (Jaringan Pusat Data): Distribusi latensi transfer jaringan untuk traffic multislice.
Latensi kolektif: Distribusi latensi kolektif menyeluruh untuk traffic multibagian.
Latensi transfer Host-ke-Perangkat: Distribusi latensi transfer host ke perangkat untuk setiap bagian data untuk traffic multislice.
Latensi transfer Perangkat-ke-Host: Distribusi perangkat ke latensi transfer host untuk setiap bagian data untuk traffic multislice.

Metrik ini terletak di 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

TPU slice versus Multislice

Tabel berikut membedakan organisasi arsitektur slice TPU dan Multislice:

	Slice TPU	Multi-slice
Interkoneksi	Workload berjalan pada satu slice TPU. Semua chip TPU dalam satu irisan terhubung dengan ICI.	Workload berjalan di beberapa slice TPU. Komunikasi dalam slice terjadi melalui ICI. Komunikasi antarbagian terjadi melalui DCN.
Kumpulan node yang didukung	Slice TPU host tunggal dan slice TPU multi-host	Grup slice TPU multi-host
Jenis workload yang direkomendasikan	IndexedJob atau JobSet	JobSet

Langkah selanjutnya

Pelajari cara Mengorkestrasi workload Multislice dengan TPU slice