Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Pengantar Jalur Pembelajaran di Cloud

Pathways adalah sistem yang dirancang untuk memungkinkan pembuatan sistem machine learning (ML) berskala besar, multi-tugas, dan diaktifkan secara tersebar. Dengan begitu, Anda dapat menggunakan ribuan atau puluhan ribu akselerator, dengan kemampuan untuk mengalokasikan berbagai jumlah komputasi secara dinamis untuk berbagai tugas berdasarkan persyaratan pemrosesannya.

Pathways menyederhanakan komputasi ML berskala besar dengan memungkinkan satu klien JAX mengorkestrasi workload di seluruh slice TPU besar, yang berpotensi mencakup ribuan chip TPU.

Google menggunakan Pathways secara internal untuk melatih model besar seperti Gemini. Pathways on Cloud memberikan manfaat yang sama kepada pelanggan Google Cloud .

Sebelum memulai

Pilih tab untuk melihat bagaimana Anda berencana menggunakan contoh di halaman ini:
gcloud

Instal Google Cloud CLI, lalu login ke gcloud CLI dengan identitas gabungan Anda. Setelah login, inisialisasi Google Cloud CLI dengan menjalankan perintah berikut:
```
gcloud init
```
REST

Untuk menggunakan contoh REST API di halaman ini dalam lingkungan pengembangan lokal, gunakan kredensial yang Anda berikan ke gcloud CLI.
Untuk mengetahui informasi selengkapnya, lihat Melakukan autentikasi untuk menggunakan REST dalam dokumentasi autentikasi Google Cloud .

Sebelum memulai

Anda dapat menjalankan Pathways di Google Kubernetes Engine (GKE) atau di cluster Cluster Director dengan Slurm. Pilih tab untuk platform Anda.

GKE

Selesaikan prasyarat berikut:

Dokumen ini memberikan ringkasan tentang cara menggunakan TPU yang dikelola Pathways di GKE untuk workload batch, real-time, dan interaktif. Panduan ini mengasumsikan bahwa Anda sudah memahami penggunaan TPU di GKE, termasuk TPU slice tunggal dan TPU multislice di GKE, serta memiliki pengalaman umum dengan TPU multislice.

Cluster Director

Selesaikan prasyarat berikut:

Dokumen ini memberikan ringkasan tentang cara menggunakan Pathways dengan TPU di cluster Cluster Director untuk workload batch dan interaktif. Anda dianggap sudah memahami Cluster Director, mesin TPU di Cluster Director, dan menjalankan tugas dengan Slurm, serta memiliki pengalaman umum dengan TPU multislice.

Untuk membuat cluster, lihat Membuat cluster Cluster Director dengan Pathways.

Pengontrol tunggal dan multi-pengontrol

Ada dua cara utama untuk mengelola dan mengorkestrasi komputasi di beberapa perangkat:

Fitur	Pengontrol tunggal (Pathways)	Multi-pengontrol (default JAX)
Kontrol	Titik kontrol tunggal: Satu program klien bertindak sebagai pengontrol pusat.	Kontrol terdistribusi: Beberapa proses berpartisipasi, masing-masing dengan instance penafsir Python-nya sendiri.
Lihat	Tampilan terpadu: Klien melihat semua perangkat sebagai satu sistem terpadu.	Tampilan yang dilokalkan: Setiap proses Python hanya melihat perangkat yang terhubung ke proses tersebut.
Pemrograman	Pemrograman yang disederhanakan: Anda berinteraksi dengan satu klien, sehingga sistem tampak seperti satu mesin besar dengan banyak akselerator lokal.	Program tunggal, beberapa data (SPMD): Terutama menggunakan paradigma SPMD, yang mengharuskan semua perangkat menjalankan program yang sama.
Fleksibilitas	Mendukung pola komputasi kompleks di luar SPMD, termasuk paralelisme pipeline asimetris dan sparsity komputasi.	Dapat kurang fleksibel dalam pengelolaan resource, terutama di berbagai slice TPU.

Komponen Pathways

Bagian berikut menguraikan komponen utama arsitektur Pathways.

Pengelola resource Pathways

Komponen ini adalah bidang kontrol pusat sistem Pathways. Layanan ini mengelola semua resource akselerator dan mengoordinasikan alokasi akselerator untuk tugas Anda. WorkManager memantau kondisi worker dan menangani penjadwalan, penjedaan, dan melanjutkan tugas. Layanan ini berfungsi sebagai titik kontak tunggal untuk error dan status sistem. Komponen ini hanya memerlukan resource CPU.

Klien Pathways

Komponen ini adalah implementasi dari Interim Framework Runtime (IFRT) yang berfungsi sebagai titik entri ke sistem Pathways. XLA menerima Operasi Tingkat Tinggi (HLO) dari program Anda. Klien Pathways berkoordinasi dengan pengelola resource Pathways untuk menentukan tempat menempatkan program yang dikompilasi untuk dieksekusi berdasarkan kode Anda. Sistem ini menyajikan tampilan terpadu sistem ke klien JAX tertentu. Komponen ini hanya memerlukan resource CPU.

Pekerja Pathways

Proses worker Pathways berjalan di mesin akselerator (VM TPU). Mereka menerima file yang dapat dieksekusi yang dikompilasi dari program Anda dari server proxy IFRT dan melakukan komputasi pada TPU. Worker Pathways mengirimkan data kembali ke program Anda melalui server proxy IFRT. Komponen ini memerlukan resource akselerator.

Klien proxy IFRT

Klien proxy IFRT adalah implementasi open source dari API Interim Framework Runtime (IFRT) yang memisahkan kode pengguna dari runtime yang mendasarinya serta meningkatkan portabilitas dan transparansi kode. JAX menggunakan implementasi ini sebagai alternatif untuk runtime multi-pengontrol default-nya. Klien proxy IFRT bertindak sebagai jembatan komunikasi antara program Anda dan komponen Pathways. Klien ini mengirimkan permintaan ke server proxy IFRT dan menerima hasil darinya. Komponen ini hanya memerlukan resource CPU.

Server proxy IFRT

Server proxy IFRT adalah server gRPC yang menerima permintaan dari klien proxy IFRT dan meneruskannya ke klien Pathways, yang menangani distribusi pekerjaan. Komponen ini hanya memerlukan resource CPU.

Server file bantuan

Server file bantuan adalah server gRPC yang ditempatkan bersama dengan worker Pathways di VM akselerator. Layanan ini menjalankan kode Python yang ditentukan pengguna langsung di VM akselerator untuk mengurangi latensi transfer data antara pengontrol dan akselerator. Server file bantuan berinteraksi dengan worker Pathways melalui protokol berversi kustom pada transport gRPC.

Diagram arsitektur yang menunjukkan cara klien JAX, proxy IFRT,
pengelola resource Pathways, dan worker Pathways berinteraksi. — Komponen Pathways

Komponen Pathways di GKE

Bagian ini memetakan komponen Pathways ke komponen GKE seperti container dan Pod.

Anda dapat menemukan image container Pathways di lokasi berikut:

Jenis penampung	Lokasi
Server proxy IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION`
Pengelola/pekerja resource Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION`

Pengelola resource Pathways

Setelah membuat cluster GKE, Anda dapat menggunakan containerSpec berikut untuk men-deploy resource manager Pathways:

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

Container menggunakan argumen berikut:

--server_port: port yang digunakan resource manager Pathways untuk berkomunikasi dengan komponen Pathways lainnya.
--node_type: jenis node. Tetapkan argumen ini ke resource_manager untuk pengelola resource Pathways; argumen ini tidak diperlukan untuk container lainnya.
--instance_count: jumlah slice TPU.
--instance_type: jenis TPU dan topologi slice, dalam format tpuTPU_TYPE:TOPOLOGY—misalnya, tpuv5e:4x4.
--gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

Server proxy IFRT

Anda dapat menggunakan containerSpec berikut untuk men-deploy server proxy IFRT:

  - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

Container menggunakan argumen berikut:

--resource_manager_address: nama host dan port yang digunakan server proxy untuk berkomunikasi dengan resource manager Pathways. Port harus cocok dengan nilai --server_port yang digunakan untuk container pengelola resource Pathways.
--server_port: port yang digunakan server proxy IFRT untuk berkomunikasi dengan klien proxy IFRT.
--gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

Pekerja Pathways

Anda dapat menggunakan containerSpec berikut untuk men-deploy worker Pathways:

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

Container menggunakan argumen berikut:

--resource_manager_address: nama host dan port yang digunakan TPU worker untuk berkomunikasi dengan resource manager Pathways. Port harus cocok dengan nilai --server_port yang digunakan untuk container pengelola resource Pathways.
--server_port: port yang digunakan worker untuk berkomunikasi dengan server proxy dan resource manager Pathways.
--gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

Resource manager Pathways, server proxy IFRT, dan worker Pathways semuanya dapat memiliki port yang berbeda, tetapi dalam contoh ini, resource manager Pathways dan worker Pathways menggunakan port yang sama.

Komponen Pathways di Cluster Director

Bagian ini memetakan komponen Pathways ke tugas Slurm di cluster Cluster Director.

Skrip Pathways menjalankan komponen Pathways dengan Enroot. Skrip setup_env.sh menggunakan image container berikut, yang cocok dengan versi JAX (0.11.1) dari lingkungan Python yang dibangunnya:

Jenis penampung	Lokasi
Server proxy IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1`
Pengelola/pekerja resource Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1`

Anda memulai komponen Pathways dengan mengirimkan skrip start_pathways.sh sebagai tugas layanan Pathways. Anda tidak menulis sendiri perintah berikut: skrip memperoleh argumen dari variabel lingkungan yang Anda tetapkan saat mengirimkan start_pathways.sh, seperti SLICE_SHAPE, MIN_SLICES, MAX_SLICES, dan GCS_BUCKET.

Pengelola resource Pathways

Tugas layanan Pathways menjalankan pengelola resource Pathways di satu node CPU. Perintah ini menjalankan perintah berikut di penampung server Pathways:

/pathways/server \
    --server_port=29001 \
    --node_type=resource_manager \
    --instance_type=tpuTPU_TYPE:TOPOLOGY \
    --instance_count=MAX_SLICES \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Perintah ini menggunakan argumen berikut:

--server_port: port yang diproses oleh resource manager Pathways. Server proxy IFRT dan worker Pathways terhubung ke port ini.
--node_type: jenis node, resource_manager. Hanya pengelola resource Pathways yang menetapkan argumen ini.
--instance_type: jenis TPU dan topologi setiap slice, dari SLICE_SHAPE. Misalnya, SLICE_SHAPE=2x4 menetapkan tpuv6e:2x4.
--instance_count: jumlah slice TPU, dari MAX_SLICES.
--gcs_scratch_location: jalur di bucket Cloud Storage Anda, dari GCS_BUCKET, yang digunakan untuk file sementara seperti cache kompilasi.
--enable_gcs_olm: apakah Pathways membuat aturan siklus proses objek Cloud Storage untuk membersihkan entri cache kompilasi lama. Skrip menonaktifkan ini karena memerlukan izin tambahan pada bucket.

Server proxy IFRT

Tugas layanan Pathways menjalankan server proxy IFRT di node CPU yang sama dengan pengelola resource Pathways. Perintah ini menjalankan perintah berikut di penampung server proxy IFRT:

/pathways/proxy_server \
    --server_port=29000 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --num_elastic_instances=MAX_SLICES

Perintah ini menggunakan argumen berikut:

--server_port: port yang diproses oleh server proxy IFRT. Program JAX Anda terhubung ke port ini.
--resource_manager_address: alamat resource manager Pathways. RM_IP adalah alamat IP internal node CPU.
--gcs_scratch_location: jalur Cloud Storage yang sama dengan resource manager Pathways.
--num_elastic_instances: jumlah slice yang dapat tidak tersedia tanpa menghentikan sesi proxy IFRT, dari MAX_SLICES. Workload Anda menerapkan MIN_SLICES.

Saat setidaknya MIN_SLICES slice berjalan, tugas layanan akan menulis alamat server proxy IFRT ke ~/pathways/cluster.env. Program JAX Anda, klien proxy IFRT, membaca file ini untuk terhubung ke Pathways. Program JAX berjalan sebagai workload Anda—misalnya, run_maxtext.sh—atau di notebook Jupyter.

Pekerja Pathways

Tugas layanan Pathways mengirimkan array tugas Slurm (run_worker_slice.sh) dengan satu tugas array untuk setiap slice TPU. Setiap tugas array menjalankan perintah berikut di penampung server Pathways pada setiap VM TPU di bagiannya:

/pathways/server \
    --server_port=29001 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Perintah ini menggunakan argumen berikut:

--server_port: port yang diproses oleh worker Pathways. Resource manager Pathways dan server proxy IFRT terhubung ke port ini.
--resource_manager_address: alamat resource manager Pathways, sama seperti untuk server proxy IFRT.
--gcs_scratch_location, --enable_gcs_olm: sama seperti untuk Resource manager Pathways.

Skrip ini juga menetapkan variabel lingkungan topologi TPU dan MEGASCALE_* untuk setiap worker dari node slice dan indeks tugas array.

Jika slice TPU di-preempt atau gagal, Slurm akan mengantrekan ulang tugas array-nya. Jika Anda memiliki sudo tanpa sandi pada node cluster, tugas layanan juga akan mematikan node slice sehingga Slurm dapat membuatnya kembali. Resource manager Pathways dan server proxy IFRT terus berjalan, sehingga program JAX Anda tidak perlu terhubung ke alamat baru.

Pengantar Jalur Pembelajaran di Cloud Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.

Sebelum memulai

gcloud

REST

Sebelum memulai

GKE

Cluster Director

Pengontrol tunggal dan multi-pengontrol

Komponen Pathways

Pengelola resource Pathways

Klien Pathways

Pekerja Pathways

Klien proxy IFRT

Server proxy IFRT

Server file bantuan

Komponen Pathways di GKE

Pengelola resource Pathways

Server proxy IFRT

Pekerja Pathways

Komponen Pathways di Cluster Director

Pengelola resource Pathways

Server proxy IFRT

Pekerja Pathways

Langkah berikutnya

Pengantar Jalur Pembelajaran di Cloud