JAX 워크로드를 Pathways로 포팅

Pathways를 사용하는 JAX의 분산 특성으로 인해 통신 오버헤드로 인해 일부 작업이 제대로 확장되지 않을 수 있습니다. Pathways는 비동기 디스패치와 같은 기능을 통해 이러한 오버헤드를 최소화하지만 JAX 워크로드를 Pathways로 포팅하거나 Pathways 워크로드로 JAX를 많은 수의 가속기로 확장할 때 유의해야 할 사항이 있습니다.

시작하기 전에

다음 사항이 필요합니다.

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

프로세스 색인

Pathways를 사용하는 JAX는 Pathways 클러스터의 모든 기기를 로컬로 취급합니다. 이렇게 하면 기기 관리가 간소화되고 JAX가 사용 가능한 모든 리소스를 활용할 수 있습니다. 실제로는 다음과 같은 의미입니다.

• jax.process_index()는 모든 기기에서 항상 0입니다.
• jax.devices() 및 jax.local_devices()는 전체 작업에서 모든 TPU 기기를 반환합니다.

하드웨어 유형 및 코로케이션

최상의 성능을 위해 모든 Pathways 구성요소와 사용자 작업을 동일한 Google Cloud 클라우드 영역에 배치합니다. IFRT 프록시 및 리소스 관리자와 같은 대형 CPU를 사용합니다. 64개의 vCPU와 256GB 메모리가 제공되는 전용 n2-standard-64 이상을 사용하는 것이 좋습니다.

PathwaysUtils

Pathways-utils는 Pathways on Cloud 아키텍처에서 JAX 워크로드의 배포 및 실행을 간소화할 수 있는 필수 유틸리티와 도구를 제공하는 Python 기반 GitHub 저장소입니다. 이 패키지는 클라우드 환경에 필요한 조정을 처리하므로 JAX 개발자는 최소한의 플랫폼별 구성으로 핵심 머신러닝 워크플로에 집중할 수 있습니다. 특히 다음을 제공합니다.

• '프록시' JAX 백엔드: 이 커스텀 백엔드를 사용하면 JAX 애플리케이션이 JAX_PLATFORMS=proxy 환경 변수를 설정하여 Pathways 인프라를 사용할 수 있습니다.
• 통합 프로파일링 유틸리티: 애플리케이션의 성능을 파악할 수 있는 프로파일링 기능입니다. jax.profiler.start_trace 및 jax.profiler.start_server와 같은 표준 JAX 프로파일링 API를 사용하면 JAX 코드뿐만 아니라 기본 Pathways 구성요소도 프로파일링하여 클라우드 환경 내 실행에 대한 전체적인 뷰를 제공할 수 있습니다.
• Orbax를 사용한 분산 체크포인트: Pathways 환경 내에서 Orbax 라이브러리를 사용할 때 분산 체크포인트를 사용하고 체크포인트를 복원할 수 있는 커스텀 Orbax 체크포인트 핸들러입니다. 이 통합은 pathwaysutils를 가져오는 한 기존 Orbax 체크포인트 코드에 변경사항을 적용하지 않고도 작동하는 것을 목표로 합니다.
• 탄력적 학습 프리미티브: Pathways를 사용하여 강력하고 확장 가능한 학습 워크플로를 빌드하는 데 사용할 수 있는 기본 탄력적 학습 프리미티브를 제공합니다. 이러한 프리미티브를 사용하면 학습 작업이 사용 가능한 리소스의 변화에 동적으로 적응하여 클라우드 환경의 효율성과 복원력을 개선할 수 있습니다.

체크포인트

Orbax는 Cloud Storage를 사용한 분산 체크포인트 및 복원을 위해 Pathways로 철저하게 테스트됩니다. 환경 변수 ENABLE_PATHWAYS_PERSISTENCE=1를 설정하고 train.py에서 import pathwaysutils; pathwaysutils.initialize()를 호출하면 IFRT 프록시를 통해 체크포인트 작업을 효율적으로 처리하는 커스텀 ArrayHandler가 등록되어 가속기의 Pathways 작업자가 데이터를 직접 저장하고 복원할 수 있습니다.

Colocated Python

Colocated Python은 사용자가 지정한 Python 코드를 TPU 또는 GPU 호스트에서 직접 실행할 수 있는 오픈소스 JAX API로, 멀티 컨트롤러 JAX에서 더 간단하게 사용할 수 있습니다. 이 방법을 사용하면 데이터 로드 및 체크포인트와 같은 컴퓨팅 집약적인 태스크를 위해 클라이언트와 TPU 머신 간에 데이터를 전송하지 않아도 됩니다. 공동 배치된 Python JAX API를 실행하도록 Pathways 클러스터를 구성하려면 샘플 사이드카 Dockerfile을 사용하여 사이드카 컨테이너 이미지를 빌드하고 푸시합니다.

데이터 로딩

학습 중에 데이터 세트에서 배치를 반복적으로 로드하여 모델에 피드합니다. 가속기의 작업량이 고갈되는 것을 방지하려면 호스트 간에 배치를 샤딩하는 효율적인 비동기 데이터 로더가 필요합니다. Pathways로 학습을 실행할 때 데이터 로더는 CPU VM에서 실행되고 (멀티 컨트롤러 설정에서 사용되는 TPU VM과 다름) TPU VM에 데이터를 전달합니다. 이로 인해 데이터를 읽는 데 지연 시간이 길어지지만 CPU 호스트에서 X개의 배치를 미리 읽고 읽은 데이터를 TPU에 비동기식으로 디스패치하여 부분적으로 완화됩니다. 이 솔루션은 중소 규모로 실행할 때 충분합니다.

GKE의 헤드 워크로드 컨테이너에서 멀티 프로세스 데이터 로더 (예: Grain)를 실행하면 컨테이너가 pathways-proxy 및 pathways-rm 사이드카와 헤드 노드를 공유합니다. 이 공유 아키텍처로 인해 다음을 구성해야 합니다.

• 명시적 작업자 수 설정: 컨테이너 내에서 os.cpu_count()는 컨테이너의 CPU 한도가 아닌 호스트 VM의 총 vCPU를 보고합니다. os.cpu_count()에서 파생하는 대신 명시적인 작업자 수 (컨테이너의 CPU 한도 이하)를 전달합니다.
• 메모리 지원 /dev/shm 볼륨 마운트: 멀티 프로세스 데이터 로더는 프로세스 간 통신에 POSIX 공유 메모리 (/dev/shm)를 사용합니다. 미리 가져온 배치가 기본 컨테이너 공유 메모리 한도를 초과하지 않도록 메모리 지원 emptyDir (medium: Memory) 볼륨이 /dev/shm에 마운트되어 있는지 확인합니다.

대규모로 최적의 성능을 얻으려면 같은 위치에 배치된 Python을 사용하여 가속기에서 직접 데이터 파이프라인을 실행하여 입력 데이터 파이프라인을 같은 위치에 배치하는 것이 좋습니다. 이렇게 하면 CPU 병목 현상이 제거되고 데이터 전송에 TPU의 빠른 상호 연결이 활용됩니다.

multihost_dataloading.py의 RemoteIterator 구현에서 TFDS 기반 입력 파이프라인을 마이그레이션하는 참조 구현을 찾을 수 있습니다. 이 구현은 멀티 컨트롤러 JAX와 Pathways 모두에서 공동 배치된 Python JAX API를 사용하여 분산 방식으로 작동합니다.

Jax 버전 관리

Pathways 출시 버전은 호환성과 안정성을 보장하기 위해 JAX 버전과 긴밀하게 결합되어 있습니다. 잠재적인 문제를 방지하려면 Pathways 아티팩트와 JAX 버전이 일치하는지 확인하세요. 각 Pathways 출시 버전은 jax-<version> 형식의 태그를 통해 호환되는 JAX 버전을 명확하게 지정합니다.

컴파일 캐시

Pathways 영구 컴파일 캐시는 Pathways 서버가 컴파일된 XLA 실행 파일을 Cloud Storage와 같은 영구 위치에 저장하여 중복 컴파일을 방지할 수 있는 기능입니다. 이 기능은 기본적으로 사용 설정되어 있습니다. 캐시 위치는 --gcs_scratch_location 플래그로 리소스 관리자 및 Pathways 워커 컨테이너에 전달됩니다. 연결된 스토리지 비용을 최소화하기 위해 캐시는 수명 주기 정책을 Cloud Storage 위치에 연결합니다. Cloud Storage 버킷당 정책은 50개로 제한됩니다. 따라서 모든 워크로드에서 공통 Cloud Storage 위치를 사용하는 것이 좋습니다.

--gcs_scratch_location 플래그를 생략하거나 비워두면 Pathways가 자동으로 인메모리 컴파일 캐시로 대체됩니다. 이 캐시는 Pathways 배포에서 컴파일된 XLA 실행 파일을 유지하지 않습니다.

이 캐시는 Pathways 워크로드에 대해 pathwaysutils.initialize()에서 사용 중지되는 JAX 컴파일 캐시와 유사합니다.

컴파일 캐시에는 다음 Cloud Storage 권한이 필요합니다.

• storage.buckets.get: 버킷 메타데이터를 검색합니다.
• storage.buckets.update: Pathways가 캐시 제거를 위해 TTL을 적용하는 객체 수명 주기 정책을 설정하는 데 필수적입니다.
• storage.objects.list: 버킷 내의 기존 캐시 객체를 나열합니다.
• storage.objects.create: 새로 컴파일된 실행 파일을 캐시에 쓰기
• storage.objects.get: 버킷에서 캐시된 실행 파일을 읽습니다.

프로파일링

JAX 프로파일러를 사용하여 JAX 프로그램의 trace를 생성할 수 있습니다. Pathways에서 지원되는 두 가지 일반적인 방법은 다음과 같습니다.

• 프로그래매틱
  • JAX 코드에서 프로그래매틱 방식으로 프로필 캡처
• 수동
  • JAX 코드에서 프로파일러 서버를 시작한 후 온디맨드로 프로필 캡처

두 경우 모두 프로필은 Cloud Storage 버킷에 기록됩니다. Cloud Storage 버킷에 여러 trace 파일이 생성되며, 타임스탬프 폴더가 다를 수 있습니다. 예를 들면 다음과 같습니다.

• 트레이스를 호출한 기본 Python 프로세스 (일반적으로 노트북 VM): <jax-client-vm-name>.xplane.pb
• Pathways IFRT 프록시: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
• Pathways 리소스 관리자: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
• Pathways 작업자: server.*<tpu-node-name>.xplane.pb

이러한 trace 파일은 다음 명령어를 실행하여 TensorBoard로 분석할 수 있습니다. 텐서보드 및 모든 프로파일링 도구에 대한 자세한 내용은 프로파일러를 사용하여 TensorFlow 성능 최적화를 참조하세요.

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

다음을 바꿉니다.

• BUCKET : trace 파일을 저장할 Cloud Storage 버킷
• PREFIX: trace 파일을 저장할 Cloud Storage 버킷 내 경로

프로그래매틱 프로필 캡처

코드 내에서 프로필을 캡처합니다. 프로필은 타임스탬프 디렉터리 아래의 gs://<bucket>/<prefix> 내에 저장됩니다.

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

수동 프로필 캡처

프로필 정보를 수동으로 캡처하려면 Python 코드에서 프로파일러 서버를 시작해야 합니다.

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

프로파일러 서버가 실행되는 동안 프로필을 캡처하고 데이터를 대상 Cloud Storage 위치로 내보낼 수 있습니다.

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

프로그램의 trace 내에서 Compile 및 Execute와 같은 IFRT 프록시 클라이언트 메서드의 타이밍 정보를 찾을 수 있습니다. 컴파일 및 실행 중에 IFRT 프록시 gRPC 서버와의 상호작용을 자세히 설명하는 이러한 이벤트는 GrpcClientSessionUserFuturesWorkQueue라는 스레드에 표시됩니다. 트레이스에서 이 스레드를 검사하면 이러한 작업의 성능에 관한 유용한 정보를 얻을 수 있습니다.

XLA 플래그

Pathways를 사용하는 경우 pathways-proxy 컨테이너에서 XLA 플래그를 설정해야 합니다.

Cluster Toolkit(gcluster job submit 명령어)을 사용하는 경우 --pathways-proxy-args를 사용하여 XLA 플래그를 설정합니다.

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

HLO 덤프

XLA 컴파일러에 제공되는 고수준 옵티마이저 (HLO) 입력을 검사하려면 --pathways-proxy-args를 사용하여 HLO를 지정된 Cloud Storage 위치에 덤프하도록 Pathways를 구성하면 됩니다.

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

다음 단계

• Pathways를 사용하여 GKE 클러스터 만들기
• Pathways를 사용한 멀티 호스트 추론
• Pathways로 일괄 워크로드 실행
• Pathways 대화형 모드
• Pathways를 사용한 복원력 있는 학습
• 클라우드에서 Pathways 문제 해결