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 berbagai jumlah komputasi secara dinamis untuk berbagai tugas berdasarkan persyaratan pemrosesannya.

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

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

Sebelum memulai

Pastikan Anda memiliki:

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

Pengontrol tunggal dan multi-pengontrol

Pada dasarnya, ada dua cara berbeda 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 dengannya.

Pemrograman

Pemrograman yang disederhanakan: Pengguna berinteraksi dengan satu klien, sehingga sistem tampak sebagai 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 selain 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

Ini adalah bidang kontrol pusat sistem Pathways. Slurm mengelola semua resource akselerator dan bertanggung jawab untuk mengoordinasikan alokasi akselerator untuk tugas pengguna. Alat ini 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

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

Pekerja Pathways

Proses ini 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

Ini adalah implementasi OSS 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. Ini adalah implementasi 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 pekerjaan yang sebenarnya. Komponen ini hanya memerlukan resource CPU.

Server file bantuan

Server gRPC ini ditempatkan bersama 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 worker Pathways melalui protokol berversi kustom pada transport gRPC.

Menunjukkan hubungan komponen Pathways.
Komponen Pathways

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 resource Pathways

us-docker.pkg.dev/cloud-tpu-v2-images/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/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: Resource manager Pathways menggunakan port ini untuk berkomunikasi dengan komponen Pathways lainnya.
  • --node_type: jenis node. Ini harus ditetapkan ke "resource_manager" untuk pengelola resource Pathways dan tidak diperlukan untuk container lainnya.
  • --instance_count: jumlah slice TPU.
  • --instance_type: jenis TPU dan topologi 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 resource manager Pathways. Port harus sama dengan nilai --server_port yang digunakan untuk container pengelola resource Pathways.
  • --server_port: server proxy IFRT menggunakan port ini 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/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 TPU worker untuk berkomunikasi dengan pengelola resource Pathways. Port harus sama dengan nilai --server_port yang digunakan untuk penampung pengelola resource Pathways
  • --server_port: worker menggunakan port ini untuk berkomunikasi dengan server proxy dan resource manager Pathways.
  • --gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

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

Langkah berikutnya