Pathways로 일괄 워크로드 실행

이 가이드에서는 Pathways 컨트롤러 구성요소와 동일한 클러스터에서 완료될 때까지 실행되는 JAX 학습 워크로드를 사용하여 일괄 워크로드를 실행하는 방법을 보여줍니다.

Google Kubernetes Engine(GKE) 클러스터 또는 Slurm이 있는 Cluster Director 클러스터에서 일괄 워크로드를 실행할 수 있습니다.

  • GKE에서 JAX 워크로드가 완료되면 Pathways 클러스터 구성요소가 종료됩니다.
  • 클러스터 디렉터에서 Pathways 컨트롤러 구성요소는 별도의 Slurm 작업으로 실행되므로 JAX 워크로드가 완료된 후에도 중지할 때까지 계속 실행됩니다.

이 페이지에는 두 플랫폼에 대한 안내가 모두 포함되어 있습니다. 각 섹션에서 GKE 또는 Cluster Director 탭을 선택합니다.

시작하기 전에

시작하기 전에

GKE

다음 기본 요건을 완료합니다.

Cluster Director

다음 기본 요건을 완료합니다.

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

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

GKE

오픈소스 GitHub 저장소에서 최신 버전의 안정적인 JAX를 사용하여 MaxText Docker 이미지를 빌드하려면 다음 명령어를 실행합니다.

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

다음을 바꿉니다.

  • PROJECT_ID:Google Cloud 프로젝트 ID입니다.
  • USER:Google Cloud 사용자 ID

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

Cluster Director

학습 이미지를 빌드할 필요가 없습니다. setup_env.sh가 /home/pathways의 Pathways 환경에 MaxText를 설치했고 run_maxtext.sh가 이를 실행합니다.

일괄 워크로드 실행

GKE

  1. gcluster 바이너리가 PATH 환경 변수에 없으면 gcluster가 포함된 디렉터리를 PATH에 추가합니다.

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    CLUSTER_TOOLKIT_PATH를 gcluster 바이너리가 포함된 디렉터리의 절대 경로(예: $HOME/cluster-toolkit)로 바꿉니다.

  2. 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=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.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"

    다음을 바꿉니다.

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

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

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

    gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

    출력은 다음과 비슷합니다.

    completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

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

    gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION

Cluster Director

~/pathways에서 클러스터의 로그인 노드에 다음 명령어를 실행합니다. 스크립트는 서로를 찾고 해당 디렉터리에 로그를 작성합니다.

  1. Pathways 서비스를 시작합니다.

    cd ~/pathways
SLICE_SHAPE=TOPOLOGY \
    MIN_SLICES=MIN_SLICES \
    MAX_SLICES=MAX_SLICES \
    GCS_BUCKET=gs://BUCKET_NAME \
    sbatch start_pathways.sh

    다음을 바꿉니다.

    • TOPOLOGY: 각 슬라이스의 TPU 토폴로지입니다. TPU 노드 세트의 acceleratorTopology와 일치해야 합니다(예: 2x4).
    • MIN_SLICES: 학습을 시작하고 계속하는 데 필요한 최소 실행 슬라이스 수입니다.
    • MAX_SLICES: 사용할 최대 슬라이스 수입니다. 일반적으로 노드 집합의 슬라이스 수입니다.
    • BUCKET_NAME: Pathways 스크래치 데이터 및 체크포인트의 Cloud Storage 버킷입니다.

    TPU 파티션이 클러스터의 기본 파티션이 아닌 경우 PARTITION을 해당 이름으로 설정합니다.

    로그 pathways-service-JOB_ID.out(여기서 JOB_ID은 Slurm 작업 ID)에 다음과 비슷한 출력이 표시되면 서비스가 준비된 것입니다.

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
 Running TPU Slices : 2 / 2 (min=2)

  2. MaxText 학습 워크로드를 제출합니다. 서비스 작업 직후에 제출할 수 있습니다. 서비스가 준비될 때까지 최대 1시간 동안 대기합니다.

    sbatch run_maxtext.sh

    MaxText 작업은 Pathways 서비스에 실행 중인 슬라이스가 MIN_SLICES개 이상 있을 때까지 기다린 후 학습을 시작합니다. 워크로드는 합성 데이터에서 Llama 3.1 8B를 학습시킵니다. 슬라이스 설정과 버킷은 서비스에서 가져오므로 이 작업에서 설정하지 마세요. 다음과 같은 선택적 변수를 설정할 수 있습니다.

    • STEPS: 학습 단계 수입니다. 기본값은 30입니다.
    • CHECKPOINT_PERIOD: 체크포인트 간 단계 수입니다. 기본값은 5입니다.
    • BASE_OUTPUT_DIR: MaxText가 체크포인트를 작성하는 Cloud Storage 경로입니다. 기본값은 gs://BUCKET_NAME/maxtext-runs입니다.
    • REDUCED_LAYERS=1: 2x2 슬라이스에 대해 전체 32개 레이어 대신 8개 레이어 변형을 학습시킵니다.

    워크로드 진행 상황을 따릅니다(JOB_ID를 MaxText 작업의 Slurm 작업 ID로 바꿈). MaxText는 로그를 표준 오류에 기록합니다.

    tail -f maxtext-llama31-JOB_ID.err

    각 학습 단계에서는 다음과 유사한 줄을 출력합니다.

    completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238

    서비스를 다시 시작하지 않고 나중에 run_maxtext.sh를 다시 제출할 수 있습니다.

  3. 완료되면 워크로드를 중지한 후 Pathways 서비스를 중지합니다.

    scancel -u ${USER} -n maxtext-llama31
scancel -u ${USER} -n pathways-service

    서비스를 취소하면 Pathways 작업자 작업도 취소됩니다. TPU VM은 정적 노드이므로 계속 실행되고 다른 Slurm 작업에 사용할 수 있습니다.

문제가 발생하면 Cloud에서 Pathways 문제 해결을 참조하세요.

다음 단계