Cloud의 학습 과정 소개

Pathways는 대규모 멀티태스크 희소 활성화 머신러닝 시스템을 만들 수 있도록 설계된 시스템입니다. 이를 통해 수천 또는 수만 개의 가속기를 사용할 수 있으며, 처리 요구사항에 따라 다양한 태스크에 다양한 양의 컴퓨팅을 동적으로 할당할 수 있습니다.

Pathways는 단일 JAX 클라이언트가 여러 대규모 TPU 슬라이스에서 워크로드를 조정할 수 있도록 지원하여 대규모 머신러닝 계산을 간소화합니다. 이 슬라이스는 수천 개의 TPU 칩에 걸쳐 있을 수 있습니다.

Pathways는 Google에서 Gemini와 같은 대규모 모델을 학습시키는 데 내부적으로 사용됩니다. Cloud의 Pathways는 Google Cloud 고객에게 동일한 이점을 제공합니다.

시작하기 전에

다음 사항이 필요합니다.

이 문서에서는 Google Kubernetes Engine (GKE)에서 Pathways 관리형 TPU를 일괄, 실시간, 대화형 워크로드에 사용하는 방법을 간략하게 설명합니다. 이 문서에서는 사용자가 Google Kubernetes Engine의 단일 및 멀티슬라이스 TPU를 포함하여 GKE에서 TPU를 사용하는 데 이미 익숙하고 멀티슬라이스 TPU에 대한 일반적인 경험이 있다고 가정합니다.

단일 컨트롤러 및 멀티 컨트롤러

여러 기기에서 계산을 관리하고 조정하는 방법에는 크게 두 가지가 있습니다.

기능

단일 컨트롤러 (Pathways)

멀티 컨트롤러 (JAX 기본값)

제어

단일 제어 지점: 단일 클라이언트 프로그램이 중앙 컨트롤러 역할을 합니다.

분산 제어: 여러 프로세스가 참여하며 각 프로세스에는 자체 Python 인터프리터 인스턴스가 있습니다.

보기

통합 뷰: 고객은 모든 기기를 하나의 통합된 시스템으로 봅니다.

현지화된 뷰: 각 Python 프로세스는 연결된 기기만 볼 수 있습니다.

프로그래밍

간소화된 프로그래밍: 사용자는 단일 클라이언트와 상호작용하여 시스템이 여러 로컬 가속기가 있는 단일 대형 머신으로 표시됩니다.

SPMD: 주로 SPMD 패러다임을 사용하므로 모든 기기에서 동일한 프로그램을 실행해야 합니다.

유연성

비대칭 파이프라인 동시 로드 및 계산 희소성을 포함하여 SPMD를 넘어서는 더 복잡한 계산 패턴을 지원합니다.

리소스 관리, 특히 여러 TPU 슬라이스에서 유연성이 떨어질 수 있습니다.

Pathways 구성요소

다음 섹션에서는 Pathways 아키텍처의 주요 구성요소를 간략하게 설명합니다.

Pathways 리소스 관리자

Pathways 시스템의 중앙 컨트롤 플레인입니다. 모든 가속기 리소스를 관리하고 사용자 작업에 대한 가속기 할당을 조정합니다. 작업자의 상태를 모니터링하고 작업 예약, 일시중지, 재개를 처리합니다. 오류 및 시스템 상태에 대한 단일 연락 지점 역할을 합니다. 이 구성요소에는 CPU 리소스만 필요합니다.

Pathways 클라이언트

이는 Pathways 시스템의 진입점 역할을 하는 임시 프레임워크 런타임(IFRT)의 구현입니다. 프로그램에서 상위 수준 작업 (HLO)을 수신합니다. Pathways 클라이언트는 Pathways 리소스 관리자와 조정하여 사용자 코드를 기반으로 실행을 위해 컴파일된 프로그램을 배치할 위치를 결정합니다. 지정된 JAX 클라이언트에 시스템의 통합 뷰를 제공합니다. 이 구성요소에는 CPU 리소스만 필요합니다.

Pathways 작업자

가속기 머신 (TPU VM)에서 실행되는 프로세스입니다. IFRT 프록시 서버에서 프로그램의 컴파일된 실행 파일을 수신하고 TPU에서 계산을 수행합니다. Pathways 작업자는 IFRT 프록시 서버를 통해 데이터를 프로그램으로 다시 전송합니다. 이 구성요소에는 가속기 리소스가 필요합니다.

IFRT 프록시 클라이언트

이는 기본 런타임에서 사용자 코드를 분리하고 코드 이식성과 투명성을 향상하는 Interim Framework Runtime (IFRT) API의 OSS 구현입니다. JAX는 이 구현을 기본 멀티 컨트롤러 런타임의 대안으로 사용합니다. IFRT 프록시 클라이언트는 프로그램과 Pathways 구성요소 간의 통신 브리지 역할을 합니다. IFRT 프록시 서버에 요청을 전송하고 결과를 수신합니다. IFRT API의 OSS 구현입니다. 이 구성요소에는 CPU 리소스만 필요합니다.

IFRT 프록시 서버

이 gRPC 서버는 IFRT 프록시 클라이언트로부터 요청을 수신하고 이를 Pathways 클라이언트로 전달하여 실제 작업 분산을 처리합니다. 이 구성요소에는 CPU 리소스만 필요합니다.

사이드카 서버

이 gRPC 서버는 가속기 VM에서 Pathways 워커와 함께 배치되어 가속기 VM에서 직접 사용자 지정 Python 코드를 실행하여 컨트롤러에서 가속기로의 데이터 전송 지연 시간을 줄입니다. 사이드카 서버는 gRPC 전송에서 커스텀 버전 프로토콜을 통해 Pathways 워커와 상호작용합니다.

Pathways 구성요소의 관계를 보여줍니다.
Pathways 구성요소

GKE의 Pathways 구성요소

이 섹션에서는 Pathways 구성요소를 컨테이너 및 포드와 같은 Google Kubernetes Engine 구성요소에 매핑합니다.

Pathways 컨테이너 이미지는 다음 위치에서 찾을 수 있습니다.

컨테이너 유형

위치

IFRT 프록시 서버

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>

Pathways 리소스 관리자/작업자

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Pathways 리소스 관리자

GKE 클러스터를 만든 후 다음 containerSpec를 사용하여 Pathways 리소스 관리자를 배포할 수 있습니다.

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

인수 설명:

  • --server_port: Pathways 리소스 관리자는 이 포트를 사용하여 다른 Pathways 구성요소와 통신합니다.
  • --node_type: 노드 유형입니다. Pathways 리소스 관리자의 경우 'resource_manager'로 설정해야 하며 다른 컨테이너에는 필요하지 않습니다.
  • --instance_count: TPU 슬라이스 수입니다.
  • --instance_type: 슬라이스의 TPU 유형 및 토폴로지입니다. tpu{TPU type}:{TPU topology} 형식입니다(예: tpuv5e:4x4).
  • --gcs_scratch_location: 임시 파일에 사용되는 Cloud Storage 버킷입니다.

IFRT 프록시 서버

다음 containerSpec를 사용하여 IFRT 프록시 서버를 배포할 수 있습니다.

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

인수 설명:

  • --resource_manager_address: 프록시 서버가 Pathways 리소스 관리자와 통신하는 데 사용하는 호스트 이름 및 포트입니다. 포트는 Pathways 리소스 관리자 컨테이너에 사용되는 --server_port 값과 동일해야 합니다.
  • --server_port: IFRT 프록시 서버는 이 포트를 사용하여 IFRT 프록시 클라이언트와 통신합니다.
  • --gcs_scratch_location: 임시 파일에 사용되는 Cloud Storage 버킷입니다.

Pathways 작업자

다음 containerSpec를 사용하여 Pathways 작업자를 배포할 수 있습니다.

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

인수 설명:

  • --resource_manager_address: TPU 워커가 Pathways 리소스 관리자와 통신하는 데 사용하는 호스트 이름 및 포트입니다. 포트는 Pathways 리소스 관리자 컨테이너에 사용되는 --server_port 값과 동일해야 합니다.
  • --server_port: 작업자는 이 포트를 사용하여 프록시 서버 및 Pathways 리소스 관리자와 통신합니다.
  • --gcs_scratch_location: 임시 파일에 사용되는 Cloud Storage 버킷입니다.

Pathways 리소스 관리자, IFRT 프록시 서버, Pathways 작업자는 모두 다른 포트를 가질 수 있지만 이 예시에서는 Pathways 리소스 관리자와 Pathways 작업자가 동일한 포트를 공유합니다.

다음 단계