Pelatihan yang tangguh dengan Pathways

Halaman ini menjelaskan cara mengonfigurasi fitur ketahanan untuk workload Pathways di Google Kubernetes Engine (GKE) sehingga tugas pelatihan berskala besar dapat pulih dari gangguan terencana dan kegagalan hardware yang tidak direncanakan.

Pathways memberikan manfaat ketahanan dengan cara berikut:

  • Penangguhan-lanjutkan: Memungkinkan workload Anda menoleransi gangguan yang direncanakan, 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.
  • Multi-Tier Checkpointing: Memungkinkan workload Anda menyimpan checkpoint lokal dan yang direplikasi peer berfrekuensi tinggi di seluruh worker TPU menggunakan Python yang ditempatkan bersama, dengan cadangan Cloud Storage asinkron untuk mengurangi waktu pemulihan.
  • Hot-swap: Memungkinkan workload JobSet GKE dengan prioritas lebih tinggi mengambil alih slice TPU cadangan atau berprioritas lebih rendah dengan cepat saat terjadi kegagalan hardware.

Sebelum memulai

Sebelum memulai, selesaikan tugas-tugas berikut:

Jika Anda menggunakan pelatihan elastis, instal paket berikut dalam image container workload Anda:

  • pathwaysutils
  • Orbax (orbax-checkpoint): >= 0.12.6 (diperlukan untuk snapshot dalam memori dengan ocp.v1.training.pathways.Snapshotter)
  • Grain (grain): wajib jika Anda menggunakan grain.experimental.ElasticIterator

Jika Anda menggunakan Multi-Tier Checkpointing, selesaikan juga prasyarat berikut:

  • Konfigurasikan cluster GKE Anda untuk Multi-Tier Checkpointing dengan mengikuti Melatih model machine learning berskala besar di GKE dengan Multi-Tier Checkpointing:
    • Aktifkan add-on Multi-Tier Checkpointing dan Cloud Storage FUSE (GcsFuseCsiDriver), bersama dengan Workload Identity Federation for GKE, di cluster Anda.
    • Buat bucket Cloud Storage dengan namespace hierarkis diaktifkan.
    • Berikan peran Identity and Access Management (IAM) roles/storage.objectUser dan roles/storage.insightsCollectorService ke akun layanan Kubernetes Multi-Tier Checkpointing.
    • Deploy resource kustom CheckpointConfiguration yang menentukan bucket Cloud Storage, pemilih node TPU, dan ukuran volume dalam memori (inMemoryVolumeSize).
  • Bangun dan kirim image container sidecar Python yang ditempatkan bersama—misalnya, dengan menggunakan Dockerfile sidecar contoh. Dalam workload dan image file bantuan, gunakan versi jax dan jaxlib yang cocok yang sesuai dengan rilis Pathways Anda dan memenuhi persyaratan versi minimum berikut:
    • Versi Orbax: >= 0.12.4
    • Versi JAX (termasuk Python yang ditempatkan bersama): >= 0.10.1

Penangguhan-pelanjutan

GKE biasanya mengirimkan pemberitahuan preemption ke Pod akselerator sebelum melakukan preemption pada Pod. Pathways mengaktifkan toleransi preemption secara default di semua deployment cloud, dan tugas akselerator Pathways memproses pemberitahuan ini.

Saat pemberitahuan preemption tiba, Pathways akan menentukan terlebih dahulu apakah dapat menyimpan dan memulihkan workload secara transparan. Jika memungkinkan, 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 dipertahankan.

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

Workload ML standar 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, bergantung pada komponen Pathways stateful dan 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 Anda menetapkan minimum_slice_count ke None atau ke jumlah total slice, pelatihan akan dijeda saat slice gagal dan dilanjutkan setelah semua slice responsif kembali. Karena jumlah slice tidak berubah di seluruh upaya, workload Anda dapat menggunakan kembali mesh perangkat yang ada dan file yang dapat dieksekusi yang telah dikompilasi.
    • Pengubahan ukuran replika: Jika Anda menetapkan minimum_slice_count ke kurang dari jumlah total slice, pelatihan akan dilanjutkan pada slice aktif yang tersisa saat slice gagal 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 ulang 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/cloud-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. Sebaliknya, ElasticIterator melacak satu posisi set data global yang dapat Anda pulihkan 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.
  • Menerapkan @elastic_manager.elastic_retry ke run_training_loop akan mengonfigurasi fungsi untuk menunggu slice yang responsif, menetapkan jax.default_device ke slice aktif, menjaga thread monitor latar belakang tetap aktif, dan mencoba lagi pada jax.errors.JaxRuntimeError.
  • Pada awal upaya pemulihan, kode Anda harus membangun ulang mesh perangkat dari elastic_manager.active_devices dan memanggil snapshotter.load untuk memulihkan snapshot terbaru ke sharding mesh yang telah diupdate. Biarkan buffer perangkat dari upaya yang gagal keluar dari cakupan sehingga JAX dapat mengosongkan memorinya.
  • 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 berjalan di awal setiap percobaan setelah pengelola memperbarui active_slice_indices dan sebelum fungsi loop pelatihan yang digabungkan berjalan.
  • on_elastic_event_callback: Callback opsional yang berjalan setelah elastic_retry menangkap slice-down jax.errors.JaxRuntimeError, 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). Dalam fungsi yang didekorasi dengan @elastic_manager.elastic_retry, elastic_retry akan 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 akan memanggil elastic.wait_for_slices sekali untuk melakukan polling kondisi slice hingga setidaknya minimum_slice_count slice aktif. Kemudian, fungsi ini akan memperbarui 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.

Checkpointing Multi-Tingkat

Tugas pelatihan berskala besar sering mengalami gangguan, dan menulis atau memulihkan checkpoint langsung ke dan dari Cloud Storage dapat menghambat progres pelatihan. Checkpoint Multitier mengurangi overhead checkpoint dan waktu pemulihan dengan mengoordinasikan tiga tingkat penyimpanan di GKE:

  • Disk RAM lokal: Worker TPU menulis checkpoint frekuensi tinggi ke volume CSI dalam memori (multitier-checkpoint.csi.storage.gke.io) yang dipasang di setiap worker node. Volume ini tetap ada saat Pod dimulai ulang di node yang sama.
  • Replikasi peer dalam cluster: Layanan replikator GKE mereplikasi shard checkpoint lokal di seluruh worker TPU peer dalam cluster.
  • Cadangan Cloud Storage: Layanan replikator mencadangkan checkpoint lokal secara asinkron ke bucket Cloud Storage dengan namespace hierarkis diaktifkan, tanpa memblokir langkah-langkah pelatihan TPU.

Tidak seperti JAX multi-pengontrol, yang setiap host TPU-nya menjalankan skrip pelatihannya sendiri, Pathways menjalankan satu pengontrol di Pod pathways-head dan mengirimkan pekerjaan ke Pod pathways-worker. Untuk membaca dan menulis shard checkpoint langsung di RAM disk lokal setiap host TPU tanpa merutekan data checkpoint melalui satu pengontrol, Pathways menggunakan Python yang ditempatkan bersama. Setiap Pod worker menjalankan container colocated-python-sidecar bersama container pathways-worker, dan kedua container memasang volume CSI Multi-Tier Checkpointing dan direktori memori bersama (/tmp/sidecar). Selama inisialisasi dan pelatihan, Orbax mengirimkan penyiapan, penyimpanan, dan pemulihan checkpoint dari pengontrol Pathways ke sidecar worker di setiap host TPU. Multi-Tier Checkpointing di Pathways juga mendukung pelatihan elastis saat memperkecil atau memperbesar jumlah slice TPU aktif.

Pathways mendukung Checkpoint Multi-Tingkat dengan MaxText. Untuk petunjuk penyiapan dan konfigurasi, lihat panduan checkpointing multi-tingkat MaxText.

Mengonfigurasi workload Pathways Anda

Anda dapat mengonfigurasi workload Pathways untuk Multi-Tier Checkpointing menggunakan Cluster Toolkit (gcluster job submit) atau dengan mengonfigurasi manifes JobSet Pathways kustom.

Menggunakan Cluster Toolkit

Saat Anda mengirimkan tugas dengan gcluster job submit --pathways, teruskan flag --gke-mtc-enabled, --gke-mtc-ramdisk-dir, dan --pathways-colocated-python-sidecar-image:

gcluster job submit \
    --pathways \
    --gke-mtc-enabled \
    --gke-mtc-ramdisk-dir="/tmp/mtc_checkpoints" \
    --pathways-colocated-python-sidecar-image="SIDECAR_IMAGE" \
    ...

Ganti SIDECAR_IMAGE dengan URI image container sidecar Python yang ditempatkan bersama.

Saat Anda menetapkan flag ini, gcluster akan otomatis mengonfigurasi hal berikut:

  • Meneruskan --sidecar_name=external ke container pathways-proxy.
  • Menjalankan container colocated-python-sidecar di setiap Pod worker dan memasang direktori /tmp/sidecar bersama (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar) di pathways-worker dan colocated-python-sidecar.
  • Memasang volume CSI multitier-checkpoint.csi.storage.gke.io di jalur yang ditentukan oleh --gke-mtc-ramdisk-dir di pathways-worker dan colocated-python-sidecar pada Pod worker, tetapi tidak di Pod pathways-head.

Menggunakan manifes JobSet Pathways kustom

Jika Anda mengelola manifes JobSet Pathways sendiri, bukan menggunakan gcluster, teruskan --sidecar_name=external ke pathways-proxy. Dalam spesifikasi Pod worker, pasang volume memori sidecar bersama dan volume CSI multitier-checkpoint.csi.storage.gke.io pada container pathways-worker dan colocated-python-sidecar:

containers:
- name: pathways-worker
  args:
  - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
- name: colocated-python-sidecar
  volumeMounts:
  - name: shared-tmp
    mountPath: /tmp/sidecar
  - name: checkpoint-mount
    mountPath: /tmp/mtc_checkpoints
volumes:
- name: shared-tmp
  emptyDir:
    medium: Memory
- name: checkpoint-mount
  csi:
    driver: multitier-checkpoint.csi.storage.gke.io

Menggunakan Orbax Multi-Tier Checkpointing dalam workload JAX kustom

Dalam skrip pelatihan JAX kustom yang berjalan di Pathways, gunakan paket orbax.checkpoint.experimental.emergency.multi_tier_checkpointing dan tetapkan use_colocated_python=True di initialize_multi_tier_checkpointing dan ReplicatorCheckpointManagerOptions:

from etils import epath
import jax
import orbax.checkpoint as ocp
from orbax.checkpoint.experimental.emergency.multi_tier_checkpointing import (
    initialization,
    replicator_checkpoint_manager,
)
import pathwaysutils

pathwaysutils.initialize()

local_checkpoint_dir = epath.Path("/tmp/mtc_checkpoints")

# Initialize multi-tier checkpointing across worker sidecars.
# When using elastic training, pass the active slice devices to `devices`
# on initial setup and whenever the active slice count changes.
initialization.initialize_multi_tier_checkpointing(
    local_checkpoint_directory=local_checkpoint_dir,
    backup_interval_minutes=10,
    num_slices=num_slices,
    run_name=run_name,
    data_parallelism=data_parallelism,
    use_colocated_python=True,
    devices=jax.devices(),
)

options = replicator_checkpoint_manager.ReplicatorCheckpointManagerOptions(
    save_interval_steps=20,
    use_colocated_python=True,
)
checkpoint_manager = (
    replicator_checkpoint_manager.ReplicatorCheckpointManager(
        local_checkpoint_dir,
        options=options,
        global_mesh=mesh,
    )
)

# Restore from the latest checkpoint if one exists.
start_step = 1
latest_step = checkpoint_manager.latest_step()
if latest_step is not None:
  restore_args = jax.tree.map(
      lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
      state,
  )
  restored = checkpoint_manager.restore(
      latest_step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeRestore(
              item=state,
              restore_args=restore_args,
          )
      ),
  )
  state = restored.state
  start_step = latest_step + 1

# Step 0 is reserved by the replicator service to indicate "no checkpoint".
for step in range(start_step, final_step + 1):
  state = jitted_train_step(state)
  checkpoint_manager.save(
      step,
      args=ocp.args.Composite(
          state=ocp.args.PyTreeSave(item=state)
      ),
  )

checkpoint_manager.wait_until_finished()
checkpoint_manager.close()

Hot-swap

Hot-swap adalah 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 mempercepat pemulihan.

Saat Anda membuat JobSet, GKE menjadwalkan workload di beberapa slice, seperti yang ditentukan dalam konfigurasi JobSet. Jika terjadi kegagalan hardware pada satu atau beberapa slice, GKE akan menandai Pod yang terpengaruh sebagai gagal. Jika cluster GKE Anda memiliki slice cadangan yang menjalankan tugas dengan prioritas lebih rendah, pengontrol JobSet akan memetakan ulang workload slice yang gagal dari tugas dengan prioritas 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 dan simpan sebagai high-prior-job.yaml:

apiVersion: scheduling.k8s.io/v1
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 workload Pathways Anda. Jika Anda mengirimkan workload menggunakan Cluster Toolkit, teruskan --priority=high-prior-job ke perintah gcluster job submit, yang menetapkan priorityClassName pada Pod pathways-head dan worker. Jika Anda menulis manifes JobSet sendiri, tetapkan priorityClassName di template Pod tugas yang direplikasi worker, yang menjalankan container pathways-worker:

replicatedJobs:
- name: worker
  template:
    spec:
      template:
        spec:
          priorityClassName: high-prior-job

Jika Anda membuat cluster menggunakan blueprint TPU Cluster Toolkit dengan Pathways diaktifkan, cluster tersebut sudah menyertakan class prioritas very-low, low, medium, high, dan very-high, sehingga Anda dapat meneruskan salah satunya ke --priority, bukan membuat sendiri.

Langkah berikutnya