Pelatihan yang tangguh dengan Pathways

Pathways memberikan manfaat ketahanan dengan cara berikut:

  • Menangguhkan-Melanjutkan: toleransi dalam menghadapi gangguan terencana seperti pemberitahuan preemption tanpa mengharuskan pengguna menulis kode penanganan preemption kustom.
  • Pelatihan Elastis: toleransi dalam menghadapi kegagalan hardware yang tidak direncanakan tanpa menyebabkan klien mengalami error, tetapi mengharuskan pengguna menulis kode pemulihan khusus model.

Sebelum memulai

Pastikan Anda memiliki:

Menangguhkan-melanjutkan

Biasanya, GKE mengirimkan pemberitahuan preemption ke pod akselerator, sebelum pod tersebut di-preempt. Toleransi preemption Pathways diaktifkan secara default pada semua deployment cloud dan tugas akselerator Pathways memproses pemberitahuan ini.

Saat pemberitahuan preemption tiba, Pathways pertama-tama akan menentukan apakah workload saat ini dapat dipulihkan—apakah Pathways dapat menyimpan dan memulihkan workload secara transparan. Jika ya, Pathways 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 persisten.

Jika workload tidak dapat dipulihkan, Pathways akan menghentikan 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 memulai ulang JobSet.

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

Pelatihan elastis

Pelatihan elastis memungkinkan tugas pelatihan Anda terus berjalan meskipun terjadi kegagalan hardware. Hal ini dicapai melalui kombinasi kemampuan sistem Pathways dan logika pemulihan model yang ditentukan pengguna:

  • Deteksi kegagalan: Jika terjadi kegagalan hardware (misalnya, pekerja TPU mengalami error), sistem Pathways akan mendeteksinya dan memberi tahu tugas pelatihan pengguna melalui pengecualian saat data yang berada di hardware tersebut diakses. Pemberitahuan ini tidak menyebabkan workload Anda mengalami error; pemberitahuan ini memungkinkan kode Anda menangani pemberitahuan dan mengonfigurasi ulang resource Anda untuk melanjutkan pemrosesan atau keluar dengan baik.
  • Penangan elastisitas yang ditentukan pengguna: Kode model pengguna harus dapat menangani pengecualian ini. Hal inilah yang menjadikannya "pemulihan khusus model".
    • Snapshot: Pendekatan yang paling umum adalah menyimpan snapshot status model Anda secara berkala. Jika terjadi kegagalan, Anda dapat memuat dari snapshot terbaru untuk melanjutkan pelatihan.
    • Konfigurasi ulang: Anda mungkin perlu mengonfigurasi ulang tugas pelatihan untuk menyesuaikan jumlah slice yang tersedia. Misalnya, jika satu slice berhenti berfungsi, Anda dapat mengurangi jumlah slice aktif sebanyak satu hingga penggantian tersedia. Untuk mengetahui informasi selengkapnya, lihat Penangan Elastis.
    • Update grafik Data/Komputasi: Kode Anda harus menangani perubahan jumlah perangkat yang tersedia untuk komputasi dengan membuat ulang grafik komputasi sesuai kebutuhan. Hal ini mungkin melibatkan pemartisian ulang data atau mengompilasi ulang model Anda.
  • Peran Pathways dalam pemulihan: Pathways menyediakan primitif untuk mendukung konfigurasi ulang yang ditentukan pengguna:
    • Penggantian slice: Jika slice yang gagal diganti, klien dapat diberi tahu setelah slice baru tersedia. Kode Anda kemudian dapat dikonfigurasi ulang untuk menggunakan slice baru ini.
    • Pemulihan transparan: Pathways menangani detail pemulihan tingkat yang lebih rendah, seperti membangun kembali koneksi ke bagian cluster yang sehat.
  • Utilitas di pathwaysutils: Kumpulan utilitas Pathways yang ditentukan di pathways-utils.

Mengimplementasikan penangan elastis

Sebagian besar kode yang harus Anda tulis akan berada di penangan elastis yang ditentukan pengguna. Penangan ini bereaksi terhadap peristiwa elastis (seperti slice TPU yang tidak tersedia) dengan membuat ulang mesh dan menginisialisasi ulang loop pelatihan.

Setiap workload bersifat unik. Kompleksitas penangan elastis dapat diskalakan dengan kompleksitas workload. Input dan output penangan harus berupa argumen dan nilai yang ditampilkan minimum yang diperlukan untuk menginisialisasi ulang loop pelatihan.

def elastic_handler(elastic_utils, *args, **kwargs):
  mesh = initialize_mesh(**kwargs["mesh_kwargs"])
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **kwargs["initialize_training_loop_kwargs"])

  step, snapshot = elastic_utils.get_next_snapshot()
  state = initial_state.replace(**snapshot)

  return state, step, mesh, jitted_train_step, other_variables

Mengupdate loop pelatihan

Anda harus melakukan perubahan berikut pada loop pelatihan:

  1. Membuat pengelola elastis
  2. Menggabungkan loop pelatihan Anda di dalam blok try-except yang menangani jax.errors.JaxRuntimeError
  3. Dalam penangan jax.errors.JaxRuntimeError, panggil maybe_reshard_down. Pengelola elastis akan melakukan resharding jika error terkait dengan peristiwa elastis atau jika tidak, akan memunculkan kembali error tersebut.
  4. Panggil maybe_snapshot dan maybe_reshard_up di akhir loop pelatihan
import pathwaysutils
from pathwaysutils.elastic import manager

pathwaysutils.initialize()

def initialize_mesh(**kwargs):
  ...


def initialize_training_loop(**kwargs):
  ...


def train_loop(
    final_step,
    elastic_manager,
    mesh_kwargs,
    initialize_training_loop_kwargs,
):
  mesh = initialize_mesh(**mesh_kwargs)
  initial_state, initial_step, jitted_train_step, other_variables =
      initialize_training_loop(mesh, **initialize_training_loop_kwargs)

  step = initial_step
  while step < final_step:
    try:
      state = jitted_train_step(state)

      elastic_manager.maybe_snapshot(step=step, snapshot=state)
      handler_returns = elastic_manager.maybe_reshard_up(
          step=step,
          snapshot=state,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns
      step += 1
    except jax.errors.JaxRuntimeError as error:
      handler_returns = elastic_manager.maybe_reshard_down(
          error=error,
          elastic_handler=elastic_handler,
          handler_args=(),
          handler_kwargs=dict(
              mesh_kwargs=mesh_kwargs,
              initialize_training_loop_kwargs=initialize_training_loop_kwargs,
          ),
      )
      if handler_returns:
        state, step, mesh, jitted_train_step, other_variables = handler_returns

  return state


def main():
  elastic_manager = manager.Manager(
      devices=jax.devices(),
      snapshot_period=10,
      snapshot_buffer_size=1,
      reshard_check_period=5,
      max_elastic_down_event_count=10,
      max_reshard_retry_count=3,
  )

  train_loop(100, elastic_manager, {}, {})

Mengonfigurasi pengelola elastis

Pengelola elastis dapat dikonfigurasi dengan beberapa cara yang berbeda. Frekuensi snapshot ditentukan oleh periode snapshot. Periode snapshot memengaruhi jumlah rata-rata langkah yang hilang karena peristiwa elastis. Periode pemeriksaan resharding menentukan frekuensi loop pelatihan Anda akan melakukan polling untuk ketersediaan slice. max_elastic_down_event_count memungkinkan Anda menetapkan jumlah peristiwa elastis karena kehilangan slice yang akan didukung oleh loop pelatihan Anda. max_reshard_retry_count menentukan jumlah percobaan ulang resharding yang harus dilakukan oleh pengelola elastis. Pengelola adalah objek singleton dan hanya boleh dibuat satu kali.

Snapshot

Berdasarkan konfigurasi pengelola elastis, fungsi ini dapat mengambil snapshot data ke memori host yang akan tersedia untuk digunakan oleh penangan elastis Anda selama peristiwa elastis.

Mengurangi sharding

Setelah menangkap jax.errors.JaxRuntimeError, Pathways akan memeriksa apakah error tersebut disebabkan oleh peristiwa elastis karena slice yang hilang. Jika ya, Pathways akan memanggil penangan elastis dalam loop hingga berhasil atau percobaan ulang maksimum. Jika error tidak disebabkan oleh peristiwa elastis, error akan muncul lagi. Nilai yang ditampilkan dari penangan elastis diteruskan ke pemanggil.

Meningkatkan sharding

Berdasarkan konfigurasi pengelola elastis dan jika ada slice yang tidak tersedia, Pathways akan memeriksa apakah slice tambahan telah tersedia. Jika ya, Pathways akan segera menyimpan snapshot (jika snapshot yang sudah ada untuk langkah saat ini belum diambil) dan memanggil penangan elastis dalam loop hingga berhasil atau jumlah percobaan ulang maksimum tercapai. Jika resharding terjadi, nilai yang ditampilkan dari penangan elastis diteruskan ke pemanggil. Jika tidak, None akan ditampilkan.

Hot-swap

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

Saat JobSet dibuat, GKE akan 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. Saat menjadwalkan ulang Jobset ini, jika Anda memilih untuk menyimpan slice cadangan di cluster GKE yang dapat digunakan untuk Tugas dengan prioritas lebih rendah, sistem JobSet akan memetakan ulang workload slice yang gagal dari tugas dengan prioritas lebih tinggi ke slice cadangan yang digunakan oleh tugas dengan prioritas lebih rendah dalam cluster GKE yang sama. Pemetaan ulang ini biasanya memerlukan waktu kurang dari satu menit.

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

  1. Mode Default: Jika slice TPU cadangan yang tidak ada aktivitas tersedia dalam cluster yang sama, penjadwal Kubernetes akan memprioritaskan penjadwalan tugas yang dimulai ulang ke slice ini, bukan menunggu slice yang gagal diperbaiki. Hal ini memberikan pemulihan yang lebih cepat.
  2. Workload Heterogen: Di cluster yang menjalankan beberapa workload dengan PriorityClass Kubernetes yang dikonfigurasi, JobSet yang dimulai ulang dapat memicu hot swap. Jika afinitas tugas yang dimulai ulang cocok dengan resource tugas dengan prioritas lebih rendah, Kubernetes akan meng-preempt tugas dengan prioritas lebih rendah, sehingga tugas dengan prioritas lebih tinggi dapat segera dimulai. Misalnya, Anda dapat mengonfigurasi pod pekerja Pathways dengan prioritas yang berbeda menggunakan PriorityClass.

Untuk menggunakan prioritas di cluster Anda, tentukan class prioritas, misalnya:

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 pekerja Pathways Anda dengan menambahkan teks berikut ke podspec Pod pathways-worker Anda.

priorityClassName: high-prior-job

Langkah berikutnya