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

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

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

시작하기 전에

다음 사항이 필요합니다.

• Cluster Toolkit 설정
• 설치된 Kubernetes 도구
• gcloud CLI 설치
• GKE API 사용 설정

정지-재개

일반적으로 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 시스템은 이를 감지하고 다음에 해당 하드웨어에 있는 데이터에 액세스할 때 예외를 통해 사용자의 학습 작업에 알립니다. 이 알림은 워크로드를 비정상 종료하지 않습니다. 코드가 알림을 처리하고 리소스를 재구성하여 처리를 계속하거나 정상적으로 종료할 수 있습니다.
• 사용자 정의 탄력성 핸들러: 사용자의 모델 코드가 이 예외를 처리할 수 있어야 합니다. 이것이 바로 '모델별 복구'입니다.
  • 스냅샷: 가장 일반적인 방법은 모델 상태의 스냅샷을 주기적으로 저장하는 것입니다. 오류가 발생하면 가장 최근 스냅샷에서 로드하여 학습을 재개할 수 있습니다.
  • 재구성: 사용 가능한 슬라이스 수에 맞게 학습 작업을 재구성해야 할 수 있습니다. 예를 들어 슬라이스 하나가 작동을 중지하면 교체할 수 있을 때까지 활성 슬라이스 수를 하나 줄일 수 있습니다. 자세한 내용은 Elastic 핸들러를 참조하세요.
  • 데이터/계산 그래프 업데이트: 코드는 필요에 따라 계산 그래프를 다시 만들어 계산에 사용할 수 있는 기기 수의 변경사항을 처리해야 합니다. 여기에는 데이터를 다시 파티셔닝하거나 모델을 다시 컴파일하는 작업이 포함될 수 있습니다.
• 복구에서 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 포드의 포드 사양에 다음 텍스트를 추가하여 새 우선순위 클래스를 Pathways 워커 작업에 연결합니다.

priorityClassName: high-prior-job

다음 단계

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