학습 과정을 통한 탄력적인 학습

이 페이지에서는 대규모 학습 작업이 계획된 중단 및 계획되지 않은 하드웨어 장애로부터 복구될 수 있도록 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 클러스터를 구성합니다.
  • 클러스터에서 GKE용 워크로드 아이덴티티 제휴와 함께 다중 계층 체크포인트 및 Cloud Storage FUSE(GcsFuseCsiDriver) 부가기능을 사용 설정합니다.
  • 계층적 네임스페이스가 사용 설정된 Cloud Storage 버킷을 만듭니다.
  • 다중 등급 체크포인트 Kubernetes 서비스 계정에 roles/storage.objectUser 및 roles/storage.insightsCollectorService Identity and Access Management(IAM) 역할을 부여합니다.
  • Cloud Storage 버킷, TPU 노드 선택기, 인메모리 볼륨 크기(inMemoryVolumeSize)를 지정하는 CheckpointConfiguration 커스텀 리소스를 배포합니다.
• 예를 들어 샘플 사이드카 Dockerfile을 사용하여 코로케이션된 Python 사이드카 컨테이너 이미지를 빌드하고 푸시합니다. 워크로드 및 사이드카 이미지에서 Pathways 출시와 일치하고 다음 최소 버전 요구사항을 충족하는 일치하는 jax 및 jaxlib 버전을 사용합니다.
  • Orbax 버전: >= 0.12.4
  • JAX 버전(공동 배치된 Python 포함): >= 0.10.1

정지-재개

GKE는 일반적으로 포드를 선점하기 전에 가속기 포드에 선점 알림을 전송합니다. Pathways는 기본적으로 모든 Cloud Deploy 배포에서 선점 허용을 사용 설정하며 Pathways 가속기 작업은 이러한 알림을 수신합니다.

선점 알림이 도착하면 Pathways는 먼저 워크로드를 투명하게 저장하고 복원할 수 있는지 확인합니다. 가능한 경우 Pathways는 GKE가 가속기 작업을 제거하기 전에 현재 상태를 Cloud Storage와 같은 영구 스토리지에 기록하여 ML 워크로드를 투명하게 정지하려고 시도합니다. 나중에 GKE가 작업을 다시 예약하면 Pathways는 유지된 상태를 다시 읽어 ML 워크로드를 재개합니다.

워크로드를 복원할 수 없는 경우 Pathways는 가속기 작업을 종료하고 탄력적 학습이 구성된 경우 워크로드에 실패를 전달합니다. 탄력적 학습이 구성되지 않은 경우 GKE는 JobSet 재시작 정책에 따라 전체 워크로드를 다시 시작합니다.

JAX를 사용하여 정의된 일반적인 ML 워크로드는 고대역폭 메모리(HBM) 스냅샷을 사용하여 복원할 수 있는 스테이트리스(Stateless) Pathways XLA 구성요소에 의존합니다. JAX 공동 배치 Python API를 사용하여 정의된 워크로드와 같은 특정 ML 워크로드는 스테이트풀(Stateful) 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는 원격 스토리지에서 읽지 않고 활성 TPU 메시에서 직접 상태를 재구성하고 다시 샤딩합니다.

Pathways 프록시 서버에서 탄력성 사용 설정

GKE의 Pathways에서 탄력적 기기 세트를 사용 설정하려면 JobSet 매니페스트의 pathways-proxy 컨테이너의 args 필드에 --num_elastic_instances 플래그를 추가하고 워크로드의 탄력적 TPU 슬라이스 수를 지정합니다.

- 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 버킷입니다.
• 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를 확인하여 정상적인 대체 슬라이스를 감지합니다. 대체 슬라이스를 사용할 수 있게 되면 코드가 이를 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)가 수행하는 슬라이스 상태 점검 사이의 간격(초)입니다. 기본값은 10입니다.
• timeout: TimeoutError를 발생시키기 전에 minimum_slice_count 슬라이스가 활성화될 때까지 기다리는 최대 시간(초)입니다. 기본값은 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): 저장된 호스트 복제본을 프로브하여 살아남은 슬라이스를 식별하고, 정상 복제본을 연결하고, abstract_state로 지정된 대상 샤딩으로 상태를 다시 샤딩합니다.
• join(): 큐에 추가된 스냅샷 작업이 완료될 때까지 차단합니다.
• latest: 가장 최근에 완료된 스냅샷의 메타데이터(step 포함)를 반환하거나 완료된 스냅샷이 없는 경우 None를 반환합니다.

슬라이스 실패 시 샤딩 줄이기

실행 중에 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 프로파일러 trace를 중지하고 on_elastic_event_callback를 실행하고 retry_policy(attempt, error)를 평가합니다.
• 다음 시도를 시작하기 전에 elastic_retry는 elastic.wait_for_slices를 한 번 호출하여 최소 minimum_slice_count개의 슬라이스가 활성화될 때까지 슬라이스 상태를 폴링합니다. 그런 다음 elastic_manager.active_slice_indices를 업데이트하고 나머지 슬라이스에서 학습 루프 함수를 실행합니다.

슬라이스가 다시 조인될 때 샤딩 증가

워크로드가 복제본 크기 조절 모드에서 축소된 슬라이스 수로 실행되면 백그라운드 모니터 스레드가 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 워커는 각 워커 노드에 마운트된 인메모리 CSI 볼륨(multitier-checkpoint.csi.storage.gke.io)에 고주파 체크포인트를 작성합니다. 이 볼륨은 동일한 노드에서 포드가 다시 시작되어도 유지됩니다.
• 클러스터 내 피어 복제: GKE 복제기 서비스는 클러스터의 피어 TPU 워커 간에 로컬 체크포인트 샤드를 복제합니다.
• Cloud Storage 백업: 복제기 서비스는 TPU 학습 단계를 차단하지 않고 계층적 네임스페이스가 사용 설정된 Cloud Storage 버킷에 로컬 체크포인트를 비동기식으로 백업합니다.

각 TPU 호스트가 자체 학습 스크립트를 실행하는 멀티 컨트롤러 JAX와 달리 Pathways는 pathways-head Pod에서 단일 컨트롤러를 실행하고 작업을 pathways-worker Pod에 전달합니다. 단일 컨트롤러를 통해 체크포인트 데이터를 라우팅하지 않고 각 TPU 호스트의 로컬 RAM 디스크에서 직접 체크포인트 샤드를 읽고 쓰기 위해 Pathways는 공동 배치된 Python을 사용합니다. 각 worker 포드는 pathways-worker 컨테이너와 함께 colocated-python-sidecar 컨테이너를 실행하며, 두 컨테이너 모두 멀티 티어 체크포인트 CSI 볼륨과 공유 메모리 디렉터리(/tmp/sidecar)를 마운트합니다. 초기화 및 학습 중에 Orbax는 Pathways 컨트롤러에서 각 TPU 호스트의 워커 사이드카로 체크포인트 설정, 저장, 복원을 전달합니다. 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 포드에서 colocated-python-sidecar 컨테이너를 실행하고 pathways-worker 및 colocated-python-sidecar 모두에 공유 /tmp/sidecar 디렉터리(--cloud_pathways_sidecar_shm_directory=/tmp/sidecar)를 마운트합니다.
• worker 포드의 pathways-worker 및 colocated-python-sidecar 모두에서 --gke-mtc-ramdisk-dir로 지정된 경로에 multitier-checkpoint.csi.storage.gke.io CSI 볼륨을 마운트하지만 pathways-head 포드에는 마운트하지 않습니다.

커스텀 Pathways JobSet 매니페스트 사용

gcluster를 사용하는 대신 자체 Pathways JobSet 매니페스트를 관리하는 경우 --sidecar_name=external를 pathways-proxy에 전달합니다. worker 포드 사양에서 공유 사이드카 메모리 볼륨과 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는 영향을 받는 포드를 실패로 표시합니다. GKE 클러스터에 우선순위가 낮은 작업을 실행하는 예비 슬라이스가 있는 경우 JobSet 컨트롤러는 실패한 슬라이스의 워크로드를 우선순위가 높은 작업에서 해당 예비 슬라이스로 다시 매핑합니다. 이 재매핑은 일반적으로 1분 미만으로 걸립니다.

JobSet가 다시 시작되면 다음과 같은 상황에서 핫 스왑이 발생할 수 있습니다.

• 기본 모드: 동일한 클러스터 내에 여분의 유휴 TPU 슬라이스가 있는 경우 Kubernetes 스케줄러는 실패한 슬라이스가 복구될 때까지 기다리지 않고 이러한 슬라이스에 다시 시작된 작업을 우선적으로 예약합니다. 이 예약 동작은 더 빠른 복구를 제공합니다.
• 이기종 워크로드: 구성된 Kubernetes PriorityClass 리소스로 여러 워크로드를 실행하는 클러스터에서 다시 시작된 JobSet가 핫 스왑을 트리거할 수 있습니다. 다시 시작된 작업의 어피니티가 우선순위가 낮은 작업의 리소스와 일치하면 Kubernetes는 우선순위가 낮은 작업을 선점하여 우선순위가 높은 작업이 슬라이스 복구를 기다리지 않고 시작할 수 있도록 합니다. 예를 들어 PriorityClass 리소스를 사용하여 Pathways 워커 포드를 다양한 우선순위로 구성할 수 있습니다.

클러스터에서 우선순위를 사용하려면 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 포드 모두에 priorityClassName가 설정됩니다. 자체 JobSet 매니페스트를 작성하는 경우 pathways-worker 컨테이너를 실행하는 worker 복제된 작업의 포드 템플릿에서 priorityClassName를 설정합니다.

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

Pathways가 사용 설정된 Cluster Toolkit TPU 청사진을 사용하여 클러스터를 만든 경우 클러스터에 이미 very-low, low, medium, high, very-high 우선순위 클래스가 포함되어 있으므로 자체 클래스를 만드는 대신 이 중 하나를 --priority에 전달할 수 있습니다.

다음 단계

• Pathways를 사용하여 GKE 클러스터 만들기
• Pathways로 일괄 워크로드 실행
• Pathways 대화형 모드
• JAX 워크로드를 Pathways로 포팅
• Cloud의 Pathways 문제 해결
• 멀티 계층 체크포인트를 사용하여 GKE에서 대규모 머신러닝 모델 학습
• MaxText 다중 등급 체크포인트 가이드
• 대규모 AI 학습 작업에 다중 계층 체크포인트 사용