Course 3：ML 基础设施

课程简介

分布式训练、GPU 管理、训练平台选型。

🎬 本课程视频：MLOps Production — 机器学习工程生产实践

一、ML 基础设施概览

1.1 基础设施的角色

ML 基础设施支撑着从数据到模型的完整生命周期。没有可靠的基础设施，再好的算法也无法转化为生产价值。

ML 基础设施的核心组件：
1. 计算资源：CPU、GPU、TPU 集群
2. 存储系统：数据湖、特征存储、模型注册表
3. 编排系统：工作流调度、资源管理
4. 网络系统：数据传输、API 网关、负载均衡
5. 监控系统：日志收集、指标监控、告警

1.2 基础设施架构演进

阶段一：单机时代
- 单台机器训练和推理
- Jupyter Notebook + 本地存储
- 适合原型开发和少量数据

阶段二：集群时代
- 多台机器组成计算集群
- 分布式存储（HDFS、S3）
- 任务调度（YARN、Kubernetes）

阶段三：云原生时代
- 容器化和微服务
- 弹性伸缩
- 托管服务（SageMaker、Vertex AI）
- MLOps 平台

二、分布式训练

2.1 为什么需要分布式训练

随着模型和数据规模的增大，单机训练变得不可行：
- 数据并行：数据太大无法装入单机内存
- 模型并行：模型太大无法装入单机显存
- 加速需求：缩短训练时间（从数周到数小时）

2.2 数据并行（Data Parallelism）

数据并行是最常见的分布式训练方式。核心思想：

将训练数据划分为 $N$ 份，每份分配给一个 worker
每个 worker 拥有完整的模型副本
每个 worker 独立计算梯度
通过梯度同步更新模型参数

同步 SGD：所有 worker 计算完梯度后同步更新
- 优点：训练稳定，与单机 SGD 行为一致
- 缺点：慢 worker（Straggler）拖慢整体速度

异步 SGD：每个 worker 独立更新参数（无需等待）
- 优点：训练速度快
- 缺点：梯度"过时"问题——一个 worker 的梯度可能基于过时的参数计算

# PyTorch 数据并行示例
import torch.nn as nn

model = nn.Linear(100, 10)
model = nn.DataParallel(model)  # 自动数据并行
output = model(input_data)  # 自动分割数据

2.3 模型并行（Model Parallelism）

当模型太大无法装入单个 GPU 时，需要模型并行：

层间并行（Pipeline Parallelism）：将模型的不同层分配到不同设备。例如，前几层在 GPU 0 上，中间层在 GPU 1 上，后几层在 GPU 2 上。数据像流水线一样依次通过各层。
层内并行（Tensor Parallelism）：将单个层的计算分配到多个设备。例如，一个大的全连接层可以被分割到两个 GPU 上各计算一半。

2.4 混合并行

现代大规模训练通常同时使用数据并行和模型并行：

在节点间使用数据并行
在节点内使用模型并行
加上梯度检查点（Gradient Checkpointing）减少显存占用

2.5 分布式训练框架

PyTorch Distributed：torch.distributed 包，支持多种后端
TensorFlow Distribution Strategy：tf.distribute.MirroredStrategy
Horovod：Uber 开源的分布式训练框架
DeepSpeed：微软的分布式训练优化库
Megatron-LM：NVIDIA 的大模型训练框架

三、GPU 管理

3.1 GPU 资源特性

GPU 在 ML 训练和推理中有不可替代的地位，但 GPU 资源管理面临挑战：

显存限制：单卡显存 16GB-80GB，大模型需要跨卡
利用率：GPU 利用率低是常见问题
碎片化：不同团队使用不同 GPU 类型
成本：GPU 是昂贵的计算资源

3.2 GPU 利用率优化

显存优化：
- 梯度累积：累积多个 batch 的梯度后统一更新
- 混合精度训练（AMP）：FP16 训练减少显存一半
- 梯度检查点：牺牲计算换取显存
- offload：将部分参数卸载到 CPU 内存

计算优化：
- 选择合适的 batch size（通常 2 的幂次效率最高）
- 使用 DataLoader 的多进程加载
- 将数据预处理与 GPU 计算流水线化

# PyTorch 混合精度训练
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():  # 自动混合精度
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

3.3 GPU 调度

Kubernetes + GPU 插件：容器化的 GPU 调度
NVIDIA GPU Operator：自动管理 GPU 驱动和运行时
SLURM：传统 HPC 集群管理

3.4 成本管理

按需实例 vs 预留实例：预留实例可节省 40-60%
抢占式实例：价格低但可能被回收
Spot 实例：大幅节省成本但可靠性差
GPU 配额管理：确保资源不被滥用

四、ML 平台

4.1 平台功能

一个完整的 ML 平台应该提供：

数据层：
- 数据湖/数据仓库
- 特征存储
- 数据版本控制

训练层：
- 训练任务调度
- 超参数调优
- 实验追踪

部署层：
- 模型注册表
- 推理服务
- 在线实验平台

监控层：
- 指标收集
- 漂移检测
- 性能监控

4.2 平台架构

典型的 ML 平台架构：

┌─────────────────────────────────┐
│         用户界面/CLI/API         │
├─────────────────────────────────┤
│  Pipeline  Orchestration(Airflow)│
├──────────┬──────────┬───────────┤
│ Data     │  Training│ Serving   │
│ Pipeline │  Pipeline│ Pipeline  │
├──────────┴──────────┴───────────┤
│     Feature Store   │  Model    │
│ (Redis / Feast)     │  Registry│
├─────────────────────────────────┤
│  Compute(K8s/Slurm) │ Storage   │
│  GPU / CPU Cluster  │ S3/HDFS  │
└─────────────────────────────────┘

4.3 平台选型

开源方案：
- MLflow：实验追踪 + 模型注册
- Kubeflow：Kubernetes 上的 ML 工作流
- Airflow：工作流编排
- Feast：特征存储

商业方案：
- AWS SageMaker：全托管 ML 平台
- GCP Vertex AI：Google 的 ML 平台
- Azure ML：微软的 ML 平台
- Databricks：数据+AI 统一平台

4.4 基础设施即代码（IaC）

使用代码管理基础设施是现代 ML 平台的最佳实践：

# Terraform 示例 - 创建 ML 基础设施
resource "aws_sagemaker_notebook_instance" "ml_dev" {
  name          = "ml-development"
  role_arn      = aws_iam_role.sagemaker.arn
  instance_type = "ml.t3.medium"
}

resource "aws_sagemaker_model" "prod_model" {
  name               = "prod-model-v1"
  execution_role_arn = aws_iam_role.sagemaker.arn
  primary_container {
    image = "${aws_ecr_repository.ml_model.repository_url}:latest"
  }
}

五、总结

ML 基础设施支撑从数据到模型的完整生命周期
分布式训练通过数据并行和模型并行加速训练
GPU 管理包括显存优化、利用率优化、成本控制
ML 平台整合数据、训练、部署、监控的全链路能力
基础设施即代码保证环境的一致性和可复现性

五、云原生 ML 基础设施

5.1 Kubernetes 与 ML

Kubernetes（K8s）已经成为现代 ML 基础设施的核心。它为 ML 工作负载提供：

资源管理：自动调度 GPU/CPU 资源
弹性伸缩：根据负载自动调整计算资源
滚动更新：无宕机地更新服务
服务发现：自动负载均衡

Kubernetes 的关键 ML 组件：
- Kubeflow：Kubernetes 上的 ML 工作流平台
- KServe：模型推理服务
- Volcano：批处理调度器（适合分布式训练）

5.2 存储选型

对象存储（S3、GCS、MinIO）：
- 存储训练数据、模型文件、日志
- 高持久性、无限扩展
- 适合数据的长期存储

文件存储（NFS、EFS）：
- 多节点共享文件系统
- 适合分布式训练中共享数据
- 性能和扩展性受限于单个文件系统

块存储（EBS、Persistent Disk）：
- 高性能单节点挂载
- 适合数据库等需要低延迟的存储
- 不能跨节点共享

5.3 网络架构

ML 系统的网络架构需要考虑：
- 数据上传/下载带宽：训练数据量大，需要高带宽网络
- 节点间通信：分布式训练中高频率梯度同步，需要低延迟的 InfiniBand 或 RoCE
- 推理服务延迟：在线推理需要在几十到几百毫秒内返回结果

5.4 成本优化策略

计算成本优化：
- 使用预留实例或承诺使用折扣
- 利用 Spot/Preemptible 实例做容错性强的训练任务
- 使用 Auto Scaling 根据负载自动调整资源

存储成本优化：
- 生命周期管理：热数据存高性能存储、冷数据存档
- 数据压缩：列式存储格式（Parquet、ORC）降低存储成本
- 定期清理：删除无用的中间数据和模型

GPU 成本优化：
- 使用 GPU 共享和分时调度
- 对小模型使用 CPU 推理
- 模型量化减少 GPU 计算量