AI 基础设施 (AI Infrastructure)

AI基础设施覆盖了从硬件到软件、从训练到推理的完整技术栈。理解这些基础设施是进行大规模AI工程的前提。

---

计算与集群

GPU 基础

现代AI训练和推理高度依赖GPU，核心概念：

CUDA Cores vs Tensor Cores：

• CUDA Cores：通用并行计算核心，每个时钟周期完成一次浮点运算（FP32），适合通用并行计算
• Tensor Cores：专为矩阵乘法设计的加速单元，每个时钟周期完成一个 \(4 \times 4\) 的矩阵乘加运算（FMA），AI训练的主力

VRAM（显存）：

显存是GPU训练的核心瓶颈。训练时显存主要被以下部分占用：

占用项	说明	大致比例
模型参数	Weight（FP16/FP32）	较小
梯度	与参数同大小	较小
优化器状态	Adam需要2份状态（均值+方差）	较大
激活值	中间层输出，反向传播需要	最大
KV Cache	推理时的注意力缓存	推理时最大

常见GPU型号对比：

GPU	VRAM	FP16 算力	典型用途
RTX 4090	24GB	330 TFLOPS	个人研究/小模型训练
A100 80GB	80GB	312 TFLOPS	主流训练/推理
H100 SXM	80GB	990 TFLOPS	大规模训练
H200	141GB	990 TFLOPS	大模型训练/长上下文推理

Multi-GPU 通信

CUDA：

• NVIDIA 的并行计算平台和编程模型
• PyTorch 通过 torch.cuda 提供高层接口
• CUDA_VISIBLE_DEVICES 环境变量控制可见GPU

NCCL (NVIDIA Collective Communications Library)：

• NVIDIA 开发的多GPU通信库
• 支持 AllReduce、AllGather、Broadcast 等集合通信操作
• 自动选择最优通信路径（NVLink > PCIe > InfiniBand）
• PyTorch 的 torch.distributed 后端默认使用 NCCL

NVLink vs PCIe：

• NVLink：GPU间直连，带宽 600-900 GB/s（H100），训练大模型必需
• PCIe Gen5：约 64 GB/s，通常是通信瓶颈

集群管理

Slurm：

• HPC（高性能计算）领域最常用的作业调度系统
• 支持资源分配、队列管理、作业调度
• 学术界和研究机构的标准选择

# 提交训练任务
sbatch --gres=gpu:4 --nodes=2 --ntasks-per-node=4 train.sh

# 查看队列
squeue -u $USER

# 交互式申请资源
srun --gres=gpu:1 --pty bash

Kubernetes (K8s)：

• 容器编排平台，工业界主流选择
• 配合 NVIDIA GPU Operator 支持GPU调度
• 优势：弹性伸缩、服务发现、滚动更新

Cloud vs On-Prem：

维度	云端 (Cloud)	自建 (On-Prem)
初始成本	低（按需付费）	高（硬件采购）
长期成本	可能更高	GPU利用率高时更划算
弹性	随时扩缩容	需要提前规划
数据安全	需要额外措施	数据不出机房
适用场景	实验探索、弹性需求	持续训练、合规要求

---

存储与数据通道

存储层次

速度高 ←──────────────────────────────→ 速度低
容量小                                   容量大

GPU VRAM → NVMe SSD → SATA SSD → NAS/NFS → Object Storage (S3)
 ~2TB/s    ~7GB/s     ~0.5GB/s   ~1-10Gb/s  ~数百MB/s

训练场景的存储选择：

• 训练数据集：大容量存储（NAS, S3），通过数据加载管线预取
• Checkpoint：高速SSD，频繁读写
• 日志与指标：Object Storage 或数据库

Data Loading 瓶颈

训练时数据加载可能成为瓶颈，导致GPU空闲等待（"starving"）：

常见优化手段：

• 多进程加载：PyTorch DataLoader 的 num_workers 参数
• 预取（Prefetch）：在GPU计算时异步加载下一批数据
• 内存映射（Memory-Mapped Files）：避免反复读取磁盘
• 数据格式优化：使用 WebDataset、TFRecord 等连续读取友好的格式
• 缓存到本地SSD：将远程数据缓存到本地 NVMe

# PyTorch DataLoader 优化示例
train_loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,          # 多进程加载
    pin_memory=True,        # 锁定内存，加速到GPU的传输
    prefetch_factor=2,      # 每个worker预取2个batch
    persistent_workers=True # 不在每个epoch重建worker
)

Checkpoint 管理

训练大模型时，Checkpoint 管理至关重要：

• 存储大小：一个 7B 参数模型的 Checkpoint（含优化器状态）约 56GB（FP32）
• 存储频率：通常每隔一定 step 或一定时间保存一次
• 存储策略：只保留最新 N 个 Checkpoint + 关键里程碑 Checkpoint
• 异步保存：使用后台线程保存 Checkpoint，不阻塞训练

---

训练基础设施

分布式训练

Data Parallel (DP) vs Distributed Data Parallel (DDP)

Data Parallel (DP)：

• PyTorch 早期的简单实现（torch.nn.DataParallel）
• 单进程多线程，受 Python GIL 限制
• 通信模式：GPU0 收集所有梯度 → 更新参数 → 广播回其他GPU
• 问题：GPU0 成为通信和显存瓶颈，已不推荐使用

Distributed Data Parallel (DDP)：

• 多进程（每个GPU一个进程），无 GIL 限制
• 通信模式：AllReduce（梯度平均），无单点瓶颈
• 梯度同步与反向传播重叠进行（Overlap Communication with Computation）
• 现在的标准做法

# DDP 基本用法
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.to(local_rank), device_ids=[local_rank])

启动命令：

torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
    --master_addr="192.168.1.1" --master_port=29500 train.py

FSDP (Fully Sharded Data Parallel)

DDP 的问题是每个GPU都要保存完整的模型参数、梯度和优化器状态。当模型太大无法放入单卡时，DDP 就不够用了。

FSDP 的做法是将参数、梯度和优化器状态分片（shard）到多个GPU：

• 前向传播时：通过 AllGather 收集需要的完整参数层
• 计算完成后：释放非本地分片的参数，只保留本地分片
• 反向传播时：收集参数 → 计算梯度 → ReduceScatter 梯度 → 释放

效果：每个GPU的显存占用从 \(O(N)\) 降到 \(O(N/P)\)，\(P\) 为GPU数量。

DeepSpeed ZeRO

微软的 DeepSpeed 库提供 ZeRO (Zero Redundancy Optimizer) 优化，思路与 FSDP 类似但更细粒度：

Stage	分片内容	显存节省
ZeRO-1	只分片优化器状态	约 4x
ZeRO-2	分片优化器状态 + 梯度	约 8x
ZeRO-3	分片优化器状态 + 梯度 + 参数	线性扩展

ZeRO-3 在功能上等价于 FSDP，但 DeepSpeed 提供了更多工程优化（如 ZeRO-Infinity 可 offload 到 CPU/NVMe）。

// DeepSpeed 配置示例 (ZeRO-2)
{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"},
        "contiguous_gradients": true,
        "overlap_comm": true
    },
    "fp16": {"enabled": true},
    "train_batch_size": 64
}

显存优化

Gradient Checkpointing（梯度检查点）

正常训练需要保存所有中间层的激活值用于反向传播。Gradient Checkpointing 只保存部分层的激活值，其余在反向传播时重新计算：

• 显存节省：从 \(O(L)\) 降到 \(O(\sqrt{L})\)，\(L\) 为层数
• 代价：增加约 30% 的计算时间（重新前向传播）
• 权衡：用计算换显存，训练大模型时几乎必用

from torch.utils.checkpoint import checkpoint

# 对特定层启用 gradient checkpointing
output = checkpoint(self.heavy_layer, input, use_reentrant=False)

Mixed Precision (AMP)

使用混合精度训练，核心思想是在不损失精度的前提下用低精度格式加速计算：

• 前向和反向传播：使用 FP16 或 BF16（减少显存、加速 Tensor Core 计算）
• 参数更新：保持 FP32 的 Master Weight（保证数值精度）
• Loss Scaling：放大 Loss 防止 FP16 梯度下溢

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for data, target in train_loader:
    optimizer.zero_grad()
    with autocast(dtype=torch.bfloat16):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

BF16 vs FP16：

• FP16：指数位5bit，精度高，但容易溢出，需要 Loss Scaling
• BF16：指数位8bit，动态范围与 FP32 一致，不需要 Loss Scaling，更稳定

Offloading

当GPU显存不够时，将部分数据卸载到CPU内存甚至NVMe：

• CPU Offloading：将优化器状态或不活跃的参数放到CPU内存
• NVMe Offloading：进一步将数据放到SSD（DeepSpeed ZeRO-Infinity）
• 代价：引入 CPU-GPU 数据传输延迟

容错与恢复

Checkpoint 策略

大规模训练可能持续数周，硬件故障是必然事件：

• 定期保存：每 N 个 step 保存一次完整 Checkpoint
• 异步保存：使用独立线程/进程保存，不阻塞训练
• 分布式保存：每个 rank 保存自己的分片，加速保存过程
• 增量保存：只保存与上次的差异（节省存储）

Elastic Training（弹性训练）

允许训练过程中动态增减节点：

• PyTorch Elastic (TorchElastic)：节点故障后自动重启，支持动态扩缩容
• 场景：抢占式实例（Spot Instance）随时可能被回收
• 机制：通过 torchrun 的 --rdzv_backend 实现节点发现和协调

# 弹性训练：最少2节点，最多4节点
torchrun --nnodes=2:4 --nproc_per_node=4 \
    --rdzv_backend=c10d --rdzv_endpoint=host:port train.py

---

推理基础设施

Serving Frameworks

框架	开发方	特点
Triton Inference Server	NVIDIA	通用推理服务，支持多种后端（TensorRT, PyTorch, ONNX等）
TGI	HuggingFace	专为 Transformer 优化，Docker 部署
vLLM	UC Berkeley	PagedAttention，高吞吐
TensorRT-LLM	NVIDIA	极致性能，需编译优化

关于 vLLM 的详细内容，参见 vLLM 与 KV Cache

Batching 策略

策略	说明	吞吐量
No Batching	逐个请求处理	最低
Static Batching	攒满一批再处理	中等
Dynamic Batching	设置等待时间窗口	较高
Continuous Batching	每个iteration动态调度	最高

Auto-Scaling（自动伸缩）

根据负载动态调整推理实例数量：

• 基于指标：GPU利用率、请求队列长度、P99延迟
• Kubernetes HPA：Horizontal Pod Autoscaler，根据指标自动调整 Pod 数量
• 预热（Warm-up）：新实例需要加载模型，通常需要数十秒到几分钟
• 缩容保护：避免频繁扩缩容导致的抖动（设置冷却期）

---

参考

• Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models", SC 2020
• Zhao et al., "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel", VLDB 2023
• DeepSpeed Documentation
• PyTorch Distributed Overview

---

评论