Men-deploy TPU Multislice di GKE


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

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

  1. Pengantar Cloud TPU
  2. Arsitektur sistem Cloud TPU
  3. Tentang TPU di GKE

Apa itu TPU Multislice

Multislice TPU adalah pengaturan arsitektur VM dalam slice TPU tempat 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 di satu slice TPU.

Ketersediaan multislice di GKE

  • Standar 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 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 multipengontrol 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 menjalankan 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 gcloud components update.

Menjalankan workload di Multislice

Bagian ini menunjukkan cara menjalankan beban kerja 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:

  1. Membuat tiga node pool slice TPU multi-host
  2. Memverifikasi status node pool

Membuat 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. Anda dapat membuat node pool 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 mengidentifikasi lokasi yang tersedia, 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, lihat Memilih versi TPU.
  • TPU_TOPOLOGY: Topologi fisik untuk slice TPU. Format topologi bergantung pada versi TPU. Untuk mempelajari topologi TPU lebih lanjut, gunakan tabel di Memilih topologi.

    Untuk mempelajari lebih lanjut, lihat Topologi.

  • 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 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.

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 node pool slice TPU yang tersedia. Untuk mempelajari pemesanan TPU lebih lanjut, silakan melihat 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.
  • --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 disediakan dan harus sama dengan produk nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip di setiap VM.

Terraform

  1. Pastikan Anda menggunakan penyedia google versi 4.84.0 atau yang lebih baru.
  2. Tambahkan blok berikut ke konfigurasi Terraform:

    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: 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 TPU node 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 yang diinginkan untuk slice TPU. Format topologi bergantung pada versi TPU yang Anda gunakan. Untuk mempelajari topologi TPU lebih lanjut, 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. Nilai ini 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.

Konsol

Untuk membuat node pool dengan TPU:

  1. Buka halaman Google Kubernetes Engine di konsol Google Cloud.

    Buka Google Kubernetes Engine

  2. Di daftar cluster, klik nama cluster yang ingin diubah.

  3. Klik Add node pool.

  4. Di bagian Node pool details, centang kotak Specify node locations.

  5. Pilih nama zona berdasarkan versi TPU yang ingin Anda gunakan. Untuk mengidentifikasi lokasi yang tersedia, lihat Ketersediaan TPU di GKE.

  6. Dari panel navigasi, klik Node.

  7. Di bagian Konfigurasi Mesin, pilih TPU.

  8. Di menu drop-down Seri, pilih salah satu opsi berikut:

    • CT3P: Untuk TPU v3.
    • CT4P: Untuk TPU v4.
    • CT5LP: Untuk TPU v5e.
  9. Di menu drop-down Jenis mesin, pilih nama mesin yang akan digunakan untuk node. Gunakan tabel Memilih versi TPU untuk mempelajari cara menentukan jenis mesin dan topologi TPU yang membuat node pool slice TPU multi-host.

  10. Di menu drop-down Topologi TPU, pilih topologi fisik untuk slice TPU.

  11. Dalam dialog Perubahan yang diperlukan, klik Buat perubahan.

  12. Pastikan Jenis boot disk adalah Persistent disk standar atau Persistent disk SSD.

  13. Secara opsional, pilih kotak centang Aktifkan node di Spot VM agar dapat menggunakan Spot VM untuk node di node pool.

  14. Klik Create.

Memverifikasi status node pool

  1. 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.
  2. Gunakan kubectl, di Cloud Shell, untuk melihat node slice TPU Anda:

    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 workload JAX yang menampilkan jumlah chip TPU global di slice TPU, lalu keluar.

Untuk menjalankan beban kerja JAX, lakukan hal berikut:

  1. 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
    

    Standard

    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 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, 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. Nilainya harus nol atau hasil kali nilai yang ditentukan dalam TPU_TOPOLOGY ({A}x{B}x{C}) dibagi dengan jumlah chip TPU 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), 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.
    • maxRestarts: 4 menunjukkan jumlah maksimum permulaan ulang JobSet oleh GKE 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 tugas saat pod gagal.
    • Nilai di bagian afinitas memastikan bahwa hanya ada satu workload Multislice TPU yang berjalan dalam 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 hak istimewa untuk mengakses TPU. Untuk mempelajari lebih lanjut, lihat Menjalankan penampung tanpa mode dengan hak istimewa.
  2. Terapkan manifes:

    kubectl apply -f tpu-multislice.yaml
    
  3. Pastikan beban kerja diizinkan:

    kubectl get jobsets
    

    Outputnya mirip dengan hal berikut ini:

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. 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.

Konfigurasi tambahan

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

Mengaktifkan hostNetwork di Pod GKE Standard

Untuk meningkatkan performa jaringan di antara slice TPU, sebaiknya aktifkan hostNetworking. Gunakan hostNetwork: true dalam spesifikasi Pod untuk melewati semua stack jaringan Kubernetes dan mengizinkan Pod Kubernetes 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 titik pemeriksaan yang sama.

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 dapat melihat log dari GKE menggunakan Logs Explorer dengan filter berikut untuk melihat log penampung 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.

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 Data Center): 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 terletak di skema penampung 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 pengaturan arsitektur slice TPU dan Multislice:

Slice TPU Multi-slice
Interkoneksi 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

Langkah selanjutnya