再来谈谈大模型的分离式推理架构

发布日期：2024-08-03 20:52:56 浏览次数： 5047

作者：zartbot

微信搜一搜，关注“zartbot”

知乎上看到方佳瑞博士的一篇文章《LLM分离式推理可能带来的软硬件变革的迷思》^[1]

恰逢这周工作上有一些和HugeCTR相关的事情, 那么就从软硬件一体化的视角来阐述一下整个架构的演进, 特别是在分离式推理架构上. 以下观点仅代表个人,和作者任职机构无关.

1. 推理系统和训练系统的区别

最简单的一句话是: 推理系统没有所谓的DP并行. 背后隐藏的一个含义是两个系统的Workload是完全不一样的.

1.1 训练系统

到达速率和服务速率为确定性分布

在训练系统中数据以Batch的方式到达, 然后计算时间也相对确定, 一方面是因为backward过程的同步需求, 另一方面是训练语料本身有长短的分布但也做了Padding, 当然可以通过一些技术对Padding进行优化提升计算效率.

1.2 推理系统

到达速率假设为泊松分布, 服务速率受实现方式和服务策略影响

推荐系统请求到达的分布假设是一个泊松分布, 另一方面input token和output token的分布则会带来服务时间有一个特定的分布, 简单的来看按泊松分布算, 或者有长尾的情况,例如Pareto分布.而Prefill-Decoder的方式也会影响这个分布, 因此在调度系统上该如何考虑是一个更值得深思的问题. 这些问题也是最近一段时间工作的一个方向.

1.3 推理系统SLO

对于不同的用户而言(例如按照充值程度分金/银/铜)SLO是不同的, 针对不同的SLO下的TTFT/TBT的整个推理系统的优先级调度和延迟隐藏也有很多可以去做的事情. 另一个非常重要的工作是对用户请求的服务时间的预测

推理系统的请求和服务时间分布相对于训练系统更加有不确定性, 因此在各个子系统内的调度编排和Locality的利用上以及时间/空间折中处理上有着巨大的创新空间.其实这也是Prefill-Decode/KV-Cache Centric Sched有收益的本质原因.

接下来我们分别从软件系统和硬件系统来谈谈.

2. 软件架构

从软件架构来看, 数据和控制平面的分离是一个非常经典的设计模式.

2.1 控制面

控制平面主要负责一些用户请求服务时间预测/调度/集群管理和高可靠/Cache及相关的分离式内存池管理的任务, 这些任务从架构上来看应该用通用的CPU进行处理.

2.2 数据面

数据面主要是用于计算的Prefill Node和Decode Node以及一些弹性内存池相关的数据搬移的工作. 从数据面来看, 其缓存结构在推荐系统中已经有很好的借鉴, 那就是HugeCTR一类的框架中的Hierarchical Parameter Server

从Embedding Cache变成了需要在LLM处理KVCache, 但是相应的软件栈和结构还是有区别的. Embedding TableLookup更多的是Hash Lookup ExatclyMatch. 而在LLM中KVCache的处理是一个Longest Prefix Match. 因此CPU Memory和SSD的软件架构还需要一些调整.

2.3 存储设计和内存缓冲池

当然系统也会面临DRAM不够的问题需要落盘. 然后既要SSD容量大又要高I/O,同时还要考虑低成本和可运维. 因此在SSD前面构建一个分布式的弹性内存池应该是必须要做的了, 这里有一篇华为《MemServe: Context Caching for Disaggregated LLM Serving with Elastic Memory Pool》^[2]

另一个业务场景是DeepSeek会将用户的上下文落盘到SSD中并保存24小时. 那么相应的软件架构应该就很好构建了.

我们可以通过Trie树或者TreeBitMap的算法来构建整个查询树.

当访问到一个叶子后, 即记录该叶子对应的KVCache所在的内存的节点ID和指针并放入处理List, 以便以后进行异步的DMA/RDMA. 这些在用户请求到达后即可进行异步执行处理并在队列适当的时候由调度器控制执行相应Prefill Node GPU VRAM的Prefetch.

Trie树也可以根据input token做一些并行化的处理方式, 例如通过一致性Hash针对前几个Token将workload分散到不同的CPU服务器上. 而这些服务器又可以构建成一个很大规模的弹性内存池.

对于弹性内存层的Cache Evict可以按照LRU的方式进行, 落盘的时候还需要考虑一些KVCache生命周期的问题. 落盘后也需要对Trie树进行更新将其相应的叶子节点指针更新到指向某个文件的特定的块.

谈到落盘,可以借鉴类似于LSM的方式合并压缩后落入磁盘, 并且还可以根据实际情况进行冷温热分层. 因此对于存储落盘到SSD的场景, 个人并不是很喜欢在GPU实例上的本地SSD, 损坏后维修带来的业务影响还是有的, 另一方面可能还会带来一些workload偏斜的影响, 因此可能更倾向于一种分布式并行存储的方式.

当然这里还有一个不同的地方是对于长时间没有更新的block也需要定期的进行GC, 例如按照DeepSeek保留24h的做法, 可以对超过24h的数据删除并在Trie树上剪枝更新.

3. 硬件架构

3.1 三网融合的推理系统

黄大年茶思屋最近也有一篇《英伟达、谷歌、Meta等5大巨头Scale-up超节点规模大比拼，揭示未来AI网络最优解》^[3]. 个人对这种不根据workload来谈互联的分析持怀疑态度, 但是有些结论是正确的. 其中最重要的一点就是“三网合一”即文章中的观点

当前的AI网络是三个独立网：前端存储网VPC（以太）、后端参数面（IB、RoCE2）和超节点（NVLink，HCCS）。三个网长期共存不合理的，一定会融合。

当然这个问题在方博士的文章中也有提及:

现在分离式架构都是用GPU训练集群做推理，节点内NVLink互联，节点间用IB或ROCE的RDMA互联。这种配置分离式架构完全是浪费，好比李云龙攻打平安县城，章明星称之为富裕仗。

有几个难题:

3.2 FrontEnd: CPU实例和GPU实例耦合

在这样的一个推理系统中, NVL72-GB200是一个非常不错的方案, ScaleUP规模大, 而CPU又直接和GPU有C2C的带宽, 这样控制面和弹性内存池的设计难度会小不少.另一个问题是泊松分布到达下GPU的调度的问题,这个涉及一些GPU本身SIMT硬件架构的缺陷, 就不展开了.

而对于国内非Grace-B/H的大概只能通过RDMA将弹性内存池和GPU连接了, 因此需要VPC上跑超大规模RDMA并且完全商用的, 现阶段全球能做的大概也就只有AWS SRD/Google Falcon和阿里云eRDMA. 而同时在这个网络上又要兼顾Prefill Node之间的SP并行带来的alltoall incast的影响, 对于其它很多客户而言, 包括英伟达自己都是有很大挑战的.

为什么需要呢, 从另一个角度来看,现在的Prefill/Decoder的转移其实都是在8卡PCIe连接的同服务器的CPU上,然后再通过Messenger转发到另一个Decoder实例, 并通过Async Load加载. 也就是说CPU的算力和内存空间和GPU是绑定的, 外置的CPU服务器如果可以直接GDR访问GPU显存来做异步的调度和prefix lookup整个的性价比还应该有所提升.

我们注意到Google Vertex AI很早就通过一些闲置的CPU实例和GPU混布来构建, 并且通过CPU实例来做参数服务器,当然还有一个是Cerebas的Swarm-X/Memory-X架构

另外Enfabrica也有一些机会,难怪NV也投资了.

3.3 ScaleOut: Prefill-Decode M:N的部署

即章明星老师提到的异构分离式架构, 需要在H800/A800和H20之间构建很高的Bi-Section带宽, 同时考虑到组网规模的问题还需要避免网络上hash冲突带来的利用率降低的问题.

3.4 ScaleUP: 是否还需要Load/Store ?

针对推理系统来看, ScaleUP主要用于TP/SP并行,例如模型规模大了或者算力无法满足SLO时或负载均衡的考虑带来的并行策略. 细粒度的Load/Store访存似乎并不需要,因此Eth-X这样的方案或许也就够用了.当然还有一些优化方式,例如英伟达的GPS/PROACT/Fine-PACK的处理方式等. 当然另一方面有现成的NVLink用用也挺好的. 等待UALink似乎又要考虑一些问题, 例如BRCM Altas4的single vendor供应的问题和InfityFabric这些IP的成熟度, 当前的国产GPU厂商可能面临二选一的决策, 但是还是请先考虑是否还需要Load/Store over Scale-UP Link.

这个问题的答案取决于业务部署模式的考虑,成本的衡量,弹性的重要程度,是否训推一体,生态的兼容程度等. 未来多模态的业务场景也是一个变数.