第1章 AI 基础设施概论

“AI is the new electricity. Infrastructure is the grid that delivers it.”

—— 改编自 Andrew Ng

1.1 本章导入：为什么 AI 基础设施突然变得如此重要

2024 年，训练 GPT-4 据报道消耗了约 6300 万美元的算力成本。而到了 2026 年，前沿模型的单次训练成本已经突破 5 亿美元。这背后不是某个算法的魔法，而是一个庞大工程体系的运转——从数千张 GPU 组成的集群，到跨越数据中心的网络互联，再到 PB 级存储系统的高速吞吐。

AI 基础设施（AI Infrastructure，简称 AI Infra）就是支撑这一切的底座。

如果你是算法工程师，理解基础设施能帮你写出更高效的训练代码；如果你是后端工程师，理解 AI 工作负载能帮你设计更可靠的推理服务；如果你是技术管理者，理解基础设施的成本结构能帮你做出更明智的投资决策。

本章将建立全书的认知框架：我们定义 AI 基础设施的边界，回顾其演进历程，明确系统设计的核心目标，并勾勒出 AI Infra 工程师的能力画像。

1.2 AI 基础设施的定义与边界

1.2.1 什么是 AI 基础设施

AI 基础设施是指支撑 AI 模型训练、部署和运营全过程所需的硬件、软件和网络资源的总和。

它包含三个层次：

┌─────────────────────────────────────────────────┐
│              应用层 (Application Layer)           │
│    ChatGPT · Copilot · 推荐系统 · 自动驾驶        │
├─────────────────────────────────────────────────┤
│           AI 平台层 (AI Platform Layer)           │
│  训练框架 · 模型仓库 · 特征存储 · 实验追踪         │
│  PyTorch · Ray · MLflow · Triton Inference Server│
├─────────────────────────────────────────────────┤
│          基础设施层 (Infrastructure Layer)         │
│  GPU/TPU 集群 · 高速网络 · 分布式存储             │
│  Kubernetes · CUDA · NCCL · RDMA · 对象存储       │
├─────────────────────────────────────────────────┤
│              物理层 (Physical Layer)               │
│  数据中心 · 电力 · 散热 · 机架 · 布线              │
└─────────────────────────────────────────────────┘

本书聚焦于基础设施层和 AI 平台层的底层部分——也就是从硬件到框架之间的那一整条技术栈。

1.2.2 边界在哪里

AI 基础设施和传统云原生基础设施有重叠，也有本质区别：

维度	传统云原生	AI 基础设施
核心资源	CPU、内存、磁盘	GPU、显存、高速网络
工作负载特征	短生命周期、无状态	长时运行、有状态、通信密集
调度粒度	容器级（毫秒级启动）	GPU 级（分钟级分配）
网络需求	万兆以太网够用	200-400Gbps InfiniBand/RoCE
存储模式	块存储/对象存储	高性能并行文件系统
故障影响	单容器重启	整个训练任务崩溃（数百万美元损失）

Tip

实践建议：不要把 AI 集群当成普通 Kubernetes 集群来运维。GPU 是稀缺且昂贵的资源，调度策略、故障恢复、监控告警都需要专门设计。一个朴素的经验法则——AI 基础设施的单位计算成本是传统云原生的 5-20 倍，容错投入应当同比例提升。

1.2.3 AI 基础设施的核心组件

完整的 AI 基础设施包含以下核心组件：

graph TB
    subgraph "AI Infrastructure Stack"
        A[计算: GPU/NPU/CPU]
        B[网络: IB/RoCE/以太网]
        C[存储: 并行FS/对象存储]
        D[调度: K8s/Slurm/YuniKorn]
        E[通信: NCCL/RCCL/MPI]
        F[框架: PyTorch/JAX]
        G[运行时: CUDA/ROCm/oneAPI]
    end

    A --> D
    B --> E
    C --> D
    D --> F
    E --> F
    G --> A

后续章节将逐一深入每个组件。

1.3 AI 系统演进：从单机训练到万卡集群

理解 AI 基础设施的演进，本质上是理解模型规模与计算需求的指数增长如何驱动系统架构的变革。

1.3.1 四个阶段

阶段 1: 单 GPU 时代 (2012-2017)
    AlexNet → ResNet → BERT-base
    1-8 张 GPU, 单机即可
    
阶段 2: 数据并行时代 (2017-2020)
    GPT-2 → BERT-large → 早期 GPT-3
    数十张 GPU, 数据并行 + 梯度聚合
    
阶段 3: 模型并行时代 (2020-2023)
    GPT-3 (175B) → PaLM → 通义千问
    数百到数千张 GPU, 张量并行 + 流水线并行
    
阶段 4: 万卡集群时代 (2023-)
    GPT-4 → Claude 3 → Gemini Ultra
    数千到数万张 GPU, 3D 并行 + 专家并行

1.3.2 规模的量变到质变

模型参数量从 2012 年的几百万（AlexNet）增长到 2024 年的万亿级别（GPT-4 估计约 1.8 万亿参数），增长了约 100 万倍。这一增长带来的是系统设计的质变：

1. 显存墙的出现

一个 175B 参数的模型，仅模型权重（FP16）就需要约 350 GB 显存——远超单张 A100 的 80GB。这迫使我们必须将模型拆分到多张卡上。

2. 通信成为瓶颈

当模型被拆分到数千张 GPU 上时，每一步前向/反向传播都需要跨节点交换数据。在万卡规模下，网络通信可能占到总训练时间的 30-50%。

3. 故障成为常态

假设单张 GPU 的 MTBF（平均无故障时间）为 6 个月，那么一个 10000 卡集群的期望故障间隔仅为约 26 分钟。故障恢复机制不再是可选项，而是必需品。

Warning

关键认知：在万卡集群中，“正常运行”反而是异常状态。你的系统必须假设硬件随时会出故障——GPU 掉卡、网络抖动、存储延迟飙升都是家常便饭。设计目标不是”消除故障”，而是”在故障中持续运转”。

1.3.3 训练成本的经济学

模型	估计 GPU 数量	训练时长	估计成本
GPT-3 (175B)	~10,000 A100	~34 天	~$4.6M
GPT-4	~25,000 A100	~100 天	~$63M
前沿模型 (2025)	~50,000+ H100	~90 天	~$500M+

成本不仅来自硬件采购，还包括：

电力：一个 50,000 H100 集群功耗约 50MW，年电费超 $50M
网络：InfiniBand 交换机和线缆约占总成本的 10-15%
存储：训练数据集和 checkpoint 需要数 PB 高速存储
运维：需要 24/7 值班的 SRE 团队

1.4 系统设计的核心目标

AI 基础设施的设计需要在多个目标之间权衡。以下是四个核心维度：

1.4.1 吞吐量（Throughput）

定义：单位时间内完成的工作量（训练时为 samples/sec 或 tokens/sec；推理时为 requests/sec）。

吞吐量直接决定了训练成本。一个 50,000 GPU 集群，如果能将整体吞吐提升 10%，相当于节省了 5,000 张 GPU——每年约 $30M 的节省。

关键公式：

\[\text{MFU (Model FLOPS Utilization)} = \frac{\text{实际达到的 FLOPS}}{\text{GPU 理论峰值 FLOPS}}\]

MFU 是衡量训练效率最重要的指标。优秀的训练系统在 GPT 类模型上能达到 45-55% 的 MFU，而未经优化的系统可能只有 15-25%。

1.4.2 延迟（Latency）

定义：从输入到输出的端到端响应时间。

延迟在推理场景中尤为关键。用户等待 LLM 生成第一个 token 的时间（Time-To-First-Token, TTFT）直接影响用户体验：

TTFT	用户感知
< 200ms	即时响应
200ms - 1s	可接受的延迟
1s - 3s	明显卡顿
> 3s	用户流失

1.4.3 成本（Cost）

成本分为两个维度：

CapEx（资本支出）：GPU 采购、数据中心建设——一次性投入
OpEx（运营支出）：电费、带宽、人力——持续投入

一个实用的成本指标是每百万 token 推理成本（Cost per Million Tokens）。2024 年行业基准约为 $0.15-3.00/M tokens，取决于模型规模和服务质量。

1.4.4 可靠性（Reliability）

对于训练任务，可靠性意味着不丢失训练进度。一个训练了 30 天的大模型，如果因为故障而从头开始，损失可能是数百万美元。

对于推理服务，可靠性意味着SLA 保障。大模型推理服务的 SLA 通常要求 99.9% 以上可用性。

# 一个简单的 checkpoint 策略示例
# 展示如何在训练循环中实现容错

import torch
import os
from datetime import datetime

class TrainingCheckpoint:
    def __init__(self, model, optimizer, save_dir="/checkpoints"):
        self.model = model
        self.optimizer = optimizer
        self.save_dir = save_dir
        self.step = 0
        self.best_loss = float('inf')

    def save(self, loss, step):
        """保存 checkpoint，保留最近 3 个版本"""
        os.makedirs(self.save_dir, exist_ok=True)

        # 只在 loss 改善时保存 best
        is_best = loss < self.best_loss
        if is_best:
            self.best_loss = loss

        checkpoint = {
            'step': step,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'loss': loss,
            'timestamp': datetime.now().isoformat(),
        }

        path = os.path.join(self.save_dir, f'checkpoint_step_{step}.pt')
        torch.save(checkpoint, path)

        # 保留最近 3 个 + best
        self._cleanup(keep=3)

        if is_best:
            best_path = os.path.join(self.save_dir, 'best_model.pt')
            torch.save(checkpoint, best_path)

        print(f"[Checkpoint] Saved step {step} (loss={loss:.4f}, best={is_best})")

    def load_latest(self):
        """加载最新的 checkpoint 用于故障恢复"""
        checkpoints = sorted(
            [f for f in os.listdir(self.save_dir)
             if f.startswith('checkpoint_step_')],
            key=lambda x: int(x.split('_')[-1].split('.')[0])
        )
        if not checkpoints:
            return 0, float('inf')

        latest = checkpoints[-1]
        path = os.path.join(self.save_dir, latest)
        ckpt = torch.load(path)
        self.model.load_state_dict(ckpt['model_state_dict'])
        self.optimizer.load_state_dict(ckpt['optimizer_state_dict'])
        print(f"[Checkpoint] Resumed from step {ckpt['step']}")
        return ckpt['step'], ckpt['loss']

    def _cleanup(self, keep=3):
        """只保留最近 N 个 checkpoint"""
        checkpoints = sorted(
            [f for f in os.listdir(self.save_dir)
             if f.startswith('checkpoint_step_')],
            key=lambda x: int(x.split('_')[-1].split('.')[0])
        )
        for old in checkpoints[:-keep]:
            os.remove(os.path.join(self.save_dir, old))

1.4.5 目标之间的权衡

这四个目标之间存在固有的张力：

graph LR
    A[吞吐量] --"冲突"--> B[延迟]
    C[成本] --"冲突"--> D[可靠性]
    A --"互补"--> C
    B --"互补"--> D

提高吞吐量（大 batch size）可能增加延迟
降低成本（使用廉价硬件）可能降低可靠性
提高可靠性（频繁 checkpoint）可能降低吞吐量

系统设计的本质就是在这些目标之间找到适合业务场景的平衡点。

1.5 AI Infra 工程师的能力模型

AI 基础设施工程是一个高度交叉的领域，要求兼具广度和深度的知识储备。

1.5.1 技能雷达

                    硬件知识
                      ★★★★★
                        │
     分布式系统 ────────┼──────── 性能优化
      ★★★★☆            │          ★★★★★
                        │
                   ─────┼─────
                        │
     软件工程 ──────────┼──────── 网络知识
      ★★★★☆            │          ★★★☆☆
                        │
                   AI/ML 基础
                      ★★★☆☆

具体来说，AI Infra 工程师需要掌握：

必备技能（P0）： - Linux 系统编程与性能分析（perf, strace, eBPF） - GPU 编程基础（CUDA, memory hierarchy） - 分布式训练原理（数据并行、张量并行、流水线并行） - 容器化与编排（Docker, Kubernetes） - Python + C++ 混合编程

进阶技能（P1）： - RDMA/InfiniBand 网络编程 - 性能剖析与调优（Nsight Systems, PyTorch Profiler） - 存储系统设计（Lustre, GPFS, 对象存储） - 云原生 AI 平台架构

加分项（P2）： - 芯片架构知识（NVIDIA Hopper/Blackwell, AMD MI300） - 编译器优化（XLA, Triton, TVM） - 数据中心运维经验（电力、散热、PUE）

1.5.2 不同角色的关注重点

角色	核心关注	典型工具
训练 infra 工程师	提高 MFU、降低故障率	NCCL, Megatron-LM, DeepSpeed
推理 infra 工程师	降低延迟、提高 QPS	TensorRT-LLM, vLLM, Triton
平台工程师	多租户隔离、资源调度	Kubernetes, Ray, YuniKorn
SRE	可观测性、故障恢复	Prometheus, Grafana, Jaeger

Tip

职业建议：如果你正在转型到 AI Infra 领域，建议从推理服务优化入手——门槛相对较低（单机即可实践），市场需求旺盛（每个部署 LLM 的公司都需要），且技能可迁移到训练场景（理解了推理的计算特征，再学训练会事半功倍）。

1.5.3 一个典型的 AI Infra 工程师的日常

09:00  查看 Grafana 仪表盘，检查昨夜训练任务的状态
09:30  分析某个训练任务的 GPU 利用率下降问题
       → 用 Nsight Systems 抓 trace
       → 发现 NCCL AllReduce 出现异常延迟尖峰
10:30  与网络团队排查 InfiniBand 链路误码率
12:00  午饭时和算法团队讨论新模型的显存需求
14:00  给推理集群升级 vLLM 到最新版本
       → 编写灰度发布方案
       → 在 1/10 流量上验证延迟指标
16:00  编写技术文档：新的 checkpoint 恢复流程
17:00  Code Review：同事提交的调度器优化 PR
18:00  On-call 开始，随时准备处理生产事故

1.6 小结

AI 基础设施是连接底层硬件与上层 AI 应用的整套技术栈，包含计算、网络、存储、调度、通信等核心组件。
AI 系统经历了从单机到万卡集群的四个演进阶段，每个阶段都带来了新的系统设计挑战。
系统设计的核心是在吞吐量、延迟、成本和可靠性之间找到适合业务的平衡点。
AI Infra 工程师是一个高度交叉的岗位，需要横跨硬件、系统软件、分布式系统和 AI/ML 的知识体系。

延伸阅读

《Designing Machine Learning Systems》 — Chip Huyen，系统化讲解 ML 系统设计
《AI Engineering》 — Chip Huyen，聚焦 AI 工程实践
NVIDIA Multi-Instance GPU (MIG) 技术白皮书 — 理解 GPU 虚拟化
Google Pathways 论文 — 理解异构计算调度思路
DeepSpeed Megatron-LM 文档 — 理解大规模分布式训练工程实现
“The Carbon Emissions of Machine Learning Training” — Patterson et al., 了解 AI 训练的能源消耗
LLM Training Cost Tracker (Epoch AI) — 持续追踪前沿模型训练成本

--- title: "第1章 AI 基础设施概论" --- # 第1章 AI 基础设施概论 > "AI is the new electricity. Infrastructure is the grid that delivers it." > > —— 改编自 Andrew Ng ## 1.1 本章导入：为什么 AI 基础设施突然变得如此重要 2024 年，训练 GPT-4 据报道消耗了约 6300 万美元的算力成本。而到了 2026 年，前沿模型的单次训练成本已经突破 5 亿美元。这背后不是某个算法的魔法，而是一个庞大工程体系的运转——从数千张 GPU 组成的集群，到跨越数据中心的网络互联，再到 PB 级存储系统的高速吞吐。 AI 基础设施（AI Infrastructure，简称 AI Infra）就是支撑这一切的底座。如果你是算法工程师，理解基础设施能帮你写出更高效的训练代码；如果你是后端工程师，理解 AI 工作负载能帮你设计更可靠的推理服务；如果你是技术管理者，理解基础设施的成本结构能帮你做出更明智的投资决策。本章将建立全书的认知框架：我们定义 AI 基础设施的边界，回顾其演进历程，明确系统设计的核心目标，并勾勒出 AI Infra 工程师的能力画像。 ## 1.2 AI 基础设施的定义与边界 ### 1.2.1 什么是 AI 基础设施 AI 基础设施是指**支撑 AI 模型训练、部署和运营全过程所需的硬件、软件和网络资源的总和**。它包含三个层次： ``` ┌─────────────────────────────────────────────────┐ │ 应用层 (Application Layer) │ │ ChatGPT · Copilot · 推荐系统 · 自动驾驶 │ ├─────────────────────────────────────────────────┤ │ AI 平台层 (AI Platform Layer) │ │ 训练框架 · 模型仓库 · 特征存储 · 实验追踪 │ │ PyTorch · Ray · MLflow · Triton Inference Server│ ├─────────────────────────────────────────────────┤ │ 基础设施层 (Infrastructure Layer) │ │ GPU/TPU 集群 · 高速网络 · 分布式存储 │ │ Kubernetes · CUDA · NCCL · RDMA · 对象存储 │ ├─────────────────────────────────────────────────┤ │ 物理层 (Physical Layer) │ │ 数据中心 · 电力 · 散热 · 机架 · 布线 │ └─────────────────────────────────────────────────┘ ``` 本书聚焦于**基础设施层和 AI 平台层的底层部分**——也就是从硬件到框架之间的那一整条技术栈。 ### 1.2.2 边界在哪里 AI 基础设施和传统云原生基础设施有重叠，也有本质区别： | 维度 | 传统云原生 | AI 基础设施 | |------|-----------|------------| | 核心资源 | CPU、内存、磁盘 | **GPU、显存、高速网络** | | 工作负载特征 | 短生命周期、无状态 | **长时运行、有状态、通信密集** | | 调度粒度 | 容器级（毫秒级启动） | **GPU 级（分钟级分配）** | | 网络需求 | 万兆以太网够用 | **200-400Gbps InfiniBand/RoCE** | | 存储模式 | 块存储/对象存储 | **高性能并行文件系统** | | 故障影响 | 单容器重启 | **整个训练任务崩溃（数百万美元损失）** | ::: callout-tip **实践建议**：不要把 AI 集群当成普通 Kubernetes 集群来运维。GPU 是稀缺且昂贵的资源，调度策略、故障恢复、监控告警都需要专门设计。一个朴素的经验法则——**AI 基础设施的单位计算成本是传统云原生的 5-20 倍**，容错投入应当同比例提升。 ::: ### 1.2.3 AI 基础设施的核心组件完整的 AI 基础设施包含以下核心组件： ```{mermaid} graph TB subgraph "AI Infrastructure Stack" A[计算: GPU/NPU/CPU] B[网络: IB/RoCE/以太网] C[存储: 并行FS/对象存储] D[调度: K8s/Slurm/YuniKorn] E[通信: NCCL/RCCL/MPI] F[框架: PyTorch/JAX] G[运行时: CUDA/ROCm/oneAPI] end A --> D B --> E C --> D D --> F E --> F G --> A ``` 后续章节将逐一深入每个组件。 ## 1.3 AI 系统演进：从单机训练到万卡集群理解 AI 基础设施的演进，本质上是理解**模型规模与计算需求的指数增长**如何驱动系统架构的变革。 ### 1.3.1 四个阶段 ``` 阶段 1: 单 GPU 时代 (2012-2017) AlexNet → ResNet → BERT-base 1-8 张 GPU, 单机即可阶段 2: 数据并行时代 (2017-2020) GPT-2 → BERT-large → 早期 GPT-3 数十张 GPU, 数据并行 + 梯度聚合阶段 3: 模型并行时代 (2020-2023) GPT-3 (175B) → PaLM → 通义千问数百到数千张 GPU, 张量并行 + 流水线并行阶段 4: 万卡集群时代 (2023-) GPT-4 → Claude 3 → Gemini Ultra 数千到数万张 GPU, 3D 并行 + 专家并行 ``` ### 1.3.2 规模的量变到质变模型参数量从 2012 年的几百万（AlexNet）增长到 2024 年的万亿级别（GPT-4 估计约 1.8 万亿参数），增长了约 **100 万倍**。这一增长带来的是系统设计的质变： **1. 显存墙的出现** 一个 175B 参数的模型，仅模型权重（FP16）就需要约 350 GB 显存——远超单张 A100 的 80GB。这迫使我们必须将模型拆分到多张卡上。 **2. 通信成为瓶颈** 当模型被拆分到数千张 GPU 上时，每一步前向/反向传播都需要跨节点交换数据。在万卡规模下，网络通信可能占到总训练时间的 **30-50%**。 **3. 故障成为常态** 假设单张 GPU 的 MTBF（平均无故障时间）为 6 个月，那么一个 10000 卡集群的期望故障间隔仅为约 **26 分钟**。故障恢复机制不再是可选项，而是必需品。 ::: callout-warning **关键认知**：在万卡集群中，"正常运行"反而是异常状态。你的系统必须假设硬件随时会出故障——GPU 掉卡、网络抖动、存储延迟飙升都是家常便饭。设计目标不是"消除故障"，而是"在故障中持续运转"。 ::: ### 1.3.3 训练成本的经济学 | 模型 | 估计 GPU 数量 | 训练时长 | 估计成本 | |------|-------------|---------|---------| | GPT-3 (175B) | ~10,000 A100 | ~34 天 | ~$4.6M | | GPT-4 | ~25,000 A100 | ~100 天 | ~$63M | | 前沿模型 (2025) | ~50,000+ H100 | ~90 天 | ~$500M+ | 成本不仅来自硬件采购，还包括： - **电力**：一个 50,000 H100 集群功耗约 50MW，年电费超 $50M - **网络**：InfiniBand 交换机和线缆约占总成本的 10-15% - **存储**：训练数据集和 checkpoint 需要数 PB 高速存储 - **运维**：需要 24/7 值班的 SRE 团队 ## 1.4 系统设计的核心目标 AI 基础设施的设计需要在多个目标之间权衡。以下是四个核心维度： ### 1.4.1 吞吐量（Throughput） **定义**：单位时间内完成的工作量（训练时为 samples/sec 或 tokens/sec；推理时为 requests/sec）。吞吐量直接决定了训练成本。一个 50,000 GPU 集群，如果能将整体吞吐提升 10%，相当于节省了 5,000 张 GPU——每年约 $30M 的节省。 **关键公式**： $$\text{MFU (Model FLOPS Utilization)} = \frac{\text{实际达到的 FLOPS}}{\text{GPU 理论峰值 FLOPS}}$$ MFU 是衡量训练效率最重要的指标。优秀的训练系统在 GPT 类模型上能达到 **45-55% 的 MFU**，而未经优化的系统可能只有 15-25%。 ### 1.4.2 延迟（Latency） **定义**：从输入到输出的端到端响应时间。延迟在推理场景中尤为关键。用户等待 LLM 生成第一个 token 的时间（Time-To-First-Token, TTFT）直接影响用户体验： | TTFT | 用户感知 | |------|---------| | < 200ms | 即时响应 | | 200ms - 1s | 可接受的延迟 | | 1s - 3s | 明显卡顿 | | > 3s | 用户流失 | ### 1.4.3 成本（Cost）成本分为两个维度： - **CapEx（资本支出）**：GPU 采购、数据中心建设——一次性投入 - **OpEx（运营支出）**：电费、带宽、人力——持续投入一个实用的成本指标是**每百万 token 推理成本**（Cost per Million Tokens）。2024 年行业基准约为 $0.15-3.00/M tokens，取决于模型规模和服务质量。 ### 1.4.4 可靠性（Reliability）对于训练任务，可靠性意味着**不丢失训练进度**。一个训练了 30 天的大模型，如果因为故障而从头开始，损失可能是数百万美元。对于推理服务，可靠性意味着**SLA 保障**。大模型推理服务的 SLA 通常要求 99.9% 以上可用性。 ```python # 一个简单的 checkpoint 策略示例 # 展示如何在训练循环中实现容错 import torch import os from datetime import datetime class TrainingCheckpoint: def __init__(self, model, optimizer, save_dir="/checkpoints"): self.model = model self.optimizer = optimizer self.save_dir = save_dir self.step = 0 self.best_loss = float('inf') def save(self, loss, step): """保存 checkpoint，保留最近 3 个版本""" os.makedirs(self.save_dir, exist_ok=True) # 只在 loss 改善时保存 best is_best = loss < self.best_loss if is_best: self.best_loss = loss checkpoint = { 'step': step, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'loss': loss, 'timestamp': datetime.now().isoformat(), } path = os.path.join(self.save_dir, f'checkpoint_step_{step}.pt') torch.save(checkpoint, path) # 保留最近 3 个 + best self._cleanup(keep=3) if is_best: best_path = os.path.join(self.save_dir, 'best_model.pt') torch.save(checkpoint, best_path) print(f"[Checkpoint] Saved step {step} (loss={loss:.4f}, best={is_best})") def load_latest(self): """加载最新的 checkpoint 用于故障恢复""" checkpoints = sorted( [f for f in os.listdir(self.save_dir) if f.startswith('checkpoint_step_')], key=lambda x: int(x.split('_')[-1].split('.')[0]) ) if not checkpoints: return 0, float('inf') latest = checkpoints[-1] path = os.path.join(self.save_dir, latest) ckpt = torch.load(path) self.model.load_state_dict(ckpt['model_state_dict']) self.optimizer.load_state_dict(ckpt['optimizer_state_dict']) print(f"[Checkpoint] Resumed from step {ckpt['step']}") return ckpt['step'], ckpt['loss'] def _cleanup(self, keep=3): """只保留最近 N 个 checkpoint""" checkpoints = sorted( [f for f in os.listdir(self.save_dir) if f.startswith('checkpoint_step_')], key=lambda x: int(x.split('_')[-1].split('.')[0]) ) for old in checkpoints[:-keep]: os.remove(os.path.join(self.save_dir, old)) ``` ### 1.4.5 目标之间的权衡这四个目标之间存在固有的张力： ```{mermaid} graph LR A[吞吐量] --"冲突"--> B[延迟] C[成本] --"冲突"--> D[可靠性] A --"互补"--> C B --"互补"--> D ``` - 提高吞吐量（大 batch size）可能增加延迟 - 降低成本（使用廉价硬件）可能降低可靠性 - 提高可靠性（频繁 checkpoint）可能降低吞吐量系统设计的本质就是**在这些目标之间找到适合业务场景的平衡点**。 ## 1.5 AI Infra 工程师的能力模型 AI 基础设施工程是一个高度交叉的领域，要求兼具**广度和深度**的知识储备。 ### 1.5.1 技能雷达 ``` 硬件知识 ★★★★★ │ 分布式系统 ────────┼──────── 性能优化 ★★★★☆ │ ★★★★★ │ ─────┼───── │ 软件工程 ──────────┼──────── 网络知识 ★★★★☆ │ ★★★☆☆ │ AI/ML 基础 ★★★☆☆ ``` 具体来说，AI Infra 工程师需要掌握： **必备技能（P0）**： - Linux 系统编程与性能分析（perf, strace, eBPF） - GPU 编程基础（CUDA, memory hierarchy） - 分布式训练原理（数据并行、张量并行、流水线并行） - 容器化与编排（Docker, Kubernetes） - Python + C++ 混合编程 **进阶技能（P1）**： - RDMA/InfiniBand 网络编程 - 性能剖析与调优（Nsight Systems, PyTorch Profiler） - 存储系统设计（Lustre, GPFS, 对象存储） - 云原生 AI 平台架构 **加分项（P2）**： - 芯片架构知识（NVIDIA Hopper/Blackwell, AMD MI300） - 编译器优化（XLA, Triton, TVM） - 数据中心运维经验（电力、散热、PUE） ### 1.5.2 不同角色的关注重点 | 角色 | 核心关注 | 典型工具 | |------|---------|---------| | 训练 infra 工程师 | 提高 MFU、降低故障率 | NCCL, Megatron-LM, DeepSpeed | | 推理 infra 工程师 | 降低延迟、提高 QPS | TensorRT-LLM, vLLM, Triton | | 平台工程师 | 多租户隔离、资源调度 | Kubernetes, Ray, YuniKorn | | SRE | 可观测性、故障恢复 | Prometheus, Grafana, Jaeger | ::: callout-tip **职业建议**：如果你正在转型到 AI Infra 领域，建议从**推理服务优化**入手——门槛相对较低（单机即可实践），市场需求旺盛（每个部署 LLM 的公司都需要），且技能可迁移到训练场景（理解了推理的计算特征，再学训练会事半功倍）。 ::: ### 1.5.3 一个典型的 AI Infra 工程师的日常 ``` 09:00 查看 Grafana 仪表盘，检查昨夜训练任务的状态 09:30 分析某个训练任务的 GPU 利用率下降问题 → 用 Nsight Systems 抓 trace → 发现 NCCL AllReduce 出现异常延迟尖峰 10:30 与网络团队排查 InfiniBand 链路误码率 12:00 午饭时和算法团队讨论新模型的显存需求 14:00 给推理集群升级 vLLM 到最新版本 → 编写灰度发布方案 → 在 1/10 流量上验证延迟指标 16:00 编写技术文档：新的 checkpoint 恢复流程 17:00 Code Review：同事提交的调度器优化 PR 18:00 On-call 开始，随时准备处理生产事故 ``` ## 1.6 小结 - **AI 基础设施**是连接底层硬件与上层 AI 应用的整套技术栈，包含计算、网络、存储、调度、通信等核心组件。 - AI 系统经历了**从单机到万卡集群**的四个演进阶段，每个阶段都带来了新的系统设计挑战。 - 系统设计的核心是在**吞吐量、延迟、成本和可靠性**之间找到适合业务的平衡点。 - AI Infra 工程师是一个高度交叉的岗位，需要横跨硬件、系统软件、分布式系统和 AI/ML 的知识体系。 ## 延伸阅读 - **《Designing Machine Learning Systems》** — Chip Huyen，系统化讲解 ML 系统设计 - **《AI Engineering》** — Chip Huyen，聚焦 AI 工程实践 - **NVIDIA Multi-Instance GPU (MIG) 技术白皮书** — 理解 GPU 虚拟化 - **Google Pathways 论文** — 理解异构计算调度思路 - **DeepSpeed Megatron-LM 文档** — 理解大规模分布式训练工程实现 - **"The Carbon Emissions of Machine Learning Training"** — Patterson et al., 了解 AI 训练的能源消耗 - **LLM Training Cost Tracker** (Epoch AI) — 持续追踪前沿模型训练成本