Google menggunakan teknologi AI untuk menerjemahkan konten ke dalam bahasa pilihan Anda. Terjemahan AI mungkin mengandung kesalahan.

Pelatihan yang tangguh dengan Pathways

Pathways memberikan manfaat ketahanan dengan cara berikut:

Suspend-Resume: memungkinkan workload Anda menoleransi gangguan terencana seperti pemberitahuan preemption tanpa memerlukan kode penanganan preemption kustom.
Pelatihan elastis: memungkinkan workload Anda menoleransi kegagalan hardware yang tidak direncanakan tanpa membuat klien error dengan menggunakan logika pemulihan dalam kode pelatihan Anda.

Sebelum memulai

Sebelum memulai, selesaikan tugas-tugas berikut:

Penangguhan-pelanjutan

Google Kubernetes Engine (GKE) biasanya mengirimkan pemberitahuan preemption ke Pod akselerator sebelum Pod di-preempt. Toleransi preemption Pathways diaktifkan secara default di semua deployment cloud dan tugas akselerator Pathways memproses pemberitahuan ini.

Saat pemberitahuan preemption tiba, Pathways akan menentukan terlebih dahulu apakah Pathways dapat menyimpan dan memulihkan workload secara transparan. Jika demikian, GKE akan mencoba menangguhkan workload ML Anda secara transparan dengan menulis statusnya saat ini ke penyimpanan persisten seperti Cloud Storage sebelum GKE mengeluarkan tugas akselerator Anda. Saat GKE menjadwalkan ulang tugas Anda nanti, Pathways akan melanjutkan workload ML Anda dengan membaca kembali statusnya yang dipertahankan.

Jika workload tidak dapat dipulihkan, Pathways akan menonaktifkan tugas akselerator dan meneruskan kegagalan ke tugas Anda jika pelatihan elastis dikonfigurasi. Jika pelatihan elastis tidak dikonfigurasi, GKE akan memulai ulang seluruh workload berdasarkan kebijakan mulai ulang JobSet.

Workload ML umum yang ditentukan menggunakan JAX mengandalkan komponen XLA Pathways stateless yang dapat dipulihkan menggunakan snapshot memori bandwidth tinggi (HBM). Workload ML tertentu seperti yang ditentukan menggunakan Python API yang ditempatkan bersama JAX mengandalkan komponen Pathways stateful; workload ini tidak dapat dipulihkan.

Pelatihan elastis

Pelatihan elastis memungkinkan tugas pelatihan Anda berlanjut saat terjadi kegagalan hardware atau preemption slice, tanpa memulai ulang seluruh workload. Dalam tugas pelatihan multi-slice, Pathways menjaga proses klien pengontrol tunggal dan slice TPU yang responsif tetap aktif saat slice individual tidak tersedia.

Pelatihan elastis menggabungkan kemampuan sistem Pathways dengan logika pemulihan dalam kode pelatihan Anda:

Deteksi kegagalan: Saat terjadi kegagalan hardware atau preemption pada slice TPU, Pathways akan membatalkan penetapan slice yang terpengaruh dan memunculkan jax.errors.JaxRuntimeError saat workload Anda mengakses data di slice tersebut pada waktu berikutnya. Karena proses klien tetap aktif, kode Anda dapat menangkap error ini, mengonfigurasi ulang resource, dan melanjutkan pelatihan.
Mode pemulihan: Class pathwaysutils.elastic.manager.Manager dan dekorator elastic_retry-nya mendukung dua strategi pemulihan:
- Jeda dan lanjutkan: Jika minimum_slice_count adalah None atau sama dengan jumlah total slice dalam tugas, pelatihan akan dijeda saat slice hilang dan menunggu hingga semua slice responsif kembali sebelum melanjutkan. Karena jumlah slice tidak berubah di seluruh percobaan, workload Anda dapat menggunakan kembali mesh perangkat dan file yang dapat dieksekusi yang telah dikompilasi.
- Pengubahan ukuran replika: Jika minimum_slice_count kurang dari jumlah total slice, pelatihan akan segera dilanjutkan pada slice aktif yang tersisa dengan membangun ulang mesh perangkat dan mengompilasi ulang untuk topologi yang lebih kecil. Thread pemantauan latar belakang secara berkala memeriksa slice yang tidak aktif dan mencatat slice yang baru responsif di elastic_manager.available_inactive_slices. Loop pelatihan Anda kemudian dapat membangun kembali mesh perangkat dengan slice pengganti, melakukan inisialisasi ulang untuk topologi yang diperluas, dan menempatkan status saat ini ke mesh baru.
Pemulihan status: Untuk melanjutkan pelatihan setelah peristiwa elastis, workload Anda memulihkan parameter model, status pengoptimal, dan status iterator data menggunakan snapshot dalam memori. Class Snapshotter di orbax.checkpoint menyalin status yang direplikasi secara asinkron dari HBM TPU ke memori CPU pinned_host di setiap slice. Selama pemulihan, Snapshotter merekonstruksi status dari replika host yang bertahan dan melakukan sharding ulang secara langsung ke mesh TPU aktif tanpa membaca dari penyimpanan jarak jauh.

Mengaktifkan elastisitas di server proxy Pathways

Untuk mengaktifkan set perangkat elastis di Pathways di GKE, tambahkan flag --num_elastic_instances ke kolom args dari container pathways-proxy dalam manifes JobSet Anda dan tentukan jumlah slice TPU elastis dalam workload Anda:

- name: pathways-proxy
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
  args:
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --server_port=29000
  - --gcs_scratch_location=gs://BUCKET_NAME
  - --num_elastic_instances=NUM_SLICES

Ganti kode berikut:

BUCKET_NAME: bucket Cloud Storage yang digunakan untuk file sementara Pathways.
NUM_SLICES: jumlah slice TPU dalam workload Anda.

Memfilter perangkat aktif dan membangun mesh perangkat

Dalam workload Pathways, jax.devices() menampilkan semua perangkat dalam alokasi awal selama masa aktif proses klien. Saat Anda menggunakan pelatihan elastis, gunakan elastic_manager.active_devices agar mesh perangkat dan ukuran batch Anda hanya menggunakan slice yang aktif dan responsif:

import jax
import numpy as np
from pathwaysutils.elastic import manager

elastic_manager = manager.Manager()


def create_mesh(
    elastic_manager: manager.Manager, devices_per_slice: int
) -> jax.sharding.Mesh:
  device_array = np.asarray(elastic_manager.active_devices).reshape(
      (elastic_manager.active_slice_count, devices_per_slice)
  )
  return jax.sharding.Mesh(device_array, axis_names=("replica", "model"))

Saat jumlah slice aktif berubah dalam mode pengubahan ukuran replika, gunakan pemuat data elastis seperti grain.experimental.ElasticIterator. Iterator data per host standar menyimpan status per host dan tidak dapat dipulihkan ke jumlah host yang berbeda, sedangkan ElasticIterator melacak satu posisi set data global yang dapat dipulihkan di sejumlah slice aktif.

Loop pelatihan elastis

Loop pelatihan elastis mengoordinasikan pemantauan slice, pemulihan kegagalan, dan snapshot dalam memori, seperti yang ditunjukkan dalam contoh kode berikut:

Fungsi train Anda membuat satu instance manager.Manager() dan Snapshotter Orbax di luar wrapper percobaan ulang sehingga pelacakan slice dan snapshot memori host tetap ada di seluruh upaya.
Mendekorasi fungsi run_training_loop dalam dengan @elastic_manager.elastic_retry akan menunggu slice yang responsif, menetapkan jax.default_device ke slice aktif, menjaga thread pemantauan latar belakang tetap aktif di seluruh langkah, dan mencoba ulang loop saat jax.errors.JaxRuntimeError terjadi.
Saat memulihkan dari kegagalan slice di awal percobaan, kode Anda harus membangun ulang mesh perangkat dari elastic_manager.active_devices dan memanggil snapshotter.load untuk memulihkan snapshot terbaru ke sharding mesh yang diupdate, sekaligus membiarkan buffer perangkat dari percobaan yang gagal keluar dari cakupan.
Selama setiap langkah loop, kode Anda secara berkala memanggil snapshotter.save untuk mencadangkan status pelatihan, dan memeriksa elastic_manager.available_inactive_slices untuk mendeteksi slice pengganti yang responsif. Saat slice pengganti tersedia, kode Anda akan menambahkannya ke elastic_manager.active_slice_indices, membangun ulang mesh, melakukan inisialisasi ulang loop pelatihan, dan menempatkan status ke mesh yang diperluas.

Contoh berikut menunjukkan cara menerapkan loop pelatihan elastis dengan menggabungkan manager.Manager, elastic_retry, dan Snapshotter dalam memori Orbax:

import jax
import orbax.checkpoint as ocp
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()


def train(
    final_step: int,
    snapshot_interval: int,  # in steps
    devices_per_slice: int,
):
  elastic_manager = manager.Manager()
  snapshotter = ocp.v1.training.pathways.Snapshotter(replica_axis_index=0)

  @elastic_manager.elastic_retry(
      minimum_slice_count=1,
      poll_interval=10,  # in seconds
      timeout=300,  # in seconds
      retry_policy=manager.ElasticRetryLimit(max_attempts=10),
  )
  def run_training_loop():
    mesh = create_mesh(elastic_manager, devices_per_slice)
    state, abstract_state, jitted_train_step, data_iterator = (
        initialize_training_loop(mesh)
    )
    step = 0

    # Restore from the latest in-memory snapshot after a slice failure.
    if snapshotter.latest is not None:
      snapshotter.join()
      restored = snapshotter.load({
          "state": abstract_state,
          "step": step,
          "iterator_state": data_iterator.get_state(),
      })
      state = restored["state"]
      step = int(restored["step"]) + 1
      data_iterator.set_state(restored["iterator_state"])

    while step < final_step:
      batch = next(data_iterator)
      state = jitted_train_step(state, batch)

      if step % snapshot_interval == 0:
        snapshotter.save(
            step,
            {
                "state": state,
                "step": step,
                "iterator_state": data_iterator.get_state(),
            },
        )

      # If replacement slices are healthy, reinitialize and place state onto the new mesh.
      if elastic_manager.available_inactive_slices:
        elastic_manager.active_slice_indices |= (
            elastic_manager.available_inactive_slices
        )
        mesh = create_mesh(elastic_manager, devices_per_slice)
        iterator_state = data_iterator.get_state()
        _, abstract_state, jitted_train_step, data_iterator = (
            initialize_training_loop(mesh)
        )
        data_iterator.set_state(iterator_state)
        state_shardings = jax.tree.map(lambda x: x.sharding, abstract_state)
        state = jax.device_put(state, state_shardings)

      step += 1

    return state

  state = run_training_loop()
  snapshotter.join()
  return state

Mengonfigurasi pengelola elastis dan perilaku percobaan ulang

Anda membuat instance manager.Manager(devices=None) satu kali per proses. Jika Anda menghilangkan devices, pengelola akan menggunakan jax.devices() dan mengelompokkan perangkat berdasarkan atribut slice_index.

Instance Manager mengekspos properti berikut untuk memeriksa topologi cluster selama pelatihan:

slice_to_devices: Pemetaan dari setiap slice_index ke urutan objek jax.Device yang termasuk dalam slice tersebut.
all_slice_indices: frozenset dari semua indeks slice dalam alokasi.
active_slice_indices: frozenset indeks slice yang berpartisipasi dalam mesh komputasi aktif.
active_devices: Daftar objek jax.Device yang termasuk dalam slice aktif, diurutkan berdasarkan slice_index.
inactive_slice_indices: frozenset indeks slice dalam alokasi yang tidak berpartisipasi dalam mesh komputasi.
available_inactive_slices: frozenset indeks slice tidak aktif yang telah dideteksi oleh thread monitor latar belakang sebagai responsif dan siap untuk bergabung kembali.
total_slice_count dan active_slice_count: Jumlah total slice yang dialokasikan dan jumlah slice aktif.
default_device: Perangkat 0 dari slice aktif, yang otomatis ditetapkan elastic_retry sebagai perangkat JAX default untuk setiap upaya.

Anda dapat menyesuaikan perilaku percobaan ulang dan penskalaan dengan meneruskan argumen berikut ke @elastic_manager.elastic_retry:

minimum_slice_count: Jumlah minimum slice yang sehat yang diperlukan sebelum memulai eksekusi awal atau upaya percobaan ulang. Jika ditetapkan ke None (default), elastic_retry akan menunggu semua slice dalam alokasi (total_slice_count), yang beroperasi dalam mode jeda dan lanjutkan. Jika ditetapkan ke bilangan bulat yang lebih kecil (seperti 1), elastic_retry akan beroperasi dalam mode pengubahan ukuran replika.
poll_interval: Interval dalam detik antara health check slice yang dilakukan oleh pengelola pelatihan elastis (Manager) selama pelatihan. Setelan defaultnya adalah 10.
timeout: Jumlah detik maksimum untuk menunggu slice minimum_slice_count menjadi aktif sebelum memunculkan TimeoutError. Defaultnya adalah None (tanpa waktu tunggu).
pre_callback: Callback opsional yang dijalankan di awal setiap percobaan setelah active_slice_indices diperbarui dan sebelum fungsi loop pelatihan yang digabungkan dijalankan.
on_elastic_event_callback: Callback opsional yang dijalankan setelah slice-down jax.errors.JaxRuntimeError tertangkap, sebelum upaya percobaan ulang berikutnya dimulai. Anda dapat menggunakan callback ini untuk membersihkan checkpoint yang tidak lengkap atau mereset iterator data.
retry_policy: (attempt: int, error: Exception) -> bool yang dapat dipanggil yang menentukan apakah akan melanjutkan percobaan ulang setelah peristiwa elastis. Gunakan manager.ElasticRetryLimit(max_attempts) untuk membatasi jumlah total upaya. Jika dihilangkan, elastic_retry akan mencoba ulang tanpa batas.

Pada peristiwa elastis, JAX mengosongkan buffer perangkat dari upaya yang gagal setelah frame stack dilepaskan dan tidak ada variabel yang mereferensikan array tersebut. Hindari menyimpan referensi ke array perangkat dari upaya yang gagal di dalam variabel global atau callback yang berjalan lama agar memori dilepaskan sebelum upaya berikutnya dimulai.

Snapshot dalam memori dengan Orbax

Class ocp.v1.training.pathways.Snapshotter (tersedia di orbax-checkpoint versi 0.12.6 dan yang lebih baru) mempertahankan cadangan dalam memori dari status pelatihan Anda di memori pinned_host di seluruh slice:

Snapshotter(replica_axis_index=0): Melakukan inisialisasi snapshotter dan menentukan sumbu mesh yang mewakili replika slice paralel data. Semua array JAX dalam status tersimpan harus direplikasi di sepanjang sumbu mesh ini sehingga slice yang bertahan mempertahankan salinan lengkap status model dan pengoptimal saat slice gagal. Jika status Anda berisi array perangkat tunggal, replikasi array tersebut di seluruh mesh sebelum memanggil save.
save(step, state): Mentransfer state ke memori pinned_host secara asinkron dan membaginya menjadi hierarki per replika di sepanjang replica_axis_index. Jika transfer snapshot sebelumnya masih berlangsung, save akan melewati permintaan baru untuk menghindari penghentian loop pelatihan.
load(abstract_state, reset_snapshot_state=True): Memeriksa replika host yang disimpan untuk mengidentifikasi slice yang bertahan, menggabungkan replika yang sehat, dan melakukan sharding ulang status ke sharding target yang ditentukan oleh abstract_state.
join(): Memblokir hingga operasi snapshot yang diantrekan selesai.
latest: Menampilkan metadata (termasuk step) untuk snapshot yang paling baru diselesaikan, atau None jika tidak ada snapshot yang selesai.

Mengurangi sharding pada kegagalan slice

Jika slice TPU gagal atau di-preempt selama eksekusi, Pathways akan memunculkan jax.errors.JaxRuntimeError (seperti DATA_LOSS, DEADLINE_EXCEEDED, NOT_FOUND, atau INTERNAL). Dekorator elastic_retry menangkap pengecualian ini dan memeriksa elastic.is_error_due_to_slice_down(error):

Jika error tidak disebabkan oleh peristiwa slice-down, elastic_retry akan segera memunculkan kembali pengecualian.
Jika error disebabkan oleh peristiwa slice-down, elastic_retry akan menghentikan trace profiler JAX yang aktif, menjalankan on_elastic_event_callback, dan mengevaluasi retry_policy(attempt, error).
Sebelum memulai upaya berikutnya, elastic_retry memanggil elastic.wait_for_slices satu kali, yang melakukan polling kondisi slice hingga setidaknya minimum_slice_count slice aktif, mengupdate elastic_manager.active_slice_indices, dan menjalankan fungsi loop pelatihan Anda pada slice yang bertahan.

Meningkatkan sharding saat slice bergabung kembali

Saat workload Anda berjalan pada jumlah slice yang dikurangi dalam mode pengubahan ukuran replika, thread monitor latar belakang secara berkala menjalankan health check pada elastic_manager.inactive_slice_indices. Saat slice yang diperbaiki atau dijadwalkan ulang lulus health check, thread monitor akan menambahkan indeks slice-nya ke elastic_manager.available_inactive_slices.

Jika elastic_manager.available_inactive_slices tidak kosong di akhir langkah pelatihan, workload akan diskalakan kembali tanpa kehilangan progres pelatihan:

Menambahkan indeks slice di elastic_manager.available_inactive_slices ke elastic_manager.active_slice_indices akan menggabungkan slice yang digabungkan kembali ke elastic_manager.active_devices dan menghapusnya dari available_inactive_slices.
Membangun ulang mesh perangkat dari elastic_manager.active_devices serta menginisialisasi ulang langkah pelatihan dan iterator data akan menyesuaikan komputasi dengan topologi yang diperluas.
Memanggil jax.device_put akan menempatkan status pelatihan saat ini ke sharding mesh baru tanpa memulihkan dari snapshot.

Hot-swap

Hot-swap mengacu pada fitur GKE JobSet API yang memungkinkan tugas berprioritas lebih tinggi mengambil alih resource dengan cepat dari tugas berprioritas lebih rendah, sehingga meminimalkan periode nonaktif dan memastikan pemulihan yang lebih cepat.

Saat JobSet dibuat, GKE menjadwalkan workload di beberapa slice, seperti yang ditentukan dalam konfigurasi JobSet. Jika terjadi kegagalan hardware pada satu atau beberapa slice, Pod yang terpengaruh akan ditandai sebagai gagal. Jika cluster GKE Anda memiliki slice cadangan yang menjalankan tugas berprioritas lebih rendah, sistem JobSet akan memetakan ulang workload slice yang gagal dari tugas berprioritas lebih tinggi ke slice cadangan tersebut. Pemetaan ulang ini biasanya memerlukan waktu kurang dari satu menit.

Saat JobSet dimulai ulang, hot-swap dapat terjadi dalam situasi berikut:

Mode default: Jika slice TPU cadangan yang tidak ada aktivitas tersedia dalam cluster yang sama, scheduler Kubernetes akan memprioritaskan penjadwalan tugas yang dimulai ulang ke slice ini, daripada menunggu slice yang gagal diperbaiki. Perilaku penjadwalan ini memberikan pemulihan yang lebih cepat.
Workload heterogen: Di cluster yang menjalankan beberapa workload dengan resource PriorityClass Kubernetes yang dikonfigurasi, JobSet yang dimulai ulang dapat memicu hot-swap. Jika afinitas tugas yang dimulai ulang cocok dengan resource tugas berprioritas lebih rendah, Kubernetes akan mendahului tugas berprioritas lebih rendah tersebut sehingga tugas berprioritas lebih tinggi dapat dimulai tanpa menunggu perbaikan slice. Misalnya, Anda dapat mengonfigurasi Pod worker Pathways dengan prioritas yang berbeda menggunakan resource PriorityClass.

Untuk menggunakan prioritas di cluster Anda, tentukan manifes PriorityClass:

kind: PriorityClass
metadata:
  name: high-prior-job
value: 2000
globalDefault: false
description: "This priority class should be used for high priority job."

Terapkan YAML ini ke cluster GKE Anda:

kubectl apply -f high-prior-job.yaml

Selanjutnya, lampirkan class prioritas baru ke tugas worker Pathways dengan menambahkan teks berikut ke podspec Pod pathways-worker Anda.

priorityClassName: high-prior-job

Pelatihan yang tangguh dengan Pathways Tetap teratur dengan koleksi Simpan dan kategorikan konten berdasarkan preferensi Anda.