William 的技术笔记

AI Infra 工程师：要懂什么，要练什么

2026-10-25 · #AI Infra#GPU#分布式训练#推理#职业

几年前说“基础设施工程师”，大家想到的是机房、K8s、CI/CD。现在说 “AI Infra”，指的是另一套东西：几千张 GPU 怎么一起训练一个模型，一个 70B 的模型怎么在线上以几十毫秒的延迟回答问题，一次训练挂了怎么在十分钟内从检查点恢复。

这个岗位的需求在过去两年涨得很快，但它的边界也很模糊：有人写 CUDA 算子，有人调 NCCL，有人做调度平台，有人做推理服务，有人搭 Agent 的沙箱，都叫 AI Infra。这篇文章想把这件事拆清楚：AI Infra 工程师在做什么、需要了解哪些知识、需要练哪些技能。

AI Infra 在做什么

一句话概括：让模型训得动、训得快、跑得起、跑得便宜。

可以把整个技术栈从下往上分成几层：

层次 关心的问题 典型技术
硬件与互联 算力、显存、带宽 GPU / TPU、HBM、NVLink、InfiniBand / RoCE
算子与编译 单卡跑多快 CUDA、Triton、FlashAttention、torch.compile
分布式训练 多卡怎么协同 FSDP / ZeRO、Megatron、NCCL、DeepSpeed
推理服务 延迟、吞吐、成本 vLLM、SGLang、TensorRT-LLM、量化
集群与平台 资源怎么分、坏了怎么办 Kubernetes、Slurm、调度器、可观测性
数据与存储 数据怎么喂、检查点怎么存 数据管道、并行文件系统、对象存储
应用侧基础设施 模型怎么被安全地用起来 沙箱、工具网关、评估平台

实际岗位通常会落在其中一两层，常见的方向有：

  • 训练 Infra：负责大规模预训练、后训练（SFT、RL）的框架和稳定性；
  • 推理 Infra：负责模型上线，压延迟、提吞吐、降成本；
  • 集群 / 平台 Infra：负责 GPU 集群的调度、监控、故障处理，给研究员提供“开箱即用”的算力；
  • 数据 Infra：负责训练数据的处理管道和存储；
  • Agent Infra：负责沙箱、工具调用、评估这些让 Agent 能跑起来的东西。前面写过的 Agent 沙箱 和 MCP 还是 CLI 都属于这一层。

不管在哪一层，这个岗位和传统后端最大的区别是：你面对的资源极其昂贵，而且瓶颈往往在硬件上。一张 H100 一小时的成本是普通 CPU 机器的几十倍，一个千卡训练任务慢 10%，一个月就能多烧掉一辆车。所以 AI Infra 的工作几乎都围绕着一个问题：这块 GPU 现在有没有在干活？没有的话，它在等什么？

需要了解的知识

1. GPU 和硬件：先知道机器长什么样

不需要会设计芯片，但必须对硬件有“物理直觉”：

  • GPU 的结构：SM（流式多处理器）、Tensor Core、寄存器 / 共享内存 / L2 / HBM 的层级。很多优化的本质就是让数据尽量待在更快的那一层。
  • 算力 vs 带宽：一张 H100 的 BF16 稠密算力约 1 PFLOPS，HBM 带宽约 3.35 TB/s。两者的比值决定了一个算子是算力瓶颈还是带宽瓶颈，这就是 Roofline 模型。大矩阵乘法通常是前者，逐元素操作、LayerNorm、以及大模型推理的 decode 阶段通常是后者。
  • 互联：机内 GPU 之间走 NVLink / NVSwitch（几百 GB/s 量级），机间走 InfiniBand 或 RoCE（每卡 400 Gb/s 量级）。机内和机间差了一个数量级，这直接决定了不同并行策略该放在哪一层。
  • 显存是硬约束：80GB 就是 80GB，放不下就是 OOM，没有“稍微慢一点”的选项。

2. 模型本身：会算账

AI Infra 工程师不一定要会设计模型，但一定要看得懂 Transformer，并且能算出它要多少资源。最常用的几笔账：

训练算力：训练一个稠密模型大约需要 6 × 参数量 × token 数 次浮点运算。比如 7B 模型训 2T token：

6 × 7e9 × 2e12 ≈ 8.4e22 FLOPs
H100 按 40% MFU 算，约 400 TFLOPS 有效算力
8.4e22 / 4e14 ≈ 2.1e8 GPU·秒 ≈ 5.8 万 GPU·小时
用 1024 张卡，大约 2.4 天

训练显存：混合精度 + Adam，每个参数大约要 16 字节（BF16 权重 2 + BF16 梯度 2 + FP32 主权重 4 + Adam 两个动量 8），还没算激活值。7B 模型就是约 112 GB，单张 80GB 卡放不下，这就是为什么需要 ZeRO / FSDP 把这些状态切到多张卡上。

推理 KV Cache：每个 token 要缓存的 K、V 大小是 2 × 层数 × KV 头数 × 头维度 × 字节数。以 Llama 3 8B（32 层、8 个 KV 头、头维度 128、BF16）为例：

2 × 32 × 8 × 128 × 2 B = 128 KB / token
8K 上下文的一个请求 ≈ 1 GB

这意味着一张卡能同时服务多少请求，往往不是由算力决定，而是由 KV Cache 能占多少显存决定。

能随手算出这几笔账，是 AI Infra 工程师最基本、也最有用的能力之一。它能让你在动手之前就知道一个方案行不行。

3. 分布式训练：并行与通信

单卡放不下、单卡太慢，就要多卡并行。主流的几种切法：

  • 数据并行（DP）：每张卡一份完整模型，各自处理不同数据，最后同步梯度。ZeRO / FSDP 是它的升级版，把优化器状态、梯度、参数也切开，用通信换显存；
  • 张量并行（TP）：把一个矩阵乘法切到多张卡上算，通信非常频繁，通常只在机内（NVLink）做；
  • 流水线并行（PP）：按层切，不同卡负责不同的层，要处理好“气泡”（有的卡在等）；
  • 序列 / 上下文并行（SP / CP）：长上下文场景下按序列维度切；
  • 专家并行（EP）：MoE 模型里把不同专家放到不同卡上，带来大量 all-to-all 通信。

真实的大规模训练通常是这几种的组合，比如“机内 TP + 机间 PP + 最外层 DP”。

这些并行策略背后都是集合通信：all-reduce、all-gather、reduce-scatter、all-to-all。你需要知道它们各自的通信量、NCCL 是怎么实现的（Ring、Tree），以及怎么让通信和计算重叠起来。

除了“跑得快”，大规模训练还有一个经常被低估的问题：稳定性。千卡任务里，每天坏几张卡、掉几个网口是常态。检查点怎么存得快、故障怎么快速定位、任务怎么自动重启，往往比再提 5% 的 MFU 更重要。

4. 推理：一个完全不同的优化问题

推理和训练的瓶颈不一样。LLM 推理分两个阶段：

  • Prefill：一次性处理整段输入，计算密集，和训练很像；
  • Decode：每次只生成一个 token，但每一步都要把全部权重从 HBM 读一遍，严重受限于显存带宽。

举个例子，8B 模型 BF16 权重约 16 GB，H100 带宽约 3.35 TB/s，那么 batch size 为 1 时，每生成一个 token 至少要约 5ms，单请求上限大约 200 token/s。想提高吞吐，就要把多个请求拼在一起，让一次权重读取服务更多 token。

围绕这个特点，推理优化的常见手段有：

  • Continuous Batching：请求随到随拼，而不是等一批凑齐；
  • PagedAttention：像操作系统的虚拟内存一样分页管理 KV Cache，减少碎片（vLLM 的核心思想）；
  • Prefix Caching：多个请求共享同样的系统提示词时，复用已经算好的 KV；
  • 量化：INT8、FP8、INT4 权重量化，KV Cache 量化，用精度换带宽和显存；
  • 投机解码（Speculative Decoding）：用小模型先猜几个 token，大模型一次验证，把串行的 decode 变成部分并行；
  • PD 分离：把 prefill 和 decode 放在不同的机器上，各自按自己的瓶颈优化。

衡量推理服务的指标也要熟悉：TTFT（首 token 延迟）、TPOT / ITL（每个输出 token 的间隔）、吞吐、以及最终的每百万 token 成本。这几个指标经常互相拉扯，没有银弹，只有权衡。

5. 算子与编译：单卡的极限

当上面的手段都用过了，剩下的性能往往藏在算子里。

  • CUDA：至少要能读懂一个 kernel，理解 thread / block / warp、共享内存、合并访存这些概念；
  • Triton：用 Python 写 GPU kernel，门槛比 CUDA 低很多，现在很多自定义算子都用它；
  • 算子融合：把多个小算子合成一个，减少 HBM 读写。FlashAttention 就是最经典的例子，它本质上没有减少计算量，而是通过分块（tiling）避免了把巨大的注意力矩阵写回显存；
  • 编译器：torch.compile、XLA 这类工具会自动做融合和图优化，知道它们能做什么、什么时候会失效。

不是每个 AI Infra 工程师都要手写 kernel，但都应该知道性能是怎么在这一层流失的。

6. 集群、调度与可观测性

有了几千张卡，就要管好它们：

  • 调度：Kubernetes 和 Slurm 是两大主流。AI 场景有一些特殊需求，比如 Gang Scheduling（一个训练任务的 512 张卡必须同时拿到，否则谁都别跑）、拓扑感知调度（尽量让一个任务的卡在同一个交换机下）、抢占和配额；
  • 故障处理：GPU 掉卡、ECC 错误、NCCL 超时、网口 flapping。要能快速判断“是代码的问题还是机器的问题”，并自动把坏节点隔离出去；
  • 可观测性：GPU 利用率、显存、功耗、温度、网络流量、训练的 loss 和吞吐，都要能看到、能告警、能回溯。

7. 数据与存储

模型训练的另一端是数据：

  • 数据管道：清洗、去重、分词、打包，大规模下通常要用 Spark、Ray 这类分布式框架；
  • 数据加载：训练时 DataLoader 不能成为瓶颈，GPU 不应该等数据；
  • 存储：训练数据和检查点的规模都是 TB 到 PB 级。并行文件系统（如 Lustre）、对象存储、本地 NVMe 缓存各有适用场景。一个 70B 模型的完整检查点（含优化器状态）有上 TB，存得慢就意味着训练停得久。

8. Agent 时代的新问题

模型能力变强之后，AI Infra 的范围还在往上延伸。Agent 需要执行代码、调用工具、访问网络，于是出现了一批新的基础设施问题：

  • 怎么给每个 Agent 一个隔离、快速启动、可快照的执行环境（见 Agent 沙箱）；
  • 怎么统一地接入和管控工具，处理凭证和权限（见 MCP 还是 CLI）；
  • 怎么在大规模上跑评估和 RL 环境，几万个并发沙箱同时执行任务。

这部分和传统的训练、推理 Infra 不太一样，更接近云计算和安全，但正在变成一个很重要的方向。Agent 本身的概念可以看 搭建 Agent 的基础知识。

需要练的技能

知识可以查，技能要靠练。下面这几项是我认为最值得刻意投入的。

1. 估算能力

前面算的那几笔账，就是这项技能。拿到一个需求，能在几分钟内估出：要多少卡、要多少显存、理论上限是多少、瓶颈大概在哪。估算能帮你在写第一行代码之前就排除掉大部分错误方案。

练习方法很简单：每看到一个模型、一篇论文、一个性能数字，就自己算一遍，看能不能对上。

2. Profiling 与性能分析

“感觉慢”不算结论。要熟练使用：

  • PyTorch Profiler：看每个算子的耗时、通信和计算是否重叠；
  • Nsight Systems（nsys）：看整个时间线，GPU 在哪里空闲、CPU 在哪里卡住；
  • Nsight Compute（ncu）：深入单个 kernel，看它是算力瓶颈还是带宽瓶颈；
  • NCCL 测试工具（nccl-tests）：测集合通信的实际带宽，判断网络是否正常。

核心习惯是：先测量，再优化；优化之后再测量。拿着 profile 结果和 Roofline 的理论值比，才知道还有多少空间。

3. 编程：Python 为主，C++ / CUDA 为辅

  • Python：绝大多数框架和工具都是 Python 写的，要写得熟练，也要知道它的性能陷阱（GIL、同步点、隐式的 CPU-GPU 拷贝）；
  • C++ / CUDA：推理引擎、通信库、自定义算子都离不开它。不一定要精通，但要能读、能改、能调试；
  • 系统编程：进程、线程、内存、文件 IO、网络，这些操作系统层面的东西在排查问题时经常用到。

4. Linux 与网络排障

AI Infra 的很多问题最后都落到操作系统和网络上：为什么这个进程卡住了，为什么这两台机器之间带宽只有一半，为什么容器里看不到 GPU。

要熟悉 nvidia-smi、dcgm、ibstat、perf、strace、htop 这一套工具，理解 RDMA、PCIe 拓扑、NUMA 亲和性这些概念。

5. 分布式系统的调试

单机程序崩了看栈就行。一个千卡任务挂了，可能是一张卡的 ECC 错误导致它慢了一点，然后它的邻居在 all-reduce 上等超时，最后全部报 NCCL timeout。看到的报错和真正的原因可能隔着好几层。

这需要的是一种系统性的排查方法：缩小范围、二分定位、对比正常节点、看日志的时间线。这项能力很难从书上学，主要靠踩坑积累。

6. 读论文和读源码

这个领域的变化速度很快，新的并行策略、推理优化、硬件特性每隔几个月就会出现。很多最好的资料就是论文和开源代码本身：

  • 论文：Megatron-LM、ZeRO、FlashAttention、PagedAttention（vLLM）、Speculative Decoding 等都值得精读；
  • 源码：vLLM、SGLang、Megatron-LM、PyTorch FSDP、NCCL，读懂一个就能理解一大片。

7. 和研究员协作

AI Infra 工程师的“用户”通常是研究员。他们关心的是“我的实验能不能快点跑完”，而不是你用了什么调度算法。

能听懂研究员的需求、把 Infra 的约束翻译成他们能理解的语言（“这个 batch size 下你的显存会爆”“换成这个并行策略能快 30%”），是一项被低估但非常重要的能力。

一个可能的学习路径

如果从后端或通用 Infra 转过来，可以大致按这个顺序：

  1. 打基础：吃透 Transformer 的结构，能手算参数量、FLOPs、显存和 KV Cache；读一遍 CUDA 编程入门，理解 GPU 的内存层级；
  2. 跑起来：用 PyTorch 在单卡上训一个小模型，再用 FSDP / DDP 扩到多卡，用 Profiler 看看时间都花在了哪里；
  3. 上推理：用 vLLM 或 SGLang 部署一个开源模型，压测 TTFT 和吞吐，试试量化、Prefix Caching 带来的变化；
  4. 挑一个方向深入：训练、推理、集群、Agent Infra 里选一个，读它的核心论文和代码，最好能给开源项目提一个 PR；
  5. 积累实战：真正的大规模问题只有在真实的集群上才会遇到，有机会就去碰真实的故障和真实的性能瓶颈。

小结

  1. AI Infra 的目标是让模型训得动、训得快、跑得起、跑得便宜，核心问题是“GPU 有没有在干活”；
  2. 知识面要广：从 GPU 硬件、模型结构、分布式训练、推理优化，到集群调度、数据存储，再到 Agent 的执行环境；
  3. 会算账是基本功：FLOPs、显存、KV Cache、带宽上限，能在动手前判断方案是否可行；
  4. 推理和训练是两种优化问题：训练多半拼算力和通信，推理的 decode 阶段多半拼显存带宽；
  5. 技能要靠练：估算、Profiling、编程、排障、分布式调试、读论文和源码、与研究员协作。

模型会换代，框架会换代，硬件也会换代。但“理解资源的物理约束、用测量代替猜测、在复杂系统里找到真正的瓶颈”，这几件事会一直是 AI Infra 工程师的核心价值。