我们一直追求更大的模型，DeepSeek-R1 有数千亿参数，使用了数十万亿 token 的训练数据，涉及算力、存储、通信等多维度的工程挑战。有了 PyTorch 深度学习框架，只是 AI 应用落地的万里长征第一步。接下来我们将讨论深度学习框架之上的模型训练的挑战。

3.1 存得下

DeepSeek-R1 模型大小=670GB，而一台 GPU 服务器有8张H20卡，提供768GB显存，足够存下一个完整的 DeepSeek 模型。那整个行业为什么还投入大量的人力物力，顶着通信延时造成的算力损耗，也要建设分布式 GPU 集群？核心原因是单台 GPU 服务器“存不下”。

3.1.1 显存刺客：中间激活

如下图所示的模型，x1/x2/x3/x4 这些中间变量就是"中间激活"。它们是神经网络前向传播（Forward）的“堆栈帧（Stack Frame）”——记录每一层处理后的数据快照，确保反向传播（Backward）可回溯梯度，根据预测误差调整模型权重，最小化损失函数。

这些中间激活为什么会成为"显存刺客"？是因为中间激活的空间复杂度是和输入数据长度正相关的，特别的，对于 LLM 来说是O（N²）正比于输入数据长度的平方，这是一个指数爆炸式增长的数字。类似函数递归不断增长的“堆栈帧”导致的内存溢出，我们遇到了 AI Infra 的 OOM（Out of Memory）挑战。

借助 PyTorch 的 profiler 工具，我们可以直观地看到这个OOM。下图是训练过程中不同阶段的显存分配，包括模型参数（Parameter）、优化器状态（Optimizer state）、中间激活（Activation）、梯度（Gradient）。在前向传播结束后出现一个显存占用（中间激活）的尖峰，远大于模型参数本身。

3.1.2 模型并行

传统后台服务使用分片（Sharding）策略解决单机存不下的问题。与之相似，AI Infra 提出“模型并行”，就是将单个大模型拆分为多个子模块，并分布到不同 GPU 上协同工作，通过通信来共享数据。有不同的“拆分模型”策略，例如按模型模块划分，按张量（Tensor）划分的，也可以将多种拆分方法结合起来一起使用。PyTorch 深度学习框架和开源方案 Megatron 都能帮助我们高效地实现模型并行。

3.2 算得快

建设分布式 GPU 集群的原因，一个是因为“单机存不下”，另外一个是提升训练速度。但简单的机器堆叠，算力不一定有线性的增长。因为分布式训练并不是简单地把原来一个 GPU 做的事情分给多个 GPU 各自做。需要协调多个 GPU 机器计算任务分配，GPU 机器之间的数据传输会引入网络IO和通信开销，降低训练速度。

3.2.1 通信计算重叠

如下图所示的常规训练时序是串联式的，存在许多网络 IO，GPU 利用率低，训练速度慢。我们希望 GPU 大部分时间都在计算，而不是花在数据传输或等待其他 GPU 的工作上。

传统后台服务我们通过多线程或异步 IO 避免阻塞 CPU 主线程，与之相似，AI Infra 提出通信计算重叠的方法论。GPU 编程模型中有流（stream）的概念，一个流表示一个 GPU 操作队列，该队列中的操作将以添加到流中的先后顺序而依次执行。不同流之间可以并行执行。那么通过令计算和通信操作加入不同的流中，可以做到二者的执行在时间上重叠。例如 TorchRec 的 训练流水线 能帮助我们实现高效的通信计算重叠。

AI 模型训练成本很高，优秀如 DeepSeek 也要烧掉500万美金，但再贵也只是一次性的。而模型推理的成本更高，因为用户越多，AI 模型推理次数越多，总成本越高。模型推理面对的挑战和传统 Infra 非常相似，主要是2个挑战：高吞吐（降本），低延时（增效）。

4.1 降低延时

现在的 AI 模型越来越多地直面终端用户，需要和用户进行实时的交互，例如文本对话和语音合成。模型推理耗时过高，会直接造成用户体验受损，用户流失与转化率下降。

传统后台服务我们使用链接复用、缓存、柔性等技术降低系统响应时间。AI Infra 也有相似的做法。

4.1.1 CUDA Graph

在 GPU 编程模型中，CPU 和 GPU 是异构的，CPU 通过 API（例如 CUDA API） 向 GPU 提交任务，然后异步等待 GPU 的计算结果返回。GPU 收到任务后，会执行内核启动、内存拷贝、计算等操作。这个过程中，涉及到 CPU 与 GPU 之间的通信、驱动程序的处理以及 GPU 任务的调度等环节，会产生一定的延迟。模型推理需要执行大量重复的 GPU 操作，每个的 GPU 操作都要重复执行上述环节，这些非核心的 GPU 开销会成倍数地放大，影响最终响应时间。

在传统后台服务，我们使用 Redis 的 Lua 脚本封装多个 Redis 操作和计算逻辑，一次提交，减少网络开销。与之相似，AI Infra 利用 CUDA Graph 技术将多个 GPU 操作转化为一个有向无环图（DAG），然后一次性提交整个 DAG 提交到 GPU 执行，由GPU自身来管理这些操作的依赖关系和执行顺序，从而减少 CPU 与 GPU 之间的交互开销。

4.1.2 KV Cache：空间换时间

LLM 大模型推理存在大量矩阵乘法运算，且高度依赖上下文信息。每次推理都需要将之前生成过的词重新输入模型进行计算。这种计算方式使得复杂度达到了 O（N²），其中必然存在大量的重复计算。

例如，给定“天气”，模型会逐个预测剩下的字，假设接下来预测的两个字为“真好“。

将”真“拼接到”天气“的后面，即新的输入为”天气真“，再预测”好“

观察到，经过多次预测后， 和  的结果上半部分都是相同的。这是由于 LLM 模型结构的特殊设计导致的。这些重复计算的结果可以缓存（即 KV Cache）下来，空间换时间，减少计算量。几乎所有的 LLM 推理框架都支持了 KV Cache，例如vLLM 。

4.1.3 流式响应

有时候模型推理延时实在避免不了，可以从工程交互上想办法。传统后台服务的 RPC 通信是一问一答方式，这种方式不太适合语音合成或者文本对话的场景。因为大模型推理需要几秒-几十秒，如果等待模型推理结束才展示结果，用户会等待较长的时间，体验很差。

流式响应就是当模型推理计算得到第一个token或者第一个音频帧的时候，立马展示或者播放给用户，同时后续的模型推理结果在已经建立的 TCP 流上继续顺序传输。工程上从关注模型推理的整体耗时，改为关注首token或首个音频帧的耗时。几乎所有的 LLM 推理框架都支持了流式响应。

4.2 提高吞吐量

提高吞吐量是程序员在传统 Infra 领域孜孜不倦的追求，因为更高的吞吐量意味着更低的机器成本。实现 AI 应用的高吞吐本质上就是提高昂贵的 GPU 的利用率，让 GPU 单位时间能完成更多的任务。

尽管模型推理需要执行万亿次浮点运算，但 GPU 有大量的计算单元（CUDA Cores），单个请求的模型推理很难令 GPU 利用率达到饱和。提高 GPU 利用率有2个方法：传统批处理和连续批处理。这里的“传统批处理”是相对于“连续批处理”这样的新型批处理方式而言的。

4.2.1 传统批处理

其实传统后台服务也大量使用了批处理，例如 Redis 的 MGet 命令，单次请求就完成所有 key 的获取，将 N 次网络往返（RTT）压缩为1次。与之相似，模型推理的批处理就是将多个输入样本打包（batch），将原本串行的 N 次轻量的推理计算，合并为 1 次重量的计算，实现单位时间内处理更多的请求，提高了 GPU 利用率。

“打包输入样本”是一个共性需求，大部分推理框架都提供该功能，例如 Triton Inference Server 的 Batcher 。

模型批量推理流程图

4.2.2 连续批处理

传统批处理类似 “固定班次的公交车”：乘客（请求）必须等待发车时间（组建一个batch），发车后所有乘客同步前进。即使有乘客提前下车（短请求完成），车辆仍需等待所有乘客到达终点（长请求完成）才能返程接新乘客。传统批处理存在着资源浪费：GPU 要等待长请求处理完，不能处理新的请求而空闲。

这个问题在 LLM 应用领域显得特别突出，因为不同用户请求 Prompt，模型的回答结果长度差异巨大，如果使用传统批处理，GPU 空闲率很高。这个本质上是个任务调度问题，传统后台服务我们使用工作窃取算法（work stealing）解决线程空闲问题，与之相似，AI Infra 提出“连续批处理”解决这个问题。

连续批处理类似“随时随地拼车的顺风车”，每辆车（GPU）在行程中可随时上/下客。新乘客（请求）直接加入当前车辆的空位（空闲计算单元），已完成的乘客立即下车（释放资源）。几乎所有的 LLM 推理框架都支持了连续批处理能力，例如 vLLM 的 Continuous Batching

连续批推理流程图

五、结语

AI Infra 面对的工程挑战，例如计算、存储、通信，大部分是新时代的老问题，我们在传统 Infra 领域都能找到对应的场景和解决思路。差异只在于战场从 CPU 转移到 GPU，传统后台工程师积累的方法论，依然可以无缝衔接到 AI Infra。