当前位置:首页 > 科技资讯 > 正文

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题

ThinkingMachines Lab近期发布了令人振奋的重大进展!

该实验室联合创始人、前OpenAI副总裁Lilian Weng近日公开透露:

Thinking Machines的首代旗舰产品正式命名为「Connection Machine」(联结机)。

此次事件源于Thinking Machines全新开设的研究博客专栏「Connectionism」(联结主义),其首篇博文《击败LLM推理中的非确定性》深入探讨了相关技术挑战。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第1张

Thinking Machines在博客中阐述道:

我们坚信科学进步源于开放共享。

Connectionism专栏将动态追踪我们的研究轨迹:从底层内核数值计算到上层提示工程。我们期待通过此平台持续分享工作进展,并与研究社区保持高频、透明的交流。

此外,博客还提及「Connectionism」这一术语可回溯至人工智能的早期阶段——在1980年代,它特指专注于神经网络及其与生物大脑相似性的研究子领域。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第2张

Lilian Weng进一步透露了更多信息,该专栏命名不仅源于历史术语,更因为其首代旗舰模型就命名为Connection Machine,这预示着后续将有更多重要发布。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第3张

这是否意味着Thinking Machines即将推出全新模型?

在新模型发布之前,让我们先深入了解他们此次展示的技术实力与研究焦点。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第4张

博客链接:https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

据博文主要作者Horace He介绍,文章核心聚焦于一个关键议题——

LLM推理中的可复现浮点数计算(Reproducible floating point numerics in LLM inference)。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第5张

LLM推理中的非确定性挑战

可复现性是推动科学发展的关键支柱。然而,在大语言模型中获取可复现的结果却面临显著困难。

例如,用户可能发现,即使向ChatGPT多次提出完全相同的问题,得到的答复也往往不同。

这本身并不意外,因为语言模型的输出过程涉及「概率采样」:

将模型输出的逻辑值转化为概率分布,并基于此随机选择一个token。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第6张

更令人困惑的是,即便将温度参数(temperature)设置为0(理论上应使采样过程完全确定),LLM API在实践中仍可能产生非确定性输出。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第7张

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第8张

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第9张

即使使用自有硬件并运行如vLLM或SGLang等开源推理库,采样过程仍可能无法保证确定性。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第10张

那么,LLM推理引擎为何会表现出非确定性?

一个普遍假设是,浮点数的非结合性与并发执行的交互,导致了基于不同计算核心完成顺序的非确定性结果。

此项研究将其称为LLM推理非确定性的「并发+浮点数」假说

例如,研究员Jiayi Yuan、Hao Li、Xinheng Ding等人近期在arXiv上发布了一篇预印本论文,其中指出:

GPU中的浮点运算存在非结合性,即 (a+b)+c ≠ a+(b+c),这是由有限精度和舍入误差导致的。

此特性直接影响Transformer架构中注意力分数和logits的计算,跨多个线程的并行操作可能因执行顺序差异而产生不同结果。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第11张

论文链接:https://arxiv.org/abs/2506.09501

尽管该假说具有一定合理性,但并未完全揭示问题的全貌。

例如,在GPU上对相同数据重复执行相同的矩阵乘法运算,通常总能获得逐位一致的结果。

我们确实在使用浮点数,GPU也确实执行大量并发计算。那么为何在此类测试中未观察到非确定性?⬇️

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第12张

要深入理解LLM推理非确定性的根本原因,我们需要进行更细致的探究。

令人意外的是,即便定义LLM推理的确定性本身也并非易事。

以下看似矛盾的陈述可能同时成立:

GPU上的某些计算核(Kernel)确实具有非确定性。

然而,语言模型前向传播所使用的所有Kernel都是确定性的。

此外,LLM推理服务(如vLLM)的前向传播也可声称是确定性的。

但从最终用户视角观察,推理结果却呈现非确定性。

此次,Thinking Machines致力于揭示LLM推理非确定性背后的真正根源,并阐述如何克服这一挑战,实现真正可复现的推理结果。

核心发现:

LLM前向传播无需原子加法;其非确定性的真实来源是「批次大小变化」而非「原子竞争」。

为确保推理服务避免非确定性、实现Transformer的批处理不变性,我们必须在计算核中内建「批处理不变性」。

幸运的是,我们可以默认每个逐点操作都具有批处理不变性。因此,仅需关注涉及归约的三个关键操作——RMSNorm、矩阵乘法和注意力机制

这三者的实现难度依次递增。每个操作都需要额外考量,才能在保证合理性能的前提下实现批处理不变性。

批处理不变的RMSNorm:数据并行RMSNorm

理想情况下,我们希望采用无需核心间通信的并行化策略。

一种实现方法是为每个计算核心分配独立的批处理元素,从而确保每个归约操作完全在单个核心内完成。

这被称为「数据并行」策略,因为它沿无需通信的维度进行并行化。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第13张

批处理不变的矩阵乘法:数据并行Matmul

与RMSNorm类似,矩阵乘法的标准并行策略也是一种「数据并行」策略,将整个归约过程保持在一个核心内。

最直观的方式是将输出张量分割为二维分块(tiles),并将每个分块分配给不同的核心。每个核心计算属于其分块的点积,再次在单个核心内完成全部归约。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第14张

与RMSNorm不同,矩阵乘法还需考虑算术强度与利用张量核心的约束,这迫使我们将二维分块而非单个输出元素作为分配单元,以实现高效的计算核。

关键在于,矩阵乘法可视为逐点操作与归约操作的组合。

确保矩阵乘法批处理不变性的最简方法是编译一个通用计算核配置,并应用于所有形状。

尽管这可能牺牲部分性能,但在大语言模型推理中通常可接受:

与cuBLAS相比,性能损失仅约20%。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第15张

批处理不变的注意力机制

在实现矩阵乘法的批处理不变性后,注意力机制引入了两个额外挑战——因其包含两个矩阵乘法。

与RMSNorm和矩阵乘法仅在特征维度归约不同,注意力机制需在特征维度和序列维度上进行归约。

此外,注意力机制还需处理各种影响序列处理的推理优化(如分块预填充、前缀缓存等)。

带KV缓存的FlashAttention可能破坏批处理不变性,根本原因在于它将「缓存KV」与「当前KV」分开计算。

不同的KV块数量 → 不同的掩码/完整块组合 → 不同的归约路径。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第16张

只需在计算核启动前统一更新KV-cache页表,确保任意时刻的KV布局一致,即可解决此问题。

大语言模型推理中的注意力形状通常需要分裂归约计算核,常称为Split-KV或FlashDecoding。

固定分裂数量的Split-KV策略(即FlashDecode),由于具体的归约策略取决于请求中处理的查询token数量,不幸地破坏了批处理不变性

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第17张

当查询长度非常小(如在解码阶段),计算核中可能缺乏足够并行性。此时,需沿归约维度——即KV维度——进行分割。典型的KV分割策略是计算所需并行度,然后均匀划分KV维度。例如,若KV长度为1000且需要4个分割,每个核心将处理250个元素。

此外,注意力中常用的分裂归约策略也对批处理不变性构成挑战。

为实现批处理不变性,我们不再固定分割数量,而是固定每个分割的大小,从而得到可变数量的分割。

通过这种方式,无论处理多少token,我们总能保证相同的归约顺序。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第18张

这实现了批处理不变性,因为归约策略不再依赖于一次性处理的查询token数量!

采用「固定块大小」Split-KV策略,注意力机制也能像RMSNorm和矩阵乘法一样实现批处理不变性,从而确保确定性推理。

实践实现与实验验证

通过集成vLLM的FlexAttention后端并利用torch.Library,Thinking Machines提供了一个在vLLM上实现确定性推理的演示。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第19张

代码仓库:https://github.com/thinking-machines-lab/batch_invariant_ops

补全结果的非确定性程度如何?

使用Qwen/Qwen3-235B-A22B-Instruct-2507模型,在温度设为0的条件下,以提示词「Tell me about Richard Feynman」(非思考模式)采样1000次补全,每次生成1000个token。

令人惊讶的是,我们获得了80种不同的补全结果,其中最常见的出现了78次。

分析差异发现,所有补全结果在前102个token上完全一致!首次分歧出现在第103个token。

所有补全均生成序列「Feynman was born on May 11, 1918, in」。

然而,992次补全后续生成「Queens, New York」,而8次补全生成「New York City」。

另一方面,当启用批处理不变计算核时,所有1000次补全完全一致。

性能表现

本研究未对批处理不变计算核进行深度性能优化,但仍通过实验测试了其性能。

实验设置:单GPU API服务器运行Qwen-3-8B模型,处理1000个序列请求,输出长度在90到110之间。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第20张

大部分性能下降源于vLLM中FlexAttention集成尚未充分优化。尽管如此,性能下降处于可接受范围。

实现真正的同策略强化学习

如研究人员指出,训练与推理间的数值计算差异,无形中将同策略强化学习转化为异策略强化学习。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第21张

相关讨论:https://fengyao.notion.site/off-policy-rl

如果两次相同的推理请求都无法保证逐位一致,那么训练与推理间的位级一致就更无从谈起。

确定性推理使我们能够同步优化训练栈,确保采样与训练在数值上逐位一致,从而实现真正的同策略RL

BigmathRLVR设定下,研究人员进行了实验:策略以Qwen 2.5-VL instruct 8B初始化,最大滚动长度为4096

未进行异策略校正(即未使用重要性加权)时,训练中期奖励出现崩溃;

加入异策略校正(重要性加权)后,训练平稳推进;

采样器与训练器实现逐位一致时,策略完全同源(KL=0),训练同样保持平稳。

同时,研究人员绘制了采样器与训练器对数概率之间的KL散度:三种设定呈现明显差异——

带重要性加权(下图红线)时,KL约在0.001附近波动,偶有尖峰;

未加权(下图绿线)时,KL在奖励崩溃前后显著飙升;

真正同策略(下图蓝线)时,KL始终为0,表明训练策略与采样策略完全一致。

需要强调的是:未使用重要性加权的实验在第318步左右出现显著损失峰值,同时KL散度同步骤升;而应用异策略校正实现「真正同策略」的两种设置,RL优化过程均持续且平稳。

Thinking Machines Lab突破性研究:彻底解决LLM推理非确定性难题 LLM推理 非确定性 批处理不变性 强化学习 第22张

精益求精,追求卓越

现代软件系统层层抽象,机器学习中的非确定性微小数值差异,常使人倾向于妥协:

既然系统本质「概率化」,多一点不确定性又何妨?

在单元测试中放宽容错阈值、将训练与采样间的对数概率差异视为「假阳性」,似乎也可行。

但我们必须抵制这种「将就」心态。只需投入更多努力,我们完全有能力定位并修复这些非确定性根源!

Thinking Machines希望本研究能为社区提供一套清晰解决推理非确定性的框架,并激励更多人深入理解自身系统

参考资料:

https://x.com/lilianweng/status/1965828743152509198

https://x.com/cHHillee/status/1965828670167331010

https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/