跳到主要内容

训练集群管理

高效的 AI 训练集群管理是大规模模型训练的基础。本章介绍如何使用 Kubernetes 等工具管理 AI 训练集群。

AI 集群的特点​

AI 训练集群与传统应用集群有显著差异：

特性传统应用AI 训练
资源类型CPU、内存为主GPU、显存为主
任务时长秒到小时级小时到天级
通信模式低频网络请求高频 GPU 通信
故障影响单实例影响小全局同步失败
资源需求相对稳定波动较大

这些特点决定了 AI 集群需要专门的管理策略。

Kubernetes 基础​

Kubernetes（K8s）已成为容器编排的事实标准，也被广泛用于 AI 工作负载管理。

核心概念​

Pod：最小部署单元，可以包含一个或多个容器

apiVersion: v1
kind: Pod
metadata:
  name: gpu-pod
spec:
  containers:
  - name: training
    image: pytorch/pytorch:latest
    resources:
      limits:
        nvidia.com/gpu: 2
    command: ["python", "train.py"]

Deployment：管理 Pod 的副本和更新策略

Job：一次性任务，完成后终止

CronJob：定时执行的任务

GPU 资源管理​

Kubernetes 通过设备插件管理 GPU 资源：

resources:
  limits:
    nvidia.com/gpu: 4  # 请求 4 张 GPU

安装 NVIDIA 设备插件：

kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.0/nvidia-device-plugin.yml

节点选择​

使用标签和选择器将任务调度到合适的节点：

apiVersion: v1
kind: Pod
metadata:
  name: training-pod
spec:
  nodeSelector:
    accelerator: nvidia-h100
  containers:
  - name: training
    image: training:latest

Kubeflow​

Kubeflow 是 Kubernetes 上的机器学习工具包，提供完整的 ML 工作流支持。

核心组件​

组件功能
Kubeflow PipelinesML 工作流编排
Training Operators分布式训练支持
Katib超参数调优
KServe模型服务
Notebook ServersJupyter 环境

分布式训练示例​

使用 PyTorchJob 运行分布式训练：

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-training
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch-training:latest
            resources:
              limits:
                nvidia.com/gpu: 8
    Worker:
      replicas: 3
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: pytorch-training:latest
            resources:
              limits:
                nvidia.com/gpu: 8

启动训练​

kubectl apply -f pytorchjob.yaml

# 查看状态
kubectl get pytorchjob

# 查看日志
kubectl logs -f pytorch-training-master-0

Volcano 调度器​

Volcano 是面向高性能计算场景的 Kubernetes 调度器，特别适合 AI 训练。

核心特性​

  • Gang Scheduling：组调度，确保所有任务同时启动
  • Fair Share：公平调度，防止资源被独占
  • Queue Management：队列管理，支持优先级
  • Preemption：抢占机制，高优先级任务优先

Gang Scheduling 示例​

apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: training-job
spec:
  minAvailable: 4  # 最少需要 4 个 Pod 同时运行
  schedulerName: volcano
  tasks:
  - replicas: 4
    name: worker
    template:
      spec:
        containers:
        - name: pytorch
          image: pytorch:latest
          resources:
            limits:
              nvidia.com/gpu: 8

为什么需要 Gang Scheduling​

分布式训练要求所有节点同时启动：

没有 Gang Scheduling：
节点 1: [运行中...] [等待其他节点] [超时失败]
节点 2: [排队中...]
节点 3: [排队中...]
节点 4: [排队中...]

有 Gang Scheduling：
节点 1: [等待]
节点 2: [等待]
节点 3: [等待]
节点 4: [等待]
        ↓ 资源就绪后同时启动
节点 1: [运行中...]
节点 2: [运行中...]
节点 3: [运行中...]
节点 4: [运行中...]

资源配额与限制​

ResourceQuota​

限制命名空间的资源总量：

apiVersion: v1
kind: ResourceQuota
metadata:
  name: gpu-quota
  namespace: ai-team
spec:
  hard:
    requests.nvidia.com/gpu: 32
    limits.nvidia.com/gpu: 32

LimitRange​

限制单个 Pod 的资源范围：

apiVersion: v1
kind: LimitRange
metadata:
  name: gpu-limits
spec:
  limits:
  - type: Container
    max:
      nvidia.com/gpu: 8
    min:
      nvidia.com/gpu: 1

存储管理​

数据集存储​

训练数据通常存储在共享文件系统中：

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: dataset-pvc
spec:
  accessModes:
  - ReadOnlyMany
  resources:
    requests:
      storage: 10Ti
  storageClassName: nfs

在 Pod 中挂载​

spec:
  volumes:
  - name: dataset
    persistentVolumeClaim:
      claimName: dataset-pvc
  containers:
  - name: training
    volumeMounts:
    - name: dataset
      mountPath: /data

存储性能优化​

  • 缓存层：使用 Alluxio 等工具缓存热点数据
  • 本地存储：将数据预加载到节点本地 SSD
  • 数据预取：训练前提前拉取数据

监控与告警​

GPU 监控​

使用 DCGM Exporter 暴露 GPU 指标：

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: dcgm-exporter
spec:
  selector:
    matchLabels:
      app: dcgm-exporter
  template:
    metadata:
      labels:
        app: dcgm-exporter
    spec:
      containers:
      - name: dcgm-exporter
        image: nvcr.io/nvidia/k8s/dcgm-exporter:3.1.7-3.1.4-ubuntu20.04
        resources:
          limits:
            nvidia.com/gpu: 1

关键监控指标​

指标说明告警阈值
GPU 利用率GPU 计算资源使用率< 50% 持续 10 分钟
显存使用率GPU 显存使用率> 95%
GPU 温度GPU 核心温度> 85°C
训练进度Epoch/Step 进度长时间无更新
任务状态Job 运行状态Failed

Prometheus 规则示例​

groups:
- name: gpu-alerts
  rules:
  - alert: LowGPUUtilization
    expr: dcgm_gpu_utilization < 50
    for: 10m
    labels:
      severity: warning
    annotations:
      summary: "GPU 利用率过低"
      description: "GPU {{ $labels.gpu }} 利用率低于 50%"

故障处理​

常见故障类型​

故障类型症状处理方法
GPU 故障CUDA 错误、训练中断隔离节点、更换 GPU
网络故障NCCL 超时、通信失败检查网络配置、重启任务
存储故障I/O 错误、数据读取失败检查存储服务、恢复数据
OOM显存不足减小批次、启用梯度检查点

检查点机制​

定期保存训练状态，支持故障恢复：

# PyTorch 检查点保存
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, f'checkpoint_{epoch}.pt')

# 恢复训练
checkpoint = torch.load('checkpoint_10.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']

自动重启策略​

spec:
  restartPolicy: OnFailure  # 失败时重启
  backoffLimit: 3           # 最大重试次数
  activeDeadlineSeconds: 86400  # 最大运行时间

多租户管理​

命名空间隔离​

apiVersion: v1
kind: Namespace
metadata:
  name: team-a
  labels:
    team: a
---
apiVersion: v1
kind: Namespace
metadata:
  name: team-b
  labels:
    team: b

RBAC 权限控制​

apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
  name: training-role
  namespace: team-a
rules:
- apiGroups: ["kubeflow.org"]
  resources: ["pytorchjobs"]
  verbs: ["create", "get", "list", "delete"]
- apiGroups: [""]
  resources: ["pods", "pods/log"]
  verbs: ["get", "list"]

最佳实践​

资源规划​

  1. 预留资源：为系统组件预留 10-15% 资源
  2. 资源池化：按团队或项目划分资源池
  3. 弹性伸缩：结合云资源实现弹性扩容

任务管理​

  1. 优先级队列：紧急任务优先调度
  2. 任务预估：预估资源需求，避免资源浪费
  3. 超时机制：设置合理的超时时间

成本优化​

  1. Spot 实例：使用竞价实例降低成本
  2. 资源回收：及时释放空闲资源
  3. 共享存储：避免数据冗余存储

小结​

Kubernetes 已成为 AI 集群管理的主流平台，配合 Kubeflow、Volcano 等工具，可以实现高效的资源调度和任务管理。理解 AI 工作负载的特点，选择合适的调度策略和监控方案，是构建稳定可靠训练集群的关键。

参考资料​