Pengantar Jalur Pembelajaran di Cloud

Pathways adalah sistem yang dirancang untuk memungkinkan pembuatan sistem machine learning berskala besar, multi-tugas, dan diaktifkan secara tersebar. Sistem ini memungkinkan penggunaan ribuan atau puluhan ribu akselerator, dengan kemampuan untuk mengalokasikan jumlah komputasi yang bervariasi secara dinamis untuk tugas yang berbeda 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. Pathways on Cloud memberikan manfaat yang sama kepada Google Cloud pelanggan.

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. Dokumen ini mengasumsikan bahwa Anda sudah memahami cara menggunakan TPU dengan GKE termasuk TPU satu 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-pengontrol (Default JAX)

Kontrol

Titik kontrol tunggal: Program klien tunggal 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 terlokalisasi: Setiap proses Python hanya melihat perangkat yang terhubung ke proses tersebut.

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 di luar SPMD, termasuk paralelisme pipeline asimetris dan sparsity komputasi.

Dapat menjadi 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. Pengelola ini mengelola semua resource akselerator dan bertanggung jawab untuk mengoordinasikan alokasi akselerator untuk tugas pengguna. Pengelola ini memantau kondisi pekerja dan menangani penjadwalan, jeda, dan melanjutkan tugas. Pengelola ini berfungsi sebagai satu titik kontak 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. Klien ini menerima Operasi Tingkat Tinggi (HLO) dari program Anda. Klien Pathways bertanggung jawab untuk berkoordinasi dengan pengelola resource Pathways untuk menentukan tempat menempatkan program yang dikompilasi untuk dieksekusi berdasarkan kode pengguna. Klien ini menyajikan tampilan sistem yang terpadu ke klien JAX tertentu. Komponen ini hanya memerlukan resource CPU.

Pekerja Pathways

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

Klien proxy IFRT

Ini adalah implementasi OSS dari Interim Framework Runtime (IFRT) API yang memisahkan kode pengguna dari runtime yang mendasarinya, dan 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 mengirim 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 tugas yang sebenarnya. Komponen ini hanya memerlukan resource CPU.

Server sidecar

Server gRPC ini ditempatkan bersama dengan pekerja 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 versi kustom pada transportasi gRPC.

Menampilkan 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 Container

Lokasi

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 Anda membuat cluster GKE, Anda dapat menggunakan containerSpec berikut untuk men-deploy pengelola resource 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: pengelola resource Pathways menggunakan port ini untuk berkomunikasi dengan komponen Pathways lainnya.
  • --node_type: jenis node. Nilai 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 pengelola resource 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 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 container pengelola resource Pathways
  • --server_port: pekerja menggunakan port ini untuk berkomunikasi dengan proxy server dan pengelola resource Pathways.
  • --gcs_scratch_location: bucket Cloud Storage yang digunakan untuk file sementara.

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

Langkah berikutnya