AI Infra 架构全景介绍

AI Infra 架构全景

一、什么是 AI Infra

AI Infra（AI 基础设施）是支撑大模型从开发到落地全过程的软件栈。它解决的核心问题是：如何让模型在有限的硬件资源上跑得更快、更大、更稳。

从抽象的视角看，整个 AI Infra 可以划分为三个层次：

┌─────────────────────────────────────────────────────────┐
│                    应用层 (Applications)                 │
│              Chatbot / Agent / RAG / 业务系统            │
├─────────────────────────────────────────────────────────┤
│                    模型服务层 (Serving)                  │
│         vLLM / TensorRT-LLM / TGI / llama.cpp           │
├─────────────────────────────────────────────────────────┤
│                    模型适配层 (Adaptation)               │
│           SWIFT / PEFT / LLaMA-Factory / Axolotl        │
├─────────────────────────────────────────────────────────┤
│                    训练加速层 (Training)                 │
│           DeepSpeed / FSDP / Megatron-LM                │
├─────────────────────────────────────────────────────────┤
│                    基础框架层 (Framework)                │
│              PyTorch / JAX / TensorFlow                 │
├─────────────────────────────────────────────────────────┤
│                    硬件层 (Hardware)                     │
│            NVIDIA GPU / AMD / 华为昇腾 / 谷歌TPU         │
└─────────────────────────────────────────────────────────┘

下面从下往上逐一拆解每一层。

---

二、硬件层

硬件是 AI 的算力基础，当前市场格局：

厂商	代表产品	特点	适用场景
NVIDIA	A100/H100/B200	生态最成熟，软件支持最好	训练、推理通用
AMD	MI250/MI300	性价比优势，生态追赶中	推理、部分训练
华为	昇腾910/310	国内自主，配套CANN软件栈	国产化要求场景
谷歌	TPU v4/v5	云端自用，不对外销售	Google内部/Cloud

关键概念：

• 显存带宽：决定数据传输速度，直接影响推理延迟
• 显存容量：决定能跑多大的模型（如H100 80GB ≈ 可单卡跑70B 4bit量化模型）
• 多卡互联：NVLink/InfiniBand，决定多卡通信效率

---

三、基础框架层

这是最底层的编程接口，负责将硬件能力抽象为张量计算。

PyTorch（当前主流）

核心概念：

• Tensor：多维数组，在GPU上的数据结构
• Autograd：自动微分，自动构建计算图并反向传播
• nn.Module：所有神经网络层的基类

# 最底层的操作
import torch
x = torch.tensor([1.0], device='cuda')  # 数据放到GPU
w = torch.nn.Linear(784, 256, device='cuda')  # 参数在GPU
y = w(x)  # 前向计算
y.backward()  # 反向传播，自动算梯度

JAX（Google生态）

特点：函数式编程、即时编译、适合研究场景。与PyTorch的关系类似于"Python vs C++"------JAX更底层但更灵活，PyTorch更易用但更厚。

这一层解决的核心问题

• 硬件抽象：写一份代码跑在CPU/GPU/TPU上
• 自动求导：不用手写反向传播公式
• 算子库：CUDA核函数已封装好，直接调用

---

四、训练加速层

当模型单卡放不下，或单卡训练太慢时，需要分布式训练框架。

DeepSpeed（微软）

核心组件：

组件	解决的问题	实现方式
ZeRO	显存不够分	将优化器状态、梯度、参数分片到多卡
流水线并行	模型层数太多	不同层放不同GPU
张量并行	单层参数太大	单层的参数切分到多卡
梯度检查点	显存不够	用计算换显存，不存中间激活值
CPU Offload	显存还不够	把部分数据放CPU内存

Megatron-LM（NVIDIA）

专注于张量并行，将Transformer的Attention和MLP层切分到多卡，通过精细的通信优化达到高效率。常与DeepSpeed配合使用：Megatron负责模型切分，DeepSpeed负责显存优化。

FSDP（PyTorch原生）

PyTorch 1.11后内置的ZeRO实现，功能类似DeepSpeed，但更"PyTorch原生"。选择哪个取决于生态：DeepSpeed功能更全，FSDP与PyTorch集成更紧。

这一层解决的核心问题

• 显存不够：模型太大，单卡装不下
• 训练太慢：单卡算不过来，需要多卡并行
• 多卡效率：通信开销要尽可能小

---

五、模型适配层（微调框架）

微调框架解决的是"如何用有限的资源让大模型适应新任务"。

SWIFT（ModelScope生态）

核心模块：

┌─────────────────────────────────────────┐
│              SWIFT 组件                  │
├─────────────────────────────────────────┤
│ Model Hub：从ModelScope下载模型           │
├─────────────────────────────────────────┤
│ Adapter：注入LoRA/QLoRA适配器             │
│   - 原模型冻结                            │
│   - 只训练小矩阵（通常<1%参数）             │
├─────────────────────────────────────────┤
│ Template：数据格式化成模型对话格式          │
│   - 自动处理system/user/assistant标记     │
├─────────────────────────────────────────┤
│ Trainer：封装训练循环                     │
│   - 混合精度、梯度累积、日志记录             │
├─────────────────────────────────────────┤
│ Export：合并LoRA权重，导出完整模型          │
└─────────────────────────────────────────┘

PEFT（HuggingFace生态）

HuggingFace的官方微调库，提供LoRA、Prefix Tuning、Adapter等参数高效微调方法。与Transformers库深度集成。

LoRA 的机制

假设原始层是 W（4096×4096），LoRA做的事情：

原始输出 = x @ W           # 计算量大，不训练

LoRA输出 = x @ W + (x @ A) @ B
                    ↑      ↑
                低秩矩阵A  低秩矩阵B
                (4096×64) (64×4096)
                ← 只训练这两个 →

• 训练参数从 16M 降到 0.5M
• 推理时可合并回 W' = W + A@B，无额外推理开销

这一层解决的核心问题

• 算力有限：消费级显卡也能微调
• 时间有限：几小时完成微调，而非几天
• 多任务：一套模型可以挂多个LoRA适配器，按需切换

---

六、模型服务层（推理框架）

推理框架解决的是"模型训练好后，如何高效响应大量用户请求"。

vLLM

核心创新点：

1. PagedAttention

传统推理：每个请求预先分配固定大小的KV Cache，存在内存浪费。

vLLM：将KV Cache按块管理（类似操作系统虚拟内存），按需分配。

传统方式（固定分配）：
请求A：[████████░░░░] 40%浪费
请求B：[██████░░░░░░] 50%浪费

PagedAttention（分页管理）：
请求A：[█][█][█][█][ ] 按需分配
请求B：[█][█][█][ ][ ]
        └── 空闲块可复用给其他请求

2. 连续批处理

传统批处理：等凑齐一批才执行，空闲时间多。

vLLM：请求来了随时加入执行中的批次，类似CPU的抢占式调度。

TensorRT-LLM（NVIDIA）

NVIDIA官方推理框架，特点：

• 算子融合：将多个小算子合并成一个，减少kernel启动开销
• 量化支持：INT4/INT8量化，显存减半
• 仅NVIDIA GPU可用，但性能最优

llama.cpp

特点：

• 纯C++实现，无PyTorch依赖
• 支持CPU推理
• GGUF量化格式（CPU友好）
• 适合边缘设备、本地部署

TGI（HuggingFace）

vLLM的竞品，功能类似，但生态更偏HuggingFace。选择时主要看：vLLM性能更好，TGI功能更全（如embedding、微调服务等）。

这一层解决的核心问题

• 延迟：首token时间要快（毫秒级）
• 吞吐：同时服务尽可能多的用户
• 显存：KV Cache占主要，要高效管理

---

七、框架选型指南

场景	推荐框架	理由
从头预训练大模型	DeepSpeed + Megatron	需要极致并行能力
单卡/双卡微调	SWIFT / PEFT	LoRA省显存，上手快
在线高并发推理	vLLM / TensorRT-LLM	PagedAttention + 连续批处理
本地CPU推理	llama.cpp	无需GPU，GGUF量化
国内合规部署	昇腾CANN + SWIFT	国产化适配，ModelScope生态
多模态模型	HuggingFace Transformers	模型最全，生态最广

---

八、各层的关系：从代码到硬件

一条训练命令经过的完整路径：

swift train --model Qwen-7B --dataset alpaca

执行链路：

1. swift 命令行 → 解析参数，调用 Swift 库函数
                ↓
2. Swift.prepare_model() → 注入 LoRA 适配器
                ↓
3. Trainer.train() → 启动训练循环
                ↓
4. DeepSpeed.initialize() → 初始化分布式环境，ZeRO 分片
                ↓
5. PyTorch Distributed → 调用 NCCL 进行多卡通信
                ↓
6. CUDA Kernel → 在 GPU 上执行矩阵运算
                ↓
7. GPU 硬件 → 实际完成计算

每一层都在解决上一层的复杂性，让用户用更少的代码完成更多的事。

---

九、总结

AI Infra 本质上是一个层层抽象的软件栈：

层	抽象	解决的问题
硬件	GPU显存/算力	物理资源
基础框架	Tensor + Autograd	用代码描述神经网络
训练加速	分布式并行	模型太大/训练太慢
模型适配	参数高效微调	算力不够也能调模型
模型服务	批处理 + 内存管理	在线服务要快、要稳

学习路径建议：

1. 先跑通一个推理服务（vLLM），理解"模型部署"是什么
2. 再做一次LoRA微调（SWIFT），理解"参数高效"的含义
3. 最后读DeepSpeed文档，理解分布式训练的难点
4. 遇到具体问题时，再深入对应框架的源码

不需要一次性掌握所有框架，先理解每一层解决什么问题，使用时再深入。