在 DwarfStar 中实现分布式 LLM 推理
高端 NVIDIA 显卡,以及运行它们所需的服务器和电力,成本都非常高昂,尤其是当你需要足够大的 VRAM 来运行超大模型时。到目前为止,替代方案一直是 Apple 硬件,或是 DGX Spark——即便它因 memory bandwidth(内存带宽)而受到严重限制,仍能以足够快的速度运行 LLM 的提示处理(prefill(预填充))。Mac Studio 提供了高达 512GB 的 unified memory(统一内存),其 memory bandwidth 虽然不算高(但比 Spark 好得多),算力也相对有限,但在当前形势下,价格总体上还算公道。
例如,使用 DwarfStar 时,配备 512GB 的 Mac Studio M3 Ultra 可以以 150 t/s 的 prefill 和约 10-13 t/s 的 decoding(解码)速度运行 DeepSeek v4 PRO,表现不算出色,但对某些用例来说已经可用。即使是 2-bit quantized(量化)的版本,DeepSeek v4 PRO 的表现依然非常坚挺,与同等量化下的 Flash 相当(今天我让 PRO 写了一个 C 编译器,视频很快就会发布)。能够以约 1.2 万美元的总花费在家中运行一个前沿模型,我认为这绝非一件小事。
人们本可以期待情况会越来越好,但前景却显得扑朔迷离。NVIDIA 方案几乎没有降价的希望,即使是小公司也难以轻松购置并维护一个用于本地推理的小型数据中心。与此同时,RAM 短缺使得我们不太可能看到搭载 M5 Ultra、拥有约 1.2T/s memory bandwidth 和更强算力的 Mac Studio(M5 Max 在算力上已经更快,并且在每个 GPU 核心内都配备了对某些模型有帮助的 Neural Accelerators)。
因此,就本地推理而言,目前最好的机器可能是一台笔记本电脑。128GB 的 M5 Max 可以以相当不错的 prefill 和 decoding 速度运行 2-bit quantized 的 DeepSeek v4 Flash 和 Mimo V2.5。我们说的是约 500 t/s 的 prefill 和约 35-40 t/s 的 decoding 速度,随着上下文长度增加,性能下降的斜率也完全可以接受。根据配置不同,6千至7千美元的售价使其成为目前最划算的选择之一。
如果是这种情况,那么对于一般的本地推理项目,尤其是对于 DwarfStar 来说,关注分布式推理就开始变得有意义了。如果我们有两台、三台、四台 MacBook M5 Max 系统能做什么?或者有两台配备 512GB RAM 的 M3 Ultra 呢?
传统上,运行 distributed inference(分布式推理)主要有两种方案。一种是通过将 50% 的 transformer layers(Transformer 层)加载到计算机 A 上,剩下 50% 加载到计算机 B 上,以串行方式运行推理,从而实现内存翻倍。在这种情况下,只需要在机器间传输 activations(激活值),概念上非常简单,并且借助一些 micro-batching(微批处理)技巧,不仅可以实现内存翻倍,理论上甚至可以大幅提升提示处理速度(但无法提升 decoding:生成单个 token 时,必须等待机器 A 上的前几层计算完成,再等待机器 B 上的剩余层,以此类推——不过至少产生的热量会更少,因此可以承受持续负载),这已经相当不错了。这意味着,例如,拥有两台 512GB Mac Studio 的幸运儿可以运行全尺寸的 DeepSeek v4 PRO(即便 2-bit quants 已经运行得非常非常好了),并且通过 micro-batching 甚至可以获得更快的 prefill。
另一种方法是利用 Apple RDMA(远程直接内存访问)在两台机器之间并行执行,本质上是一种垂直切分。例如,可以尝试在机器 A 和 B 上加载相同的 2-bit quants,使两者都能容纳,并且每一侧都拥有 *全部* routed experts(路由专家)。然后,对于每一层,我们可以尝试进行必要的协调,以便在一台机器上执行一半专家、在另一台机器上执行另一半,以此类推(注意,两台机器都拥有全部专家,因此无论路由器如何决策,我们都可以将 50% 的计算发送到另一台机器,而 activations 非常小)。这对于拥有更大 routed experts 的 PRO 来说更为可行,因此通信开销的影响较小。但这种方法能否很好地奏效,还有待观察。
你可能也在想,还有 tensor parallelism(张量并行),对吧?但我敢打赌,以两台 Apple 电脑、两台 DGX Spark 之间现有的通信速度,这完全不可行(去查一下 NVLink 的速度就知道了)。上述两种方案的神奇之处在于,你只需要传输非常少量的数据。
好吧,到这里我敢打赌你在想,这不过是人人都知道的关于并行运行 LLM 的老生常谈,确实如此。但这篇文章正是为了引出这一点。如果我们能以一种完全不同的方式将两台 Mac 或 DGX 并行化,会怎样呢?开源权重模型如今正处于黄金时代,我们拥有大量且非常强大的模型。在 128GB 的 2-bit quants 类别中,有许多有趣的选择:Minimax M2.7、Mimo V2.5、DeepSeek v4 Flash 等等。与此同时,最近有人指出,LLMs ensemble(集成)(https://arxiv.org/abs/2502.18036)是一个研究不足的可能性,它允许两个模型在两台不同的机器上以完全无共享的方式运行,最后只需合并 logits(对数几率)或选择最佳续写。有多种方法可以实现这一点,即使两个模型的词表不同也可行:你可以选择 perplexity(困惑度)更低的续写(也就是说,选择更确信的那个模型:这就像一个路由隐含的双专家 MoE(混合专家模型)),甚至还可以合并 logits(由于词表不同会带来一些复杂性)并从中采样。更新的论文表明,混合这两种技术是最佳方法。无论如何:这些技术似乎确实有效,模型的表现似乎比单独运行时更好。就好像它们的知识得到了提升,因为每个模型都带来了自己对接下来该说什么的视角。
也许除了前两种方案之外,这是最合乎逻辑的第三种尝试方向。我真的希望在接下来的几个月里能找到时间多尝试一下这些。
随机一篇博客