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

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

• 일시중지-재개: 사용자가 커스텀 선점 처리 코드를 작성하지 않아도 선점 알림과 같은 계획된 중단에 대한 허용 오차를 제공합니다.
• 탄력적 학습: 계획되지 않은 하드웨어 오류에 대한 허용 오차를 제공하여 클라이언트가 비정상 종료되지 않지만 사용자가 모델 별 복구 코드를 작성해야 합니다.

시작하기 전에

다음 사항이 필요합니다.

• XPK 설치
• Kubernetes 도구 설치
• gcloud CLI 설치
• GKE API 사용 설정

일시중지-재개

일반적으로 GKE는 포드가 선점되기 전에 가속기 포드에 선점 알림을 보냅니다. Pathways 선점 허용 오차는 모든 클라우드 배포에서 기본적으로 사용 설정되며 Pathways 가속기 작업은 이러한 알림을 수신 대기합니다.

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

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

JAX를 사용하여 정의된 일반적인 ML 워크로드는 고대역폭 메모리 (HBM) 스냅샷을 사용하여 복원할 수 있는 상태 비저장 Pathways XLA 구성요소에 의존합니다. JAX 공동 배치 Python API를 사용하여 정의된 워크로드와 같은 특정 ML 워크로드는 상태 저장 Pathways 구성요소에 의존합니다. 이러한 구성요소는 복원할 수 없습니다.

탄력적 학습

탄력적 학습을 사용하면 하드웨어 오류가 발생하더라도 학습 작업을 계속할 수 있습니다. 이는 Pathways 시스템 기능과 사용자 정의 모델 복구 로직의 조합을 통해 이루어집니다.

• 오류 감지: 하드웨어 오류가 발생하면 (예: TPU 작업자가 비정상 종료됨) Pathways 시스템이 이를 감지하고 해당 하드웨어에 있는 데이터에 다음에 액세스할 때 예외를 통해 사용자 학습 작업에 알립니다. 이 알림은 워크로드를 비정상 종료하지 않습니다. 코드가 알림을 처리하고 리소스를 재구성하여 계속 처리하거나 정상적으로 종료할 수 있도록 합니다.
• 사용자 정의 탄력성 핸들러: 사용자의 모델 코드는 이 예외를 처리할 수 있어야 합니다. 이것이 '모델별 복구'를 만드는 것입니다.
  • 스냅샷: 가장 일반적인 접근 방식은 모델 상태의 스냅샷을 주기적으로 저장하는 것입니다. 오류가 발생하면 가장 최근 스냅샷에서 로드하여 학습을 재개할 수 있습니다.
  • 재구성: 사용 가능한 슬라이스 수에 맞게 학습 작업을 재구성해야 할 수 있습니다. 예를 들어 슬라이스 하나가 작동을 멈추면 대체 슬라이스를 사용할 수 있을 때까지 활성 슬라이스 수를 하나 줄일 수 있습니다. 자세한 내용은 탄력적 핸들러를 참고하세요.
  • 데이터/계산 그래프 업데이트: 코드는 필요에 따라 계산 그래프를 다시 만들어 계산에 사용할 수 있는 기기 수의 변경사항을 처리해야 합니다. 여기에는 데이터 재분할 또는 모델 재컴파일이 포함될 수 있습니다.
• 복구에서 Pathways의 역할: Pathways는 사용자 정의 재구성 을 지원하는 기본 요소를 제공합니다.
  • 슬라이스 대체: 실패한 슬라이스가 대체되면 새 슬라이스를 사용할 수 있게 되면 클라이언트에 알릴 수 있습니다. 그러면 코드가 이 새 슬라이스를 사용하도록 재구성할 수 있습니다.
  • 투명한 복구: Pathways는 클러스터의 정상 부분에 대한 연결을 다시 설정하는 것과 같은 복구의 하위 수준 세부정보를 처리합니다.
• pathwaysutils의 유틸리티: pathways-utils에 정의된 Pathways 유틸리티 집합입니다.

탄력적 핸들러 구현

작성해야 하는 대부분의 코드는 사용자 정의 탄력적 핸들러에 있습니다. 이 핸들러는 메시를 다시 만들고 학습 루프를 다시 초기화하여 탄력적 이벤트 (예: TPU 슬라이스를 사용할 수 없게 됨)에 반응합니다.

각 워크로드는 고유합니다. 탄력적 핸들러의 복잡성은 워크로드의 복잡성에 따라 확장될 수 있습니다. 핸들러의 입력 및 출력은 학습 루프를 다시 초기화하는 데 필요한 최소 인수 및 반환 값이어야 합니다.

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

학습 루프 업데이트

학습 루프를 다음과 같이 변경해야 합니다.

1. 탄력적 관리자 만들기
2. jax.errors.JaxRuntimeError를 처리하는 try-except 블록 내에 학습 루프 래핑
3. jax.errors.JaxRuntimeError 핸들러 내에서 maybe_reshard_down을 호출합니다. 오류가 탄력적 이벤트와 관련된 경우 탄력적 관리자가 다시 분할하거나 다시 발생시킵니다.
4. 학습 루프 끝에서 maybe_snapshot 및 maybe_reshard_up 호출

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, {}, {})

탄력적 관리자 구성

탄력적 관리자는 몇 가지 다른 방법으로 구성할 수 있습니다. 스냅샷 생성 빈도는 스냅샷 기간에 따라 결정됩니다. 스냅샷 기간은 탄력적 이벤트로 인해 손실된 평균 단계 수에 영향을 미칩니다. 다시 분할 확인 기간은 학습 루프가 슬라이스 가용성을 폴링하는 빈도를 결정합니다. max_elastic_down_event_count를 사용하면 학습 루프가 지원할 슬라이스 손실로 인한 탄력적 이벤트 수를 설정할 수 있습니다. max_reshard_retry_count는 탄력적 관리자가 다시 분할을 재시도해야 하는 횟수를 지정합니다. 관리자는 싱글톤 객체이며 한 번만 만들어야 합니다.

스냅샷

탄력적 관리자 구성에 따라 함수는 탄력적 이벤트 중에 탄력적 핸들러에서 사용할 수 있는 호스트 메모리에 데이터를 스냅샷할 수 있습니다.

분할 감소

jax.errors.JaxRuntimeError를 포착한 후 Pathways는 오류가 손실된 슬라이스로 인한 탄력적 이벤트로 인한 것인지 확인합니다. 그렇다면 성공하거나 최대 재시도 횟수에 도달할 때까지 루프에서 탄력적 핸들러를 호출합니다. 오류가 탄력적 이벤트로 인한 것이 아니면 오류가 다시 발생합니다. 탄력적 핸들러의 반환 값은 호출자에게 전달됩니다.

분할 증가

탄력적 관리자 구성에 따라 사용할 수 없는 슬라이스가 있는 경우 Pathways는 추가 슬라이스를 사용할 수 있게 되었는지 확인합니다. 그렇다면 현재 단계의 기존 스냅샷이 아직 생성되지 않은 경우 스냅샷을 즉시 저장하고 성공하거나 최대 재시도 횟수에 도달할 때까지 루프에서 탄력적 핸들러를 호출합니다. 다시 분할이 발생하면 탄력적 핸들러의 반환 값이 호출자에게 전달됩니다. 그렇지 않으면 None이 반환됩니다.

핫스왑

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

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

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

1. 기본 모드: 동일한 클러스터 내에서 예비 유휴 TPU 슬라이스를 사용할 수 있는 경우 Kubernetes 스케줄러는 실패한 슬라이스가 복구될 때까지 기다리는 대신 이러한 슬라이스에 다시 시작된 작업을 예약하는 것을 우선시합니다. 이렇게 하면 복구가 더 빨라집니다.
2. 이종 워크로드: 구성된 Kubernetes PriorityClass로 여러 워크로드를 실행하는 클러스터에서 다시 시작된 JobSet은 핫스왑을 트리거할 수 있습니다. 다시 시작된 작업의 선호도가 우선순위가 낮은 작업의 리소스와 일치하면 Kubernetes는 우선순위가 낮은 작업을 선점하여 우선순위가 높은 작업을 즉시 시작할 수 있도록 합니다. 예를 들어 PriorityClass를 사용하여 우선순위가 다른 Pathways 작업자 포드를 구성할 수 있습니다.

클러스터에서 우선순위를 사용하려면 우선순위 클래스를 정의합니다. 예를 들면 다음과 같습니다.

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

priorityClassName: high-prior-job

다음 단계

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