当 GPU 集群接近满载时,真正拖慢训练的往往不只是算力不足,而是调度失序:团队在聊天频道里争抢资源,空闲 GPU 散落在不同节点,大型训练任务即使看到“总量足够”,也迟迟无法启动。
AI21 在共享 GKE 集群中集中使用数千个 Google Cloud A3 与 A3 Ultra 实例,承载 Jamba 系列模型以及其他训练和智能体优化工作负载。引入 AI Hypercomputer、GKE 和 Kueue 后,其高优先级任务最长等待时间从 72 小时降到 12 小时,人工调度干预从每周约 20 次降到零,集群碎片率也从 15% 降到 8%。
接近 100% 利用率时,问题不只是“谁先用”
AI21 最初通过 Slack 的 #gpu-resources 频道协调容量。集群尚有余量时,这种方式足够灵活;但保留算力长期接近 100% 利用率后,每次申请都会变成跨团队谈判。
这里其实有两个不同问题。
1. 资源竞争:下一个轮到谁
训练团队、调试任务和推理服务都需要 GPU。如果没有统一队列、优先级和配额规则,团队负责人只能人工判断哪个项目更紧急。
这种方式难以审计,也容易出现两个极端:大型关键任务长期饥饿,或者高优先级任务频繁打断其他团队的工作。
2. 资源碎片:总量够,但形状不对
假设集群共有 8 张空闲 GPU,但分别以 1 + 1 + 4 + 2 的形式散落在四台机器上。对一个要求单节点 8 卡的任务来说,可用容量仍然是零。
多节点训练还会放大这个问题。如果任务的部分 Pod 先占住资源,其余 Pod 却无法调度,就可能形成“占着资源但无法工作”的僵局。人工重新分配优先级并不能解决这种装箱问题,调度系统需要理解任务的整体资源需求和集群拓扑。
为什么选择 Kueue,而不是替换 Kubernetes 调度器
AI21 评估了 Apache YuniKorn、Volcano 和 Kueue。最终选择 Kueue,核心原因不是它功能最多,而是它能够在保留标准 Kubernetes 调度链路的前提下增加批处理准入能力:
- 不需要替换 Kubernetes 核心调度组件;
- 不必重写现有 Job 规格;
- 可以把一个训练任务作为整体排队和准入;
- 能与 GKE 以及 Spot VM、Dynamic Workload Scheduler 等容量类型结合;
- 可以通过队列、配额和优先级统一管理共享 GPU 池。
这里要区分“准入”和“放置”。Kueue 判断一个工作负载何时可以获得配额并进入集群,Kubernetes 调度器负责把具体 Pod 放到节点上。两者配合后,可以避免大型任务在资源尚未凑齐时过早占住部分 GPU。
AI21 还使用了两个重要能力:
- Admission Fair Sharing(AFS):根据团队历史资源使用情况调整准入顺序,倾向于让使用较少的团队先获得机会,同时不抢占已经运行的任务。
- Topology Aware Scheduling:让准入决策考虑节点和物理拓扑,避免批准一个实际上无法按所需形状落地的任务。
AFS 的价值在于把“公平”与“强制抢占”拆开。对昂贵且运行时间很长的训练任务来说,为了公平而中断已运行作业,可能浪费大量计算;调整尚未启动任务的顺序通常更稳妥。
可以这样实践:为共享 GPU 集群建立准入队列
下面是一个最小化示例,演示如何用 Kueue 管理 Kubernetes Job。它不是 AI21 生产配置的复刻,而是一套可以据此改造的起点。
运行前请确认:
- 集群已经安装 Kueue;
- GPU 节点已经安装 NVIDIA device plugin,能够报告
nvidia.com/gpu; - 你有权限创建
ResourceFlavor、ClusterQueue和PriorityClass; - 将示例中的 CPU、内存和 GPU 配额改成集群真实容量。
kubectl create namespace ai-training --dry-run=client -o yaml | kubectl apply -f -
kubectl apply -f - <<'EOF'
apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
name: critical-training
value: 100000
preemptionPolicy: Never
globalDefault: false
description: High-priority model training without Kubernetes preemption
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
name: gpu-reserved
spec: {}
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
name: shared-gpu
spec:
namespaceSelector:
matchLabels:
kueue-access: shared-gpu
resourceGroups:
- coveredResources:
- cpu
- memory
- nvidia.com/gpu
flavors:
- name: gpu-reserved
resources:
- name: cpu
nominalQuota: 256
- name: memory
nominalQuota: 2Ti
- name: nvidia.com/gpu
nominalQuota: 32
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
name: training
namespace: ai-training
spec:
clusterQueue: shared-gpu
EOF
kubectl label namespace ai-training kueue-access=shared-gpu --overwrite
接着提交一个需要两个 Pod、每个 Pod 使用 8 张 GPU 的训练任务。suspend: true 很关键:任务先进入 Kueue 队列,只有获得整体配额后才会被解除挂起。
kubectl apply -f - <<'EOF'
apiVersion: batch/v1
kind: Job
metadata:
name: distributed-training
namespace: ai-training
labels:
kueue.x-k8s.io/queue-name: training
spec:
suspend: true
parallelism: 2
completions: 2
completionMode: Indexed
template:
spec:
priorityClassName: critical-training
restartPolicy: Never
containers:
- name: trainer
image: nvidia/cuda:12.3.2-base-ubuntu22.04
command:
- bash
- -lc
- |
echo "Starting worker ${JOB_COMPLETION_INDEX}"
nvidia-smi
sleep 300
resources:
requests:
cpu: "8"
memory: 32Gi
nvidia.com/gpu: "8"
limits:
cpu: "8"
memory: 32Gi
nvidia.com/gpu: "8"
EOF
可以用以下命令观察队列和工作负载状态:
kubectl get clusterqueue shared-gpu
kubectl get localqueue -n ai-training
kubectl get workloads -n ai-training
kubectl describe workload -n ai-training
kubectl get jobs,pods -n ai-training -w
如果任务一直处于等待状态,优先检查三件事:配额是否足够、是否存在满足单 Pod GPU 请求的节点,以及节点标签、污点和亲和性规则是否把可用节点排除在外。
生产环境通常还应按团队创建不同的 LocalQueue,统一指向共享的 ClusterQueue,而不是给每个团队永久切割一块 GPU。这样既保留组织边界,也能让暂时空闲的容量服务其他团队。若要启用公平共享和拓扑感知,应根据正在使用的 Kueue 与 GKE 版本配置对应特性;这些 API 和开关可能随版本变化,不宜直接复制未经验证的旧配置。
成功指标不应只看 GPU 利用率
AI21 的总成本没有因为这次改造而下降。其保留算力原本就被有意维持在接近 100% 利用率,优化目标不是少买 GPU,而是减少等待和协调成本。
更有意义的指标包括:
- 高优先级任务从提交到准入的 P50、P95 和最长等待时间;
- 每周人工调整队列、删除 Pod或重新分配容量的次数;
- 已分配但没有产生有效训练进度的 GPU 时间;
- 因单节点或拓扑限制无法使用的碎片比例;
- 各团队实际用量与其公平份额之间的偏差;
- 任务因抢占、失败或部分调度而浪费的 GPU 小时数。
只观察平均利用率,可能掩盖严重问题:集群可以显示 98% 忙碌,但关键训练任务仍等待数天,甚至有一部分 GPU 被无法完成调度的 Pod 占住。
落地前的检查清单
将共享 GPU 集群从人工协调迁移到自动准入时,可以按以下顺序推进:
- 先统一入口:要求训练 Job 进入明确的 Kueue 队列,避免保留绕过队列的隐蔽通道。
- 区分优先级与公平性:业务关键程度决定优先级,团队历史用量决定公平顺序,两者不要混成一个字段。
- 默认采用整体准入:尤其是多 Pod、多节点训练,避免部分任务先占住 GPU。
- 把拓扑当成容量的一部分:32 张分散的 GPU 不一定等价于四台完整的 8-GPU 节点。
- 谨慎使用抢占:中断长时间训练的代价很高,可以先采用不抢占运行中任务的公平准入。
- 保留可观测性:记录排队原因、准入时间、配额消耗和失败事件,避免把 Slack 黑箱换成调度器黑箱。
- 再连接弹性容量:当保留容量用满后,再把相同准入策略扩展到 Spot VM 或动态申请的容量。
这次实践最值得借鉴的地方,并不是某个单独的调度功能,而是把 GPU 分配从“人与人协商”变成“策略驱动的整体准入”。当集群越接近满载,队列、公平性和拓扑意识就越不是附加功能,而是让昂贵算力真正转化为实验速度的基础设施。