Pathways로 일괄 워크로드 실행

이 문서의 목적에 따라 배치 워크로드는 완료될 때까지 실행되고 Pathways 컨트롤러 구성요소 (IFRT 프록시 서버 및 Pathways 리소스 관리자)와 함께 특히 Pathways 클러스터와 동일한 GKE 클러스터 내에 배포되는 JAX 워크로드로 정의됩니다. JAX 워크로드가 완료되면 Pathways 클러스터 구성요소가 종료됩니다. 이 가이드에서는 JAX 학습 워크로드를 사용하여 이를 보여줍니다.

시작하기 전에

다음 사항이 필요합니다.

Maxtext를 사용하여 학습 이미지 빌드

MaxText는 Google에서 개발한 오픈소스 대규모 언어 모델 (LLM) 프로젝트입니다. JAX로 작성되었으며 Google Cloud TPU 및 GPU에서 효율적으로 실행되도록 고성능 및 확장성을 갖추도록 설계되었습니다.

OSS GitHub 저장소의 최신 안정화 JAX 버전을 사용하여 MaxText Docker 이미지를 빌드하려면 다음 명령어를 실행하세요.

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext/dependencies/scripts
gcloud config set project PROJECT_ID
bash ./docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
bash ./docker_upload_runner.sh CLOUD_IMAGE_NAME=USER_runner # This script needs bash version >= 4.2 to execute.

이 명령어는 MaxText Kubernetes 이미지를 gcr.io/$PROJECT_ID/${USER}_runner에 푸시합니다. 이 Docker 이미지를 사용하여 Pathways 백엔드를 통해 TPU에서 학습을 실행할 수 있습니다.

Cluster Toolkit으로 일괄 워크로드 실행

gcluster job submit 명령어를 사용하여 사전 빌드된 MaxText Docker 이미지를 제출합니다.

gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m MaxText.train /deps/src/MaxText/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

작업 제출 옵션에 대한 자세한 내용은 Cluster Toolkit 작업 제출 가이드를 참고하세요.

다음을 바꿉니다.

  • WORKLOAD: 워크로드를 식별하는 고유한 이름입니다. DNS 라벨 제한으로 인해 22자 이하여야 합니다.
  • CLUSTER: GKE 클러스터의 이름
  • WORKLOAD_NODEPOOL_COUNT: TPU 슬라이스 노드 풀 수
  • COMPUTE_TYPE: TPU 머신 유형 (예: ct6e-standard-4t). 각 TPU 버전에서 지원되는 TPU 유형에 대한 자세한 내용은 TPU 버전을 참고하세요.
  • TOPOLOGY: TPU 배치 토폴로지 (예: 2x4)
  • PROJECT_ID: Google Cloud 프로젝트 ID
  • ZONE: 워크로드를 실행할 계획인 영역
  • USER: Google Cloud 사용자 ID
  • BUCKET_NAME: 출력을 위한 Cloud Storage 버킷
  • RUN_NAME: 워크플로 실행을 식별하는 사용자 할당 이름

gcluster job logs 명령어를 사용하여 워크로드의 진행 상황을 확인합니다.

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --main-only=false

completed step: 1, seconds: 0.484, TFLOP/s/device: 87.349, Tokens/s/device: 2117.382, total_weights: 2945, loss: 10.888
completed step: 2, seconds: 0.407, TFLOP/s/device: 103.699, Tokens/s/device: 2513.735, total_weights: 3253, loss: 9.697
completed step: 3, seconds: 0.248, TFLOP/s/device: 170.300, Tokens/s/device: 4128.167, total_weights: 3154, loss: 9.641
completed step: 4, seconds: 0.216, TFLOP/s/device: 195.122, Tokens/s/device: 4729.880, total_weights: 3119, loss: 9.547
completed step: 5, seconds: 0.272, TFLOP/s/device: 155.298, Tokens/s/device: 3764.512, total_weights: 2837, loss: 10.179
completed step: 6, seconds: 0.472, TFLOP/s/device: 89.489, Tokens/s/device: 2169.266, total_weights: 3069, loss: 9.776

워크로드가 완료되기 전에 취소하려면 gcluster job cancel 명령어를 사용합니다.

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

다음 단계