Menskalakan workload ML menggunakan Ray

Dokumen ini memberikan detail tentang cara menjalankan beban kerja machine learning (ML) dengan Ray dan JAX di TPU. Ada dua mode berbeda untuk menggunakan TPU dengan Ray: Mode yang berfokus pada perangkat (PyTorch/XLA) dan Mode yang berfokus pada host (JAX).

Dokumen ini mengasumsikan bahwa Anda telah menyiapkan lingkungan TPU. Untuk informasi selengkapnya, lihat referensi berikut:

• Cloud TPU: Menyiapkan lingkungan Cloud TPU dan Mengelola resource TPU
• Google Kubernetes Engine (GKE): Men-deploy workload TPU di GKE Autopilot atau Men-deploy workload TPU di GKE Standard

Mode yang berfokus pada perangkat (PyTorch/XLA)

Mode yang berfokus pada perangkat mempertahankan sebagian besar gaya terprogram PyTorch klasik. Dalam mode ini, Anda menambahkan jenis perangkat XLA baru, yang berfungsi seperti perangkat PyTorch lainnya. Setiap proses berinteraksi dengan satu perangkat XLA.

Mode ini ideal jika Anda sudah terbiasa dengan PyTorch dengan GPU dan ingin menggunakan abstraksi coding yang serupa.

Bagian berikut menjelaskan cara menjalankan beban kerja PyTorch/XLA di satu atau beberapa perangkat tanpa menggunakan Ray, lalu cara menjalankan beban kerja yang sama di beberapa host menggunakan Ray.

Membuat TPU

1. Buat variabel lingkungan untuk parameter pembuatan TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-b
   export ACCELERATOR_TYPE=v5p-8
   export RUNTIME_VERSION=v2-alpha-tpuv5

   Deskripsi variabel lingkungan

   Variabel	Deskripsi
   PROJECT_ID	Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
   TPU_NAME	Nama TPU.
   ZONE	Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
   ACCELERATOR_TYPE	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
   RUNTIME_VERSION	Versi software Cloud TPU.

2. Gunakan perintah berikut untuk membuat VM TPU v5p dengan 8 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Hubungkan ke VM TPU menggunakan perintah berikut:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

Jika Anda menggunakan GKE, lihat panduan KubeRay di GKE untuk mengetahui informasi penyiapan.

Persyaratan penginstalan

Jalankan perintah berikut di VM TPU untuk menginstal dependensi yang diperlukan:

1. Simpan kode berikut ke file. Misalnya, requirements.txt.

   --find-links https://storage.googleapis.com/libtpu-releases/index.html
   --find-links https://storage.googleapis.com/libtpu-wheels/index.html
   torch~=2.6.0
   torch_xla[tpu]~=2.6.0
   ray[default]==2.40.0
   

2. Untuk menginstal dependensi yang diperlukan, jalankan:

   pip install -r requirements.txt
   

Jika Anda menjalankan workload di GKE, sebaiknya buat Dockerfile yang menginstal dependensi yang diperlukan. Sebagai contoh, lihat Menjalankan workload Anda di node slice TPU dalam dokumentasi GKE.

Menjalankan workload PyTorch/XLA di satu perangkat

Contoh berikut menunjukkan cara membuat tensor XLA di satu perangkat, yaitu chip TPU. Hal ini mirip dengan cara PyTorch menangani jenis perangkat lainnya.

1. Simpan cuplikan kode berikut ke file. Misalnya, workload.py.

   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   
   t = torch.randn(2, 2, device=xm.xla_device())
   print(t.device)
   print(t)
   

   Pernyataan impor import torch_xla melakukan inisialisasi PyTorch/XLA, dan fungsi xm.xla_device() menampilkan perangkat XLA saat ini, yaitu chip TPU.

2. Tetapkan variabel lingkungan PJRT_DEVICE ke TPU.

   export PJRT_DEVICE=TPU
   

3. Jalankan skrip.

   python workload.py
   

   Output-nya terlihat mirip dengan ini: Pastikan output menunjukkan bahwa perangkat XLA ditemukan.

   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   

Menjalankan PyTorch/XLA di beberapa perangkat

1. Perbarui cuplikan kode dari bagian sebelumnya agar dapat berjalan di beberapa perangkat.

   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   
   def _mp_fn(index):
       t = torch.randn(2, 2, device=xm.xla_device())
       print(t.device)
       print(t)
   
   if __name__ == '__main__':
       torch_xla.launch(_mp_fn, args=())
   

2. Jalankan skrip.

   python workload.py
   

   Jika Anda menjalankan cuplikan kode di TPU v5p-8, output-nya akan terlihat mirip dengan berikut:

   xla:0
   xla:0
   xla:0
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   xla:0
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   

torch_xla.launch() menggunakan dua argumen: fungsi dan daftar parameter. Fungsi ini membuat proses untuk setiap perangkat XLA yang tersedia dan memanggil fungsi yang ditentukan dalam argumen. Dalam contoh ini, ada 4 perangkat TPU yang tersedia, sehingga torch_xla.launch() membuat 4 proses dan memanggil _mp_fn() di setiap perangkat. Setiap proses hanya memiliki akses ke satu perangkat, sehingga setiap perangkat memiliki indeks 0, dan xla:0 dicetak untuk semua proses.

Menjalankan PyTorch/XLA di beberapa host dengan Ray

Bagian berikut menunjukkan cara menjalankan cuplikan kode yang sama di slice TPU multi-host yang lebih besar. Untuk informasi selengkapnya tentang arsitektur TPU multi-host, lihat Arsitektur sistem.

Dalam contoh ini, Anda menyiapkan Ray secara manual. Jika sudah memahami cara menyiapkan Ray, Anda dapat langsung ke bagian terakhir, Menjalankan workload Ray. Untuk informasi selengkapnya tentang cara menyiapkan Ray untuk lingkungan produksi, lihat referensi berikut:

• Meluncurkan Cluster Ray
• Memulai Google Cloud cluster GKE dengan TPU untuk KubeRay

Membuat VM TPU multi-host

1. Buat variabel lingkungan untuk parameter pembuatan TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-b
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=v2-alpha-tpuv5

   Deskripsi variabel lingkungan

   Variabel	Deskripsi
   PROJECT_ID	Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
   TPU_NAME	Nama TPU.
   ZONE	Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
   ACCELERATOR_TYPE	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
   RUNTIME_VERSION	Versi software Cloud TPU.

2. Buat TPU v5p multi-host dengan 2 host (v5p-16, dengan 4 chip TPU di setiap host) menggunakan perintah berikut:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --version=$RUNTIME_VERSION

Menyiapkan Ray

TPU v5p-16 memiliki 2 host TPU, masing-masing dengan 4 chip TPU. Dalam contoh ini, Anda akan memulai node head Ray di satu host dan menambahkan host kedua sebagai node pekerja ke cluster Ray.

1. Hubungkan ke host pertama menggunakan SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0

2. Instal dependensi dengan file persyaratan yang sama seperti di bagian Instal persyaratan.

   pip install -r requirements.txt
   

3. Mulai proses Ray.

   ray start --head --port=6379
   

   Catatan: Jika mendapatkan error seperti Command 'ray' not found, Anda mungkin perlu menambahkan direktori yang berisi file yang dapat dieksekusi ray ke variabel lingkungan PATH sistem Anda. Anda juga dapat menentukan jalur lengkap saat menjalankan perintah ray (misalnya, ~/.local/bin/ray).
   Catatan: Perintah ray start meminta Anda untuk mengaktifkan atau menonaktifkan pengumpulan statistik penggunaan. Untuk informasi selengkapnya, lihat Pengumpulan statistik penggunaan dalam dokumentasi Ray.

   Outputnya terlihat mirip dengan yang berikut ini:

   Enable usage stats collection? This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup. Confirm [Y/n]: y
   Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.
   
   Local node IP: 10.130.0.76
   
   --------------------
   Ray runtime started.
   --------------------
   
   Next steps
   To add another node to this Ray cluster, run
       ray start --address='10.130.0.76:6379'
   
   To connect to this Ray cluster:
       import ray
       ray.init()
   
   To terminate the Ray runtime, run
       ray stop
   
   To view the status of the cluster, use
       ray status
   

   Host TPU ini sekarang menjadi node head Ray. Catat baris yang menunjukkan cara menambahkan node lain ke cluster Ray, mirip dengan baris berikut:

   To add another node to this Ray cluster, run
       ray start --address='10.130.0.76:6379'
   

   Anda akan menggunakan perintah ini di langkah berikutnya.

4. Periksa status cluster Ray:

   ray status
   

   Outputnya terlihat mirip dengan yang berikut ini:

   ======== Autoscaler status: 2025-01-14 22:03:39.385610 ========
   Node status
   ---------------------------------------------------------------
   Active:
   1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
   Pending:
   (no pending nodes)
   Recent failures:
   (no failures)
   
   Resources
   ---------------------------------------------------------------
   Usage:
   0.0/208.0 CPU
   0.0/4.0 TPU
   0.0/1.0 TPU-v5p-16-head
   0B/268.44GiB memory
   0B/119.04GiB object_store_memory
   0.0/1.0 your-tpu-name
   
   Demands:
   (no resource demands)
   

   Cluster hanya berisi 4 TPU (0.0/4.0 TPU) karena Anda baru saja menambahkan node head.

   Setelah node head berjalan, Anda dapat menambahkan host kedua ke cluster.

5. Hubungkan ke host kedua menggunakan SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1

6. Instal dependensi dengan file persyaratan yang sama seperti di bagian Persyaratan penginstalan.

   pip install -r requirements.txt
   

7. Mulai proses Ray. Untuk menambahkan node ini ke cluster Ray yang ada, gunakan perintah dari output perintah ray start. Pastikan untuk mengganti alamat IP dan port dalam perintah berikut:

   ray start --address='10.130.0.76:6379'

   Outputnya terlihat mirip dengan yang berikut ini:

   Local node IP: 10.130.0.80
   [2025-01-14 22:30:07,397 W 75572 75572] global_state_accessor.cc:463: Retrying to get node with node ID 35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
   
   --------------------
   Ray runtime started.
   --------------------
   
   To terminate the Ray runtime, run
   ray stop
   

8. Periksa kembali status Ray:

   ray status
   

   Outputnya terlihat mirip dengan yang berikut ini:

   ======== Autoscaler status: 2025-01-14 22:45:21.485617 ========
   Node status
   ---------------------------------------------------------------
   Active:
   1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
   1 node_35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
   Pending:
   (no pending nodes)
   Recent failures:
   (no failures)
   
   Resources
   ---------------------------------------------------------------
   Usage:
   0.0/416.0 CPU
   0.0/8.0 TPU
   0.0/1.0 TPU-v5p-16-head
   0B/546.83GiB memory
   0B/238.35GiB object_store_memory
   0.0/2.0 your-tpu-name
   
   Demands:
   (no resource demands)
   

   Host TPU kedua kini menjadi node di cluster. Daftar resource yang tersedia sekarang menampilkan 8 TPU (0.0/8.0 TPU).

   Catatan: Anda harus memulai proses Ray di setiap VM TPU sebelum menjalankan beban kerja.

Menjalankan workload Ray

1. Perbarui cuplikan kode untuk dijalankan di cluster Ray:

   import os
   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   import ray
   
   import torch.distributed as dist
   import torch_xla.runtime as xr
   from torch_xla._internal import pjrt
   
   # Defines the local PJRT world size, the number of processes per host.
   LOCAL_WORLD_SIZE = 4
   # Defines the number of hosts in the Ray cluster.
   NUM_OF_HOSTS = 4
   GLOBAL_WORLD_SIZE = LOCAL_WORLD_SIZE * NUM_OF_HOSTS
   
   def init_env():
       local_rank = int(os.environ['TPU_VISIBLE_CHIPS'])
   
       pjrt.initialize_multiprocess(local_rank, LOCAL_WORLD_SIZE)
       xr._init_world_size_ordinal()
   
   # This decorator signals to Ray that the `print_tensor()` function should be run on a single TPU chip.
   @ray.remote(resources={"TPU": 1})
   def print_tensor():
       # Initializes the runtime environment on each Ray worker. Equivalent to
       # the `torch_xla.launch call` in the Run PyTorch/XLA on multiple devices section.
       init_env()
   
       t = torch.randn(2, 2, device=xm.xla_device())
       print(t.device)
       print(t)
   
   ray.init()
   
   # Uses Ray to dispatch the function call across available nodes in the cluster.
   tasks = [print_tensor.remote() for _ in range(GLOBAL_WORLD_SIZE)]
   ray.get(tasks)
   
   ray.shutdown()
   

2. Jalankan skrip di node head Ray. Ganti ray-workload.py dengan jalur ke skrip Anda.

   python ray-workload.py

   Outputnya terlihat mirip dengan yang berikut ini:

   WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
   xla:0
   xla:0
   xla:0
   xla:0
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   xla:0
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   

   Output menunjukkan bahwa fungsi berhasil dipanggil di setiap perangkat XLA (8 perangkat dalam contoh ini) di slice TPU multi-host.

Mode yang berfokus pada host (JAX)

Bagian berikut menjelaskan mode yang berfokus pada host dengan JAX. JAX menggunakan paradigma pemrograman fungsional dan mendukung semantik program tunggal dan beberapa data (SPMD) tingkat tinggi. Alih-alih membuat setiap proses berinteraksi dengan satu perangkat XLA, kode JAX dirancang untuk beroperasi di beberapa perangkat pada satu host secara serentak.

JAX dirancang untuk komputasi berperforma tinggi dan dapat menggunakan TPU secara efisien untuk pelatihan dan inferensi skala besar. Mode ini ideal jika Anda sudah terbiasa dengan konsep pemrograman fungsional sehingga Anda dapat memanfaatkan potensi penuh JAX.

Petunjuk ini mengasumsikan bahwa Anda telah menyiapkan lingkungan Ray dan TPU, termasuk lingkungan software yang menyertakan JAX dan paket terkait lainnya. Untuk membuat cluster TPU Ray, ikuti petunjuk di Memulai Google Cloud cluster GKE dengan TPU untuk KubeRay. Untuk informasi selengkapnya tentang cara menggunakan TPU dengan KubeRay, lihat Menggunakan TPU dengan KubeRay.

Menjalankan workload JAX di TPU host tunggal

Contoh skrip berikut menunjukkan cara menjalankan fungsi JAX di cluster Ray dengan TPU satu host, seperti v6e-4. Jika Anda memiliki TPU multi-host, skrip ini akan berhenti merespons karena model eksekusi multi-pengontrol JAX. Untuk mengetahui informasi selengkapnya tentang cara menjalankan Ray di TPU multi-host, lihat Menjalankan workload JAX di TPU multi-host.

1. Buat variabel lingkungan untuk parameter pembuatan TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-a
   export ACCELERATOR_TYPE=v6e-4
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Deskripsi variabel lingkungan

   Variabel	Deskripsi
   PROJECT_ID	Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
   TPU_NAME	Nama TPU.
   ZONE	Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
   ACCELERATOR_TYPE	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
   RUNTIME_VERSION	Versi software Cloud TPU.

2. Gunakan perintah berikut untuk membuat VM TPU v6e dengan 4 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Hubungkan ke VM TPU menggunakan perintah berikut:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

4. Instal JAX dan Ray di TPU Anda.

   pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

5. Simpan kode berikut ke file. Misalnya, ray-jax-single-host.py.

   import ray
   import jax
   
   @ray.remote(resources={"TPU": 4})
   def my_function() -> int:
       return jax.device_count()
   
   h = my_function.remote()
   print(ray.get(h)) # => 4
   

   Jika Anda terbiasa menjalankan Ray dengan GPU, ada beberapa perbedaan utama saat menggunakan TPU:

   • Daripada menetapkan num_gpus, tentukan TPU sebagai resource kustom dan tetapkan jumlah chip TPU.
   • Tentukan TPU menggunakan jumlah chip per node pekerja Ray. Misalnya, jika Anda menggunakan v6e-4, menjalankan fungsi jarak jauh dengan TPU yang disetel ke 4 akan menggunakan seluruh host TPU.
   • Hal ini berbeda dengan cara GPU biasanya berjalan, dengan satu proses per host. Menetapkan TPU ke angka yang bukan 4 tidak direkomendasikan.
     • Pengecualian: Jika memiliki v6e-8 atau v5litepod-8 satu host, Anda harus menetapkan nilai ini ke 8.

6. Jalankan skrip.

   python ray-jax-single-host.py

Menjalankan workload JAX di TPU multi-host

Contoh skrip berikut menunjukkan cara menjalankan fungsi JAX di cluster Ray dengan TPU multi-host. Contoh skrip menggunakan v6e-16.

1. Buat variabel lingkungan untuk parameter pembuatan TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-a
   export ACCELERATOR_TYPE=v6e-16
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Deskripsi variabel lingkungan

   Variabel	Deskripsi
   PROJECT_ID	Project ID Google Cloud Anda. Gunakan project yang ada atau buat project baru.
   TPU_NAME	Nama TPU.
   ZONE	Zona tempat VM TPU akan dibuat. Untuk mengetahui informasi selengkapnya tentang zona yang didukung, lihat Region dan zona TPU.
   ACCELERATOR_TYPE	Jenis akselerator menentukan versi dan ukuran Cloud TPU yang ingin Anda buat. Untuk mengetahui informasi selengkapnya tentang jenis akselerator yang didukung untuk setiap versi TPU, lihat versi TPU.
   RUNTIME_VERSION	Versi software Cloud TPU.

2. Gunakan perintah berikut untuk membuat VM TPU v6e dengan 16 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Instal JAX dan Ray di semua pekerja TPU.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --zone=$ZONE \
      --worker=all \
      --command="pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

4. Simpan kode berikut ke file. Misalnya, ray-jax-multi-host.py.

   import ray
   import jax
   
   @ray.remote(resources={"TPU": 4})
   def my_function() -> int:
       return jax.device_count()
   
   ray.init()
   num_tpus = ray.available_resources()["TPU"]
   num_hosts = int(num_tpus) # 4
   h = [my_function.remote() for _ in range(num_hosts)]
   print(ray.get(h)) # [16, 16, 16, 16]
   

   Jika Anda terbiasa menjalankan Ray dengan GPU, ada beberapa perbedaan utama saat menggunakan TPU:

   • Mirip dengan workload PyTorch di GPU:
     • Beban kerja JAX di TPU berjalan dalam mode multi-pengontrol, satu program, beberapa data (SPMD).
     • Kolektif antarperangkat ditangani oleh framework machine learning.
   • Tidak seperti workload PyTorch di GPU, JAX memiliki tampilan global tentang perangkat yang tersedia di cluster.

5. Salin skrip ke semua pekerja TPU.

   gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all

6. Jalankan skrip.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --zone=$ZONE \
      --worker=all \
      --command="python ray-jax-multi-host.py"

Menjalankan workload JAX Multislice

Multislice memungkinkan Anda menjalankan workload yang mencakup beberapa slice TPU dalam satu Pod TPU atau di beberapa pod melalui jaringan pusat data.

Anda dapat menggunakan paket ray-tpu untuk menyederhanakan interaksi Ray dengan slice TPU.

Instal ray-tpu menggunakan pip.

pip install ray-tpu

Untuk informasi selengkapnya tentang cara menggunakan paket ray-tpu, lihat Memulai di repositori GitHub. Untuk contoh penggunaan Multislice, lihat Menjalankan di Multislice.

Melakukan orkestrasi workload menggunakan Ray dan MaxText

Untuk informasi selengkapnya tentang cara menggunakan Ray dengan MaxText, lihat Menjalankan tugas pelatihan dengan MaxText.

Resource TPU dan Ray

Ray memperlakukan TPU secara berbeda dengan GPU untuk mengakomodasi perbedaan penggunaan. Dalam contoh berikut, ada total sembilan node Ray:

• Node head Ray berjalan di VM n1-standard-16.
• Node pekerja Ray berjalan di dua TPU v6e-16. Setiap TPU terdiri dari empat pekerja.

$ ray status
======== Autoscaler status: 2024-10-17 09:30:00.854415 ========
Node status
---------------------------------------------------------------
Active:
 1 node_e54a65b81456cee40fcab16ce7b96f85406637eeb314517d9572dab2
 1 node_9a8931136f8d2ab905b07d23375768f41f27cc42f348e9f228dcb1a2
 1 node_c865cf8c0f7d03d4d6cae12781c68a840e113c6c9b8e26daeac23d63
 1 node_435b1f8f1fbcd6a4649c09690915b692a5bac468598e9049a2fac9f1
 1 node_3ed19176e9ecc2ac240c818eeb3bd4888fbc0812afebabd2d32f0a91
 1 node_6a88fe1b74f252a332b08da229781c3c62d8bf00a5ec2b90c0d9b867
 1 node_5ead13d0d60befd3a7081ef8b03ca0920834e5c25c376822b6307393
 1 node_b93cb79c06943c1beb155d421bbd895e161ba13bccf32128a9be901a
 1 node_9072795b8604ead901c5268ffcc8cc8602c662116ac0a0272a7c4e04
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/727.0 CPU
 0.0/32.0 TPU
 0.0/2.0 TPU-v6e-16-head
 0B/5.13TiB memory
 0B/1.47TiB object_store_memory
 0.0/4.0 tpu-group-0
 0.0/4.0 tpu-group-1

Demands:
 (no resource demands)

Deskripsi kolom penggunaan resource:

• CPU: Jumlah total CPU yang tersedia di cluster.
• TPU: Jumlah TPU chip dalam cluster.
• TPU-v6e-16-head: ID khusus untuk resource yang sesuai dengan pekerja 0 dari slice TPU. Hal ini penting untuk mengakses setiap slice TPU.
• memory: Memori heap pekerja yang digunakan oleh aplikasi Anda.
• object_store_memory: Memori yang digunakan saat aplikasi Anda membuat objek di penyimpanan objek menggunakan ray.put dan saat menampilkan nilai dari fungsi jarak jauh.
• tpu-group-0 dan tpu-group-1: ID unik untuk setiap slice TPU. Hal ini penting untuk menjalankan tugas di slice. Kolom ini ditetapkan ke 4 karena ada 4 host per slice TPU di v6e-16.

Catatan: Semua resource ini terdeteksi secara otomatis saat runtime Ray dimulai (ray start), sehingga Anda tidak perlu menetapkannya secara eksplisit.