AI Infra 框架体系介绍
层次核心问题关键框架硬件提供算力和显存GPU/昇腾/TPU基础框架硬件抽象 + 自动求导训练加速显存不够、训练太慢模型适配资源受限下的模型定制SWIFT/PEFT推理服务低延迟、高吞吐应用业务逻辑自研学习时不必一次性掌握所有框架,先理解每一层解决什么问题,使用时再深入对应框架的细节。
AI Infra 框架体系详解
快速术语表
阅读正文前,建议先熟悉以下核心术语:
| 术语 | 英文 | 一句话解释 |
|---|---|---|
| ZeRO | Zero Redundancy Optimizer | 将优化器状态、梯度、参数分片到多卡,消除冗余 |
| LoRA | Low-Rank Adaptation | 训练低秩矩阵替代全参数微调,参数量减少99% |
| KV Cache | Key-Value Cache | 推理时缓存Attention的K/V矩阵,避免重复计算 |
| PagedAttention | Paged Attention | vLLM的内存管理机制,将KV Cache分页存储 |
| 连续批处理 | Continuous Batching | 动态调度请求,无需等待批次填满 |
| 张量并行 | Tensor Parallelism | 将单层参数切分到多张GPU |
| 流水线并行 | Pipeline Parallelism | 按层切分模型,不同层在不同GPU |
| 算子融合 | Operator Fusion | 合并多个小算子为一个kernel,减少启动开销 |
| GGUF | GGUF Format | llama.cpp使用的量化格式,CPU友好 |
| 混合精度 | Mixed Precision | 同时使用FP16和FP32,平衡精度与速度 |
一、为什么需要这些框架?
在没有框架的情况下,直接用 PyTorch 实现大模型训练或推理的复杂度:
| 场景 | 裸写 PyTorch 的复杂度 | 框架带来的价值 |
|---|---|---|
| 单卡训练 | 手写训练循环、优化器调度、检查点保存(约100行) | 框架封装为10行配置 |
| 多卡训练 | 手写 all_reduce 通信、进程管理、数据分片(500+行) | DDP/DeepSpeed 一行开启 |
| 模型太大 | 单卡显存放不下,无法运行 | ZeRO 自动分片到多卡 |
| 线上推理 | 需自己实现批处理队列、KV Cache 管理、内存回收(1000+行) | vLLM 开箱即用 |
核心结论:框架 = 别人写好的、经过验证的、高性能的函数集合,让你专注于模型本身而非工程细节。
二、分层架构总览
AI Infra 自下而上分为六个层次:
┌─────────────────────────────────────────────────────────────┐
│ 应用层 │
│ Chatbot / Agent / RAG / 业务系统 │
│ 关注:业务逻辑、用户体验 │
├─────────────────────────────────────────────────────────────┤
│ 推理服务层 │
│ vLLM / TensorRT-LLM / TGI / llama.cpp │
│ 关注:延迟、吞吐、显存效率 │
├─────────────────────────────────────────────────────────────┤
│ 模型适配层 │
│ SWIFT / PEFT / LLaMA-Factory │
│ 关注:资源受限下的模型定制、快速迭代 │
├─────────────────────────────────────────────────────────────┤
│ 训练加速层 │
│ DeepSpeed / FSDP / Megatron-LM │
│ 关注:显存不足、训练太慢、多卡效率 │
├─────────────────────────────────────────────────────────────┤
│ 基础框架层 │
│ PyTorch / JAX / TensorFlow │
│ 关注:硬件抽象、自动求导、张量计算 │
├─────────────────────────────────────────────────────────────┤
│ 硬件层 │
│ NVIDIA GPU / 华为昇腾 / AMD / 谷歌 TPU │
│ 关注:算力、显存、带宽 │
└─────────────────────────────────────────────────────────────┘
三、训练加速层:DeepSpeed
DeepSpeed 是微软开源的分布式训练框架,核心解决大模型训练中的显存瓶颈和通信效率问题。
3.1 核心组件
┌─────────────────────────────────────────────────┐
│ DeepSpeed 架构 │
├─────────────────────────────────────────────────┤
│ 1. DataLoader │
│ ├── 数据分片与多卡并行加载 │
│ └── 保证各 GPU 数据不重叠 │
├─────────────────────────────────────────────────┤
│ 2. Model Parallelism Engine │
│ ├── 张量并行 (Tensor Parallelism) │
│ ├── 流水线并行 (Pipeline Parallelism) │
│ └── 专家并行 (Expert Parallelism) │
├─────────────────────────────────────────────────┤
│ 3. ZeRO Optimizer │
│ ├── ZeRO-1: 优化器状态分片 │
│ ├── ZeRO-2: 梯度分片 │
│ ├── ZeRO-3: 模型参数分片 │
│ └── ZeRO-Offload: CPU 内存卸载 │
├─────────────────────────────────────────────────┤
│ 4. Communication Backend │
│ ├── NCCL 通信库封装 │
│ └── 梯度同步与参数广播 │
├─────────────────────────────────────────────────┤
│ 5. Checkpoint Manager │
│ ├── 分布式检查点保存 │
│ └── 断点续训支持 │
└─────────────────────────────────────────────────┘
3.2 ZeRO 工作机制
ZeRO(Zero Redundancy Optimizer)通过分片消除显存冗余:
| 阶段 | 分片内容 | 显存节省 | 通信开销 | 适用场景 |
|---|---|---|---|---|
| ZeRO-1 | 优化器状态 | 4倍 | 无增加 | 中等规模训练 |
| ZeRO-2 | 优化器状态 + 梯度 | 8倍 | 轻微增加 | 大规模训练 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | 64倍 | 显著增加 | 超大模型训练 |
| ZeRO-Offload | 分片 + CPU卸载 | >100倍 | CPU-GPU传输 | 显存极度受限 |
3.3 训练框架对比
| 框架 | 定位 | 核心能力 | 适用场景 |
|---|---|---|---|
| DeepSpeed | 通用训练加速 | ZeRO 显存优化,支持万亿参数 | 从零预训练、大规模微调 |
| Megatron-LM | 张量并行优化 | 精细化张量切分,NVIDIA 生态深度优化 | 超大模型预训练 |
| FSDP | PyTorch 原生 | PyTorch 内置的 ZeRO 实现 | 不想引入额外依赖的场景 |
选择建议:
- 需要极致显存优化 → DeepSpeed
- 使用 NVIDIA 超大规模集群 → Megatron-LM
- 追求简洁和原生生态 → FSDP
四、模型适配层:SWIFT
SWIFT 是 ModelScope 生态下的微调框架,专注于参数高效微调(PEFT),让消费级显卡也能微调大模型。
4.1 核心组件
┌─────────────────────────────────────────────────┐
│ SWIFT 架构 │
├─────────────────────────────────────────────────┤
│ 1. Model Hub │
│ ├── 自动模型下载与缓存 │
│ └── 支持 Qwen/Llama/ChatGLM 等主流模型 │
├─────────────────────────────────────────────────┤
│ 2. Adapter Manager ← 核心 │
│ ├── LoRA/QLoRA 适配器注入 │
│ ├── 原模型参数冻结 │
│ └── 仅训练增量参数 (<1%) │
├─────────────────────────────────────────────────┤
│ 3. Template Engine │
│ ├── 对话模板格式化 │
│ └── 支持多种模型的消息格式 │
├─────────────────────────────────────────────────┤
│ 4. Trainer │
│ ├── 训练循环封装 │
│ ├── 混合精度自动管理 │
│ └── 梯度累积支持 │
├─────────────────────────────────────────────────┤
│ 5. Exporter │
│ ├── LoRA 权重合并 │
│ └── ONNX/GGUF 格式导出 │
└─────────────────────────────────────────────────┘
4.2 LoRA 数学原理
假设原始线性层为 W ∈ R^{d×k},LoRA 引入两个低秩矩阵:
原始计算: y = x @ W
LoRA 计算: y = x @ W + (x @ A) @ B
其中 A ∈ R^{d×r}, B ∈ R^{r×k}, r << min(d,k)
参数量变化:
- 原层: d × k
- LoRA: r × (d + k)
- 当 d=k=4096, r=8 时,参数量从 16M 降至 65K(减少 99.6%)
关键特性:
- 训练时只更新 A 和 B,原模型 W 冻结
- 推理时可合并为
W' = W + A@B,无额外开销 - 支持多任务:一套基座模型挂多个 LoRA 适配器,按需切换
4.3 微调框架对比
| 框架 | 生态 | 特点 | 适用场景 |
|---|---|---|---|
| SWIFT | ModelScope | 国内模型支持好,开箱即用 | 使用 Qwen/ChatGLM 等国产模型 |
| PEFT | HuggingFace | LoRA 变体最丰富,社区活跃 | 使用 LLaMA/Mistral 等国际模型 |
| LLaMA-Factory | 独立 | Web UI 友好,适合非开发者 | 快速实验、教学演示 |
五、推理服务层:vLLM
vLLM 是加州大学伯克利分校开源的推理框架,核心创新是 PagedAttention,显著提升推理吞吐量。
5.1 核心组件
┌─────────────────────────────────────────────────┐
│ vLLM 架构 │
├─────────────────────────────────────────────────┤
│ 1. Model Loader │
│ ├── HuggingFace 模型加载 │
│ └── AWQ/GPTQ 量化模型支持 │
├─────────────────────────────────────────────────┤
│ 2. Scheduler (连续批处理) ← 核心 │
│ ├── 动态请求调度 │
│ ├── 无需等待批次填满 │
│ └── 抢占式调度策略 │
├─────────────────────────────────────────────────┤
│ 3. PagedAttention Engine ← 核心 │
│ ├── KV Cache 分页管理 │
│ ├── 按需分配,消除内存碎片 │
│ └── 跨请求内存共享 (相同 prompt) │
├─────────────────────────────────────────────────┤
│ 4. Decoding Engine │
│ ├── Transformer 前向计算 │
│ └── CUDA kernel 优化 │
├─────────────────────────────────────────────────┤
│ 5. API Server │
│ ├── OpenAI 兼容接口 │
│ └── /v1/chat/completions 端点 │
└─────────────────────────────────────────────────┘
5.2 PagedAttention 内存管理
传统推理框架为每个请求预先分配固定长度的 KV Cache,存在严重的内存浪费:
| 维度 | 传统方式 | PagedAttention |
|---|---|---|
| 内存分配 | 静态预留最大长度 | 动态按需分配 |
| 内存碎片 | 存在内部碎片(预留但未使用) | 消除碎片 |
| 内存复用 | 不支持 | 支持 prefix 共享 |
| 显存利用率 | 通常 40-60% | 可达 80-90% |
工作原理类比:
- 传统方式:每个客人预订固定大小的房间,住不满也空着
- PagedAttention:将房间拆成"页",需要几页给几页,满页再分配新页
5.3 推理框架对比
| 框架 | 定位 | 核心能力 | 适用场景 |
|---|---|---|---|
| vLLM | 高吞吐推理 | PagedAttention,连续批处理 | 通用高并发服务 |
| TensorRT-LLM | NVIDIA 极致优化 | 算子融合,INT4/INT8 量化 | 追求极致性能,仅 NVIDIA |
| TGI | HuggingFace 生态 | 功能完整,持续更新 | HuggingFace 深度用户 |
| llama.cpp | CPU 推理 | GGUF 量化,无 GPU 依赖 | 本地部署、边缘设备 |
选择建议:
- 通用场景 → vLLM
- 追求 NVIDIA 极致性能 → TensorRT-LLM
- 无 GPU 环境 → llama.cpp
六、横向对比与选型
6.1 功能维度对比
| 维度 | 训练框架 (DeepSpeed) | 微调框架 (SWIFT) | 推理框架 (vLLM) |
|---|---|---|---|
| 核心目标 | 大规模预训练 | 高效领域适配 | 低延迟服务 |
| 参数更新 | 全参数更新 | 增量更新 (<1%) | 无更新 |
| 显存优化 | ZeRO 分片 | LoRA 低秩 | PagedAttention |
| 并行策略 | 3D 并行 | 单卡/数据并行 | 张量并行 |
| 关键指标 | samples/sec | 显存占用 / 收敛速度 | latency / throughput |
| 硬件要求 | 多机多卡 A100/H100 | 单卡 RTX 4090 起 | 单卡/多卡,显存优先 |
6.2 框架组合关系
训练阶段: PyTorch + DeepSpeed/Megatron → 产出基座模型
↓
微调阶段: PyTorch + SWIFT/PEFT → 产出 LoRA 适配器
↓
推理阶段: vLLM/TensorRT-LLM → 对外提供服务
关键说明:
- 这些框架不是互斥的,可以在不同阶段组合使用
- DeepSpeed 和 FSDP 是替代关系,选其一即可
- vLLM 和 TensorRT-LLM 是竞争关系,根据硬件和性能需求选择
6.3 选型决策树
开始
│
├─ 你的目标是什么?
│ │
│ ├─ 从零预训练大模型
│ │ └─ DeepSpeed + Megatron-LM
│ │
│ ├─ 微调已有模型
│ │ ├─ 有 A100/H100(显存充足)
│ │ │ ├─ 全参微调 → DeepSpeed
│ │ │ └─ 参数高效微调 → SWIFT/PEFT
│ │ └─ 只有消费级显卡(RTX 3090/4090)
│ │ └─ LoRA 微调 → SWIFT/PEFT + QLoRA
│ │
│ └─ 部署推理服务
│ ├─ 需要高并发、低延迟
│ │ ├─ NVIDIA GPU → vLLM 或 TensorRT-LLM
│ │ └─ 国产昇腾 → MindIE
│ └─ CPU 推理
│ └─ llama.cpp
七、核心概念深化
7.1 大模型显存占用拆解
以 70B 参数模型为例,FP16 精度:
| 组成部分 | 计算公式 | 显存占用 |
|---|---|---|
| 模型参数 | 70B × 2 bytes | 140 GB |
| 梯度 | 同参数大小 | 140 GB |
| 优化器状态 (Adam) | 参数 × 12 bytes | 420 GB |
| 激活值 | 取决于 batch size 和序列长度 | 动态(通常 20-50 GB) |
| 总计(训练) | ~720 GB |
各框架的优化方向:
- ZeRO:将优化器状态、梯度、参数分片到多卡
- LoRA:只训练增量参数,不存储完整梯度
- PagedAttention:管理 KV Cache(推理阶段,非训练)
7.2 量化技术体系
| 量化类型 | 精度 | 显存节省 | 质量损失 | 适用场景 | 代表框架 |
|---|---|---|---|---|---|
| FP16/BF16 | 16-bit | 2倍(对比FP32) | 基本无损 | 训练、推理 | PyTorch 原生 |
| INT8 | 8-bit | 4倍 | 轻微 | 推理 | TensorRT, vLLM |
| INT4/GPTQ | 4-bit | 8倍 | 可控 | 推理 | vLLM, AutoGPTQ |
| NF4 | 4-bit | 8倍 | 可控 | QLoRA 微调 | bitsandbytes |
| GGUF | 可变 | 4-8倍 | 可控 | CPU 推理 | llama.cpp |
量化时机:
- 训练时:通常用 FP16/BF16,一般不用 INT 量化
- 微调时:QLoRA 在加载时量化,训练时部分反量化
- 推理时:INT4/INT8 量化,减少显存、加速推理
7.3 并行策略一览
| 并行策略 | 原理 | 通信开销 | 代表框架 | 适用场景 |
|---|---|---|---|---|
| 数据并行 | 每卡完整模型,切分数据 | 低(梯度同步) | DDP, DeepSpeed | 模型能单卡装下 |
| 张量并行 | 单层参数切分到多卡 | 高(每层通信) | Megatron, DeepSpeed | 单层太大 |
| 流水线并行 | 按层切分,不同层在不同卡 | 中(边界通信) | DeepSpeed, Megatron | 层数多 |
| 专家并行 | MoE 模型的专家分布 | 中(路由通信) | DeepSpeed | MoE 架构 |
组合使用:大模型训练通常使用"3D 并行" = 数据并行 + 张量并行 + 流水线并行。
八、国产化生态(华为昇腾)
对于国内部署场景,需要了解华为昇腾生态:
| 层级 | 华为昇腾生态 | 对应 NVIDIA 生态 |
|---|---|---|
| 硬件 | 昇腾 910B / 310 | A100 / H100 / T4 |
| 基础软件 | CANN (Compute Architecture for Neural Networks) | CUDA |
| 框架适配 | torch_npu (PyTorch 昇腾版) | PyTorch CUDA |
| 训练加速 | 昇腾版 DeepSpeed | DeepSpeed |
| 微调 | SWIFT (已原生支持昇腾) | PEFT |
| 推理 | MindIE / vLLM 昇腾版 | vLLM / TensorRT-LLM |
关键差异:
- CANN 替代 CUDA,API 语义不同,需要代码适配
- 模型需要格式转换(如
.onnx→.om) - 生态成熟度仍有差距,但国内合规项目必须考虑
九、学习路径建议
9.1 实践顺序
-
推理入门(最快获得反馈)
pip install vllm vllm serve Qwen/Qwen2-7B-Instruct --port 8000 curl http://localhost:8000/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "Qwen/Qwen2-7B-Instruct", "messages": [{"role": "user", "content": "Hello"}] }' -
微调实践(理解参数高效)
- 用 SWIFT 在公开数据集上执行 LoRA 微调
- 对比微调前后的输出差异
- 调整
r参数观察显存和效果变化
-
训练原理(理论储备)
- 阅读 DeepSpeed ZeRO 论文
- 理解 ZeRO-1/2/3 的通信与显存权衡
- 了解何时需要分布式训练
9.2 常见误区澄清
| 误区 | 真相 |
|---|---|
| LoRA 会显著降低模型效果 | 合适 rank(如 64)下效果接近全参微调 |
| vLLM 只能跑 NVIDIA GPU | 已支持 AMD ROCm,社区版支持昇腾 |
| DeepSpeed 和 FSDP 可以一起用 | 是替代关系,选其一即可 |
| 量化后模型质量一定下降 | INT8 基本无损,INT4 有轻微但可控下降 |
| 推理框架也能做训练 | 不能,推理框架只做前向传播 |
| 微调必须用多卡 | LoRA 可在单卡 RTX 4090 微调 70B 模型 |
十、术语索引
| 术语 | 英文 | 章节 |
|---|---|---|
| ZeRO | Zero Redundancy Optimizer | 3.2 |
| LoRA | Low-Rank Adaptation | 4.2 |
| QLoRA | Quantized LoRA | 4.1 |
| KV Cache | Key-Value Cache | 5.1 |
| PagedAttention | Paged Attention | 5.2 |
| 连续批处理 | Continuous Batching | 5.1 |
| 张量并行 | Tensor Parallelism | 7.3 |
| 流水线并行 | Pipeline Parallelism | 7.3 |
| 算子融合 | Operator Fusion | 5.3 |
| GGUF | GGUF Format | 5.3 |
| 混合精度 | Mixed Precision | 7.2 |
| 数据并行 | Data Parallelism | 7.3 |
| DDP | Distributed Data Parallel | 3.3 |
| NCCL | NVIDIA Collective Communications Library | 3.1 |
| MoE | Mixture of Experts | 3.1 |
总结
AI Infra 是一个层层抽象的软件栈,每一层解决特定的问题:
| 层次 | 核心问题 | 关键框架 |
|---|---|---|
| 硬件 | 提供算力和显存 | GPU/昇腾/TPU |
| 基础框架 | 硬件抽象 + 自动求导 | PyTorch/JAX |
| 训练加速 | 显存不够、训练太慢 | DeepSpeed/FSDP |
| 模型适配 | 资源受限下的模型定制 | SWIFT/PEFT |
| 推理服务 | 低延迟、高吞吐 | vLLM/TensorRT-LLM |
| 应用 | 业务逻辑 | 自研 |
学习时不必一次性掌握所有框架,先理解每一层解决什么问题,使用时再深入对应框架的细节。
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐



所有评论(0)