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
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/${USER}_runner에 푸시합니다. 이 Docker 이미지를 사용하여 Pathways 백엔드를 통해 TPU에서 학습을 실행할 수 있습니다.

PathwaysJob API를 사용하여 일괄 워크로드 실행

다음 매니페스트는 Pathways 구성요소를 배포하고 PathwaysJob API를 사용하여 MaxText 워크로드를 실행합니다. 워크로드는 main 컨테이너에 캡슐화되고 train.py를 실행합니다.

다음 YAML을 pathways-job-batch-training.yaml이라는 파일에 복사하고 수정 가능한 값을 업데이트합니다.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  workers:
  - type: TPU_MACHINE_TYPE
    topology: TOPOLOGY
    numSlices: WORKLOAD_NODEPOOL_COUNT
  pathwaysDir: gs://BUCKET_NAME
  controller:
    deploymentMode: default
    template:
      spec:
        containers:
        - name: main
          image: gcr.io/PROJECT/USER_runner
          command:
          - bash
          - -c
          - |
            python3 -m MaxText.train /deps/src/MaxText/configs/base.yml \
            base_output_directory=gs://BUCKET_NAME \
            run_name=RUN_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

다음을 바꿉니다.

  • USER : Google Cloud 사용자 ID
  • MAX_RESTARTS : 작업을 다시 시작할 수 있는 최대 횟수입니다.
  • TPU_MACHINE_TYPE : TPU 머신 유형
  • TOPOLOGY : TPU v4 이상 토폴로지 TPU 버전 및 지원되는 토폴로지에 관한 자세한 내용은 TPU 버전을 참고하세요.
  • WORKLOAD_NODEPOOL_COUNT : Pathways 워크로드에서 사용되는 노드 풀 수
  • BUCKET_NAME : 임시 파일을 저장하는 Cloud Storage 버킷
  • PROJECT : Google Cloud 프로젝트 ID
  • RUN_NAME : 워크플로 실행을 식별하기 위해 사용자가 할당한 이름입니다.

다음과 같이 PathwaysJob YAML을 배포할 수 있습니다.

kubectl apply -f pathways-job-batch-training.yaml

이전 명령어로 생성된 PathwaysJob 인스턴스를 보려면 다음을 사용합니다.

kubectl get pathwaysjob

출력은 다음과 같습니다.

NAME             AGE
pathways-trial   9s

PathwaysJob 인스턴스의 속성을 수정하려면 PathwaysJob 인스턴스를 삭제하고 YAML을 수정한 후 이를 적용하여 새 PathwaysJob 인스턴스를 만듭니다.

컨테이너 이름 필터에서 main를 선택하여 JAX 컨테이너의 로그 탐색기로 이동하면 워크로드의 진행 상황을 확인할 수 있습니다.

학습이 진행 중임을 나타내는 다음과 같은 로그가 표시됩니다. 워크로드는 30단계 후에 완료됩니다.

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

PathwaysJob 인스턴스를 삭제하려면 다음 명령어를 사용하면 됩니다.

kubectl delete -f pathways-job-batch-training.yaml

XPK를 사용하여 일괄 워크로드 실행

이제 이전에 사용한 것과 동일한 명령어를 사용하여 XPK로 사전 빌드된 Maxtext Docker 이미지를 제출할 수 있습니다.

xpk workload create-pathways \
--workload=WORKLOAD \
--cluster=CLUSTER \
--num-slices=WORKLOAD_NODEPOOL_COUNT \
--tpu-type=TPU_TYPE \
--project=PROJECT \
--zone=ZONE \
--docker-image='gcr.io/PROJECT/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"

다음을 바꿉니다.

  • WORKLOAD: 워크로드를 식별하는 고유한 이름
  • CLUSTER: GKE 클러스터의 이름
  • WORKLOAD_NODEPOOL_COUNT : 작업을 다시 시작할 수 있는 최대 횟수입니다.
  • TPU_TYPE: TPU 유형은 만들려는 Cloud TPU의 버전과 크기를 지정합니다. TPU 버전마다 지원되는 TPU 유형에 대한 자세한 내용은 TPU 버전을 참고하세요.
  • PROJECT : Google Cloud 프로젝트 ID
  • ZONE: 워크로드를 실행할 영역
  • USER : Google Cloud 사용자 ID
  • RUN_NAME : 워크플로 실행을 식별하기 위해 사용자가 할당한 이름입니다.

다음과 같은 출력이 표시되어야 합니다.

[XPK] Follow your Pathways workload and other resources here : https://console.cloud.google.com/logs/query;query=resource.type%3D"k8s_container"%0Aresource.labels.project_id%3D"<project-name>"%0Aresource.labels.location%3D"<your-zone>"%0Aresource.labels.cluster_name%3D"<your-cluster-name>"%0Aresource.labels.pod_name:"<your-pod-name>"%0Aseverity>%3DDEFAULT

이전 XPK 명령어의 출력에 있는 링크를 사용하여 워크로드의 진행 상황을 추적합니다. 컨테이너 이름 필터에서 jax-tpu를 선택하여 JAX 컨테이너의 로그를 필터링할 수 있습니다.

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

지정된 단계 수가 지나면 워크로드가 완료됩니다. 조기에 종료하려면 다음 명령어를 사용합니다.

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

다음 단계