Mengirimkan workload JAX ke Pathways

Karena sifat JAX dengan Pathways yang terdistribusi, beberapa operasi mungkin tidak diskalakan dengan baik karena overhead komunikasi. Meskipun Pathways meminimalkan overhead ini dengan fitur seperti pengiriman asinkron, ada beberapa hal yang perlu Anda perhatikan saat mem-porting workload JAX ke Pathways atau menskalakan workload JAX dengan Pathways ke sejumlah besar akselerator.

Sebelum memulai

Pastikan Anda memiliki:

Indeks proses

JAX dengan Pathways memperlakukan semua perangkat di seluruh cluster Pathways Anda sebagai lokal. Hal ini menyederhanakan pengelolaan perangkat dan memungkinkan JAX memanfaatkan semua resource yang tersedia. Dalam praktiknya, hal ini berarti:

  • jax.process_index() selalu 0 untuk semua perangkat.
  • jax.devices() dan jax.local_devices() menampilkan semua perangkat TPU di seluruh tugas.

Jenis hardware dan kolokasi

Untuk performa terbaik, tempatkan semua komponen Pathways dan tugas pengguna di zona cloud Google Cloud yang sama. Gunakan CPU besar seperti proxy IFRT dan pengelola resource. Sebaiknya gunakan setidaknya n2-standard-64 khusus yang dilengkapi dengan 64 vCPU dan memori 256 GB.

PathwaysUtils

Pathways-utils adalah repositori GitHub berbasis Python yang menyediakan utilitas dan alat penting yang memungkinkan Anda menyederhanakan deployment dan eksekusi workload JAX di arsitektur Pathways on Cloud. Paket ini menangani adaptasi yang diperlukan untuk lingkungan cloud, sehingga developer JAX dapat berfokus pada alur kerja machine learning inti mereka dengan konfigurasi khusus platform yang minimal. Secara khusus, alat ini menawarkan:

  • Backend JAX "proxy": backend kustom ini memungkinkan aplikasi JAX Anda menggunakan infrastruktur Pathways dengan menetapkan variabel lingkungan JAX_PLATFORMS=proxy.
  • Utilitas Profiling Terintegrasi: kemampuan pembuatan profil yang memungkinkan Anda memahami performa aplikasi. Dengan menggunakan API pembuatan profil JAX standar seperti jax.profiler.start_trace dan jax.profiler.start_server, Anda dapat membuat profil tidak hanya kode JAX, tetapi juga komponen Pathways yang mendasarinya, sehingga memberikan tampilan eksekusi yang menyeluruh dalam lingkungan cloud.
  • Checkpointing Terdistribusi dengan Orbax: pengendali checkpoint Orbax kustom yang memungkinkan Anda menggunakan checkpoint terdistribusi dan memulihkan checkpoint saat menggunakan library Orbax dalam lingkungan Pathways. Integrasi ini bertujuan untuk berfungsi tanpa memerlukan perubahan apa pun pada kode checkpoint Orbax yang ada selama kode tersebut mengimpor pathwaysutils.
  • Primitif Pelatihan Elastis: menyediakan primitif pelatihan elastis dasar yang dapat Anda gunakan untuk membangun alur kerja pelatihan yang tangguh dan skalabel menggunakan Pathways. Primitif ini memungkinkan tugas pelatihan Anda beradaptasi secara dinamis terhadap perubahan resource yang tersedia, sehingga meningkatkan efisiensi dan ketahanan di lingkungan cloud.

Checkpointing

Orbax diuji secara menyeluruh dengan Pathways untuk checkpointing dan pemulihan terdistribusi dengan Cloud Storage. Saat Anda menetapkan variabel lingkungan ENABLE_PATHWAYS_PERSISTENCE=1 dan memanggil import pathwaysutils; pathwaysutils.initialize() di train.py, ArrayHandler kustom akan didaftarkan yang secara efisien menangani operasi checkpoint melalui proxy IFRT, sehingga memungkinkan worker Pathways di akselerator untuk menyimpan dan memulihkan data secara langsung.

Python yang Ditempatkan Bersama

Colocated Python adalah JAX API open source yang memungkinkan Anda menjalankan kode Python yang ditentukan pengguna langsung di host TPU atau GPU, yang lebih mudah di JAX multi-pengontrol. Dengan menggunakan pendekatan ini, Anda dapat menghindari transfer data antara mesin klien dan TPU untuk tugas yang memerlukan komputasi intensif, seperti pemuatan data dan checkpointing. Untuk mengonfigurasi cluster Pathways agar menjalankan Python JAX API yang ditempatkan bersama, bangun dan kirim image container sidecar menggunakan contoh Dockerfile sidecar.

Pemuatan data

Selama pelatihan, kita berulang kali memuat batch dari set data untuk dimasukkan ke dalam model. Memiliki loader data asinkron yang efisien dan melakukan sharding batch di seluruh host sangat penting untuk menghindari kekurangan pekerjaan pada akselerator. Saat menjalankan pelatihan dengan Pathways, loader data berjalan di VM CPU (tidak seperti VM TPU yang digunakan pada penyiapan multi-pengontrol) dan mengirimkan data ke VM TPU. Hal ini menimbulkan latensi yang lebih tinggi dalam membaca data, tetapi hal tersebut dapat dimitigasi sebagian dengan membaca terlebih dahulu sejumlah X batch pada host CPU dan mengirimkan data yang dibaca secara asinkron ke TPU. Solusi ini sudah cukup saat dijalankan dalam skala kecil hingga menengah.

Saat Anda menjalankan loader data multiproses (seperti Grain) di container workload head di GKE, container tersebut akan berbagi node head dengan sidecar pathways-proxy dan pathways-rm. Karena arsitektur bersama ini, Anda harus mengonfigurasi hal berikut:

  • Menetapkan jumlah worker eksplisit: Di dalam container, os.cpu_count() melaporkan total vCPU VM host, bukan batas CPU container. Teruskan jumlah worker eksplisit (tidak lebih dari batas CPU container), bukan mengambilnya dari os.cpu_count().
  • Memasang volume /dev/shm yang didukung memori: Loader data multiproses menggunakan memori bersama POSIX (/dev/shm) untuk komunikasi antarproses. Pastikan volume emptyDir (medium: Memory) yang didukung memori dipasang di /dev/shm sehingga batch yang diambil sebelumnya tidak melebihi batas memori bersama container default.

Untuk performa optimal dalam skala besar, sebaiknya tempatkan pipeline data input Anda bersama-sama menggunakan Python yang ditempatkan bersama untuk menjalankan pipeline data Anda langsung di akselerator. Hal ini menghilangkan bottleneck CPU dan memanfaatkan interkoneksi cepat TPU untuk transfer data.

Anda dapat menemukan implementasi referensi untuk memigrasikan pipeline input berbasis TFDS dalam implementasi RemoteIterator di multihost_dataloading.py. Implementasi ini berfungsi pada JAX multi-pengontrol dan Pathways secara terdistribusi menggunakan API JAX Python yang ditempatkan bersama.

Pembuatan Versi Jax

Rilis Pathways sangat terkait dengan versi JAX untuk memastikan kompatibilitas dan stabilitas. Untuk menghindari potensi masalah, pastikan artefak Pathways dan versi JAX Anda selaras. Setiap rilis Pathways menentukan dengan jelas versi JAX yang kompatibel melalui tag dalam bentuk jax-<version>.

Cache Kompilasi

Cache kompilasi persisten Pathways adalah fitur yang memungkinkan server Pathways menyimpan file yang dapat dieksekusi XLA yang dikompilasi di lokasi persisten, seperti Cloud Storage, untuk menghindari kompilasi yang redundan. Fitur ini diaktifkan secara default. Lokasi cache diteruskan sebagai flag --gcs_scratch_location ke resource manager dan container worker Pathways. Untuk meminimalkan biaya penyimpanan terkait, cache akan melampirkan kebijakan siklus proses ke lokasi Cloud Storage. Ada batas 50 kebijakan per bucket Cloud Storage. Oleh karena itu, sebaiknya gunakan lokasi Cloud Storage yang sama di semua workload.

Jika Anda menghilangkan flag --gcs_scratch_location atau membiarkannya kosong, Pathways akan otomatis kembali ke cache kompilasi dalam memori. Cache ini tidak mempertahankan file yang dapat dieksekusi XLA yang dikompilasi di seluruh deployment Pathways.

Cache ini mirip dengan cache kompilasi JAX yang dinonaktifkan oleh pathwaysutils.initialize() untuk workload Pathways.

Izin Cloud Storage berikut diperlukan untuk cache kompilasi:

  • storage.buckets.get: Untuk mengambil metadata bucket.
  • storage.buckets.update: Penting bagi Pathways untuk menyiapkan kebijakan siklus proses objek guna menerapkan TTL untuk penghapusan cache.
  • storage.objects.list: Untuk mencantumkan objek cache yang ada dalam bucket.
  • storage.objects.create: Untuk menulis file yang dapat dieksekusi yang baru dikompilasi ke cache.
  • storage.objects.get: Untuk membaca file yang dapat dieksekusi yang di-cache dari bucket.

Pembuatan profil

Anda dapat menggunakan profiler JAX untuk membuat trace program JAX. Ada dua cara umum yang didukung dengan Pathways:

  • Terprogram
    • Merekam profil secara terprogram dari kode JAX Anda
  • Manual
    • Merekam profil sesuai permintaan setelah memulai server profiler dari kode JAX Anda

Dalam kedua kasus tersebut, profil ditulis ke bucket Cloud Storage. Akan ada beberapa file rekaman aktivitas yang dibuat di bucket Cloud Storage, yang berpotensi berada di folder stempel waktu yang berbeda, misalnya:

  • Proses Python utama yang memanggil trace (biasanya VM notebook Anda): <jax-client-vm-name>.xplane.pb
  • Proxy IFRT Pathways: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Resource Manager Pathways: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Worker Pathways: server.*<tpu-node-name>.xplane.pb

File rekaman aktivitas ini dapat dianalisis dengan TensorBoard dengan menjalankan perintah berikut. Untuk mengetahui informasi selengkapnya tentang TensorBoard dan semua alat pembuatan profilnya, lihat Mengoptimalkan performa TensorFlow menggunakan Profiler.

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

Ganti kode berikut:

  • BUCKET : bucket Cloud Storage untuk menyimpan file rekaman aktivitas
  • PREFIX: jalur dalam bucket Cloud Storage Anda untuk menyimpan file rekaman aktivitas

Pengambilan profil terprogram

Mengambil profil dari dalam kode Anda. Profil disimpan di dalam gs://<bucket>/<prefix> di direktori stempel waktu

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Pengambilan profil manual

Untuk mengambil informasi profil secara manual, Anda harus memulai server profiler dari kode Python:

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

Saat server profiler berjalan, Anda dapat mengambil profil dan mengekspor data ke lokasi Cloud Storage target:

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

Anda dapat menemukan informasi pengaturan waktu untuk metode klien proxy IFRT seperti Compile dan Execute dalam trace program Anda. Peristiwa ini, yang menjelaskan interaksi dengan server gRPC Proxy IFRT selama kompilasi dan eksekusi, muncul di thread bernama GrpcClientSessionUserFuturesWorkQueue. Dengan memeriksa thread ini dalam rekaman aktivitas, Anda dapat memperoleh insight tentang performa operasi ini.

Flag XLA

Saat menggunakan Pathways, Anda harus menetapkan flag XLA di container pathways-proxy.

Saat Anda menggunakan Cluster Toolkit (perintah gcluster job submit), tetapkan flag XLA menggunakan --pathways-proxy-args:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

Dump HLO

Untuk memeriksa input High Level Optimizer (HLO) yang diberikan ke compiler XLA, Anda dapat mengonfigurasi Pathways untuk membuang HLO ke lokasi Cloud Storage yang ditentukan menggunakan --pathways-proxy-args:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

Langkah berikutnya