Pathways로 일괄 워크로드 실행

이 문서의 목적에 맞게 일괄 워크로드는 완료될 때까지 실행되고 Pathways 클러스터와 동일한 GKE 클러스터 내, 특히 Pathways 컨트롤러 구성요소 (IFRT 프록시 서버 및 Pathways 리소스 관리자)와 함께 배포되는 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에서 학습을 실행할 수 있습니다.

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 : you Google Cloud 프로젝트 ID
  • ZONE: 워크로드를 실행할 영역
  • USER : 사용자 ID Google Cloud
  • 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

다음 단계