在 DwarfStar 上实现分布式 LLM 推理
原文由 Salvatore Sanfilippo 于 发布,订阅该博客
高端 NVIDIA 显卡,以及运行它们所需的服务器和电力,成本都很高,尤其是当你需要足够大的显存来运行超大模型时。到目前为止,替代方案一直是苹果硬件,或是 DGX Spark——即便受限于内存带宽,它依然能让 LLM 的提示词处理(prefill)达到足够快的速度。Mac Studio 提供了最高 512GB 的统一内存,内存带宽不算高(但比 Spark 好得多),算力也够用,以当下的行情来看,价格还算相对合理。
比如,在 DwarfStar 上,配备 M3 Ultra、512GB 内存的 Mac Studio 可以以约 150 tokens/秒的速度进行 DeepSeek v4 PRO 的 prefill,解码速度约为 10-13 tokens/秒,不算出色,但在某些场景下已经堪用。即便是 2-bit 量化,DeepSeek v4 PRO 的表现依然非常坚挺,和同等量化下的 Flash 差不多(今天我让 PRO 写了一个 C 编译器,很快会发布视频)。能在家里用约 1.2 万美元的总投入跑起一个前沿模型,我认为这绝非小事。
本以为情况会越来越好,但前景看起来却有些阴云密布。NVIDIA 方案几乎不可能变得更便宜,即便是小公司,也很难轻松购置并运维一个用于本地推理的小型数据中心。与此同时,内存短缺也让 Mac Studio 推出 M5 Ultra 版本的前景变得不太明朗——原本或许能带来 1.2TB/s 的内存带宽和更强的算力(仅从算力上看,M5 Max 就已经更快了,而且每个 GPU 核心内都集成了对某些模型很有帮助的神经加速器)。
所以就本地推理的现状而言,最好的机器很可能是一台笔记本。128GB 的 M5 Max 可以以相当不错的 prefill 和解码速度运行 2-bit 量化的 DeepSeek v4 Flash 和 Mimo V2.5。我们说的是约 500 tokens/秒的 prefill 和约 35-40 tokens/秒的解码速度,而且随着上下文长度增加,性能下降的曲线也完全可以接受。按配置不同,花费约 6000 至 7000 美元,这在当下是最划算的选择之一。
如果现状如此,那么对于一般的本地推理项目,尤其是对 DwarfStar 来说,关注分布式推理就开始变得有意义了。如果我们有两台、三台、四台 M5 Max 的 MacBook 能做什么?或者有两台 512GB 内存的 M3 Ultra 呢?
传统上有两种主要的分布式推理方案。一种是通过内存叠加,将 50% 的 Transformer 层加载到 A 机器上,剩下 50% 加载到 B 机器上,然后以串行方式执行推理。这种情况下只需要在机器间传输激活值,概念上非常简单,而且通过一些微批处理的技巧,不仅能实现内存叠加,理论上还能大幅提升提示词处理速度(但解码不行:生成单个 token 时,你必须等 A 机器跑完前半部分层,再等 B 机器跑完后半部分层,以此类推——不过至少发热量会降低,因此可以承受持续负载),这已经相当不错了。这意味着,例如那些拥有两台 512GB Mac Studio 的幸运儿,就可以运行全尺寸的 DeepSeek v4 PRO(即便 2-bit 量化版已经跑得非常非常好了),而且通过微批处理还能享受到更快的 prefill。
另一种方案是利用 Apple RDMA,在两台机器间并行执行,基本上是一种纵向切分。比如,可以尝试在 A、B 两台机器上都加载同样的 2-bit 量化模型,这样两边都能放下,且每一边都拥有*全部*的路由专家(routed experts)。然后在每一层尝试进行必要的协调,让一半专家在 A 机器上执行,一半在 B 机器上执行,以此类推(注意两台机器都有全部专家,所以无论路由器如何调度,都可以把 50% 的计算量分给另一台机器,而且激活值非常小)。这种方式对 PRO 更可行,因为它的路由专家更大,所以通信开销相对不那么敏感。但能否真正很好地跑起来,还有待观察。
你可能也在想张量并行,对吧?但我敢说,以两台苹果电脑、两台 DGX Spark 之间的通信速度,这条路完全行不通(去查查 NVLink 的速度就知道了)。上面两种方案的神奇之处就在于需要传输的数据非常少。
好吧,说到这里,我猜你在想,这不就是大家都知道的那些并行跑 LLM 的老套路吗,确实如此。但这篇文章正是为了引出这一点。如果我们能以一种完全不同的方式把两台 Mac 或 DGX 并行起来,会怎么样?开源权重模型如今正处在黄金时代,数量众多且很多都非常强大。在 128GB、2-bit 量化这个档位里,就有不少有意思的模型:Minimax M2.7、Mimo V2.5、DeepSeek v4 Flash 等等。与此同时,最近有人指出,LLM 集成(https://arxiv.org/abs/2502.18036)是一个被低估的可能性,它能让两个模型在两台不同的机器上以完全无共享(shared-nothing)的方式运行,最后只需合并 logits 或选出最佳续写。实现方式有好几种,而且即便两个模型的词表不同也依然有效:你可以选择困惑度更低的那个续写(也就是选更确信的那个模型:这就像一个隐式路由的双专家 MoE),甚至还可以合并 logits(由于词表不同会带来一些复杂性)再从中采样。更新的论文表明,将这两种技术混合起来效果最好。无论如何:这些技术似乎真的有效,模型合在一起的表现似乎比单独运行时更好。就好像它们的知识得到了提升,因为每个模型都为“接下来该说什么”贡献了自己的视角。
也许除了前两种方案之外,这是最合乎逻辑的第三种尝试路径。真希望在接下来的几个月里能找到时间,好好折腾一下这些想法。
随机一篇博客
评论
登录后参与讨论