使用 Pathways 运行批量工作负载

本指南使用一个 JAX 训练工作负载，该工作负载在与 Pathways 控制器组件相同的集群中运行至完成，以演示如何运行批处理工作负载。

您可以在 Google Kubernetes Engine (GKE) 集群或具有 Slurm 的 Cluster Director 集群上运行批处理工作负载：

  • 在 GKE 上，JAX 工作负载完成后，Pathways 集群组件会终止。
  • 在 Cluster Director 上，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 训练工作负载。您可以在服务作业之后立即提交该作业：它最多会等待一小时，直到服务准备就绪。

    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 切片训练 8 层变体，而不是完整的 32 层。

    跟踪工作负载的进度（将 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 虚拟机是静态节点，因此它们会继续运行，并可供其他 Slurm 作业使用。

如果您遇到问题，请参阅排查 Pathways on Cloud 问题。

后续步骤