简介
简介(未完成)
程序员必须了解的AI系统设计与挑战知识 硬件演进 ==> 软件演进 ==> 训练挑战 ==> 推理挑战。
硬件演进
从CPU为中心到GPU为中心。传统基础设施以 CPU 为核心,通过多线程和微服务构建分布式系统,处理高并发请求(如 Web 服务)。这些都有成熟的方法论了(如”海量服务之道”)。主要工作是逻辑事务的处理,瓶颈在网络 I/O 和 CPU 核心数量。而 AI Infra 以 GPU 为核心,其设计目标从逻辑事务处理转向高吞吐浮点计算。此时CPU 多线程被 GPU 并行计算替代,内存被显存替代。为什么 GPU 会成为核心?是因为 LLM 大模型每次生成一个 token,都需要读取全量的模型参数。传统的 CPU + 内存的算力和带宽无法满足如此恐怖的计算密度,计算和通信都必须转移(offload)到 GPU 内完成。CPU 成为数据搬运工和“辅助处理器”。为了更直观地理解这个计算密度,我们做一个简单的计算。不考虑计算的延时,LLM 大模型生成一个 token 的耗时公式计算为:计算耗时 = 模型参数 * 数据精度 / 显存带宽。以 DeepSeek-R1-671B-A37B-FP8 模型为例,计算一个 token 耗时,H20 为 37B × 1byte ÷ 4000GB/s = 9ms,如果是 CPU 则为 37B × 1byte ÷ 64GB/s = 578ms。传统 Infra 的分布式理念貌似在 AI 时代失效了。传统 Infra 追求横向扩展,而 AI Infra 呈现 “AI 大型机”特性,是因为传统后台服务的可以容忍毫秒级延迟,但 AI 集群不行,GPU 的算力是 CPU 的数百倍,微秒级的延时等待也会造成很大的算力损耗,需要硬件的高度集成。
软件演进
相比传统后台应用的增删查改,AI 应用的新范式是模型训练和推理。模型训练是指通过海量数据拟合出一个复杂的神经网络模型,推理就是利用训练好的神经网络模型进行运算,输入的新数据来获得新的结论。
- 工欲善其事,必先利其器。传统后台应用依赖 tRPC 或 Spring 等微服务框架,帮助我们屏蔽负载均衡、网络通信等底层细节,我们可以把精力放在业务实现上。与之相似,AI 应用则依赖深度学习框架。如果没有深度学习框架(比如PyTorch),我们就可能陷入在茫茫的数学深渊中,挣扎于痛苦的 GPU 编程泥潭里。得益于动态计算图、自动微分和丰富的 Tensor 操作算子,PyTorch 能帮助我们快速实现模型设计。只需要描述模型结构+待学习的网络参数,不需要关心数学计算和 GPU 编程的细节。
- GPU 编程。绝大部分的 AI 应用,的确不需要我们手写数学计算的 GPU 代码。但为了满足模型创新的需求,有必要学习 GPU 编程。例如 Meta 发布的 HSTU 生成式推荐模型,核心的 hstu_attn 计算,如果直接用 PyTorch 框架算子组合实现,则时间复杂度为 O(M * N²) ,其中 M 和 N 是一个数量级,相当于O(N³) 。但是通过自定义内核,可以优化到 O(N²)。
模型训练的挑战
我们一直追求更大的模型,DeepSeek-R1 有数千亿参数,使用了数十万亿 token 的训练数据,涉及算力、存储、通信等多维度的工程挑战。有了 PyTorch 深度学习框架,只是 AI 应用落地的万里长征第一步。接下来我们将讨论深度学习框架之上的模型训练的挑战,即建设分布式 GPU 集群的原因。
- 存得下。
- 显存刺客:中间激活。在前向传播结束后出现一个显存占用(中间激活)的尖峰,远大于模型参数本身。
- 传统后台服务使用分片(Sharding)策略解决单机存不下的问题。与之相似,AI Infra 提出“模型并行”,就是将单个大模型拆分为多个子模块,并分布到不同 GPU 上协同工作,通过通信来共享数据。
- 算得快。简单的机器堆叠,算力不一定有线性的增长。因为分布式训练并不是简单地把原来一个 GPU 做的事情分给多个 GPU 各自做。需要协调多个 GPU 机器计算任务分配,GPU 机器之间的数据传输会引入网络IO和通信开销,降低训练速度。
- 通信计算重叠。传统后台服务我们通过多线程或异步 IO 避免阻塞 CPU 主线程,与之相似,AI Infra 提出通信计算重叠的方法论。GPU 编程模型中有流(stream)的概念,一个流表示一个 GPU 操作队列,该队列中的操作将以添加到流中的先后顺序而依次执行。不同流之间可以并行执行。那么通过令计算和通信操作加入不同的流中,可以做到二者的执行在时间上重叠。PS: 一个wrap阻塞了就跑另一个wrap?
模型推理的挑战
主要是2个挑战:高吞吐(降本),低延时(增效,用户体验)。
- 降低延时。传统后台服务我们使用链接复用、缓存、柔性等技术降低系统响应时间。AI Infra 也有相似的做法。
- 在 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 之间的交互开销。
- KV Cache:空间换时间。
- 流式响应
- 提高吞吐量。实现 AI 应用的高吞吐本质上就是提高昂贵的 GPU 的利用率,让 GPU 单位时间能完成更多的任务。
- 传统批处理。其实传统后台服务也大量使用了批处理,例如 Redis 的 MGet 命令,单次请求就完成所有 key 的获取,将 N 次网络往返(RTT)压缩为1次。与之相似,模型推理的批处理就是将多个输入样本打包(batch),将原本串行的 N 次轻量的推理计算,合并为 1 次重量的计算,实现单位时间内处理更多的请求,提高了 GPU 利用率。传统批处理类似 “固定班次的公交车”:乘客(请求)必须等待发车时间(组建一个batch),发车后所有乘客同步前进。即使有乘客提前下车(短请求完成),车辆仍需等待所有乘客到达终点(长请求完成)才能返程接新乘客。传统批处理存在着资源浪费:GPU 要等待长请求处理完,不能处理新的请求而空闲。这个问题在 LLM 应用领域显得特别突出,因为不同用户请求 Prompt,模型的回答结果长度差异巨大,如果使用传统批处理,GPU 空闲率很高。这个本质上是个任务调度问题,传统后台服务我们使用工作窃取算法(work stealing)解决线程空闲问题,与之相似,AI Infra 提出“连续批处理”解决这个问题。
- 连续批处理。连续批处理类似“随时随地拼车的顺风车”,每辆车(GPU)在行程中可随时上/下客。新乘客(请求)直接加入当前车辆的空位(空闲计算单元),已完成的乘客立即下车(释放资源)。