摘要
在科学探索的宏伟殿堂中,"可复现性"是支撑其穹顶的基石。然而,在当前人工智能的前沿——大型语言模型(LLM)领域,我们却面临着一个令人困惑的挑战:即使在理论上应完全确定的设置下(例如,将采样温度设为0),模型的输出仍然像幽灵一样难以捉摸,呈现出"非确定性"的特质。这一现象不仅困扰着研究人员,也对依赖模型一致性输出的应用构成了障碍。一个广为流传的假说将其归咎于GPU并行计算中固有的"浮点数非结合性"与线程执行顺序的随机性。然而,这一解释虽不无道理,却未能触及问题的全貌,如同只见树木,不见森林。
本文将以我的第一视角,带领读者踏上一场深入GPU内核与推理引擎架构的探索之旅。我们将首先解构浮点数运算的"原罪"——非结合性,并通过一个交互式动画直观展示其微妙而关键的影响。随后,我们将揭示为何"原子加操作"这一常见的并发"嫌疑人",在LLM的前向传播中其实是无辜的。真正的"罪魁祸首"远比这更隐蔽,它潜藏于推理服务器处理并发请求的机制之中——即"批处理大小(batch-size)"的动态变化。我们发现,许多底层计算核心(kernel)缺乏"批处理不变性"(batch invariance),导致一个请求的计算结果会因与之同批处理的其他请求数量而发生微小但可累积的偏差。
在定位了问题的根源后,我们将系统性地提出一套解决方案,旨在"锻造"确定性。我们将逐一剖析Transformer模型中的三个关键组件——RMSNorm、矩阵乘法(Matmul)和注意力机制(Attention),并为它们量身打造具备"批处理不变性"的内核实现策略。这趟旅程不仅涉及算法设计的权衡,还包含了对现代GPU硬件特性的深刻理解。最终,我们将通过实验数据证明,通过实施这些批处理不变的内核,我们能够在不严重牺牲性能的前提下,成功驯服非确定性的"猛兽",实现真正可复现的LLM推理。这不仅为实现真正的"在策略(on-policy)"强化学习铺平了道路,更为构建更可靠、更可信的AI系统提供了坚实的基础。
第一章:原罪——浮点数的"非结合律"幽灵
大家好,我是Horace。今天,我想和大家聊一个看似深奥,却又与我们每次和AI交互都息息相关的话题:为什么即使我把语言模型的"创造力开关"(也就是 `temperature` 参数)调到0,它每次给我的回答还是可能不一样?
这就像一个顶级厨师,严格按照同一份食谱、用同样的食材烹饪,但每次端上来的菜品味道都有微乎其微的差别。这背后,藏着一个计算机科学的"原罪"——浮点数的非结合性。
在数学世界里,我们从小就学习加法结合律:\( (a+b)+c = a+(b+c) \)。但在计算机的浮点数世界里,这个定律却不总是成立。为什么呢?
生活化类比:精密的尺子与巨大的误差
想象一下,你有一把超级精密的尺子,但它只能记录6位有效数字。现在让你测量三个长度:一个是"一粒沙"的直径(比如0.000001米),一个是"一栋楼"的高度(比如100.000米),另一个是"负一栋楼"的高度(-100.000米)。
如果你先算"楼 + (-楼)",结果是0,再加"沙",最终结果就是0.000001米。
但如果你先算"楼 + 沙",因为尺子精度有限,100.000米加上0.000001米,结果太长,尺子只能四舍五入记为100.000米,那粒"沙"的长度就被"吃掉"了。然后再减去"楼",结果就变成了0。
仅仅是计算顺序的不同,导致了"一粒沙"的神秘消失!这就是浮点数非结合性的本质:当数量级差异巨大的数字进行加减时,精度会丢失。在LLM复杂的计算中,亿万次这样的浮点数加法正在发生。不同的计算顺序,就会累积成最终输出文本的差异。下面的动画将让你亲手体验这个"数字魔法"。
动画一:浮点数求和顺序的魔力
类比:尝试拖动下面的数字牌,改变它们的求和顺序,观察在计算机有限的精度下,最终结果如何发生变化。这揭示了 \( (a+b)+c \ne a+(b+c) \) 的现象。
当前顺序计算结果: ...
第二章:被误解的"凶手"——并发与原子加
既然我们知道了计算顺序是关键,那么一个很自然的推论就是:现代GPU拥有成千上万个并行核心,它们像无数个忙碌的工人,执行任务完成的顺序自然有随机性。如果它们都向同一个"账本"上累加数字(这个操作叫"原子加"),那么累加的顺序就会随机,从而导致结果的非确定性。
这个"并发+浮点数"假说听起来非常完美,也流传甚广。但事实是,它在LLM前向推理的场景下,只是一个"红鲱鱼"——一个看似重要却误导人的线索。
生活化类比:有序的交响乐团
将GPU的一次计算想象成一场交响乐。如果指挥家(程序)让所有乐手(核心)随意演奏自己的部分,然后把声音混在一起,那将是一片混乱。但实际上,指挥家会给每个声部(例如,小提琴组)分配独立的乐谱和区域。小提琴组内部的演奏可以高度并行,但他们只对自己声部的总音量负责。最后,指挥家再将各个声部的完美成品组合起来。现代的深度学习库,就像这位高明的指挥家。在LLM前向传播(即生成文本的过程)中,绝大多数计算都被精心设计,以避免使用那种混乱的、依赖完成顺序的"原子加"。
例如,计算一个向量的和,与其让100个核心抢着往一个地方加,不如先让它们分成5组,每组20个核心各自算出一个局部和,最后再将这5个局部和相加。这种"分而治之"的树状归约(Tree Reduction)策略,在保证并行效率的同时,也固定了计算顺序,从而变得"确定"。
所以,尽管某些特定GPU操作(尤其是在模型训练的反向传播中)确实存在非确定性,但对于我们关心的"推理"过程,几乎所有的计算内核(Kernel)本身都是"运行到运行确定"的(run-to-run deterministic)。也就是说,如果你用完全相同的输入数据,在同一台机器上连续运行两次同一个内核,你会得到比特级别完全相同的结果。
那么,谜题又回来了。如果每个"乐器"(计算内核)的演奏都是确定的,为什么最终的"交响乐章"(LLM的输出)还是变幻莫测呢?
图一:两种并行求和策略
原子加(左)像多人同时向一个篮子扔球,顺序不定。树状归约(右)则像分级比赛,层层汇总,顺序固定,保证了确定性。
第三章:真正的罪魁祸首——批处理不变性的缺失
答案,隐藏在一个我们每天都在使用却很少深思的系统行为中:批处理(Batching)。
为了提升效率,LLM推理服务器不会一个一个地处理用户请求。它会像公交车一样,等待一小段时间,把多个用户的请求(连同它们正在生成的后续词元)打包成一个"批次"(batch),然后一次性扔给GPU处理。这个批次的大小是动态变化的,取决于服务器当前的负载情况。有时可能只有你的一个请求,有时可能同时有几十上百个。
问题的关键在于:许多底层的GPU计算内核,特别是矩阵乘法,它们虽然对单个批次是确定的,但其计算结果却不具备"批处理不变性"(Batch Invariance)。
生活化类比:专业咖啡师的困境
想象一位技艺精湛的咖啡师,他有一套冲煮单杯咖啡的完美流程,每次都能做到分秒不差。这是"运行到运行确定"。
现在,为了效率,老板要求他同时冲煮多杯咖啡。当他只冲煮一杯时,他可以全神贯注,水温、研磨、冲泡时间都控制得尽善尽美。但当他同时冲煮10杯时,尽管他尽力维持标准,但为整个批次准备的热水、磨豆机的工作状态,都可能与只做一杯时有微小差异。这些微小的系统性偏差,会传递到你手中的那一杯咖啡,让它的风味与单独冲煮时略有不同。这就是LLM推理正在发生的事情。你的一次请求,其计算结果会受到"拼车"的其他请求数量的影响。从你的视角来看,你无法控制也无法预知和你同批处理的请求有哪些,因此,这个批处理大小就成了一个非确定性的外部因素。
当一个缺乏"批处理不变性"的内核,与一个"非确定性"的批处理大小相遇,整个系统就对外表现出了非确定性。即使内核本身是"确定"的!
动画二:矩阵乘法的"批处理"效应
类比:下方模拟了两种矩阵乘法。左侧是单独计算(批大小为1),右侧是与许多其他数据一起批处理计算,我们只观察第一行的结果。由于GPU为优化大批量计算而采用不同策略,导致即使是同一行输入,结果也可能出现微小差异。
单独计算 (Batch=1) 结果: N/A
批处理 (Batch=N) 首行结果: N/A
结果是否比特一致: 待开始
第四章:铸造确定性——实现批处理不变的内核
既然找到了病根,我们就可以对症下药了。我们的目标是:改造Transformer模型中的关键计算内核,让它们无论批处理大小是多少,都能对同一份输入数据产生完全相同的结果。这就像要求我们的咖啡师,无论同时冲1杯还是10杯,每一杯的味道都必须与单独冲煮时完全一致。
我们需要关注三个核心操作:RMSNorm、矩阵乘法(Matmul)和注意力机制(Attention)。
4.1 RMSNorm的改造
RMSNorm需要对每个序列中的每个词元(token)的特征向量进行归一化,这涉及到一次求和归约。当批处理大小足够大时,GPU可以为每个词元分配一个核心,并行处理,这天然地保证了批处理不变性。但当批次很小时(比如只有一个请求),为了不让GPU核心闲置,高性能内核会"分裂"单个词元的计算任务,让多个核心协同完成。这就改变了计算顺序,破坏了不变性。
解决方案:最简单的方法是"以不变应万变"。我们强制内核始终使用同一种并行策略,比如,总是采用那种为小批次设计的、需要"分裂"任务的策略。这样,虽然在大批次时牺牲了一点点性能(因为并行度可能过高),但我们换来了宝贵的确定性。
动画三:RMSNorm 的并行策略
类比:观察在不同批处理大小下,GPU如何分配计算任务。在"性能优先"模式下,策略会动态改变,导致计算顺序变化。在"确定性优先"模式下,策略被固定,保证了结果的一致性。
当前模式: 性能优先
批大小 > 4 时计算顺序: 数据并行
批大小 <= 4 时计算顺序: 分裂归约
4.2 矩阵乘法的挑战
矩阵乘法是LLM的心脏。它的问题与RMSNorm类似,但更复杂。为了极致性能,现代GPU使用专门的"张量核心"(Tensor Cores),它们一次处理的是一个"瓦片"(Tile)状的数据块。高性能的库(如cuBLAS)会根据输入矩阵的形状,动态选择最优的瓦片大小和并行策略(如Split-K),这几乎必然导致批处理不变性的丧失。
解决方案:同样是"固定策略"。我们编译一个"一刀切"的内核配置,为所有形状的矩阵乘法都使用相同的瓦片大小和并行方案。幸运的是,在LLM推理中,矩阵的很多维度(如模型的隐藏层大小)是固定的,这使得固定策略的性能损失通常在可接受的范围内(大约10-20%)。
图二:矩阵乘法的瓦片化(Tiling)
为了在GPU上高效计算,大矩阵被切分成小"瓦片"。每个核心负责计算一个输出瓦片。不同的并行策略(如Data-Parallel vs. Split-K)会改变归约维度(K)的计算顺序。
4.3 注意力机制的终极挑战
注意力机制是最棘手的部分,因为它不仅涉及到矩阵乘法,还引入了序列维度上的归约。LLM推理为了优化长文本处理,会使用KV缓存(KV Cache)、分页注意力(PagedAttention)等技术。这意味着一个序列的计算可能会被切分成好几块(比如,处理提示词的prefill阶段和逐词生成的decode阶段),这给固定计算顺序带来了新的麻烦。
解决方案:我们需要一种更聪明的固定策略。对于decode阶段这种并行度严重不足的场景,我们必须使用"分裂归约"(Split-KV)。但为了保证不变性,我们不能像传统方法那样"固定分裂的份数",而是要"固定每一份的大小"。
比如,对于一个长度为1000的序列,传统方法可能会根据需要的并行度,将其分成4份,每份250。但如果下次序列长度变为900,它就会被分成4份,每份225。计算顺序就变了。我们的新策略是,固定每份的大小为256,那么长度1000的序列就会被分成三份256和一份232。下次长度为900时,它会被分成三份256和一份132。虽然最后一份的大小变了,但前三份的计算过程是完全一致的,从而保证了结果的确定性。
动画四:Attention 的 Split-KV 策略
类比:序列(K/V)是长面包,GPU核心是切刀。传统"固定份数"策略(左)会根据面包总长调整每块厚度。我们的"固定块大小"策略(右)则始终切出标准厚度的面包片,只有最后一块是边角料,保证了前面所有标准片的"加工过程"完全一致。
当前序列长度: 1000
固定份数策略 (4份) 切分: 4 x 250
固定块大小策略 (256) 切分: 3 x 256 + 1 x 232
第五章:和谐的流场——确定性之美
经过上述一系列精密的改造,我们终于驯服了非确定性的幽灵。数据在GPU中流动,不再是随机碰撞的溪流,而更像一片遵循着无形而深刻规律的星云。
下面的这个动画,虽然与LLM的具体计算无关,但它在美学上捕捉了我们追求的最终目标:一个由简单规则驱动,却能产生无限复杂、和谐而又完全确定的动态系统。每一个粒子的运动都由其所在位置的柏林噪声场(Perlin Noise Field)决定,这就像我们的内核,其行为由输入数据唯一确定。这片流动的星海,便是我心中"确定性"最美的写照。
高级动画:粒子流场
生活化类比:想象无数微小的尘埃,在空中随一阵看不见却又和谐有序的风飘动,形成了优雅的涡流和线条。这是一个完全由算法驱动的、确定性的复杂系统。
结论:拒绝失败主义,拥抱完全理解
现代软件系统充满了层层抽象。在机器学习领域,当我们遇到非确定性和微小的数值差异时,很容易选择"视而不见"。毕竟,模型本身就是"概率性"的,多一点非确定性又何妨?把单元测试的容忍度(atol/rtol)调高一点不就行了?
我们拒绝这种技术上的失败主义。通过一点点努力,我们完全可以理解我们系统中非确定性的根源,并彻底解决它们!我希望这篇深度剖析,能为社区提供一个解决推理系统非确定性的清晰路线图,并激励更多人去追求对自己系统的完全理解。
实现确定性推理,不仅仅是满足洁癖工程师的强迫症。它对于需要严格一致性的应用(如内容审核、代码生成)至关重要,更是实现真正"在策略"(on-policy)强化学习等前沿研究的基石。当我们能够精确复现每一次计算,我们才算真正掌握了手中的工具,才能在此基础上,构建更加宏伟和可靠的智能大厦。
专题对话:探索大语言模型的"非确定性"之谜
以下内容整理自一次围绕技术文章《Defeating Nondeterminism in LLM Inference》的深度对话,旨在梳理为何大型语言模型(LLM)面对相同问题会给出不同答案,以及如何系统性地消除这种"非确定性"。
一、初始问题:这篇文章说了什么?
文章指出,LLM 非确定性的根源不仅仅是 GPU 上浮点数的非结合性与并行顺序差异,更关键的是许多推理内核缺乏批次不变性(Batch Invariance):同一请求在不同批大小下的计算路径与顺序会变化,从而导致结果差异。作者提出对关键算子进行"批不变"改造的方案,并用实验验证其必要性与可行性。
二、深入核心:什么是"批次不变性"?
理想状态下,模型的输出只应取决于输入本身,与服务器上其他并发请求无关。批次不变性指:无论一个请求与多少其他请求合批处理,结果都应完全一致。然而实际中,为了性能,内核会随批大小调整策略,进而改变计算顺序与数值路径,造成外显的不确定性。
用户洞察: 也就是说四舍五入不一样,由于顺序不一样,导致了不同的结果是吧?由于浮点数运算的特性,这会为批次中的每个人产生略微不同的结果。
要点归纳:
• 浮点问题:硬件层面的精度与非结合性(如 \((a+b)+c \ne a+(b+c)\))。
• 批次顺序问题:批大小变化触发不同并行/归约策略,点燃"浮点火药"。
三、解决方案:如何实现"确定性"?
目标是在关键算子层面固定计算策略,使其对批大小不敏感,确保相同输入得到比特级一致输出。
- 1) RMSNorm:小批时常切换到"分裂归约"等不同策略。方案:强制使用统一策略(即使牺牲少量性能)。
- 2) 矩阵乘法(Matmul):因 Tensor Cores 与指令/瓦片选择在不同形状与批量下切换。方案:固定内核配置与瓦片大小/指令集,避免 Split-K 等策略漂移。
- 3) 注意力机制(Attention):同时受批大小与序列长度变化影响。方案:统一 KV 缓存更新路径;长序列按固定大小切分,不随需求动态改变切分份数。
四、实验验证:效果与代价
- 实验一:不确定性规模:在理论应确定的设置下重复 1000 次提问,标准模型出现约 80 个不同答案;确定性模型 1000 次完全一致。
- 实验二:性能代价:确定性带来速度下降,但代价可控且仍有优化空间。
- 实验三:特别用途:在"在策略(on-policy)"强化学习中,确定性模型避免训练崩溃,实现真正意义上的在策略学习。
五、进一步洞察
用户总结:LLM 的训练如果是确定的,就不那么容易崩溃,数据利用效率更高,梯度下降更有效。
六、最终结论
这并非"看起来没用"的细节问题。通过为 RMSNorm、Matmul 与 Attention 落实批次不变性,我们既能约束模型行为、提升训练稳定性,也能为在策略强化学习等应用奠定确定性与可复现性的基础,具有明显的必要性与可行性。
附录:技术细节深潜
为了实现上述的批处理不变内核,我们主要依赖于对底层硬件和编程模型(如CUDA和Triton)的深入理解。以下是一些关键的技术要点:
1. 归约策略的选择与固定
如前文所述,避免动态选择归约策略是核心。在Triton这样的高级内核编程语言中,这意味着我们要避免编写依赖于输入张量维度来改变计算逻辑的代码。例如,一个常规的归约内核可能会这样写(伪代码):
if (REDUCTION_SIZE > 1024 and NUM_THREADS > 128):
use_split_reduction_strategy()
else:
use_single_threadblock_strategy()
为了确定性,我们必须移除这种条件分支,选择一种策略并始终坚持。这通常需要对目标硬件(如H100, A100)的特性进行分析,找到一个在常见负载下性能"足够好"的"最大公约数"策略。
2. Tensor Core与指令选择
NVIDIA GPU的Tensor Core有多种MMA(Matrix Multiply-Accumulate)指令,如`wgmma`、`mma.sync`等,它们处理不同形状(M, N, K)的瓦片。高性能库会自动选择最高效的指令。例如,对于一个批大小为1的向量-矩阵乘法,使用处理大瓦片的`m64n128k16`指令会浪费大量计算资源,此时库可能会降级为不使用Tensor Core的普通CUDA核心计算。
为了批处理不变性,我们必须固定使用的MMA指令。这意味着,即使在批大小很小时,我们可能仍然会使用为大批次设计的指令,并通过"填充"(Padding)将小输入伪装成大输入的形状,以满足指令的要求。这会带来性能开销,但保证了计算方式的同一性。
3. 注意力机制中的KV缓存布局
在vLLM等现代推理引擎中,KV缓存并非连续存储。它通过"分页"(Paging)机制,将不同序列的缓存块映射到物理内存中。为了实现批处理不变的注意力计算,一个关键步骤是,在调用注意力内核之前,就必须将当前批次中所有token的Key和Value值更新到这个分页的KV缓存结构中。
这样,注意力内核看到的总是一个完整的、包含了历史信息和当前信息的Key/Value张量视图。它无需区分哪些来自缓存、哪些是新计算的,从而可以使用统一的、不依赖于prefill/decode状态的归约逻辑。这避免了因处理边界条件(如缓存块的末尾)而引入的计算顺序变化。
4. 实现框架
我们的演示是在vLLM之上,通过其`FlexAttention`后端以及`torch.Library`机制实现的。`torch.Library`允许我们在不修改PyTorch或vLLM源代码的情况下,用我们自己编写的、具备批处理不变性的内核"替换掉"默认的PyTorch算子(如`torch.matmul`)。你可以在我们的代码仓库 `thinking-machines-lab/batch-invariant-ops` 找到这些内核的实现和在vLLM中启用"确定性模式"的示例。