使用 Pathways 运行批量工作负载

在本文档中，批处理工作负载是指执行到完成的 JAX 工作负载，这些工作负载部署在与 Pathways 集群相同的 GKE 集群中，特别是与 Pathways 控制器组件（IFRT 代理服务器和 Pathways 资源管理器）一起部署。 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
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

此命令会将 MaxText Kubernetes 映像推送到 gcr.io/$PROJECT_ID/${USER}_runner。 您可以使用此 Docker 映像，通过 Pathways 后端在 TPU 上运行训练。

使用 Cluster Toolkit 运行批量工作负载

如果 gcluster 二进制文件不在 PATH 环境变量中，请将包含 gcluster 的目录添加到 PATH：

export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

将 CLUSTER_TOOLKIT_PATH 替换为包含 gcluster 二进制文件的目录的绝对路径，例如 $HOME/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=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"

如需详细了解作业提交选项，请参阅 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
  • LOCATION：GKE 集群所在的区域或可用区
  • USER：您的 Google Cloud 用户 ID
  • BUCKET_NAME：用于输出的 Cloud Storage 存储桶
  • RUN_NAME：用户分配的名称，用于标识工作流运行

使用 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

如需在工作负载完成之前将其取消，请使用 gcluster job cancel 命令：

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

后续步骤