03 - K8s 上的分布式训练

分布式训练概述

当单卡 GPU 无法满足训练需求时，需要多卡/多机分布式训练：

┌─────────────────────────────────────────────────────────────────┐
│                    分布式训练策略                                 │
│                                                                 │
│  ┌─────────────────────────┐  ┌─────────────────────────┐      │
│  │   数据并行              │  │   模型并行              │      │
│  │   (Data Parallelism)    │  │   (Model Parallelism)   │      │
│  │                         │  │                         │      │
│  │  每个 GPU 有完整模型      │  │  模型分片到不同 GPU      │      │
│  │  数据分片到不同 GPU       │  │  适用于超大模型          │      │
│  │  最常用                  │  │  如 GPT/LLM             │      │
│  └─────────────────────────┘  └─────────────────────────┘      │
│                                                                 │
│  数据并行示例：                                                  │
│  ┌───────┐ ┌───────┐ ┌───────┐ ┌───────┐                      │
│  │Worker0│ │Worker1│ │Worker2│ │Worker3│                      │
│  │Model  │ │Model  │ │Model  │ │Model  │                      │
│  │Data/4 │ │Data/4 │ │Data/4 │ │Data/4 │                      │
│  │GPU 0  │ │GPU 1  │ │GPU 2  │ │GPU 3  │                      │
│  └───┬───┘ └───┬───┘ └───┬───┘ └───┬───┘                      │
│      └─────────┴─────────┴─────────┘                           │
│              AllReduce (梯度同步)                                │
└─────────────────────────────────────────────────────────────────┘

---

Kubeflow Training Operator

Kubeflow 提供了 Training Operator，支持在 K8s 上运行分布式训练任务：

CRD	框架
PyTorchJob	PyTorch DDP
TFJob	TensorFlow
MPIJob	Horovod / MPI
XGBoostJob	XGBoost

注意：Kubeflow Training Operator 正在开发 v2 API（trainer.kubeflow.org/v2alpha1），使用 TrainingRuntime 替代 PyTorchJob。目前 v1 API（kubeflow.org/v1）仍是稳定版本，生产环境推荐使用。

安装 Training Operator

# 安装
kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone"

# 验证
kubectl get pods -n kubeflow

---

PyTorchJob 示例

单机多卡

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-ddp-single-node
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
          - name: pytorch
            image: my-training:latest
            command:
            - python
            - -m
            - torch.distributed.run
            - --nproc_per_node=4
            - train.py
            - --epochs=10
            - --batch-size=256
            resources:
              limits:
                nvidia.com/gpu: 4
            volumeMounts:
            - name: data
              mountPath: /data
            - name: output
              mountPath: /output
          volumes:
          - name: data
            persistentVolumeClaim:
              claimName: training-data     # NAS 共享存储
          - name: output
            persistentVolumeClaim:
              claimName: training-output

多机多卡

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: pytorch-ddp-multi-node
spec:
  elasticPolicy:
    rdzvBackend: c10d
    minReplicas: 1
    maxReplicas: 4
  pytorchReplicaSpecs:
    Worker:
      replicas: 4                    # 4 个 Worker 节点
      restartPolicy: OnFailure
      template:
        spec:
          tolerations:
          - key: "nvidia.com/gpu"
            operator: "Exists"
            effect: "NoSchedule"
          containers:
          - name: pytorch
            image: my-training:latest
            command:
            - python
            - train.py
            - --epochs=100
            - --batch-size=1024
            # Training Operator 自动注入 MASTER_ADDR, MASTER_PORT,
            # WORLD_SIZE, RANK 等环境变量，无需手动指定 --nnodes
            env:
            - name: NCCL_DEBUG
              value: INFO
            resources:
              limits:
                nvidia.com/gpu: 4
                rdma/hca: 1          # RDMA 网络（高性能通信）
            volumeMounts:
            - name: data
              mountPath: /data
            - name: shm                # 共享内存，NCCL 需要
              mountPath: /dev/shm
          volumes:
          - name: data
            persistentVolumeClaim:
              claimName: training-data
          - name: shm
            emptyDir:
              medium: Memory
              sizeLimit: 16Gi

---

数据管理

训练数据存储方案

方案	类型	访问模式	适用场景
阿里云 NAS	文件存储	RWX（多读多写）	共享训练数据、checkpoint
阿里云 OSS	对象存储	通过 CSI 或 SDK	海量数据集
云盘	块存储	RWO（单机读写）	单机临时缓存

NAS 存储配置

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: training-data
spec:
  accessModes:
  - ReadWriteMany               # 多机共享读写
  storageClassName: alicloud-nas
  resources:
    requests:
      storage: 500Gi

---

Arena 工具

阿里云提供 Arena CLI 简化训练任务管理：

# 安装 Arena
wget https://github.com/kubeflow/arena/releases/latest/download/arena-installer.sh
chmod +x arena-installer.sh && ./arena-installer.sh

# 提交 PyTorch 分布式训练
arena submit pytorchjob \
  --name=my-training \
  --gpus=4 \
  --workers=4 \
  --image=my-training:latest \
  --data=training-data:/data \
  "python -m torch.distributed.run --nproc_per_node=4 train.py"

# 查看训练任务
arena list
arena get my-training

# 查看日志
arena logs my-training

# 查看 GPU 使用情况
arena top node
arena top job

# 删除任务
arena delete my-training

---

训练流程完整示例

1. 准备训练镜像

FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-devel

WORKDIR /workspace

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

CMD ["python", "train.py"]

2. 推送到阿里云镜像仓库

# 登录 ACR
docker login --username=<user> registry.cn-hangzhou.aliyuncs.com

# 构建并推送
docker build -t registry.cn-hangzhou.aliyuncs.com/my-ns/training:v1 .
docker push registry.cn-hangzhou.aliyuncs.com/my-ns/training:v1

3. 上传训练数据到 NAS

# 通过临时 Pod 上传数据
kubectl run data-upload --rm -it \
  --image=busybox \
  --overrides='{"spec":{"containers":[{"name":"upload","image":"busybox","command":["sh"],"volumeMounts":[{"name":"data","mountPath":"/data"}]}],"volumes":[{"name":"data","persistentVolumeClaim":{"claimName":"training-data"}}]}}' -- sh

# 在 Pod 内使用 wget/curl 下载数据到 /data

4. 提交训练任务

kubectl apply -f pytorch-job.yaml

5. 监控训练

# 查看 Job 状态
kubectl get pytorchjob my-training

# 查看 Pod
kubectl get pods -l training.kubeflow.org/job-name=my-training

# 查看日志
kubectl logs -f my-training-worker-0

# TensorBoard（可选）
kubectl port-forward svc/tensorboard 6006:6006

---

常见问题与优化

NCCL 通信优化

env:
- name: NCCL_IB_DISABLE
  value: "0"               # 启用 InfiniBand
- name: NCCL_NET_GDR_LEVEL
  value: "5"               # GPU Direct RDMA
- name: NCCL_DEBUG
  value: "WARN"

训练容错

• 使用 checkpoint 定期保存模型
• PyTorchJob 支持 restartPolicy: OnFailure 自动重启
• 弹性训练（Elastic Training）支持 Worker 动态伸缩

成本优化

• 使用抢占式实例（价格低 70-90%）
• 训练完成后自动缩容节点
• 使用 Spot Instance + checkpoint 机制

---

总结

到此，你已经掌握了从 Docker 基础到在阿里云 ACK 上进行分布式训练的完整知识链：

Docker 基础 → Docker 进阶 → K8s 架构 → 核心资源 → 网络 → 存储
    → 调度 → 安全 → 监控 → 生产实践 → Helm → ACK → GPU → 分布式训练

后续可以深入的方向：

• Operator 开发：用 Go 开发自定义 K8s Controller
• Service Mesh（Istio）：微服务治理
• GitOps（ArgoCD）：声明式 CD
• 混合云多集群管理

---

下一步

→ 04 - 阿里云 ACR 容器镜像服务