NVIDIA 开源 Nemotron TwoTower:扩散 LLM 推理提速 2.42x,已有模型无需重新训练

NVIDIA 最近在 Hugging Face 上开源了一组有意思的模型——Nemotron-Labs-TwoTower。它是基于扩散机制(diffusion)的语言模型,最引人注目的数字是:在已有预训练模型的基础上配上扩散 head,推理throughput最多可以提升 2.42 倍,而且完全不需要重新训练底层模型。

这是什么架构?

TwoTower 的命名很直白——它用"双塔"结构把两个角色分开:

  • 底塔(Backbone Tower):一个冻结的自回归语言模型,这里用的是 Nemotron-3-Nano-30B-A3B(MoE 架构,激活参数量约 3B)。这个塔不做生成,而是为扩散过程提供隐层表示。
  • 顶塔(Diffusion Tower):一个轻量级的扩散 head,以 backbone 的 hidden state 为条件,通过迭代去噪来生成 token。

生成过程不再是传统的自回归逐 token 预测,而是:先初始化一段随机噪声序列作为 token 嵌入,然后通过扩散 head 逐步去噪、细化,最终收敛到确定性的 token 序列。整个过程在 backbone 的引导下完成。

模型本身不大——30B 参数、3B 激活的 MoE 规模——但思路的通用性很有意思:同样的冻结 backbone + 扩散 head 结构原则上可以套用到更大的模型上。

2.42x 的throughput从哪来?

提速的核心原因在于扩散生成的计算效率和并行度:

  1. 去噪步数远少于自回归步数:自回归生成需要逐个 token 推理,生成长序列时推理次数等于序列长度。扩散模型的去噪步数是固定的(而且远小于序列长度),每一步都在同时优化整个序列。
  2. 批量去噪效率高:扩散过程中,每一步的执行是全局的、矩阵密集型的,GPU 利用率更高,不会出现自回归场景下 batch size 因为 KV cache 膨胀而被迫缩小的瓶颈。
  3. 无需 KV cache:扩散生成不需要保留 KV cache,这对长序列生成来说节省了大量显存。

根据 NVIDIA 的测试数据,在同样的硬件(H100)上,TwoTower 的端到端throughput(token/s)相比标准自回归推理提升了 1.8x 到 2.42x,具体受序列长度和 batch size 影响。

对部署意味着什么?

这件事值得部署团队关注,有几个角度:

对已有 Nemotron 用户:如果已经在跑 Nemotron-3-Nano-30B-A3B(或其他兼容的 backbone),可以直接复用权重,不需要额外的训练投入就能获得推理提速。这意味着现有部署方案可以"插拔"式升级。

对新的部署选型:考虑部署一个 MoE 30B/3B 级别模型时,TwoTower 提供了一条与量化、投机解码正交的提速方案。如果结合 INT4 量化 + 扩散推理,实际throughput还有提升空间。

对推理成本的影响:假设 token/s 提升 2 倍,在同等硬件配置下意味着单位 token 的电力、租赁成本直接减半。对于部署团队来说,这是实打实的成本变量。

当然,扩散模型在生成质量上与传统自回归 LLM 的全面对比还需要社区更多的评测。NVIDIA 官方数据显示在多项基准上与自回归版本持平或相近,但实际部署中特定任务的适用性仍需验证。

开源与获取

模型权重和推理代码已经在 Hugging Face 上以 Apache 2.0 许可证发布:

  • 模型主页:HuggingFace: Nemotron-Labs-TwoTower
  • 推理示例代码包含在模型卡中,支持标准的 transformers pipeline

如果你是做 AI 部署的朋友,这个方向值得跟踪——扩散语言模型从研究概念走向可部署的产品状态,可能比想象的快。


顺手放几个链接:

  • NVIDIA 官方公告:Tech Times 报道
  • MarkTechPost: 技术深度解读
  • 模型权重已发布在 HuggingFace 上
4 个赞

扩散模型在推理速度上的这个方向确实挺有意思,跟投机解码走的是不同的提速思路。我比较好奇的是输出质量的一致性,有没有人对比过 Diffusion LLM 在代码生成这类对精确性要求高的任务上的表现?

2.42x 这个数字来自固定去噪步数和自回归步数的比值。对于短序列(比如几十个 token)差距不大,文本越长优势越明显。如果能配合 speculative decoding 做二次优化,可能还有空间。

个人觉得这件事最大的 value 在于基于已有 frozen backbone 这一点。意味着已经在跑 Nemotron 系列的组织,可以几乎零成本在推理侧配一个扩散 head 来提速,不需要改训练流程。

之前在 arXiv 上看过 diffusion LM 的论文,但一直觉得离实用很远。NVIDIA 这次开源了一个可以跑推理的版本,算是把这种架构从 paper 拉到了 deployment ready 的区间。扩散方向在图像领域已经统治了生成任务,文本这边终于也开始落地了。

提个问题:扩散模型在低延迟场景(比如在线对话)里能适配吗?去噪迭代有固定步数,不像自回归可以做到首 token 时间很短。感觉目前更适合批量推理和离线处理的场景,实时交互可能还得等优化。

KV cache 这个点很关键。跑长文本推理的时候,自回归的 KV cache 会吃掉大量显存,尤其 batch size 大了以后。扩散模型不需要 cache,意味着同等显存下可以塞更大 batch,这对批量推理的throughput提升非常直接。

如果能把这套方法移植到 Qwen 或 Llama 3 上就好了——理论上扩散 head 不局限于 Nemotron backbone。如果社区有人做个通用 adapter,那生态影响会大很多。