刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文

Mira Murati 的新公司终于发声了!

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-2

Thinking Machines Lab 今天 正式推出了他们的研究博客 Connectionism,第一篇文章就直接瞄准了 LLM 推理中让人头疼的「 非确定性 」问题。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-3

这家由前 OpenAI CTO(及前临时 CEO) Mira Murati 创立的公司,在今年 7 月刚完成了约 20 亿美元 的种子轮融资,估值达到 120 亿美元 。投资方包括 Andreessen Horowitz(领投)、Nvidia、AMD、Cisco 等科技巨头。

值得注意的是: 在拿到如此巨额融资之前,公司还没发布任何产品。

LLM「不确定性」的真相

这第一篇博文《击败大语言模型推理中的非确定性》直击要害。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-4

如果你是算法相关从业者,你应该有发现: 同样的输入,LLM 有时会给出不同的输出 。

即使设置了相同的随机种子,结果还是会变化。

很多人以为是 GPU 并发执行和浮点数运算的锅。

但 Thinking Machines 的研究发现: 真正的罪魁祸首是批次不变性缺失 。

什么意思?当你向 LLM 发送请求时,服务器会根据当前负载情况,把你的请求和其他请求打包成不同大小的批次处理。问题就出在这里——

相同的输入在不同批次大小下会产生不同的结果 。

这就像你去餐厅点菜,你点的菜味道竟然会因为厨房同时在做几道菜而改变。

听起来很荒谬?

但这就是现在 LLM 推理系统的现状。

浮点数的「蝴蝶效应」

根本原因在于浮点数的非结合性: (a+b)+c ≠ a+(b+c) 。不同的加法顺序会产生微小差异,这种差异在深度神经网络中层层放大。

具体到 LLM 推理中,矩阵乘法、RMSNorm、注意力机制等核心操作,在不同批次大小下会采用不同的约简策略。 你的请求结果竟然依赖于服务器同时在处理多少其他请求 ——

这太魔幻了。

让内核「批次不变」

Thinking Machines 提出的解决方案很直接: 实现批次不变的内核。

RMSNorm :采用数据并行策略,避免分割约简。

矩阵乘法 :使用固定内核配置,避免 Split-K 策略。

注意力机制 :采用固定分割大小策略,确保约简顺序一致。

实验结果可谓是令人惊讶:

在 1000 次采样中,原本会产生 80 个不同的完成结果 。

但在启用批次不变内核后, 所有结果完全一致 。

当然,这种确定性是有代价的。未优化版本性能下降约 2 倍,但经过改进后性能损失已经可以接受。

Connectionism:不只是一个名字

有意思的是,博客名「Connectionism」是 1980 年代研究神经网络与生物大脑相似性的 AI 子领域名称。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-5

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-6

Mira Murati(@miramurati)表示:

Thinking Machines 使命的重要部分是提高人们对 AI 的科学理解,并与更广泛的研究社区合作。今天推出 Connectionism 来分享我们的一些科学见解。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-7

联合创始人 Lilian Weng(@lilianweng)补充了一个有趣的历史细节:

除了 Connectionism 与 AI 领域早期的联系,以及强调神经网络与人类大脑的相似性这个有趣的事实外,第一代 Thinking Machines 的旗舰产品就叫 Connection Machine。

豪华团队阵容

除了 Mira Murati,核心团队还包括 OpenAI 联合创始人 John Schulman、前研究 VP Barret Zoph、前 AI 安全与机器人 VP Lilian Weng 等人。

而 Andrew Tulloch 甚至拒绝了 Zuckerberg 15 亿美元回 Meta 的邀请,选择继续与 Murati 一起创业。

团队约 30 人,其中三分之二来自 OpenAI。技术岗位年薪高达 45-50 万美元 。

Bob McGrew 和 Alec Radford 等 OpenAI 核心研究者担任顾问。

Thomas Ip(@_thomasip)精辟总结道:

LLM 推理非确定性不只是浮点数非结合性或 GPU 并发执行,核心罪魁祸首是批次方差,服务器负载不可预测地改变了数值计算。批次不变内核解锁了真正的可重复性,终于让强化学习『在线策略』变得可行。

这项工作的意义不仅在于解决了一个技术难题,更重要的是为 LLM 的可重复性和可靠性提供了科学方法。尤其是对强化学习等对一致性要求极高的应用场景,该文具有重要价值。

科学确实在分享中变得更好。

下为全文

击败大语言模型推理中的非确定性

来源 [1] : https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

发布时间: 2025年9月10日

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-8

目录

引言

原罪:浮点数非结合性

为什么内核不总是按相同顺序相加数字?

什么时候需要原子加法?

批次不变性和"确定性"

如何使内核批次不变?

批次不变的RMSNorm

批次不变的矩阵乘法

批次不变的注意力机制

实现

实验

生成结果的非确定性程度如何?

性能

真正的在线策略强化学习

结论

引用

可重复性是科学进步的基石。 然而,从大语言模型中获得可重复的结果极其困难。

例如,你可能观察到向ChatGPT多次提出同一个问题会得到不同的结果。这本身并不令人惊讶,因为从语言模型获得结果涉及"采样"过程——

将语言模型的输出转换为概率分布并概率性地选择一个token。

更令人惊讶的是,即使我们将温度调整到0(这意味着LLM总是选择概率最高的token,称为贪婪采样),使采样理论上变为确定性,LLM API在实践中仍然 不是 确定性的(见过往讨论 这里 [2] 、 这里 [3] 或 这里 [4] )。即使使用vLLM或SGLang等开源推理库在自己的硬件上运行推理,采样仍然不是确定性的(见 这里 [5] 或 这里 [6] )。

但是为什么LLM推理引擎不是确定性的呢?

一个常见的假设是浮点数非结合性和并发执行的某种组合导致了基于并发核心谁先完成的非确定性。我们将这称为LLM推理非确定性的"并发+浮点数"假设。例如, 最近的一篇arXiv预印本 [7] 写道:

GPU中的浮点运算表现出非结合性,意味着(a+b)+c≠a+(b+c),这是由于有限精度和舍入误差造成的。这一特性直接影响transformer架构中注意力分数和logits的计算,其中多个线程的并行操作可能基于执行顺序产生不同的结果。

你也可以在其他地方找到重复的"并发+浮点数"假设,比如 这里 [8] ( 有速度权衡,为了使端点快速,使用了GPU,它们进行并行[非确定性]计算。任何现代GPU神经网络计算都会受到这些影响。 ),或 这里 (https://x.com/hosseeb/status/1773146428594090473 非确定性计算。任何现代GPU神经网络计算都会受到这些影响。) , 或 这里 ( 因为GPU高度并行化,每次执行时加法或乘法的顺序可能不同,这可能级联为输出的微小差异。 ) 。

虽然这个假设并非完全错误,但它没有揭示全貌。例如,即使在GPU上,对相同数据重复运行相同的矩阵乘法总是会提供按位相等的结果。我们确实在使用浮点数。我们的GPU确实有很多并发性。为什么在这个测试中我们没有看到非确定性?

A = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)B = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)ref = torch.mm(A, B)for _ in range(1000):    assert (torch.mm(A, B) - ref).abs().max().item() == 0

要理解LLM推理非确定性的真正原因,我们必须深入研究。

不幸的是,即使 定义 LLM推理确定性的含义也很困难。或许令人困惑的是,以下陈述都同时为真:

GPU上的某些内核是 非确定性的 。

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

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

尽管如此,从使用推理服务器的任何人的角度来看,结果是 非确定性的 。

在这篇文章中,我们将解释为什么"并发+浮点数"假设没有抓住要点,揭露LLM推理非确定性背后的真正罪魁祸首,并解释如何击败非确定性并在LLM推理中获得真正可重复的结果。

原罪:浮点数非结合性

在讨论非确定性之前,解释为什么会有数值差异是有用的。毕竟,我们通常认为机器学习模型是遵循结构性规则(如交换律或结合律)的数学函数。难道不应该有一个"数学上正确"的结果,我们的机器学习库应该提供给我们吗?

罪魁祸首是 浮点数非结合性 。也就是说,对于浮点数:

(a+b)+c≠a+(b+c)

(0.1 + 1e20) - 1e20>>> 00.1 + (1e20 - 1e20)>>> 0.1

讽刺的是,破坏结合律正是使浮点数有用的原因。

浮点数有用是因为它们允许"动态"的精度级别。为了解释起见,我们将使用十进制(而不是二进制),其中浮点数的格式为尾数×10^指数。我们还将为尾数使用3位数字,为指数使用1位数字。

例如,对于值3450,我们可以精确地表示为3.45×10³。我们也可以表示更小的值,如0.486为4.86×10⁻¹。通过这种方式,浮点数允许我们表示非常小和非常大的值。在科学中,我们可能说浮点数允许我们保持恒定数量的"有效数字"。

如果你将两个具有相同指数的浮点数相加,它看起来类似于整数加法。例如,123(1.23×10²)+ 456(4.56×10²)结果为579(5.79×10²)。

但是当我们将两个具有不同指数的浮点数相加时会发生什么,比如1230和23.4?在这种情况下,精确结果是1253.4。然而,我们一次只能保持3位数字的精度。因此浮点加法会 丢弃 最后2位数字并获得值1.25×10³(或1250)。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-9

我们需要3位精度来表示1230,需要3位精度来表示23.4。然而,将这两个数字相加的结果需要5位精度来表示(1253.4)。我们的浮点格式必须丢弃末尾的34。在某种意义上,我们在相加之前实际上将原始的23.4舍入为20.0。

然而,此时我们已经破坏了信息。注意,每次我们将两个具有不同"尺度"(即不同指数)的浮点数相加时,这种情况都可能发生。实际上,将具有不同指数的浮点数相加一直在发生。事实上,如果我们能保证永远不需要不同的指数,我们就可以只使用整数!

换句话说,每次我们以不同的顺序将浮点数相加时,我们可能得到完全不同的结果。举一个极端例子,根据顺序的不同,对这个数组求和有102种不同的可能结果。

import randomvals = [1e-10, 1e-5, 1e-2, 1]vals = vals + [-v for v in vals]results = []random.seed(42)for _ in range(10000):    random.shuffle(vals)    results.append(sum(vals))results = sorted(set(results))print(f"There are {len(results)} unique results: {results}")# 输出:# There are 102 unique results: [-8.326672684688674e-17, -7.45931094670027e-17, ..., 8.326672684688674e-17]

虽然这是非相同输出的根本原因,但它并没有直接回答非确定性来自哪里。它没有帮助我们理解为什么浮点值会以不同的顺序相加,何时发生这种情况,以及如何避免。

答案在于内核是如何实现的。

为什么内核不总是按相同顺序相加数字?

如上所述,对于内核为什么以不同顺序相加数字的一个常见解释是"并发+浮点数"假设。该假设声称,如果并发线程完成的顺序是非确定性的,并且累积顺序依赖于并发线程完成的顺序(比如原子加法),我们的累积顺序也将是非确定性的。

令人困惑的是,虽然这可能导致非确定性内核,但并发性(和原子加法)最终完全不涉及LLM推理非确定性!为了解释真正的罪魁祸首是什么,让我们首先理解为什么现代GPU内核很少需要原子加法。

什么时候需要原子加法?

通常,GPU会在许多"核心"(即SM)上并发启动程序。由于核心之间没有固有的同步,如果核心需要彼此通信,这就带来了挑战。例如,如果所有核心都必须累积到同一个元素,你可以使用"原子加法"(有时称为" fetch-and-add [9] ")。原子加法是"非确定性的"——结果累积的顺序纯粹取决于哪个核心先完成。

例如,想象你正在用100个核心约简一个100元素的向量(例如 torch.sum() )。虽然你可以并行加载所有100个元素,但我们最终必须约简到单个元素。实现这一点的一种方法是使用某种"原子加法"原语,硬件保证所有加法都会被处理,但不保证顺序。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-10

原子加法确保每个核心的贡献都会反映在最终和中。然而,它不保证贡献被相加的 顺序 。顺序完全取决于哪个核心先完成,这是一个非确定性属性。因此,多次执行相同的并行程序可能导致非确定性输出。

这通常是人们所说的"非确定性"——你用完全相同的输入执行同一个内核两次,得到不同的结果。这被称为 运行到运行的非确定性 ,即你用完全相同的依赖项运行同一个python脚本两次,但得到不同的结果。

虽然并发原子加法 确实 使内核非确定性,但 绝大多数内核都不需要原子加法 。实际上,在LLM的典型前向传播中,通常 没有一个原子加法存在 。

考虑到并行化约简可以从原子加法中受益,这可能令人惊讶。原子加法最终不被需要有两个主要原因。

沿着"批次"维度通常有足够的并行性,我们不需要沿着约简维度并行化。例如,假设我们不是约简单个100维向量,而是并行约简500个向量。在这种情况下,我们可以在每个核心中约简整个向量,并允许每个核心操作不同的向量。

随着时间的推移,大多数神经网络库都采用了各种策略来在不牺牲性能的情况下实现确定性。例如,我们可以执行"分割"(或树)约简,将100元素约简分割为五个20元素约简(从而实现五路并行性)。然后,为了组合剩余的五个元素,我们可以执行单独的"清理"约简(不并行化,但操作的元素足够少以保持廉价)或利用信号量(确保每个并发线程块将以确定性顺序累积)。

由于这两个因素,对于绝大多数神经网络操作来说,避免原子加法的性能损失是微不足道的。

仍然有几个常见操作避免原子加法会有显著的性能损失。例如,PyTorch中的 scatter_add ( a[b] += c )。然而,在LLM中常用的只有FlashAttention反向传播。

然而,LLM的前向传播 不涉及需要原子加法的操作 。因此,LLM中的前向传播实际上是"运行到运行确定性的"。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-11

从推理服务器的角度来看,它 是 确定性的。给定完全相同的用户请求,它总是提供相同的确定性输出。

维基百科写道,"确定性算法是一种算法,给定特定输入,总是产生相同输出。"在这种情况下,给定完全相同的输入(即推理服务器正在处理的确切请求),前向传播总是产生完全相同的输出。

然而,前向传播本身是"确定性的"并不足以确保包含它的系统是确定性的。例如,如果我们请求的输出依赖于并行用户请求(例如批量归一化)怎么办?由于每个单独的请求无法知道并行请求将是什么,从他们的角度来看,我们的整体LLM推理也是非确定性的!

事实证明,我们请求的输出 确实 依赖于并行用户请求。不是因为我们以某种方式在批次间泄漏信息——而是因为我们的前向传播缺乏"批次不变性",导致我们请求的输出依赖于前向传播的 批次大小 。

批次不变性和"确定性"

为了解释批次不变性,让我们简化系统,仅查看矩阵乘法。你可以假设所有矩阵乘法实现都是"运行到运行确定性的"。然而,它们不是"批次不变的"。换句话说,当批次大小改变时,批次中的每个元素可能得到不同的结果。

从数学角度来看,这是一个相当不寻常的属性。矩阵乘法应该沿着批次中的每个元素"独立"——批次中的其他元素或批次有多大都不应该影响批次中特定元素的计算结果。

然而,正如我们可以实验性观察到的,这并不是真的。

import torchtorch.set_default_device('cuda') B = 2048D = 4096a = torch.linspace(-1000, 1000, B*D).reshape(B, D)b = torch.linspace(-1000, 1000, D*D).reshape(D, D)# 通过取批次的第一个元素进行矩阵向量乘法out1 = torch.mm(a[:1], b)# 进行矩阵矩阵乘法然后取批次的第一个元素out2 = torch.mm(a, b)[:1]print((out1 - out2).abs().max()) # tensor(1669.2500, device='cuda:0')

注意这 是 "运行到运行确定性的"。如果你多次运行脚本,它将确定性地返回相同的结果。

然而,当非批次不变内核用作更大推理系统的一部分时,系统可能变得非确定性。当你向推理端点发出查询时,服务器承受的负载量从用户的角度来看实际上是"非确定性的"。负载决定了内核运行的批次大小,从而改变每个单独请求的最终结果!

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-12

虽然推理服务器本身可以声称是"确定性的",但对于单个用户来说情况不同。从单个用户的角度来看,其他并发用户不是系统的"输入",而是系统的非确定性属性。这使得LLM推理从每个用户的角度来看都是"非确定性的"。

如果你将内核不具有不变性的某些属性(即批次大小)与该属性的非确定性(即服务器承受的负载)组合起来,你就得到了一个非确定性系统。

换句话说, 几乎所有LLM推理端点都是非确定性的主要原因是负载(因此批次大小)非确定性地变化! 这种非确定性并不是GPU独有的——从CPU或TPU提供的LLM推理端点也会有这种非确定性来源。

所以,如果我们想在推理服务器中避免非确定性,我们必须在内核中实现批次不变性。为了理解如何实现这一点,让我们首先看看为什么内核一开始就没有批次不变性。

如何使内核批次不变?

为了使transformer实现批次不变,我们必须使每个内核批次不变。幸运的是,我们可以假设每个逐点操作都是批次不变的。因此,我们只需要担心涉及约简的3个操作——RMSNorm、矩阵乘法和注意力机制。

方便的是,这些也按难度递增排序。每一个都需要一些额外的考虑来以合理的性能实现批次不变性。让我们首先讨论RMSNorm。

批次不变的RMSNorm

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-13

数据并行RMSNorm 理想情况下,我们希望避免在并行化策略中核心之间的通信。实现这一点的一种方法是将一个批次元素分配给每个核心,从而保证每个约简完全在单个核心内完成。这就是所谓的"数据并行"策略,因为我们只是沿着不需要通信的维度并行化。在这个例子中,我们有四行和四个核心,饱和了我们的核心。

RMSNorm可以实现为:

# x: [batch_size, hidden_dim]# weight: [hidden_dim]def rms_norm(x, weight):    return x * torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) * weight

批次不变性的要求是 每个元素的约简顺序必须固定,无论内核的批次大小如何 。注意这并不意味着我们必须总是使用相同的约简策略。例如,如果我们改变要约简的元素数量,即使约简策略发生变化,我们仍然可以是批次不变的。

因此,我们只有在批次大小影响约简策略时才会破坏批次不变性。

让我们看看RMSNorm的标准并行化策略。通常,并行算法受益于最小化核心间的通信。对于本讨论的目的,你可以假设当我们提到"核心"时,我们指的是SM。更具体地说,这里重要的属性是我们内核启动的线程块数量大于SM的数量。所以,我们可以开始的一个策略是将每个批次元素分配给一个核心,如上图所示。

增加批次大小不会影响我们的约简策略;如果批次大小为200为我们的内核提供了足够的并行性,那么批次大小为2000 肯定 会提供足够的并行性。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-14

更大批次的数据并行RMSNorm 将数据并行策略扩展到更大的批次相当简单——不是让每个核心处理一行,而是允许每个核心顺序处理不同的行。这 保持了批次不变性 ,因为每个批次元素的约简策略保持相同。

另一方面,减少批次大小可能带来挑战。因为我们将每个批次元素分配给一个核心,减少批次大小最终会导致核心数量多于批次元素,使一些核心闲置。

遇到这种情况时,一个好的内核工程师会使用前一节提到的解决方案之一(原子加法或分割约简),保持良好的并行性,从而获得良好的性能。不幸的是,这改变了约简策略,阻止了这个内核成为批次不变的。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-15

分割约简RMSNorm 如果我们有小的批次大小,我们的数据并行策略可能不再有足够的并行性来饱和我们的核心。在这种情况下,在多个核心之间"分割"约简可能更有效,允许我们充分利用GPU。然而,这 失去了 批次不变性,因为我们不再以相同的顺序约简每个元素。

最简单的解决方案是完全忽略这些情况。这并非完全 不合理 ——小批次大小意味着内核可能执行得很快,所以减速可能不是灾难性的。

如果我们 被迫 优化这种用例,一种方法是始终使用一个约简策略,即使对于非常小的批次大小也有足够的并行性。这样的约简策略会导致较大批次大小的过量并行性,但允许我们在整个大小范围内实现不错(但不是峰值)的性能。

批次不变的矩阵乘法

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-16

数据并行矩阵乘法 类似于RMSNorm,矩阵乘法的标准并行化策略是"数据并行"策略,将整个约简保持在一个核心中。最直接的思考方式是将输出张量分割为2D块,并将每个块分配给不同的核心。然后每个核心计算属于该块的点积,再次在一个核心内执行整个约简。

从本质上讲,你也可以将矩阵乘法视为逐点操作后跟约简。然后,如果我们通过将 输出 分块来并行化矩阵乘法,我们就有了一个类似的"数据并行"内核策略,将每个约简保持在一个核心内。

同样类似于RMSNorm,我们的"批次"维度(M和N)可能变得太小,迫使我们沿着约简维度(K)分割。尽管有两个"批次"维度,矩阵乘法也要求我们每个核心有更多的"工作",以便有效地利用张量核心。例如,如果你有一个[1024, K] x [K, 1024]矩阵乘法和标准的2D块大小[128, 128],数据并行策略只能将这个矩阵乘法分割为64个核心,不足以饱和GPU。

在矩阵乘法中沿约简维度分割被称为 Split-K矩阵乘法 [10] 。就像RMSNorm一样,使用这种策略会破坏批次不变性。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-17

Split-K矩阵乘法 如果我们的批次维度相当小,我们可能没有足够的并行性,需要split-k矩阵乘法。在这个例子中,我们将每个约简分割到两个核心上,它们会分别累积然后在最后合并结果。然而,将每个约简分割到两个核心上允许我们仍然利用八个核心。

然而,还有一个额外的复杂性——张量核心指令。与约简不同,我们可以简单地一次操作一行,高效的矩阵乘法内核必须一次操作整个"块"。

每个张量核心指令(比如 `wgmma.mma_async.sync.aligned.m64n128k16` [11] )内部可能有不同的约简顺序。使用不同张量核心指令的一个原因可能是批次大小非常小。例如,如果我们使用操作长度为256的块的张量核心PTX指令,但批次大小只有32,我们几乎浪费了所有计算!在批次大小为1时,最快的内核通常根本不使用张量核心。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-18

填充的张量核心指令 如果批次大小太小,我们可能处于无法在输出中放入甚至一个2D块的情况。在这种情况下,切换到较小的张量核心指令或完全放弃张量核心是最有效的!然而,这两个选项都阻止了我们的内核成为批次不变的。

所以,确保矩阵乘法批次不变性的最简单方法是编译一个内核配置并将其用于所有形状。虽然我们会失去一些性能,但这在LLM推理中通常不是灾难性的。特别是,当 M和N都 很小时最需要split-k,而幸运的是,在我们的情况下,N(即模型维度)通常相当大!

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-19

尽管获得了批次不变性,我们与cuBLAS相比只损失了大约20%的性能。注意这也不是优化的Triton内核(例如没有TMA)。然而,性能中的一些模式说明了我们的批次不变要求在哪里损失性能。首先,注意我们在非常小的批次大小时由于过大的指令和不足的并行性损失了大量性能。其次,随着批次大小的增加,有一个"拼图"模式,这是由通常通过改变块大小来改善的量化效应(块和波)引起的。

批次不变的注意力机制

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-20

FlashAttention2策略 我们沿Q并行化,同时沿K/V约简。这意味着我们的整个约简可以保持在单个核心内,使其成为另一个数据并行策略。

在为矩阵乘法获得批次不变性后,注意力机制引入了两个额外的复杂性——恰如其分,因为它包含两个矩阵乘法。

与只沿特征维度约简的RMSNorm和矩阵乘法不同,我们现在沿特征维度 和 序列维度约简。

由于上述原因,注意力机制必须处理影响序列处理方式的各种推理优化(分块预填充、前缀缓存等)。

因此,为了在LLM推理中实现确定性,我们的数值必须对一次处理多少请求 和 每个请求在推理引擎中如何分片都保持不变。

让我们首先介绍注意力机制的标准并行化策略,首次在FlashAttention2中引入。类似于RMSNorm和矩阵乘法,默认策略是"数据并行"策略。由于我们沿着key/value张量约简,数据并行策略只能沿着query张量并行化。

例如,根据推理引擎的选择,一个序列可能被分几部分处理(如在分块预填充中)或一次全部处理(如果预填充没有分割)。为了实现"批次不变性",必须确保 给定token的约简顺序不依赖于其序列中同时处理多少其他token 。如果你分别约简KV缓存中的K/V值和当前处理的token中的K/V值(如vLLM的 Triton注意力内核 [12] ),这无法实现。例如,在处理序列中的第1000个查询token时,无论KV缓存中有0个token(预填充)还是999个token(解码),约简顺序必须相同。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-21

带KV缓存的FlashAttention 为什么显式地将KV缓存与当前KV值分开处理会破坏批次不变性有点微妙,这与"边界条件"有关。特别是,想象你的块大小是32,但我们当前在KV缓存中有80个元素。然后我们计算另外48个未缓存的元素。在这种情况下,我们需要三个块(两个完整的和一个掩码的)来计算"P缓存",另外两个块(一个完整的和一个掩码的)来计算"P"。因此,当我们只有四个总块(即128)元素要计算时,这是五个总块来计算我们的约简,这肯定会改变我们的约简顺序。

为了解决这个问题,我们可以在注意力内核本身之前更新KV缓存和页表,确保我们的键和值总是一致地布局,无论正在处理多少token。

有了这个额外的细节(以及前一节提到的所有内容,如一致的块大小),我们能够实现批次不变的注意力实现!

然而,这里有一个重大问题。与矩阵乘法不同,我们在LLM推理中看到的注意力形状通常确实需要分割约简内核,通常称为Split-KV或FlashDecoding。这是因为如果我们不沿约简并行化,我们只能沿批次维度、头维度和"查询长度"维度并行化。在注意力的解码阶段,查询长度非常小,所以除非我们有非常大的批次大小,否则我们通常无法饱和GPU。

不幸的是,忽略这种情况不像RMSNorm和矩阵乘法那样容易。例如,如果你有非常长的KV缓存,尽管只处理一个请求,注意力内核可能需要很长时间。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-22

固定 #Split -KV策略(即FlashDecode) 如果我们的查询长度变得非常小(如在解码期间),我们可能最终处于内核中几乎没有并行性的情况。在这些情况下,我们需要再次沿约简维度分割——这次是KV维度。沿KV维度分割的典型策略是确定我们需要多少并行性,然后平均分割KV维度。例如,如果我们的KV长度是1000,需要4个分割,每个核心将处理250个元素。

这不幸地也破坏了批次不变性,因为我们的精确约简策略取决于我们在任何给定请求中从序列处理多少查询token。

此外,注意力机制常用的分割约简策略也对批次不变性构成挑战。例如,FlashInfer的"平衡调度算法"选择仍能饱和GPU所有核心的最大分割大小,从而使约简策略不是"批次不变的"。然而,与RMSNorm/矩阵乘法不同,仅选择固定数量的分割而不考虑批次大小是不够的。

相反,为了实现批次不变性,我们必须采用"固定分割大小"策略。换句话说,不是固定分割的数量,我们固定每个分割的大小,然后得到不同数量的分割。通过这种方式,我们可以保证无论我们处理多少token,我们总是执行相同的约简顺序。这需要一些内部FlexAttention更改,这些更改不包含在我们的代码发布中。我们将在不久的将来上游它们!

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-23

固定大小Split-KV策略 这个策略与前一个策略的唯一区别是我们的分割现在是"固定大小的"。例如,如果我们的KV长度是1000,不是将其分割为四个长度为250的均匀分割,我们会将其分割为三个固定大小长度为256的分割和一个长度为232的分割。这允许我们 保持 批次不变性,因为我们的约简策略不再依赖于我们一次处理多少查询token!

实现

我们通过利用vLLM的FlexAttention后端以及torch.Library,在vLLM之上提供了确定性推理的演示。通过torch.Library,我们能够以非侵入式的方式替换大部分相关的PyTorch操作符。你可以在 thinking-machines-lab/batch-invariant-ops [13] 找到"批次不变"内核库,以及在"确定性"模式下运行的vLLM示例。

实验

生成结果的非确定性程度如何?

我们使用 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设置运行Qwen-3-8B的API服务器,并请求1000个输出长度在90到110之间的序列。

配置

时间(秒)

vLLM默认

26

未优化的确定性vLLM

55

+ 改进的注意力内核

42

大部分减速来自vLLM中的FlexAttention集成尚未进行大量优化。尽管如此,我们看到性能不是 灾难性的 。

真正的在线策略强化学习

正如 研究人员所注意到的 [14] ,训练和推理之间的不同数值隐含地将我们的在线策略RL转变为离线策略RL。

当然,如果我们甚至无法从两个相同的推理请求中获得按位相同的结果,就不可能在训练和推理之间获得按位相同的结果。然后,确定性推理使我们也能够修改我们的训练堆栈,以在采样和训练之间获得按位相同的结果,从而产生真正的在线策略RL。

我们在 Bigmath [15] 上的RLVR设置中运行实验,RL策略从Qwen 2.5-VL指令8B初始化,最大展开长度为4096。

如果我们在没有离线策略校正(即重要性加权)的情况下训练,我们的奖励在训练过程中崩溃,而添加离线策略校正项允许训练顺利进行。但是,如果我们在采样器和训练器之间实现按位相同的结果,我们是完全在线策略的(即0 KL散度),也可以顺利训练。

我们还可以绘制采样器和训练器之间logprobs的KL散度,其中所有3次运行都有明显不同的行为。使用重要性加权运行时,它保持在0.001左右,偶尔有峰值。然而, 不使用 重要性加权运行最终会在奖励崩溃的同时导致KL散度峰值。当然,运行"真正的在线策略RL"时,我们的KL散度保持在0的平线,表明训练策略和采样策略之间 没有 散度。

刚刚OpenAI前CTOMiraMurati公司ThinkingMachinesLab发文-24

注意没有重要性加权的运行在步骤318左右有显著的损失峰值,这伴随着logprobs的相应KL散度峰值。同时,使用离线策略校正或运行"真正的在线策略"都允许RL继续顺利进行。显示"真正在线策略"的蓝线不是错误——它只是在0处的平线。

结论

现代软件系统包含许多抽象层。在机器学习中,当我们遇到非确定性和微妙的数值差异时,往往很容易掩盖它们。毕竟,我们的系统已经是"概率性的",那么再多一点非确定性有什么问题呢?在失败的单元测试上提高atol/rtol有什么问题?训练器和采样器之间logprobs的差异可能不是真正的错误,对吧?

我们拒绝这种失败主义。通过一点工作,我们 可以 理解非确定性的根本原因,甚至解决它们!我们希望这篇博客文章为社区提供了如何解决推理系统中非确定性的坚实理解,并启发其他人获得对他们系统的完全理解。

引用

请引用这项工作为:

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference", Thinking Machines Lab: Connectionism, Sep 2025.

或使用BibTeX引用:

@article{he2025nondeterminism,  author = {Horace He and Thinking Machines Lab},  title = {Defeating Nondeterminism in LLM Inference},  journal = {Thinking Machines Lab: Connectionism},  year = {2025},  note = {https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/},  doi = {10.64434/tml.20250910}}

[1]

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

[2]

这里: https://152334h.github.io/blog/non-determinism-in-gpt-4/

[3]

这里: https://community.openai.com/t/a-question-on-determinism/8185/2

[4]

这里: https://cookbook.openai.com/examples/reproducible_outputs_with_the_seed_parameter

[5]

这里: https://docs.vllm.ai/en/v0.7.0/getting_started/faq.html

[6]

这里: https://docs.sglang.ai/references/faq.html

[7]

最近的一篇arXiv预印本: https://arxiv.org/abs/2506.09501

[8]

这里: https://community.openai.com/t/a-question-on-determinism/8185

[9]

fetch-and-add: https://en.wikipedia.org/wiki/Fetch-and-add

[10]

Split-K矩阵乘法: https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/efficient_gemm.md #parallelized -reductions

[11]

wgmma.mma_async.sync.aligned.m64n128k16 : https://docs.nvidia.com/cuda/parallel-thread-execution/ #asynchronous -warpgroup-level-matrix-instructions-wgmma-mma

[12]

版权声明:
作者:shadowrocket
链接:https://www.shadowrocket9.top/95.html
来源:Shadowrocket官网
文章版权归作者所有,未经允许请勿转载。

THE END
分享
二维码
< <上一篇
下一篇>>