在实际分布式机器学习项目中资源调度和异构硬件支持一直是影响迭代效率的关键因素。传统 GPU 集群虽然成熟但在特定模型训练和推理场景下TPU 的专用架构能带来显著的性能优势。然而跨多主机的 TPU 资源编排、任务分发和故障恢复往往需要复杂的自定义脚本和运维介入。Ray 2.55 版本正式宣布支持 Google Cloud TPU并通过 KubeRay 实现了多主机 TPU 切片的自动编排。这意味着开发者现在可以直接在 Kubernetes 集群中声明式地申请 TPU 资源由 KubeRay 自动完成节点发现、网络配置、资源分配和任务调度无需手动管理 TPU 工作节点的连接和生命周期。对于正在构建多智能体系统、大模型训练流水线或需要高性能张量计算的应用来说这一集成显著降低了异构硬件集群的管理门槛。本文将以一个实际的分布式训练任务为例演示如何从零搭建环境通过 KubeRay 在 GCP 上申请 TPU 资源并运行一个支持多主机切片的 Ray 任务。我们将重点解释 TPU 资源声明方式、KubeRay 配置参数、任务分发机制和常见排错点帮助你在生产环境中安全、高效地使用这一新特性。1. 理解 Ray 与 TPU 集成的核心价值1.1 为什么需要 TPU 支持Google Cloud TPUTensor Processing Unit是专为机器学习工作负载设计的张量计算加速器。与通用 GPU 相比TPU 在矩阵乘法、卷积运算等典型 AI 计算任务上具有更高的计算密度和能效。尤其是在训练超大规模模型或运行批量推理时TPU 的 pod 模式允许将计算任务分布到数百个核心上实现近乎线性的扩展。然而TPU 的使用门槛一直较高。开发者需要手动配置 TPU 节点和客户端节点的网络连通性。处理多主机环境下的数据并行和模型并行逻辑。管理 TPU 资源的创建、释放和故障转移。Ray 的分布式任务调度框架天然适合解决这些问题。通过将 TPU 资源抽象为 Ray 集群中的一种设备资源Ray 可以自动将计算任务分发到 TPU 节点并处理节点间通信、容错和负载均衡。1.2 KubeRay 在多主机编排中的作用KubeRay 是 Ray 在 Kubernetes 上的官方操作器Operator负责管理 Ray 集群的生命周期。在 2.55 版本之前KubeRay 主要支持 CPU 和 GPU 资源调度。现在通过集成 GCP TPU 驱动和资源插件KubeRay 可以自动识别集群中的 TPU 资源类型和数量。根据 Ray 集群配置创建带有 TPU 标签的工作节点组。为 Ray 任务分配具体的 TPU 设备句柄。监控 TPU 节点健康状态并自动恢复。多主机切片Multi-host Slice是指一个计算任务需要跨多个物理主机上的 TPU 核心协同工作。KubeRay 通过协调多个 Ray 工作节点使它们能够作为一个逻辑上的 TPU pod 对外提供服务。2. 环境准备与依赖配置2.1 GCP 项目与权限准备在开始之前你需要一个启用了 TPU 服务的 GCP 项目。确保当前用户或服务账号具有以下权限tpu.admin创建和管理 TPU 节点。compute.admin管理计算实例和网络配置。kubernetes.admin创建和管理 GKE 集群。可以通过以下命令检查当前项目是否已启用必要服务gcloud services list --enabled --filtername:tpu.googleapis.com gcloud services list --enabled --filtername:container.googleapis.com如果未启用使用以下命令启用gcloud services enable tpu.googleapis.com gcloud services enable container.googleapis.com2.2 创建支持 TPU 的 GKE 集群TPU 资源需要特定的 GKE 版本和节点配置。以下命令创建一个带有 TPU 驱动支持的集群gcloud container clusters create ray-tpu-cluster \ --zoneus-central1-a \ --cluster-version1.28 \ --machine-typen1-standard-4 \ --num-nodes1 \ --enable-tpu \ --scopescloud-platform关键参数说明--zone必须选择支持 TPU 的区域如us-central1-a、europe-west4-a等。--cluster-versionGKE 版本需要与 TPU 驱动兼容1.28 是当前稳定版本。--enable-tpu启用 TPU 支持自动安装必要的驱动和插件。--scopescloud-platform授予节点访问其他 GCP 服务的权限。创建完成后配置 kubectl 连接到该集群gcloud container clusters get-credentials ray-tpu-cluster --zoneus-central1-a2.3 安装 KubeRay Operator使用 Helm 安装最新版本的 KubeRayhelm repo add kuberay https://ray-project.github.io/kuberay-helm/ helm repo update helm install kuberay-operator kuberay/kuberay-operator --namespace kuberay-system --create-namespace验证 Operator 是否正常运行kubectl get pods -n kuberay-system -l app.kubernetes.io/namekuberay-operator预期看到kuberay-operatorpod 状态为Running。3. 配置 Ray 集群支持 TPU 资源3.1 创建 Ray 集群配置文件创建一个名为ray-cluster-tpu.yaml的 Kubernetes 自定义资源文件定义带有 TPU 支持的 Ray 集群apiVersion: ray.io/v1alpha1 kind: RayCluster metadata: name: ray-tpu-cluster namespace: default spec: headGroupSpec: template: spec: containers: - name: ray-head image: rayproject/ray:2.55.0 ports: - containerPort: 6379 name: gcs - containerPort: 8265 name: dashboard - containerPort: 10001 name: client resources: requests: cpu: 1 memory: 2Gi limits: cpu: 2 memory: 4Gi env: - name: RAY_DISABLE_IMPORT_WARNING value: 1 workerGroupSpecs: - replicas: 1 minReplicas: 1 maxReplicas: 4 groupName: tpu-worker-group template: spec: containers: - name: ray-worker image: rayproject/ray:2.55.0 resources: requests: cpu: 4 memory: 16Gi google.com/tpu: 8 limits: cpu: 8 memory: 32Gi google.com/tpu: 8 env: - name: RAY_DISABLE_IMPORT_WARNING value: 1 rayStartParams: num-cpus: 8 num-tpus: 8关键配置说明workerGroupSpecs[].template.spec.containers[].resources通过google.com/tpu: 8请求 8 个 TPU 核心。rayStartParamsnum-tpus: 8告诉 Ray 每个工作节点有 8 个 TPU 设备可用。镜像版本必须为rayproject/ray:2.55.0或更高这是首个正式支持 TPU 的版本。3.2 部署 Ray 集群并验证 TPU 资源应用配置文件创建 Ray 集群kubectl apply -f ray-cluster-tpu.yaml等待所有 pod 进入 Running 状态kubectl get pods -l ray.io/clusterray-tpu-cluster检查 TPU 资源是否被正确识别kubectl describe node | grep -A5 -B5 google.com/tpu应该看到节点分配了 TPU 资源。进入 Ray 工作节点容器验证 TPU 设备# 获取 worker pod 名称 WORKER_POD$(kubectl get pods -l ray.io/node-typeworker -o name | head -1) # 进入容器执行验证命令 kubectl exec -it $WORKER_POD -- python -c import ray ray.init(auto) print(TPU resources available:, ray.available_resources().get(TPU, 0)) 预期输出显示可用的 TPU 资源数量为 8。4. 编写和运行 TPU 加速的 Ray 任务4.1 创建分布式训练示例下面是一个使用 Ray 和 TPU 进行数据并行训练的完整示例。创建一个名为tpu_training.py的文件import ray import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset # 初始化 Ray自动连接集群 ray.init() # 定义简单的神经网络模型 class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.linear1 nn.Linear(10, 50) self.relu nn.ReLU() self.linear2 nn.Linear(50, 1) def forward(self, x): x self.linear1(x) x self.relu(x) x self.linear2(x) return x # 远程函数在 TPU 设备上执行训练步骤 ray.remote(num_cpus1, num_tpus1) def train_on_tpu(worker_id, model_state, data_shard, num_epochs10): # 将模型移动到 TPU device torch.device(xla) model SimpleModel() model.load_state_dict(model_state) model.to(device) # 准备数据 dataset TensorDataset(data_shard) dataloader DataLoader(dataset, batch_size32, shuffleTrue) optimizer optim.Adam(model.parameters(), lr0.001) criterion nn.MSELoss() # 训练循环 for epoch in range(num_epochs): total_loss 0.0 for batch in dataloader: inputs batch[0].to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, inputs.sum(keepdimTrue, dim1)) loss.backward() optimizer.step() total_loss loss.item() if epoch % 5 0: print(fWorker {worker_id}, Epoch {epoch}, Loss: {total_loss/len(dataloader):.4f}) # 返回更新后的模型参数 return model.cpu().state_dict() # 主训练函数 def distributed_training(): # 准备模拟数据 num_samples 1000 input_data torch.randn(num_samples, 10) # 将数据分片到不同 worker data_shards torch.chunk(input_data, chunks8, dim0) # 初始化模型 model SimpleModel() initial_state model.state_dict() # 在每个 TPU worker 上启动训练任务 futures [] for i in range(8): # 对应 8 个 TPU 核心 future train_on_tpu.remote(i, initial_state, data_shards[i]) futures.append(future) # 收集所有 worker 的训练结果 updated_states ray.get(futures) # 平均模型参数简化版的模型平均 averaged_state {} for key in initial_state.keys(): averaged_state[key] torch.stack([state[key] for state in updated_states]).mean(dim0) print(Distributed training completed) return averaged_state if __name__ __main__: final_model distributed_training() print(Training finished, model ready for use)4.2 在集群中运行训练任务将训练脚本复制到 Ray 集群中执行# 获取 head pod 名称 HEAD_POD$(kubectl get pods -l ray.io/node-typehead -o name) # 复制脚本到 head 节点 kubectl cp tpu_training.py $HEAD_POD:/tmp/ # 在 head 节点中执行训练 kubectl exec -it $HEAD_POD -- python /tmp/tpu_training.py4.3 验证 TPU 使用情况在执行过程中可以监控 TPU 资源的使用情况# 查看 Ray 集群资源状态 kubectl exec -it $HEAD_POD -- python -c import ray ray.init(auto) print(Cluster resources:, ray.available_resources()) # 查看 TPU 设备活动 kubectl exec -it $WORKER_POD -- python -c import torch_xla.core.xla_model as xm devices xm.get_xla_supported_devices() print(Available XLA devices:, devices) 5. 多主机切片配置与高级用法5.1 配置跨主机的 TPU Pod当需要超过单个主机支持的 TPU 核心数时可以配置多主机 TPU pod。修改 Ray 集群配置workerGroupSpecs: - replicas: 2 # 两个工作节点 groupName: tpu-pod-group template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: v3-8 # 指定 TPU 类型 containers: - name: ray-worker image: rayproject/ray:2.55.0 resources: requests: cpu: 8 memory: 32Gi google.com/tpu: 8 limits: cpu: 16 memory: 64Gi google.com/tpu: 8 rayStartParams: num-cpus: 8 num-tpus: 8 block: true # 等待所有节点就绪这种配置下KubeRay 会创建两个工作节点每个节点拥有 8 个 TPU 核心共同组成一个 16 核心的 TPU pod。5.2 使用 Ray Train 进行高级分布式训练对于生产级训练任务建议使用 Ray Train 抽象层。以下示例展示如何结合 TPU 使用 Ray Trainfrom ray.train import ScalingConfig from ray.train.torch import TorchTrainer import torch.optim as optim def train_loop(config): # 自动处理设备分配 from ray.train.torch import get_device device get_device() model config[model] dataset config[dataset] model.to(device) optimizer optim.Adam(model.parameters(), lr0.001) # 训练逻辑 for epoch in range(config[epochs]): for batch in dataset: inputs, targets batch inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() outputs model(inputs) loss torch.nn.functional.mse_loss(outputs, targets) loss.backward() optimizer.step() # 报告指标 from ray.train import report report({loss: loss.item(), epoch: epoch}) # 配置分布式训练 scaling_config ScalingConfig( num_workers2, # 两个工作节点 use_gpuFalse, resources_per_worker{TPU: 4} # 每个 worker 使用 4 个 TPU 核心 ) trainer TorchTrainer( train_loop, scaling_configscaling_config, train_loop_config{ model: SimpleModel(), dataset: train_dataset, epochs: 100 } ) result trainer.fit()6. 常见问题排查与性能优化6.1 TPU 资源分配失败现象Ray 任务一直处于 pending 状态日志显示 TPU 资源不可用。排查步骤检查节点 TPU 资源状态kubectl describe node | grep -A10 Capacity验证 TPU 驱动是否正常加载kubectl exec -it $WORKER_POD -- ls /dev/accel*检查 Ray 集群资源发现kubectl exec -it $HEAD_POD -- python -c import ray ray.init(auto) print(Resources:, ray.cluster_resources()) 解决方案确认 GKE 集群创建时指定了--enable-tpu。检查 TPU 配额是否充足gcloud compute regions describe us-central1 --formatvalue(quotas)。验证 Ray 工作节点配置中的google.com/tpu请求值是否超过节点容量。6.2 多主机通信超时现象跨节点的 Ray 任务通信失败出现RayTimeoutError。排查步骤检查节点网络连通性kubectl get pods -l ray.io/node-typeworker -o wide验证防火墙规则gcloud compute firewall-rules list --filtername~gke-ray-tpu-cluster检查 Ray 日志中的网络错误kubectl logs $WORKER_POD | grep -i timeout\|connection解决方案确保所有工作节点在同一个网络子网内。配置正确的防火墙规则允许 Ray 节点间通信端口 6379、10001-10020 等。增加 Ray 超时配置ray.init(auto, _redis_passwordxxx, _node_ip_addressxxx)。6.3 TPU 性能调优建议低效使用模式频繁在 CPU 和 TPU 之间移动数据。使用不适合 TPU 架构的小批量大小。未充分利用 TPU 的矩阵计算特性。优化建议数据批量优化# 不推荐批量大小过小 dataloader DataLoader(dataset, batch_size8) # 推荐使用适合 TPU 的批量大小128 dataloader DataLoader(dataset, batch_size128)设备驻留优化# 不推荐频繁设备切换 for batch in dataloader: inputs batch[0].to(xla) # 每次迭代都移动数据 # ...训练逻辑... results results.cpu() # 移回 CPU # 推荐保持数据在 TPU 上 inputs inputs.to(xla) for batch in dataloader: # 直接在 TPU 上处理 # ...训练逻辑...使用 XLA 优化编译import torch_xla.core.xla_model as xm # 在训练循环中启用 XLA 图编译 for epoch in range(epochs): for batch in dataloader: # ...前向传播、损失计算、反向传播... xm.optimizer_step(optimizer) # XLA 优化的参数更新 xm.mark_step() # 插入编译边界7. 生产环境最佳实践7.1 资源管理与成本控制TPU 资源成本较高需要精细化管理资源请求策略resources: requests: google.com/tpu: 4 # 按实际需求申请不要过度配置 limits: google.com/tpu: 4自动扩缩容配置autoscaling: minReplicas: 1 maxReplicas: 8 metrics: - type: Resource resource: name: google.com/tpu target: type: Utilization averageUtilization: 707.2 监控与日志收集建立完整的监控体系Prometheus 监控配置apiVersion: v1 kind: ConfigMap metadata: name: ray-monitoring-config data: prometheus.yml: | global: scrape_interval: 15s scrape_configs: - job_name: ray-tpu static_configs: - targets: [ray-dashboard:8265]日志收集示例# 查看 TPU 相关日志 kubectl logs -f $WORKER_POD | grep -i tpu # 监控资源使用情况 kubectl top pod -l ray.io/node-typeworker7.3 安全与权限管理最小权限原则为 Ray 服务账号分配精确的权限范围。使用 Kubernetes RBAC 限制 pod 权限。隔离生产环境和开发环境集群。网络安全配置networkPolicy: ingress: - from: - podSelector: matchLabels: ray.io/cluster: ray-tpu-cluster ports: - protocol: TCP port: 6379 - protocol: TCP port: 10001-10020Ray 2.55 对 Google Cloud TPU 的正式支持结合 KubeRay 的自动编排能力显著简化了多主机 TPU 集群的管理复杂度。在实际项目中建议从单节点 TPU 任务开始验证基础功能再逐步扩展到多主机切片场景。重点关注资源配额管理、网络配置和性能调优确保 TPU 资源得到高效利用。对于需要长期运行的生产任务建立完善的监控、告警和成本控制机制至关重要。