透過 Pathways 進行彈性訓練

本頁說明如何在 Google Kubernetes Engine (GKE) 設定 Pathways 工作負載的韌性功能，讓大規模訓練工作能從預定中斷和非預期硬體故障中復原。

Pathways 透過下列方式提供復原力優勢：

• 暫停與恢復：讓工作負載容許預定的中斷，例如搶占通知，而不需要自訂搶占處理程式碼。
• 彈性訓練：在訓練程式碼中使用復原邏輯，讓工作負載容許非預期的硬體故障，而不會導致用戶端當機。
• 多層級查核點：工作負載可使用共置的 Python，在 TPU 工作站之間儲存高頻率的本機和對等複製查核點，並透過非同步 Cloud Storage 備份縮短復原時間。
• 熱交換：發生硬體故障時，優先順序較高的 GKE JobSet 工作負載可快速接管備用或優先順序較低的 TPU 磁碟區。

事前準備

開始前，請先完成下列工作：

• 設定 Cluster Toolkit
• 安裝 Kubernetes 工具
• 安裝 Google Cloud CLI
• 啟用 Kubernetes Engine API

如果您使用彈性訓練，請在工作負載容器映像檔中安裝下列套件：

• pathwaysutils
• Orbax (orbax-checkpoint)：>= 0.12.6 (使用 ocp.v1.training.pathways.Snapshotter 建立記憶體內快照時必須提供)
• Grain (grain)：如果使用 grain.experimental.ElasticIterator，則為必填

如果使用多層檢查點，也請完成下列必要條件：

• 按照「在 GKE 上使用多層檢查點訓練大型機器學習模型」一文，為 GKE 叢集設定多層檢查點：
  • 在叢集上啟用多層檢查點和 Cloud Storage FUSE (GcsFuseCsiDriver) 外掛程式，以及 Workload Identity Federation for GKE。
  • 建立已啟用階層命名空間的 Cloud Storage bucket。
  • 將 roles/storage.objectUser 和 roles/storage.insightsCollectorService Identity and Access Management (IAM) 角色授予多層檢查點 Kubernetes 服務帳戶。
  • 部署 CheckpointConfiguration 自訂資源，指定 Cloud Storage bucket、TPU 節點選取器和記憶體內磁碟區大小 (inMemoryVolumeSize)。
• 建構並推送共置的 Python 邊車容器映像檔，例如使用範例邊車 Dockerfile。在工作負載和 sidecar 映像檔中，使用與 Pathways 版本一致的 jax 和 jaxlib 版本，並符合下列最低版本需求：
  • Orbax 版本：>= 0.12.4
  • JAX 版本 (包括共置的 Python)：>= 0.10.1

暫停/繼續

GKE 通常會在先占加速器 Pod 之前，向該 Pod 傳送先占通知。Pathways 預設會對所有 Cloud Deploy 部署啟用先占容許，而 Pathways 加速器工作會監聽這些通知。

收到先占通知時，Pathways 會先判斷是否能以透明方式儲存及還原工作負載。如果可以，Pathways 會嘗試將目前的狀態寫入永久儲存空間 (例如 Cloud Storage)，以透明的方式暫停機器學習工作負載，然後 GKE 才會淘汰加速器工作。當 GKE 稍後重新排定工作時，Pathways 會讀回其永久狀態，繼續執行機器學習工作負載。

如果工作負載無法還原，Pathways 會關閉加速器工作，並在設定彈性訓練時，將失敗訊息轉送至工作負載。如果未設定彈性訓練，GKE 會根據 JobSet 重新啟動政策，重新啟動整個工作負載。

使用 JAX 定義的一般機器學習工作負載，會依賴無狀態 Pathways XLA 元件，這些元件可使用高頻寬記憶體 (HBM) 快照還原。某些機器學習工作負載 (例如使用 JAX 共置 Python API 定義的工作負載) 依賴有狀態的 Pathways 元件，因此無法還原。

彈性訓練

彈性訓練可讓訓練工作在硬體故障或切片先占時繼續執行，不必重新啟動整個工作負載。在多配量訓練工作中，當個別配量無法使用時，Pathways 會保留單一控制器用戶端程序和健康的 TPU 配量。

彈性訓練結合 Pathways 系統功能與訓練程式碼中的復原邏輯：

• 故障偵測：當 TPU 切片發生硬體故障或先占時，Pathways 會取消指派受影響的切片，並在工作負載下次存取該切片上的資料時，引發 jax.errors.JaxRuntimeError。由於用戶端程序仍會保持運作，因此程式碼可以擷取這項錯誤、重新設定資源，並繼續訓練。
• 復原模式：pathwaysutils.elastic.manager.Manager 類別及其 elastic_retry 裝飾器支援兩種復原策略：
  • 暫停和繼續：如果將 minimum_slice_count 設為 None 或總切片數量，訓練會在切片失敗時暫停，並在所有切片恢復正常後繼續。由於切片數量不會因嘗試次數而改變，工作負載可以重複使用現有的裝置網格和編譯的可執行檔。
  • 副本大小調整：如果將 minimum_slice_count 設為小於總切片數，當切片失敗時，系統會重建裝置網格，並針對較小的拓撲重新編譯，繼續在剩餘的使用中切片上訓練。背景監控執行緒會定期檢查閒置切片，並在 elastic_manager.available_inactive_slices 中記錄新恢復健康的切片。接著，訓練迴圈可以使用替換切片重建裝置網格、重新初始化擴充的拓撲，並將目前狀態放置到新的網格上。
• 狀態復原：如要在彈性事件發生後繼續訓練，工作負載會使用記憶體內快照，還原模型參數、最佳化工具狀態和資料疊代器狀態。orbax.checkpoint 中的 Snapshotter 類別會以非同步方式，將複製的狀態從 TPU HBM 複製到每個切片的 pinned_host CPU 記憶體。在復原期間，Snapshotter 會從存留的主機副本重建狀態，並直接重新分 shard 至運作中的 TPU 網格，不需從遠端儲存空間讀取。

啟用 Pathways Proxy 伺服器的彈性

如要在 GKE 的 Pathways 啟用彈性裝置集，請在 JobSet 資訊清單的 pathways-proxy 容器 args 欄位中新增 --num_elastic_instances 旗標，並指定工作負載中的彈性 TPU slice 數量：

- 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

更改下列內容：

• BUCKET_NAME：用於 Pathways 暫存檔案的 Cloud Storage bucket。
• NUM_SLICES：工作負載中的 TPU 磁碟區數量。

篩選使用中的裝置並建立裝置網格

在 Pathways 工作負載中，jax.devices() 會在用戶端程序生命週期內，傳回初始分配中的所有裝置。使用彈性訓練時，請使用 elastic_manager.active_devices，確保裝置網格和批次大小只使用健康、活躍的切片：

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"))

在副本大小調整模式中，如果運作中的切片數量有所變更，請使用彈性資料載入器，例如 grain.experimental.ElasticIterator。標準每主機資料疊代器會儲存每主機的狀態，且無法還原至不同數量的主機。相較之下，ElasticIterator 會追蹤單一全域資料集位置，您可以在任意數量的有效切片中還原該位置。

彈性訓練迴圈

彈性訓練迴圈會協調切片監控、故障復原和記憶體內快照，如以下程式碼範例所示：

• train 函式會在重試包裝函式外，建立單一 manager.Manager() 例項和 Orbax Snapshotter，因此切片追蹤和主機記憶體快照會在多次嘗試中保留。
• 將 @elastic_manager.elastic_retry 套用至 run_training_loop，可將函式設為等待健康切片、將 jax.default_device 設為有效切片、保持背景監控執行緒有效，並在 jax.errors.JaxRuntimeError 上重試。
• 開始嘗試復原時，程式碼必須從 elastic_manager.active_devices 重建裝置網格，並呼叫 snapshotter.load，將最新快照還原至更新後的網格分區。讓失敗嘗試的裝置緩衝區超出範圍，以便 JAX 釋出記憶體。
• 在迴圈的每個步驟中，程式碼會定期呼叫 snapshotter.save 來備份訓練狀態，並檢查 elastic_manager.available_inactive_slices 來偵測健康狀態良好的替換 slice。當有可用的替換切片時，程式碼會將其新增至 elastic_manager.active_slice_indices、重建網格、重新初始化訓練迴圈，並將狀態置於擴充的網格上。

以下範例說明如何結合 manager.Manager、elastic_retry 和 Orbax 的記憶體內 Snapshotter，實作彈性訓練迴圈：

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

設定彈性管理工具和重試行為

每個程序會例項化 manager.Manager(devices=None) 一次。如果省略 devices，管理工具會使用 jax.devices()，並依 slice_index 屬性將裝置分組。

Manager 執行個體會公開下列屬性，以便在訓練期間檢查叢集拓撲：

• slice_to_devices：從每個 slice_index 到屬於該切片的 jax.Device 物件序列的對應。
• all_slice_indices：分配中所有切片索引的 frozenset。
• active_slice_indices：參與有效運算網格的切片索引 frozenset。
• active_devices：屬於有效切片的一系列 jax.Device 物件，依 slice_index 排序。
• inactive_slice_indices：分配中未參與運算網格的切片索引 frozenset。
• available_inactive_slices：背景監控執行緒偵測到健康狀態良好且可重新加入的非作用中切片索引 frozenset。
• total_slice_count 和 active_slice_count：已分配的切片總數和有效切片數量。
• default_device：來自有效切片的裝置 0，elastic_retry會自動設為每次嘗試的預設 JAX 裝置。

您可以將下列引數傳遞至 @elastic_manager.elastic_retry，自訂重試和縮放行為：

• minimum_slice_count：開始初始執行或重試前，所需健康切片的數量下限。如果設為 None (預設值)，elastic_retry 會等待分配中的所有切片 (total_slice_count)，並在暫停和繼續模式下運作。如果設為較小的整數 (例如 1)，elastic_retry 會以副本大小調整模式運作。
• poll_interval：訓練期間，彈性訓練管理工具 (Manager) 執行 slice 健康狀態檢查的間隔時間 (以秒為單位)。預設值為 10。
• timeout：等待 minimum_slice_count 切片變為有效狀態的時間上限 (以秒為單位)，超過這個時間就會引發 TimeoutError。預設值為 None (無逾時)。
• pre_callback：選用回呼，在管理工具更新 active_slice_indices 後，以及包裝的訓練迴圈函式執行前，在每次嘗試開始時執行。
• on_elastic_event_callback：選用回呼，在 elastic_retry 擷取到切片 jax.errors.JaxRuntimeError 後，以及下次重試前執行。您可以使用這個回呼清除不完整的檢查點，或重設資料疊代器。
• retry_policy：可呼叫的 (attempt: int, error: Exception) -> bool，用於判斷是否要在彈性事件後繼續重試。使用 manager.ElasticRetryLimit(max_attempts) 限制嘗試總次數。如果省略，elastic_retry 會無限次重試。

發生彈性事件時，一旦堆疊框架展開，且沒有變數參照這些陣列，JAX 就會釋出失敗嘗試的裝置緩衝區。請避免在全域變數或長期回呼中保留失敗嘗試的裝置陣列參照，以便在下次嘗試開始前釋出記憶體。

使用 Orbax 建立記憶體內快照

ocp.v1.training.pathways.Snapshotter 類別 (適用於 orbax-checkpoint 0.12.6 以上版本) 會在 pinned_host 記憶體中，跨切片維護訓練狀態的記憶體內備份：

• Snapshotter(replica_axis_index=0)：初始化快照工具，並指定哪個網格軸代表資料平行切片副本。儲存狀態中的所有 JAX 陣列都必須沿著這個網格軸複製，這樣當某個切片失敗時，其餘切片仍會保留模型和最佳化器狀態的完整副本。如果狀態包含單一裝置陣列，請先在網格中複製這些陣列，再呼叫 save。
• save(step, state)：以非同步方式將 state 轉移至 pinned_host 記憶體，並沿著 replica_axis_index 將其分割為每個副本樹。如果先前的快照移轉作業仍在進行中，save 會略過新要求，以免訓練迴圈停滯。
• load(abstract_state, reset_snapshot_state=True)：探測已儲存的主機副本，找出存留的切片，串連健康狀態良好的副本，並將狀態重新分 shard 至 abstract_state 指定的目標 shard。
• join()：封鎖，直到任何已排入佇列的快照作業完成。
• latest：傳回最近完成的快照中繼資料 (包括 step)，如果沒有已完成的快照，則傳回 None。

減少切片失敗時的 sharding

如果 TPU 切片在執行期間失敗或遭到先占，Pathways 會引發 jax.errors.JaxRuntimeError (例如 DATA_LOSS、DEADLINE_EXCEEDED、NOT_FOUND 或 INTERNAL)。在以 @elastic_manager.elastic_retry 裝飾的函式中，elastic_retry 會擷取這個例外狀況並檢查 elastic.is_error_due_to_slice_down(error)：

• 如果錯誤不是由縮減事件引起，elastic_retry 會立即重新引發例外狀況。
• 如果錯誤是由切片事件引起，elastic_retry 會停止所有執行中的 JAX Profiler 追蹤、執行 on_elastic_event_callback，並評估 retry_policy(attempt, error)。
• 在開始下一次嘗試之前，elastic_retry 會呼叫 elastic.wait_for_slices 一次，輪詢切片健康狀態，直到至少有 minimum_slice_count 個切片處於活動狀態。接著，系統會更新 elastic_manager.active_slice_indices，並在存留的切片上執行訓練迴圈函式。

在切片重新加入時增加 sharding

工作負載在副本調整大小模式下，以較少的 slice 數量執行時，背景監控執行緒會定期對 elastic_manager.inactive_slice_indices 執行健康狀態檢查。當修復或重新排程的分塊通過健康狀態檢查時，監控執行緒會將其分塊索引新增至 elastic_manager.available_inactive_slices。

如果 elastic_manager.available_inactive_slices 在訓練步驟結束時並非空白，工作負載就會在原地擴充，不會遺失訓練進度：

• 將 elastic_manager.available_inactive_slices 中的切片索引新增至 elastic_manager.active_slice_indices，會將重新加入的切片納入 elastic_manager.active_devices，並從 available_inactive_slices 中清除。
• 從 elastic_manager.active_devices 重建裝置網格，並重新初始化訓練步驟和資料疊代器，即可調整運算，以適應擴充的拓撲。
• 呼叫 jax.device_put 會將目前的訓練狀態置於新的網格分區，不需從快照還原。

多層級查核點

大規模訓練工作經常中斷，直接在 Cloud Storage 寫入或還原檢查點，可能會造成訓練進度瓶頸。多層級查核點機制會協調 GKE 上的三個儲存層，減少查核點作業負擔和復原時間：

• 本機 RAM 磁碟：TPU 工作站會將高頻率查核點寫入掛接在各 worker 節點的記憶體內 CSI 磁碟區 (multitier-checkpoint.csi.storage.gke.io)。在同一個節點上重新啟動 Pod 時，這個磁碟區會保留下來。
• 叢集內對等複製：GKE 複製器服務會在叢集中的對等 TPU 工作站之間複製本機檢查點 shard。
• Cloud Storage 備份：複製器服務會以非同步方式，將本機檢查點備份至已啟用階層命名空間的 Cloud Storage bucket，不會阻斷 TPU 訓練步驟。

與多控制器 JAX 不同，每個 TPU 主機都會執行自己的訓練指令碼，而 Pathways 會在 pathways-head Pod 上執行單一控制器，並將工作調度至 pathways-worker Pod。為了在每個 TPU 主機的本機 RAM 磁碟上直接讀取及寫入檢查點 shard，而不透過單一控制器轉送檢查點資料，Pathways 會使用共置 Python。每個 worker Pod 都會與 pathways-worker 容器一起執行 colocated-python-sidecar 容器，且兩個容器都會掛接多層級檢查點 CSI 磁碟區和共用記憶體目錄 (/tmp/sidecar)。在初始化和訓練期間，Orbax 會將檢查點設定、儲存和還原作業，從 Pathways 控制器調度至每個 TPU 主機上的 worker sidecar。在 Pathways 上，多層級檢查點也支援彈性訓練，可縮減或擴充使用中的 TPU 磁碟區數量。

Pathways 支援搭配 MaxText 使用多層檢查點。如需設定和操作說明，請參閱 MaxText 多層檢查點指南。

設定 Pathways 工作負載

您可以使用 Cluster Toolkit (gcluster job submit) 或設定自訂 Pathways JobSet 資訊清單，為 Pathways 工作負載設定多層檢查點。

使用 Cluster Toolkit

使用 gcluster job submit --pathways 提交工作時，請傳遞 --gke-mtc-enabled、--gke-mtc-ramdisk-dir 和 --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" \
    ...

將 SIDECAR_IMAGE 替換為共置 Python 邊車容器映像檔的 URI。

設定這些旗標後，gcluster 會自動設定下列項目：

• 將 --sidecar_name=external 傳遞至 pathways-proxy 容器。
• 在每個 worker Pod 中執行 colocated-python-sidecar 容器，並在 pathways-worker 和 colocated-python-sidecar 上掛接共用的 /tmp/sidecar 目錄 (--cloud_pathways_sidecar_shm_directory=/tmp/sidecar)。
• 在 worker Pod 的 pathways-worker 和 colocated-python-sidecar 中，將 multitier-checkpoint.csi.storage.gke.io CSI 磁碟區掛接在 --gke-mtc-ramdisk-dir 指定的路徑，但不會在 pathways-head Pod 上掛接。

使用自訂 Pathways JobSet 資訊清單

如果您自行管理 Pathways JobSet 資訊清單，而非使用 gcluster，請將 --sidecar_name=external 傳遞至 pathways-proxy。在 worker Pod 規格中，將共用的 Sidecar 記憶體磁碟區和 multitier-checkpoint.csi.storage.gke.io CSI 磁碟區掛接至 pathways-worker 和 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

在自訂 JAX 工作負載中使用 Orbax 多層檢查點

在 Pathways 上執行的自訂 JAX 訓練指令碼中，使用 orbax.checkpoint.experimental.emergency.multi_tier_checkpointing 套件，並在 initialize_multi_tier_checkpointing 和 ReplicatorCheckpointManagerOptions 中設定 use_colocated_python=True：

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()

熱插拔

熱交換是 GKE JobSet API 的一項功能，可讓優先順序較高的工作快速接管優先順序較低工作的資源，盡可能縮短停機時間並加快復原速度。

建立 JobSet 時，GKE 會按照 JobSet 設定，在多個切片中排定工作負載。如果一或多個切片發生硬體故障，GKE 會將受影響的 Pod 標示為失敗。如果 GKE 叢集有備用 slice 正在執行優先順序較低的工作，JobSet 控制器會將失敗 slice 的工作負載，從優先順序較高的工作重新對應至該備用 slice。這項重新對應作業通常會在 1 分鐘內完成。

JobSet 重新啟動時，下列情況可能會發生熱插拔：

• 預設模式：如果同一個叢集內有閒置的備用 TPU 配量，Kubernetes 排程器會優先將重新啟動的工作排程至這些配量，而不是等待故障的配量修復。這種排程行為可加快復原速度。
• 異質工作負載：在執行多個工作負載且已設定 Kubernetes PriorityClass 資源的叢集中，重新啟動 JobSet 可能會觸發熱插拔。如果重新啟動的工作親和性與優先順序較低的工作資源相符，Kubernetes 會優先處理優先順序較低的工作，這樣優先順序較高的工作就能啟動，不必等待晶片修復。例如，您可以使用 PriorityClass 資源，為 Pathways worker Pod 設定不同的優先順序。

如要在叢集中使用優先順序，請定義 PriorityClass 資訊清單，並儲存為 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."

將這個 YAML 套用至 GKE 叢集：

kubectl apply -f high-prior-job.yaml

接著，將新的優先順序類別附加至 Pathways 工作負載。如果您使用 Cluster Toolkit 提交工作負載，請將 --priority=high-prior-job 傳遞至 gcluster job submit 指令，這會為 pathways-head 和 worker Pod 設定 priorityClassName。如果您自行編寫 JobSet 資訊清單，請在 worker 複製工作的 Pod 範本中設定 priorityClassName，以執行 pathways-worker 容器：

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

如果您是使用 Cluster Toolkit TPU 藍圖建立叢集，並啟用 Pathways，叢集就會包含 very-low、low、medium、high 和 very-high 優先順序類別，因此您可以將其中一個類別傳遞至 --priority，不必自行建立。

後續步驟

• 使用 Pathways 建立 GKE 叢集
• 使用 Pathways 執行批次工作負載
• Pathways 互動模式
• 將 JAX 工作負載移植到 Pathways
• Troubleshoot Pathways on Cloud
• 使用多層級檢查點功能，在 GKE 上訓練大規模機器學習模型
• MaxText 多層檢查點指南
• 使用多層查核點設定大型 AI 訓練工作