Pengantar Jalur Pembelajaran di Cloud

Pathways adalah sistem yang dirancang untuk memungkinkan pembuatan sistem machine learning berskala besar, multi-tugas, dan diaktifkan secara tersebar. Hal ini memungkinkan penggunaan ribuan atau puluhan ribu akselerator, dengan kemampuan untuk mengalokasikan jumlah komputasi yang bervariasi secara dinamis untuk berbagai tugas berdasarkan persyaratan pemrosesannya.

Pathways menyederhanakan komputasi machine learning berskala besar dengan memungkinkan satu klien JAX mengatur workload di beberapa slice TPU besar, yang berpotensi mencakup ribuan chip TPU.

Pathways digunakan secara internal di Google untuk melatih model besar seperti Gemini. Jalur di Cloud memberikan manfaat yang sama bagi pelanggan Google Cloud .

Sebelum memulai

Pastikan Anda memiliki:

Dokumen ini memberikan ringkasan tentang cara menggunakan TPU yang dikelola Pathways di Google Kubernetes Engine (GKE) untuk workload batch, real-time, dan interaktif. Anda dianggap sudah memahami cara menggunakan TPU dengan GKE termasuk TPU slice tunggal dan multi-slice di Google Kubernetes Engine, serta pengalaman umum dengan TPU multi-slice

Pengontrol tunggal dan multi-pengontrol

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

Fitur	Pengontrol tunggal (Pathways)	Multi-controller (Default JAX)
Kontrol	Titik kontrol tunggal: Satu program klien bertindak sebagai pengontrol pusat.	Kontrol terdistribusi: Beberapa proses berpartisipasi, masing-masing dengan instance interpreter 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: Pengguna berinteraksi dengan satu klien, sehingga sistem tampak seperti satu mesin besar dengan banyak akselerator lokal.	SPMD: Terutama menggunakan paradigma SPMD, yang mengharuskan semua perangkat menjalankan program yang sama.
Fleksibilitas	Mendukung pola komputasi yang lebih kompleks di luar SPMD, termasuk paralelisme pipeline asimetris dan kepadatan komputasi.	Dapat kurang fleksibel dalam pengelolaan resource, terutama di berbagai slice TPU.

Komponen jalur

Bagian berikut menguraikan komponen utama arsitektur Pathways.

Pengelola resource jalur

Ini adalah bidang kontrol pusat sistem Pathways. Layanan ini mengelola semua resource akselerator dan bertanggung jawab untuk mengoordinasikan alokasi akselerator untuk tugas pengguna. Layanan ini memantau kondisi worker dan menangani penjadwalan, menjeda, dan melanjutkan tugas. Layanan ini berfungsi sebagai satu titik kontak untuk error dan status sistem. Komponen ini hanya memerlukan resource CPU.

Klien Pathways

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

Pekerja Pathways

Ini adalah proses yang berjalan di mesin akselerator (VM TPU). Mereka menerima file yang dapat dieksekusi yang dikompilasi dari program Anda dari server proxy IFRT dan mereka melakukan komputasi di TPU. Pekerja Pathways mengirimkan data kembali ke program Anda melalui server proxy IFRT. Komponen ini memerlukan resource akselerator.

Klien proxy IFRT

Ini adalah penerapan OSS dari Interim Framework Runtime (IFRT) API yang memisahkan kode pengguna dari runtime yang mendasarinya, serta meningkatkan portabilitas dan transparansi kode. JAX menggunakan implementasi ini sebagai alternatif untuk runtime multi-pengontrol defaultnya. Klien proxy IFRT bertindak sebagai jembatan komunikasi antara program Anda dan komponen Pathways. Aplikasi ini mengirim permintaan ke server proxy IFRT dan menerima hasilnya. Ini adalah penerapan OSS dari IFRT API. Komponen ini hanya memerlukan resource CPU.

Server proxy IFRT

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

Server sidecar

Server gRPC ini ditempatkan bersama dengan worker Pathways di VM akselerator untuk menjalankan kode Python yang ditentukan pengguna di VM akselerator secara langsung guna mengurangi latensi transfer data dari pengontrol ke akselerator. Server sidecar berinteraksi dengan pekerja Pathways melalui protokol berversi kustom pada transportasi gRPC.

Menampilkan hubungan komponen Pathways. — Komponen jalur

PathwaysJob API

PathwaysJob API adalah OSS Kubernetes-native API yang Anda gunakan untuk men-deploy workload pelatihan ML dan inferensi batch. Pengontrol untuk PathwaysJob memanfaatkan JobSet API untuk mengelola siklus proses dan koordinasi semua komponen Pathways. Definisi Resource Kustom (CRD) ini memberi Anda antarmuka tingkat tinggi untuk menentukan beban kerja Pathways, sehingga tidak perlu mengelola spesifikasi Pod satu per satu secara langsung untuk skenario umum. Untuk mengetahui daftar lengkap semua parameter dan artinya yang spesifik, lihat dokumentasi PathwaysJob API di GitHub.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

Tabel berikut menjelaskan setelan untuk API PathwaysJob:

Atribut	Deskripsi
`apiVersion`	Menentukan versi API untuk Definisi Resource Kustom (CRD) `PathwaysJob`: pathways-job.pathways.domain/v1.
`kind`	Mengidentifikasi objek Kubernetes sebagai `PathwaysJob`.
`metadata.name`	Nama objek `PathwaysJob` di Kubernetes, biasanya mengikuti pola pathways-.
`spec`	Menentukan status dan konfigurasi yang diinginkan untuk `PathwaysJob`.
`spec.maxRestarts`	Jumlah maksimum `PathwaysJob` dapat dimulai ulang secara otomatis oleh sistem jika mengalami kegagalan.
`spec.pathwaysVersion`	(Opsional) Menentukan versi framework JAX yang diinginkan untuk digunakan dalam lingkungan Pathways untuk tugas ini (misalnya, jax-0.5.3).
`spec.workers`	Array yang menentukan konfigurasi untuk kumpulan pekerja `PathwaysJob`, biasanya menggunakan resource TPU.
`spec.workers[].type`	Jenis mesin TPU yang akan digunakan untuk worker node (misalnya, $TPU_MACHINE_TYPE dapat berupa ct6e-standard-4t)
`spec.workers[].topology`	Topologi slice TPU yang dialokasikan untuk pekerja (misalnya, $TOPOLOGY bisa berupa 2x2, 4x4, 2x2x2).
`spec.workers[].numSlices`	Jumlah slice TPU yang akan disediakan untuk kumpulan pekerja (misalnya, $WORKLOAD_NODEPOOL_COUNT bisa menjadi 2).
`spec.workers[].maxSliceRestarts`	(Opsional) Jumlah maksimum upaya memulai ulang setiap pekerja dalam slice jika gagal.
`spec.customComponents`	(Opsional) Array yang memungkinkan Anda menentukan dan men-deploy komponen kustom (seperti server proxy, server Pathways, atau pekerja tambahan) bersama dengan tugas utama.
`spec.customComponents[].componentType`	Menentukan jenis komponen kustom yang sedang ditentukan (misalnya, proxy_server, pathways_server, worker, colocated_python_sidecar).
`spec.customComponents[].image`	Image Docker yang akan digunakan untuk container komponen kustom ini.
`spec.customComponents[].customFlags`	Array tanda command line kustom yang akan diteruskan ke container saat dimulai.
`spec.customComponents[].customEnv`	Array variabel lingkungan kustom yang akan ditetapkan dalam container. Setiap elemen memiliki nama dan nilai.
`spec.pathwaysDir`	Bucket Cloud Storage yang digunakan oleh `PathwaysJob` untuk menyimpan artefak kompilasi dan data sementara lainnya. Bucket ini harus dibuat sebelum menjalankan beban kerja Anda.
`spec.controller`	Setelan konfigurasi pengontrol jalur, yang mengelola keseluruhan eksekusi tugas.
`spec.controller.deploymentMode`	Menentukan cara resource CPU pengontrol Pathways (pengelola resource Pathways dan Server Proxy) di-deploy. Mode default men-deploy-nya di node CPU khusus, sedangkan colocate_head_with_workers men-deploy-nya bersama pekerja TPU.
`spec.controller.elasticSlices`	(Opsional) Jumlah maksimum slice TPU yang dapat menjadi tidak tersedia selama eksekusi tugas sebelum dianggap tidak responsif.
`spec.controller.template`	(Opsional) Menentukan template Pod untuk tugas pengguna. Hal ini diperlukan untuk workload batch, tetapi tidak untuk workload interaktif.
`spec.controller.template.spec`	Spesifikasi Pod untuk tugas pengguna.
`spec.controller.template.spec.containers`	Array yang menentukan container yang akan berjalan dalam tugas pengguna
`spec.controller.template.spec.containers[].name`	Nama penampung dalam tugas pengguna (dalam contoh ini, adalah main).
`spec.controller.template.spec.containers[].image`	Image Docker yang akan digunakan untuk container di container utama (dalam contoh ini, adalah python:3.11).
`spec.controller.template.spec.containers[].command`	Perintah yang akan dijalankan saat container utama dimulai. Dalam contoh ini, `pathwaysutils` diinstal, Pathways diinisialisasi, dan perangkat JAX dicetak.

Komponen Pathways di GKE

Bagian ini memetakan komponen Pathways ke komponen Google Kubernetes Engine seperti container dan pod.

Anda dapat menemukan image container Pathways di lokasi berikut.

Jenis Penampung	Location
Server proxy IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>`
Pengelola/pekerja sumber daya Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>`

Pengelola resource jalur

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

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Deskripsi argumen:

--server_port: pengelola resource Pathways menggunakan port ini untuk berkomunikasi dengan komponen jalur lainnya.
--node_type: jenis node. Nilai ini harus ditetapkan ke "resource_manager" untuk pengelola resource Pathways dan tidak diperlukan untuk penampung lainnya.
--instance_count: jumlah slice TPU.
--instance_type: jenis dan topologi TPU slice. Dalam format tpu{TPU type}:{TPU 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/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

Deskripsi argumen:

--resource_manager_address: nama host dan port yang digunakan server proxy untuk berkomunikasi dengan pengelola resource Pathways. Port harus sama dengan nilai --server_port yang digunakan untuk penampung pengelola resource Pathways.
--server_port: server proxy IFRT menggunakan port ini untuk berkomunikasi dengan klien proxy IFRT.
--gcs_scratch_location: bucket GCS yang digunakan untuk file sementara.

Pekerja Pathways

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

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/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"

Deskripsi argumen:

--resource_manager_address: nama host dan port yang digunakan pekerja TPU untuk berkomunikasi dengan pengelola resource Pathways. Port harus sama dengan nilai --server_port yang digunakan untuk penampung pengelola resource Pathways
--server_port: pekerja menggunakan port ini untuk berkomunikasi dengan server proxy dan pengelola resource Pathways.
--gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

Pengelola resource Pathways, server proxy IFRT, dan pekerja Pathways semuanya dapat memiliki port yang berbeda, tetapi dalam contoh ini, pengelola resource Pathways dan pekerja Pathways menggunakan port yang sama.

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

Sebelum memulai

Pengontrol tunggal dan multi-pengontrol

Komponen jalur

Pengelola resource jalur

Klien Pathways

Pekerja Pathways

Klien proxy IFRT

Server proxy IFRT

Server sidecar

PathwaysJob API

Komponen Pathways di GKE

Pengelola resource jalur

Server proxy IFRT

Pekerja Pathways

Langkah berikutnya

Pengantar Jalur Pembelajaran di Cloud