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

Pathways는 다음과 같은 방식으로 복원력 이점을 제공합니다.

• 정지-재개: 커스텀 선점 처리 코드가 없어도 워크로드가 선점 알림과 같은 계획된 중단을 허용할 수 있습니다.
• 탄력적 학습: 학습 코드의 복구 로직을 사용하여 클라이언트가 비정상 종료되지 않고 워크로드가 계획되지 않은 하드웨어 장애를 허용할 수 있습니다.

시작하기 전에

시작하기 전에 다음 작업을 완료하세요.

• Cluster Toolkit 설정
• Kubernetes 도구 설치
• Google Cloud CLI 설치
• Google Kubernetes Engine API 사용 설정

정지-재개

Google Kubernetes Engine(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/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를 호출하여 업데이트된 메시 샤딩에 최신 스냅샷을 복원하는 동시에 실패한 시도의 기기 버퍼가 범위를 벗어나도록 해야 합니다.
• 루프의 각 단계에서 코드는 주기적으로 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: 슬라이스 다운 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_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를 호출하면 스냅샷에서 복원하지 않고 현재 학습 상태가 새 메시 샤딩에 배치됩니다.

핫스왑

핫스왑은 우선순위가 높은 작업이 우선순위가 낮은 작업의 리소스를 빠르게 인계받아 다운타임을 최소화하고 더 빠른 복구를 보장하는 GKE JobSet API의 기능을 의미합니다.

JobSet가 생성되면 GKE는 JobSet 구성에 지정된 대로 여러 슬라이스에 워크로드를 예약합니다. 하나 이상의 슬라이스에서 하드웨어 오류가 발생하면 영향을 받는 포드가 실패로 표시됩니다. GKE 클러스터에 우선순위가 낮은 작업을 실행하는 예비 슬라이스가 있는 경우 JobSet 시스템은 우선순위가 높은 작업에서 실패한 슬라이스의 워크로드를 해당 예비 슬라이스로 다시 매핑합니다. 이 재매핑은 일반적으로 1분 이내에 완료됩니다.

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

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

클러스터에서 우선순위를 사용하려면 PriorityClass 매니페스트를 정의합니다.

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-worker 포드의 포드 사양에 다음 텍스트를 추가하여 새 우선순위 클래스를 Pathways 워커 작업에 연결합니다.

priorityClassName: high-prior-job

다음 단계

• Pathways를 사용하여 GKE 클러스터 만들기
• Pathways로 일괄 워크로드 실행
• Pathways 대화형 모드
• JAX 워크로드를 Pathways로 포팅
• 클라우드에서 Pathways 문제 해결