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:
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 topologi2x2x1
atauct5lp-hightpu-4t
dengan topologi2x2
, 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
.
- 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 memadai untuk Cloud TPU di GKE.
- Instal JobSet v0.2.3 atau yang lebih baru.
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:
- Membuat tiga node pool slice TPU multi-host
- 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 dalamTPU_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, jikaTPU_TOPOLOGY
Anda adalah2x4x4
(TPU v4 dengan empat chip di setiap VM), makaNUM_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 dalamTPU_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:
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, jikaTPU_TOPOLOGY
adalah4x8
, maka ada 32 chip, yang berarti harus ada 8NUM_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 mengisiRESERVATION_LABEL_VALUES
di kolomreservation_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 dalamTPU_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:
Buka halaman Google Kubernetes Engine di konsol Google Cloud.
Di daftar cluster, klik nama cluster yang ingin diubah.
Klik add_box 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:
- CT3P: Untuk TPU v3.
- CT4P: Untuk TPU v4.
- CT5LP: Untuk TPU v5e.
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.
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 node pool
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.
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
, atautpu-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:
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 dengan3
.ACCELERATOR_TYPE
: Jenis akselerator TPU yang Anda gunakan saat membuat node pool. Misalnya,tpu-v4-podslice
,tpu-v5-lite-device
, atautpu-v5-lite-podslice
.TPU_TOPOLOGY
: Topologi fisik untuk slice TPU. Misalnya4x4x4
atau2x2
bergantung pada versi TPU.NUM_NODES
: Jumlah node di node pool. Nilainya harus nol atau hasil kali nilai yang ditentukan dalamTPU_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, jikaTPU_TOPOLOGY
Anda adalah2x4x4
(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
dancompletions
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 MXLAcontainerPort: 8431
adalah port untuk mengekspor metrik penggunaan TPUsecurityContext: 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.
Terapkan manifes:
kubectl apply -f tpu-multislice.yaml
Pastikan beban kerja diizinkan:
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.
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
- Pelajari cara Mengorkestrasi workload Multislice dengan slice TPU