AI Infra 架构全景介绍

AI Infra 架构全景一、什么是 AI InfraAI InfraAI 基础设施是支撑大模型从开发到落地全过程的软件栈。它解决的核心问题是如何让模型在有限的硬件资源上跑得更快、更大、更稳。从抽象的视角看整个 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 的算力基础当前市场格局厂商代表产品特点适用场景NVIDIAA100/H100/B200生态最成熟软件支持最好训练、推理通用AMDMI250/MI300性价比优势生态追赶中推理、部分训练华为昇腾910/310国内自主配套CANN软件栈国产化要求场景谷歌TPU v4/v5云端自用不对外销售Google内部/Cloud关键概念显存带宽决定数据传输速度直接影响推理延迟显存容量决定能跑多大的模型如H100 80GB ≈ 可单卡跑70B 4bit量化模型多卡互联NVLink/InfiniBand决定多卡通信效率三、基础框架层这是最底层的编程接口负责将硬件能力抽象为张量计算。PyTorch当前主流核心概念Tensor多维数组在GPU上的数据结构Autograd自动微分自动构建计算图并反向传播nn.Module所有神经网络层的基类# 最底层的操作importtorch xtorch.tensor([1.0],devicecuda)# 数据放到GPUwtorch.nn.Linear(784,256,devicecuda)# 参数在GPUyw(x)# 前向计算y.backward()# 反向传播自动算梯度JAXGoogle生态特点函数式编程、即时编译、适合研究场景。与PyTorch的关系类似于Python vs C——JAX更底层但更灵活PyTorch更易用但更厚。这一层解决的核心问题硬件抽象写一份代码跑在CPU/GPU/TPU上自动求导不用手写反向传播公式算子库CUDA核函数已封装好直接调用四、训练加速层当模型单卡放不下或单卡训练太慢时需要分布式训练框架。DeepSpeed微软核心组件组件解决的问题实现方式ZeRO显存不够分将优化器状态、梯度、参数分片到多卡流水线并行模型层数太多不同层放不同GPU张量并行单层参数太大单层的参数切分到多卡梯度检查点显存不够用计算换显存不存中间激活值CPU Offload显存还不够把部分数据放CPU内存Megatron-LMNVIDIA专注于张量并行将Transformer的Attention和MLP层切分到多卡通过精细的通信优化达到高效率。常与DeepSpeed配合使用Megatron负责模型切分DeepSpeed负责显存优化。FSDPPyTorch原生PyTorch 1.11后内置的ZeRO实现功能类似DeepSpeed但更PyTorch原生。选择哪个取决于生态DeepSpeed功能更全FSDP与PyTorch集成更紧。这一层解决的核心问题显存不够模型太大单卡装不下训练太慢单卡算不过来需要多卡并行多卡效率通信开销要尽可能小五、模型适配层微调框架微调框架解决的是如何用有限的资源让大模型适应新任务。SWIFTModelScope生态核心模块┌─────────────────────────────────────────┐ │ SWIFT 组件 │ ├─────────────────────────────────────────┤ │ Model Hub从ModelScope下载模型 │ ├─────────────────────────────────────────┤ │ Adapter注入LoRA/QLoRA适配器 │ │ - 原模型冻结 │ │ - 只训练小矩阵通常1%参数 │ ├─────────────────────────────────────────┤ │ Template数据格式化成模型对话格式 │ │ - 自动处理system/user/assistant标记 │ ├─────────────────────────────────────────┤ │ Trainer封装训练循环 │ │ - 混合精度、梯度累积、日志记录 │ ├─────────────────────────────────────────┤ │ Export合并LoRA权重导出完整模型 │ └─────────────────────────────────────────┘PEFTHuggingFace生态HuggingFace的官方微调库提供LoRA、Prefix Tuning、Adapter等参数高效微调方法。与Transformers库深度集成。LoRA 的机制假设原始层是W4096×4096LoRA做的事情原始输出 x W # 计算量大不训练 LoRA输出 x W (x A) B ↑ ↑ 低秩矩阵A 低秩矩阵B (4096×64) (64×4096) ← 只训练这两个 →训练参数从 16M 降到 0.5M推理时可合并回W W AB无额外推理开销这一层解决的核心问题算力有限消费级显卡也能微调时间有限几小时完成微调而非几天多任务一套模型可以挂多个LoRA适配器按需切换六、模型服务层推理框架推理框架解决的是模型训练好后如何高效响应大量用户请求。vLLM核心创新点1. PagedAttention传统推理每个请求预先分配固定大小的KV Cache存在内存浪费。vLLM将KV Cache按块管理类似操作系统虚拟内存按需分配。传统方式固定分配 请求A[████████░░░░] 40%浪费 请求B[██████░░░░░░] 50%浪费 PagedAttention分页管理 请求A[█][█][█][█][ ] 按需分配 请求B[█][█][█][ ][ ] └── 空闲块可复用给其他请求2. 连续批处理传统批处理等凑齐一批才执行空闲时间多。vLLM请求来了随时加入执行中的批次类似CPU的抢占式调度。TensorRT-LLMNVIDIANVIDIA官方推理框架特点算子融合将多个小算子合并成一个减少kernel启动开销量化支持INT4/INT8量化显存减半仅NVIDIA GPU可用但性能最优llama.cpp特点纯C实现无PyTorch依赖支持CPU推理GGUF量化格式CPU友好适合边缘设备、本地部署TGIHuggingFacevLLM的竞品功能类似但生态更偏HuggingFace。选择时主要看vLLM性能更好TGI功能更全如embedding、微调服务等。这一层解决的核心问题延迟首token时间要快毫秒级吞吐同时服务尽可能多的用户显存KV Cache占主要要高效管理七、框架选型指南场景推荐框架理由从头预训练大模型DeepSpeed Megatron需要极致并行能力单卡/双卡微调SWIFT / PEFTLoRA省显存上手快在线高并发推理vLLM / TensorRT-LLMPagedAttention 连续批处理本地CPU推理llama.cpp无需GPUGGUF量化国内合规部署昇腾CANN SWIFT国产化适配ModelScope生态多模态模型HuggingFace Transformers模型最全生态最广八、各层的关系从代码到硬件一条训练命令经过的完整路径swift train--modelQwen-7B--datasetalpaca执行链路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用代码描述神经网络训练加速分布式并行模型太大/训练太慢模型适配参数高效微调算力不够也能调模型模型服务批处理 内存管理在线服务要快、要稳学习路径建议先跑通一个推理服务vLLM理解模型部署是什么再做一次LoRA微调SWIFT理解参数高效的含义最后读DeepSpeed文档理解分布式训练的难点遇到具体问题时再深入对应框架的源码不需要一次性掌握所有框架先理解每一层解决什么问题使用时再深入。