第一次接触 ai infra 时,我看到的往往不是一个完整系统,而是一串缩写:gpu、hbm、nvlink、rdma、nccl、ddp、zero……每个词单独查都能找到解释,放到一起却还是不知道它们如何协作。
ai infra 可以简单理解为支撑大模型训练和推理的基础设施。它不只是购买 gpu,还包括服务器内部的连接、服务器之间的网络、分布式通信、任务调度、数据存储和模型服务。
这篇文章不按硬件、网络、软件等技术分类堆词,而是沿着一项工作的真实流程展开:
准备一张 gpu
↓
让一台服务器里的多张 gpu 协作
↓
让多台 gpu 服务器通信
↓
把大模型拆到多张 gpu
↓
通过 kubernetes 分配和管理资源
↓
读取训练数据并保存模型
↓
把训练好的模型部署成 api
为了方便记忆,可以把整个系统想象成一座 ai 工厂:gpu 是计算工人,显存是工作台,网络是运输道路,nccl 是物流调度员,kubernetes 是任务管理中心,存储则是不同距离和规模的仓库。
先让一张 gpu 工作
gpu:负责计算的工人
gpu 是擅长大量并行计算的处理器。
可以把 cpu 想成少数能力全面的管理员,把 gpu 想成许多同时处理相似计算的工人。大模型训练和推理包含大量矩阵计算,因此很适合交给 gpu。
cpu 并没有消失。它仍然负责启动程序、读取和预处理数据、调度任务,再把适合并行计算的部分交给 gpu。
cpu 准备数据
↓
gpu 执行并行计算
↓
cpu 继续处理结果或安排下一批工作
显存:gpu 的工作台
显存,也叫 vram,是 gpu 自己使用的高速内存。模型参数、计算过程中的临时数据以及推理缓存,都需要占用显存。
可以这样理解:
gpu 是工人,显存是工人的工作台。工人速度再快,工作台放不下模型,也无法开始工作。
假设一个模型运行时需要 100 gb 显存,而一张 gpu 只有 80 gb,就需要减少显存占用,或者让多张 gpu 一起保存和计算模型。
hbm 是 gpu 上常见的高带宽显存技术。这里的高带宽,是指单位时间可以搬运更多数据。gpu 计算速度很快,如果显存供料太慢,工人还是要停下来等材料。
驱动、cuda 和 cudnn:让软件使用 gpu
有了 gpu,还需要一套软件把它用起来。
- gpu driver 是驱动,让操作系统能够识别和控制 gpu;
- cuda 是 nvidia 提供的 gpu 计算平台,让 pytorch 等软件可以调用 nvidia gpu;
- cudnn 是针对深度学习计算优化的软件库,提供经过优化的常用计算能力。
可以把它们理解为:
gpu driver:让机器能够被操作系统控制
cuda :让程序能够向 gpu 下达计算任务
cudnn :提供优化好的深度学习工具
cuda 不是 gpu,cudnn 也不是独立的训练框架。它们处在应用程序与硬件之间,帮助上层框架使用 gpu。
一台服务器里的多张 gpu 如何协作
安装了 gpu 的服务器通常叫 gpu node,也就是 gpu 节点。一台服务器安装 8 张 gpu,就可以称为一个 8 卡节点。
把一台 gpu 服务器想成一间厂房,接下来要解决的问题是:厂房里的 8 名工人怎样快速交换材料?
pcie:服务器内部的通用道路
pcie 是服务器内部连接 cpu、gpu、网卡和 nvme 硬盘的通用高速通道。它像厂房里的公共道路,很多设备都可以使用。
pcie p2p 中的 p2p 是 peer-to-peer,表示两台设备直接通信。两张 gpu 使用 pcie p2p 时,可以尽量直接交换数据,减少绕经 cpu 内存的中转:
普通方式:gpu 0 → cpu 内存 → gpu 1
pcie p2p:gpu 0 ─────────→ gpu 1
这就像两个工人直接传递材料,不再先交给管理员。
nvlink:gpu 之间的专用高速道路
nvlink 是 nvidia 提供的 gpu 高速互联技术,通常比普通 pcie 更适合 gpu 之间的大量数据传输。
- pcie 是许多设备共用的道路;
- nvlink 是 gpu 之间的专用高速道路。
多张 gpu 需要频繁交换模型参数或计算结果时,更快的道路可以减少等待时间。
nvswitch:连接多条 nvlink 的立交桥
nvswitch 是连接多张 gpu 的高速交换设备。它把多条 nvlink 组织起来,让更多 gpu 能够高速互通。
如果 nvlink 是一条条高速公路,nvswitch 就是连接这些公路的立交桥。
所以三者的关系可以记成:
| 技术 | 作用 | 类比 |
|---|---|---|
| pcie | 连接服务器内的多种设备 | 通用道路 |
| nvlink | 让 gpu 高速交换数据 | gpu 专用高速路 |
| nvswitch | 连接多条 nvlink | 高速立交桥 |
topology:先看懂道路地图
topology 就是拓扑,也就是设备之间怎样连接。它会告诉我们:
- 哪些 gpu 之间有 nvlink;
- gpu 连接到哪个 cpu;
- gpu 距离哪张网卡更近;
- 数据需要经过几层 pcie 设备。
拓扑就像工厂的道路地图。两张 gpu 虽然在同一台服务器里,通信距离和可用带宽也可能不同。
常用的查看命令是:
nvidia-smi topo -m
numa:设备也有远近之分
一台服务器可能安装多个 cpu,每个 cpu 都有距离自己更近的内存、gpu 和网卡。这种设备存在本地和远端访问差异的结构叫 numa。
它像一座有多个仓库的工厂:从身边的仓库取材料很快,跨过整个厂区取材料就会更慢。假如 gpu 靠近 cpu 0,却总让 cpu 1 准备数据,程序仍然可以运行,但性能可能受到影响。
多台 gpu 服务器如何通信
现在把规模扩大到两台 8 卡服务器,共 16 张 gpu。
同一台服务器里的 gpu 可以通过 pcie 或 nvlink 通信,但跨服务器时,还需要网卡、交换机和高速网络。
nic 和 hca:厂房的货物出口
nic 就是网卡,负责让服务器连接网络。
hca 是面向高性能网络和 rdma 的网卡。普通 nic 像常见的快递站,hca 更像专门搬运大批货物的高速物流中心。ai 集群中常见 100g、200g 或 400g 的高速网卡,这里的数字表示网络速率。
理解网络性能时,经常会看到两个指标:
- bandwidth,也就是带宽,表示单位时间能传输多少数据,相当于道路有多宽;
- latency,也就是延迟,表示一次数据从出发到到达要等待多久,相当于路程需要多长时间。
搬运大批模型数据时很看重带宽;频繁交换小块数据时,延迟也非常重要。
infiniband:跨服务器的专用高速公路
infiniband,简称 ib,是面向高性能计算和 ai 集群的高速网络。它具有高带宽、低延迟等特点,并支持 rdma。
可以把它理解为不同 gpu 厂房之间的专用高速公路。
另一种常见选择是 ethernet,也就是以太网。以太网是数据中心里更通用的道路,设备普遍,维护也较熟悉。
rdma:减少中转的直达运输
rdma 是 remote direct memory access,也就是远程直接内存访问。它让一台服务器更直接地访问另一台服务器的内存,减少 cpu 参与和数据复制。
普通网络传输可能需要 cpu 和操作系统多次处理。rdma 则像给货物开通直达通道,减少沿途登记和搬运。
gpudirect rdma 又向前走了一步:高性能网卡可以更直接地访问 gpu 显存。
普通方式:gpu 显存 → cpu 内存 → 网卡 → 网络
直达方式:gpu 显存 ───────→ hca → 网络
还记得显存吗?它是 gpu 的工作台。gpudirect rdma 相当于让工作台上的货物直接装上高速货车,不必先搬到 cpu 的仓库。
roce:在以太网上使用 rdma
roce 是在以太网上实现 rdma 的技术。它像是在通用公路上开辟一条高性能直达运输通道。
因此,跨服务器 gpu 通信常见两条路线:
- infiniband:使用专用高速网络;
- roce:在高速以太网上运行 rdma。
roce 对网络拥塞和丢包比较敏感,因此通常还会配置接近无损的网络。所谓无损,不是绝对不会丢包,而是通过流量控制和拥塞管理,尽量避免因拥塞造成丢包。
pfc 会在前方拥塞时暂停某一类流量,像交警暂时关闭入口;ecn 会提前给数据包标记拥塞,像高速公路提示“前方拥堵,请减速”。
fabric 和 leaf-spine:看完整张网络
fabric 指由网卡、交换机、链路和拓扑共同组成的整个高速网络。单条链路是一条路,fabric 则是整座城市的道路系统。
leaf-spine 是数据中心常见的网络结构:服务器连接 leaf 交换机,leaf 再连接 spine 交换机。leaf 像小区道路,spine 像连接不同小区的城市主干道。
设计网络时还要留意 oversubscription,也就是带宽超卖。例如 8 台服务器各有一张 400g 网卡,接入总带宽为 3200g,但交换机向上的总带宽只有 800g。所有服务器同时通信时,就会争抢这条较窄的出口。
nccl 如何指挥 gpu 通信
底层道路准备好以后,还需要一个了解这些道路的通信软件。
nccl 是 nvidia 提供的多 gpu 通信库。它不负责修建道路,而是根据硬件连接选择通信方式,并组织多张 gpu 共同完成数据交换。
可以把 nccl 想成物流调度员:
- 同一台服务器内,可能使用 nvlink 或 pcie;
- 不同服务器之间,可能使用 infiniband 或 roce;
- 程序提出“汇总所有结果”等任务,nccl 负责安排具体运输。
nccl 不是网卡,也不是 nvlink。nvlink 和 infiniband 是道路,nccl 是使用道路的软件。
rank:每个训练进程的工号
分布式训练通常会启动多个进程,每个进程有一个全局编号,叫 rank。一般情况下,一个进程控制一张 gpu。
两台 8 卡服务器启动 16 个进程时:
- world size 是总进程数,也就是 16;
- rank 是全局工号,范围是 0 到 15;
- local rank 是当前服务器内的编号,每台机器都是 0 到 7。
可以把 rank 理解为全公司的员工号,把 local rank 理解为当前厂房里的座位号。
process group 是一组需要相互通信的进程,相当于从 16 名工人中划出不同工作小组。
collective communication:一组 gpu 共同通信
多张 gpu 共同参加的通信操作叫 collective communication,也就是集合通信。它不是 gpu 0 单独给 gpu 1 发送一份数据,而是一组 gpu 一起完成某个任务。
常见操作包括:
| 操作 | 做什么 | 类比 |
|---|---|---|
| broadcast | 一张 gpu 把相同数据发给其他 gpu | 老师给全班发同一份资料 |
| allreduce | 汇总所有 gpu 的数据,再把结果发给所有 gpu | 汇总每个人的统计结果,再通知全员 |
| allgather | 收集各 gpu 的不同数据,让所有 gpu 得到完整数据 | 每人复印自己的一页,最后拿到完整文件 |
| reducescatter | 汇总全部数据,再把不同部分分给各 gpu | 汇总整份资料,再分章节保管 |
allreduce 在数据并行训练中非常常见。每张 gpu 先计算自己那部分数据的梯度,nccl 再汇总梯度,并把相同的最终结果发回所有 gpu。
nccl 还会组织具体的传输路径。ring 会让 gpu 组成逻辑环,材料依次传给相邻工人,适合充分利用带宽;tree 会按树形逐级汇总和分发,步骤较少,通常更适合对延迟敏感的通信。
模型怎样拆到多张 gpu
道路和调度员已经就位,接下来要决定工作本身怎样分配。
batch 是一批同时处理的训练数据或推理请求。一次让 gpu 处理 32 条数据,batch size 就是 32。它像让货车装一批包裹再出发,而不是一次只送一个。
data parallel:模型相同,数据不同
data parallel,也就是数据并行,让每张 gpu 保存一份完整模型,但处理不同的数据。
gpu 0:完整模型 + 第 1 批数据
gpu 1:完整模型 + 第 2 批数据
gpu 2:完整模型 + 第 3 批数据
每张 gpu 计算完成后,再通过 allreduce 汇总梯度。还记得 allreduce 吗?它负责汇总所有人的结果,再把最终结果发给所有人。
ddp 是 pytorch 常用的数据并行实现,通常一张 gpu 对应一个进程,每个进程保存完整模型,并通过 nccl 同步梯度。
tensor parallel:一起加工同一个零件
tensor parallel,也就是张量并行,会把模型同一层中的大矩阵拆到多张 gpu 上共同计算。
它像一个零件太大,一名工人无法独立加工,只能由几名工人同时完成。由于大家在处理同一层,计算过程中需要频繁交换数据,因此很依赖 nvlink 和 nccl。
pipeline parallel:每个工位负责不同步骤
pipeline parallel,也就是流水线并行,会把模型的不同层放到不同 gpu。
gpu 0:模型前部
↓
gpu 1:模型中部
↓
gpu 2:模型后部
它和工厂流水线很像,不同工位负责不同加工步骤。实际训练还需要通过微批次减少某些工位等待造成的空闲。
moe 和 expert parallel:把任务送给不同专家
moe 是 mixture of experts,也就是混合专家模型。模型中有多个专家模块,不同输入会被分配给不同专家处理。
expert parallel,也就是专家并行,会把不同专家放到不同 gpu。它像医院设置多个科室,再根据病人的问题分诊。由于输入需要在 gpu 之间重新分发,专家并行也会带来大量通信。
zero:大家分别保管模型数据
普通数据并行中,每张 gpu 都保存完整的模型参数、梯度和优化器状态,显存压力很大。
zero 会把这些数据拆分到不同 gpu 保存。它像原来每名工人都有一整套厚重资料,现在改成每人保管一部分。需要完整数据时,可以通过前面介绍的 allgather 收集;计算后再用 reducescatter 汇总并重新拆分。
这些方法不是简单的替代关系。实际系统可能同时使用数据并行、张量并行、流水线并行或专家并行,但组合越复杂,对网络、拓扑和调度的要求也越高。
kubernetes 如何管理 gpu
前面解决的是 gpu 怎样计算和通信。放到共享集群以后,还要回答:哪项任务可以使用哪些 gpu,应该运行在哪台服务器上?
kubernetes 负责管理容器化工作负载,container runtime 则真正把容器启动起来。可以把 kubernetes 看成下达任务的管理中心,把 containerd 等运行时看成执行启动命令的工作人员。
nvidia container toolkit:让容器看到 gpu
nvidia container toolkit 让容器能够访问宿主机上的 nvidia gpu 和驱动。没有它,即使服务器本身能够看到 gpu,容器里的程序也可能无法使用。
device plugin:把 gpu 报告给 kubernetes
device plugin 会发现节点上的 gpu,并把可调度资源报告给 kubernetes。pod 可以通过资源限制申请 gpu:
resources:
limits:
nvidia.com/gpu: 1
这段配置表示 pod 需要 1 张 gpu。device plugin 负责登记资源,但不会自动把普通程序改造成 gpu 程序。
gpu operator:自动维护 gpu 软件栈
gpu operator 帮助 kubernetes 部署和管理 gpu 驱动、nvidia container toolkit、device plugin 以及监控组件等。
它像 gpu 节点的自动化维护人员,但自动化并不代表可以忽略版本兼容。驱动、容器工具、kubernetes 和工作负载仍然需要经过验证。
dcgm 是 nvidia 提供的 gpu 监控和健康检查工具,可以采集利用率、显存、温度、功耗和硬件错误等信息。
scheduler:决定任务去哪台服务器
scheduler 是 kubernetes 的调度器,负责决定 pod 运行在哪个节点。它像工厂的人事调度员,选择把任务交给哪间厂房。
普通调度只找到空闲 gpu 还不够。分布式任务还经常需要:
- gang scheduling:等所需资源全部准备好后一起启动,避免只启动一部分进程后长期互相等待;
- topology-aware scheduling:调度时考虑 gpu、cpu、网卡和交换机之间的连接关系。
还记得 topology 吗?它是设备的道路地图。拓扑感知调度不只找空闲工人,还会尽量把需要密切协作的工人安排在道路更近的位置。
mig 和 mps:多人共享一张 gpu
mig 可以把支持该功能的一张 gpu 硬件隔离成多个较小的 gpu 实例,像用实体墙把大办公室分成独立小办公室。
mps 则让多个 cuda 进程更高效地共享一张 gpu,像多人共用同一间开放办公室。
- mig 的隔离更强,资源边界更清晰;
- mps 的共享更灵活,但隔离通常较弱。
选择哪一种,要看模型大小、隔离要求、并发模式和硬件支持,而不是只比较切分后的数量。
训练数据放在哪里
gpu 计算很快,但如果数据供应不上,工人还是会停下来等仓库送货。
nvme:服务器旁边的高速小仓库
nvme 通常指服务器内的高速 ssd,适合缓存热点训练数据、临时文件或中间结果。它距离计算节点近、速度快,但容量和跨节点共享能力通常有限。
object storage:容量更大的中央仓库
object storage,也就是对象存储,适合长期保存数据集、模型文件、训练结果和 checkpoint。amazon s3 和 minio 都是常见例子。
对象存储像远处的大型中央仓库,容量大、扩展方便,但读取路径比本地 nvme 更长。
parallel file system:拥有多个装卸口的仓库
parallel file system,也就是并行文件系统,允许多台服务器同时高速读取文件。lustre 和 ibm spectrum scale 是常见例子。
普通仓库可能只有一个门,并行文件系统则像拥有许多装卸口,可以同时服务大量计算节点。
checkpoint:训练过程中的存档
checkpoint 是训练过程中保存的模型状态。例如每训练 1000 步保存一次,机器故障后就可以从最近的 checkpoint 继续,而不必从头开始。
它就像游戏存档,但一次 checkpoint 可能包含大量参数和训练状态,保存过于频繁也会影响训练性能。
如果存储供料速度跟不上 gpu 计算速度,就会出现 i/o bottleneck,也就是输入输出瓶颈。常见现象是 gpu 利用率忽高忽低,gpu 经常等待,而磁盘或存储网络已经跑满。
模型训练完成后如何提供服务
model serving 是把训练好的模型部署成 api 服务。用户调用聊天接口,后端使用 gpu 生成回答。
训练关注如何高效地更新模型参数,推理则更关注请求吞吐量、响应延迟和显存利用率。
kv cache:记住已经计算过的上下文
kv cache 保存大模型推理过程中已经计算出的注意力状态。继续生成下一个 token 时,可以复用前面的结果,不必重新计算全部上下文。
它像服务员记住了顾客之前点过什么。优点是生成更快,代价是占用显存,而且对话越长,kv cache 通常越大。
batching:一车多装几个包裹
推理中的 batching 会把多个用户请求组合成一批交给 gpu 处理。还记得 batch 吗?它是一批同时处理的数据或请求。
如果每个请求都单独执行,gpu 可能无法充分发挥并行能力。把多个请求放在一起,就像快递车装满一批包裹再出发。
continuous batching:有空位就接收新请求
普通静态批处理可能要等同一批中的所有请求结束,才能开始下一批。但不同用户生成的内容长短不一,较短请求完成后,对应位置就会闲置。
continuous batching,也就是连续批处理,会在生成过程中动态移除已完成的请求,并加入等待中的新请求。
时间 →
a a a 完
d d d d d
b b b b 完
e e e e e
c c c c c c c c c
它像公交车有人下车后马上让新乘客上车,不必等所有乘客同时到站。连续批处理可以提高 gpu 利用率,但同时塞入太多请求也可能增加单个请求的延迟,需要在吞吐量和响应速度之间取舍。
用一个完整场景串起来
现在假设我们有两台 gpu 服务器,每台 8 张 gpu,共 16 张 gpu,用它们完成一次分布式训练,再部署推理服务。
1. 准备数据
训练数据长期保存在 object storage,也就是中央大仓库。训练前,把热点数据缓存到本地 nvme,也就是服务器旁边的高速小仓库。
cpu 读取和预处理数据,再交给 gpu 执行并行计算。
2. kubernetes 分配 gpu
device plugin 把每台服务器的 gpu 数量报告给 kubernetes。scheduler 选择合适节点,gang scheduling 等 16 张 gpu 全部可用后再一起启动任务。
topology-aware scheduling 根据道路地图,尽量选择连接关系更好的 gpu 和网卡。
3. 启动训练进程
系统启动 16 个进程,每个进程控制一张 gpu:
- world size,也就是总进程数,为 16;
- rank,也就是全局工号,为 0 到 15;
- local rank,也就是单机座位号,为 0 到 7。
4. 在单台服务器内通信
8 张 gpu 优先通过 nvlink 这条 gpu 专用高速路传输数据,nvswitch 这座立交桥让它们高速互通。如果没有合适的 nvlink 路径,则可能使用 pcie p2p。
5. 在两台服务器之间通信
hca 把服务器接入 infiniband 高速网络。rdma 减少 cpu 中转,gpudirect rdma 进一步让 gpu 显存中的数据更直接地进入高性能网卡。
6. nccl 同步训练结果
每张 gpu 使用相同模型处理不同数据,这叫数据并行。完成一轮计算后,nccl 这个物流调度员执行 allreduce:
收集 16 张 gpu 的梯度
↓
汇总所有梯度
↓
把最终结果发回 16 张 gpu
nccl 会根据拓扑选择底层道路:单机内可能走 nvlink,跨服务器则可能走 infiniband 和 gpudirect rdma。
7. 模型太大时拆分工作
如果一张 gpu 放不下模型,可以组合使用:
- zero:让多张 gpu 分别保管参数、梯度和优化器状态;
- tensor parallel:多张 gpu 共同计算同一层;
- pipeline parallel:不同 gpu 负责不同模型层;
- expert parallel:不同 gpu 保存不同 moe 专家。
拆分解决了单卡显存不足的问题,也增加了通信量。模型怎样拆,必须和节点内外的网络拓扑一起考虑。
8. 保存并部署模型
训练过程中定期保存 checkpoint,也就是训练存档。训练结束后,通过 model serving 把模型部署成 api。
推理服务使用 kv cache 复用历史计算,使用 batching 同时处理多个请求,再通过 continuous batching 动态填补空闲位置,提高 gpu 利用率。
最容易记住的版本
如果暂时记不住所有细节,可以先保留下面这张地图:
| 概念 | 最短解释 |
|---|---|
| gpu | 负责并行计算的工人 |
| 显存 | gpu 的工作台 |
| pcie | 服务器内部的通用道路 |
| nvlink | gpu 之间的专用高速道路 |
| nvswitch | 连接多条 nvlink 的立交桥 |
| hca | 面向高性能网络的网卡 |
| infiniband | 服务器之间的专用高速网络 |
| rdma | 减少 cpu 中转的远程内存访问 |
| gpudirect rdma | 让 gpu 数据更直接地进入高性能网卡 |
| nccl | 指挥多张 gpu 通信的物流调度员 |
| allreduce | 汇总所有 gpu 的结果,再发给所有 gpu |
| rank | 分布式进程的全局工号 |
| data parallel | 模型相同,每张 gpu 处理不同数据 |
| tensor parallel | 多张 gpu 共同计算模型同一层 |
| pipeline parallel | 不同 gpu 负责不同模型层 |
| zero | 多张 gpu 分别保存模型训练状态 |
| device plugin | 把 gpu 资源报告给 kubernetes |
| gang scheduling | 等所需资源凑齐后一起启动 |
| checkpoint | 训练过程中的模型存档 |
| kv cache | 保存推理中已计算的上下文状态 |
| continuous batching | 有计算空位就动态加入新请求 |
总结
ai infra 不是一堆彼此孤立的缩写,而是一条连续的工作链路:
gpu 和显存负责计算
↓
pcie、nvlink 和 nvswitch 连接单机内的设备
↓
infiniband、roce 和 rdma 连接不同服务器
↓
nccl 组织多张 gpu 交换数据
↓
数据并行、张量并行等方法拆分训练工作
↓
kubernetes 分配和管理 gpu 资源
↓
nvme、对象存储和并行文件系统供应数据
↓
模型服务通过缓存和批处理提高推理效率
以后再遇到陌生术语,可以先问三个问题:
- 它是在负责计算、修路、调度,还是保存数据?
- 它工作在一张 gpu、一台服务器,还是多台服务器之间?
- 它解决的是模型放不下、通信太慢、资源难调度,还是 gpu 经常空闲?
只要把新概念放回这条工作流程里,就不需要孤立地背诵每个缩写。