Cloud의 학습 과정 소개

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

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

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

시작하기 전에

다음 사항이 필요합니다.

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

단일 컨트롤러 및 다중 컨트롤러

여러 기기에서 계산을 관리하고 오케스트레이션하는 방법에는 주로 두 가지가 있습니다.

기능	단일 컨트롤러 (학습 과정)	다중 컨트롤러 (JAX 기본값)
제어	단일 제어 지점: 단일 클라이언트 프로그램이 중앙 컨트롤러 역할을 합니다.	분산 제어: 여러 프로세스가 참여하며 각 프로세스에는 자체 Python 인터프리터 인스턴스가 있습니다.
보기	통합 뷰: 클라이언트가 모든 기기를 하나의 통합 시스템으로 인식합니다.	현지화된 뷰: 각 Python 프로세스는 연결된 기기만 표시합니다.
프로그래밍	간소화된 프로그래밍: 사용자는 단일 클라이언트와 상호작용하므로 시스템이 로컬 가속기가 많은 단일 대형 머신으로 표시됩니다.	SPMD: 주로 SPMD 패러다임을 사용하여 모든 기기에서 동일한 프로그램을 실행해야 합니다.
유연성	비대칭 파이프라인 병렬 처리 및 계산 희소성을 비롯하여 SPMD를 넘어선 더 복잡한 계산 패턴을 지원합니다.	특히 여러 TPU 슬라이드에서 리소스 관리가 덜 유연할 수 있습니다.

학습 과정 구성요소

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

Pathways 리소스 관리자

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

경로 클라이언트

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

경로 작업자

가속기 머신 (TPU VM)에서 실행되는 프로세스입니다. 이러한 작업자는 IFRT 프록시 서버에서 프로그램의 컴파일된 실행 파일을 수신하고 TPU에서 연산을 실행합니다. 경로 작업자는 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 작업자와 상호작용합니다.

PathwaysJob API

PathwaysJob API는 ML 학습 및 일괄 추론 워크로드를 배포하는 데 사용하는 OSS Kubernetes 네이티브 API입니다. PathwaysJob 컨트롤러는 JobSet API를 활용하여 모든 Pathways 구성요소의 수명 주기와 조정을 관리합니다. 이 커스텀 리소스 정의 (CRD)는 일반적인 시나리오에서 개별 포드 사양을 직접 관리할 필요성을 추상화하여 경로 워크로드를 정의할 수 있는 상위 수준 인터페이스를 제공합니다. 모든 매개변수와 그 구체적인 의미의 전체 목록은 GitHub의 PathwaysJob API 문서를 참고하세요.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

다음 표에서는 PathwaysJob API의 설정을 설명합니다.

속성	설명
`apiVersion`	`PathwaysJob` 커스텀 리소스 정의 (CRD): pathways-job.pathways.domain/v1의 API 버전을 지정합니다.
`kind`	Kubernetes 객체를 `PathwaysJob`로 식별합니다.
`metadata.name`	Kubernetes의 `PathwaysJob` 객체 이름입니다. 일반적으로 pathways-. 패턴을 따릅니다.
`spec`	`PathwaysJob`의 원하는 상태와 구성을 정의합니다.
`spec.maxRestarts`	실패가 발생할 경우 시스템에서 `PathwaysJob`를 자동으로 다시 시작할 수 있는 최대 횟수입니다.
`spec.pathwaysVersion`	(선택사항) 이 작업의 Pathways 환경 내에서 사용할 원하는 JAX 프레임워크 버전을 지정합니다 (예: jax-0.5.3).
`spec.workers`	`PathwaysJob`의 작업자 풀 구성을 정의하는 배열입니다. 일반적으로 TPU 리소스를 활용합니다.
`spec.workers[].type`	작업자 노드에 사용할 TPU 머신 유형입니다 (예: $TPU_MACHINE_TYPE은 ct6e-standard-4t일 수 있음).
`spec.workers[].topology`	작업자에게 할당된 TPU 슬라이스의 토폴로지입니다 (예: $TOPOLOGY는 2x2, 4x4, 2x2x2일 수 있음).
`spec.workers[].numSlices`	작업자 풀에 프로비저닝할 TPU 슬라이스 수입니다 (예: $WORKLOAD_NODEPOOL_COUNT는 2일 수 있음).
`spec.workers[].maxSliceRestarts`	(선택사항) 슬라이스 내 개별 작업자가 실패할 경우 다시 시작할 수 있는 최대 횟수입니다.
`spec.customComponents`	(선택사항) 기본 작업과 함께 프록시 서버, Pathways 서버 또는 추가 작업자와 같은 맞춤 구성요소를 정의하고 배포할 수 있는 배열입니다.
`spec.customComponents[].componentType`	정의되는 맞춤 구성요소의 유형을 지정합니다 (예: proxy_server, pathways_server, worker, colocated_python_sidecar).
`spec.customComponents[].image`	이 맞춤 구성요소의 컨테이너에 사용할 Docker 이미지입니다.
`spec.customComponents[].customFlags`	컨테이너가 시작될 때 컨테이너에 전달되는 맞춤 명령줄 플래그의 배열입니다.
`spec.customComponents[].customEnv`	컨테이너 내에 설정할 맞춤 환경 변수의 배열입니다. 각 요소에는 이름과 값이 있습니다.
`spec.pathwaysDir`	컴파일 아티팩트 및 기타 임시 데이터를 저장하기 위해 `PathwaysJob`에서 사용하는 Cloud Storage 버킷입니다. 이 버킷은 워크로드를 실행하기 전에 만들어야 합니다.
`spec.controller`	전체 작업 실행을 관리하는 경로 컨트롤러 구성 설정
`spec.controller.deploymentMode`	경로 컨트롤러의 CPU 리소스 (경로 리소스 관리자 및 프록시 서버)가 배포되는 방식을 지정합니다. 기본 모드에서는 전용 CPU 노드에 배포하고 colocate_head_with_workers는 TPU 작업자와 함께 배포합니다.
`spec.controller.elasticSlices`	(선택사항) 비정상으로 간주되기 전에 작업 실행 중에 사용할 수 없게 될 수 있는 최대 TPU 슬라이스 수입니다.
`spec.controller.template`	(선택사항) 사용자 작업의 포드 템플릿을 정의합니다. 이는 배치 워크로드에는 필요하지만 대화형 워크로드에는 필요하지 않습니다.
`spec.controller.template.spec`	사용자 작업의 포드 사양입니다.
`spec.controller.template.spec.containers`	사용자 작업 내에서 실행될 컨테이너를 정의하는 배열입니다.
`spec.controller.template.spec.containers[].name`	사용자 작업 내 컨테이너의 이름입니다 (이 샘플에서는 main).
`spec.controller.template.spec.containers[].image`	기본 컨테이너의 컨테이너에 사용할 Docker 이미지입니다 (이 샘플에서는 python:3.11).
`spec.controller.template.spec.containers[].command`	기본 컨테이너가 시작될 때 실행되는 명령어입니다. 이 샘플에서는 `pathwaysutils`를 설치하고, Pathways를 초기화하고, JAX 기기를 출력합니다.

GKE의 Pathways 구성요소

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

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

컨테이너 유형	위치
IFRT 프록시 서버	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>`
경로 리소스 관리자/작업자	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>`

Pathways 리소스 관리자

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

  - 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: 노드 유형입니다. 경로 리소스 관리자의 경우 '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: 임시 파일에 사용되는 GCS 버킷입니다.

경로 작업자

다음 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 버킷입니다.

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

Cloud의 학습 과정 소개 컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.

시작하기 전에

단일 컨트롤러 및 다중 컨트롤러

학습 과정 구성요소

Pathways 리소스 관리자

경로 클라이언트

경로 작업자

IFRT 프록시 클라이언트

IFRT 프록시 서버

사이드카 서버

PathwaysJob API

GKE의 Pathways 구성요소

Pathways 리소스 관리자

IFRT 프록시 서버

경로 작업자

다음 단계

Cloud의 학습 과정 소개