使用 Pathways 执行多主机推理

多主机推理是一种运行模型推理的方法，可将模型分布在多个加速器主机上。这样一来，您就可以推理无法在单个主机上运行的大型模型。Pathways 可用于批量和实时多主机推理应用场景。

准备工作

请确保您已备妥：

• 创建了使用 Trillium 芯片 (v6e-16) 的 GKE 集群。
• 已安装的 Kubernetes 工具
• 已启用 Google Kubernetes Engine API

使用 JetStream 运行批量推理

JetStream 是一种吞吐量和内存优化引擎，用于在 XLA 设备（主要是张量处理单元 [TPU]，以 JAX 编写）上进行大语言模型 (LLM) 推理。

您可以使用预构建的 JetStream Docker 映像来运行批处理推理工作负载，如以下 YAML 所示。此容器基于 OSS JetStream 项目构建。 如需详细了解 MaxText-JetStream 标志，请参阅 JetStream MaxText 服务器标志。以下示例使用 Trillium 芯片 (v6e-16) 加载 Llama3.1-405b int8 检查点并对其执行推理。此示例假定您已拥有一个 GKE 集群，其中包含至少一个 v6e-16 节点池。

启动模型服务器和 Pathways

1. 获取集群的凭据并将其添加到本地 kubectl 上下文中。

         gcloud container clusters get-credentials $CLUSTER \
         --zone=$ZONE \
         --project=$PROJECT \
         && kubectl config set-context --current --namespace=default
       

2. 部署 LeaderWorkerSet (LWS) API。

         VERSION=v0.4.0
         kubectl apply --server-side -f "https://github.com/kubernetes-sigs/lws/releases/download/${VERSION}/manifests.yaml"
       

3. 将以下 YAML 复制并粘贴到名为 pathways-job.yaml 的文件中： 此 YAML 已针对 v6e-16 切片形状进行了优化。如需详细了解如何将 Meta 检查点转换为与 JAX 兼容的检查点，请按照创建推理检查点中的检查点创建指南进行操作。 例如，此处提供了 Llama3.1-405B 的相关说明：Llama3.1-405B 的检查点转换。

       apiVersion: leaderworkerset.x-k8s.io/v1
       kind: LeaderWorkerSet
       metadata:
         name: jetstream-pathways
         annotations:
           leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
       spec:
         replicas: 1
         leaderWorkerTemplate:
           leaderTemplate:
             metadata:
               labels:
                 app: jetstream-pathways
             spec:
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 4x4
               tolerations:
               - key: "google.com/tpu"
                 operator: "Exists"
                 effect: "NoSchedule"
               containers:
               - name: pathways-proxy
                 image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
                 args:
                 - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
                 - --server_port=38681
                 - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
                 imagePullPolicy: Always
                 ports:
                 - containerPort: 38681
               - name: pathways-rm
                 env:
                 - name: HOST_ADDRESS
                   value: "$(LWS_LEADER_ADDRESS)"
                 - name: TPU_SKIP_MDS_QUERY
                   value: "true"
                 image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
                 args:
                 - --server_port=38677
                 - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
                 - --node_type=resource_manager
                 - --instance_count=1
                 - --instance_type=tpuv6e:TPU_TOPOLOGY # Example: 4x4
                 - --temporary_flags_for_debugging=temporary_flag_for_debugging_worker_expected_tpu_chip_config=megachip_tccontrol
                 imagePullPolicy: Always
                 ports:
                 - containerPort: 38677
               - name: jax-tpu
                 image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:jax-0.8.0
                 # Optimized settings used to serve Llama3.1-405b.
                 args:
                 - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
                 - model_name=llama3.1-405b
                 - load_parameters_path=CHECKPOINT_PATH
                 - max_prefill_predict_length=1024
                 - max_target_length=2048
                 - async_checkpointing=false
                 - steps=1
                 - ici_fsdp_parallelism=1
                 - ici_autoregressive_parallelism=2
                 - ici_tensor_parallelism=8
                 - scan_layers=false
                 - weight_dtype=bfloat16
                 - per_device_batch_size=10
                 - enable_single_controller=true
                 - quantization=int8
                 - quantize_kvcache=true
                 - checkpoint_is_quantized=true
                 - enable_model_warmup=true
                 imagePullPolicy: Always
                 ports:
                 - containerPort: 9000
                 startupProbe:
                   httpGet:
                     path: /healthcheck
                     port: 8000
                     scheme: HTTP
                   periodSeconds: 1
                   initialDelaySeconds: 900
                   failureThreshold: 10000
                 livenessProbe:
                   httpGet:
                     path: /healthcheck
                     port: 8000
                     scheme: HTTP
                   periodSeconds: 60
                   failureThreshold: 10
                 readinessProbe:
                   httpGet:
                     path: /healthcheck
                     port: 8000
                     scheme: HTTP
                   periodSeconds: 60
                   failureThreshold: 10
               - name: jetstream-http
                 image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.4
                 imagePullPolicy: Always
                 ports:
                 - containerPort: 8000
           # The size variable defines the number of worker nodes to be created.
           # It must be equal to the number of hosts + 1 (for the leader node).
           size: 5
           workerTemplate:
             spec:
               nodeSelector:
                 cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR_TYPE # Example: tpu-v6e-slice
                 cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY # Example: 4x4
               tolerations:
               - key: "google.com/tpu"
                 operator: "Exists"
                 effect: "NoSchedule"
               containers:
               - name: worker
                 args:
                 - --server_port=38679
                 - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
                 - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
                 image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
                 imagePullPolicy: Always
                 ports:
                 - containerPort: 38679
                 resources:
                   limits:
                     google.com/tpu: "4"
       ---
       apiVersion: v1
       kind: Service
       metadata:
         name: jetstream-svc
       spec:
         selector:
           app: jetstream-pathways
         ports:
         - protocol: TCP
           name: jetstream-http
           port: 8000
           targetPort: 8000
       

   替换以下内容：
   • TPU_ACCELERATOR_TYPE：TPU 加速器类型。例如 tpu-v6e-slice。
   • TPU_TOPOLOGY：TPU 拓扑。例如 2x4。
   • CHECKPOINT_PATH：检查点的 Cloud Storage 路径。
   应用此 YAML。等待 LeaderWorkerSet 被调度。设置完成后，模型服务器可能需要一些时间才能恢复检查点。 对于 405B 模型，此过程大约需要 7 分钟。
4. 查看 Kubernetes 日志，了解 JetStream 模型服务器是否已准备就绪： 在之前的 YAML 中，工作负载名为 `jetstream-pathways`，而 `0` 是头节点。

         kubectl logs -f jetstream-pathways-0 -c