Thinking  Machines Lab終于放大招了！

剛剛，聯(lián)合創(chuàng)始人、OpenAI前副總裁Lilian Weng透露：

Thinking Machines的第一代旗艦產(chǎn)品名為「Connection Machine」（聯(lián)結(jié)機(jī)）。

事情是這樣的：今天，Thinking Machines開辟了研究博客專欄「Connectionism」（聯(lián)結(jié)主義），發(fā)表了第一篇博客文章「Defeating Nondeterminism in LLM Inference」（擊敗LLM推理中的非確定性）。

Thinking Machines介紹說:

我們相信科學(xué)因分享而更美好。  

Connectionism專欄將隨著我們的研究變化：從內(nèi)核數(shù)值計(jì)算到提示工程。在這里，我們分享我們的工作進(jìn)展，并與研究社區(qū)頻繁而開放地交流。

此外，還補(bǔ)充道，「Connectionism」這一名稱可追溯至早期的AI時(shí)代——在1980年代，該術(shù)語指代專門研究神經(jīng)網(wǎng)絡(luò)及其與生物大腦相似性的子領(lǐng)域。

圖片

而Lilian Weng爆出了更大的料，專欄之所以叫這個(gè)名稱，還有一個(gè)原因：第一代旗艦?zāi)Ｐ途徒蠧onnection Machine，不光是這篇博客文章，而且還有更多好東西要來了！

圖片

莫非Thinking Machines馬上要發(fā)布新模型了？

在期待新的LLM之前，我們先看看這次Thinking Machines到底有哪些絕活，他們到底關(guān)注哪些研究領(lǐng)域。

圖片

傳送門：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

據(jù)博文的主要作者Horace He介紹，這次的博文主要關(guān)于他心中的重要話題——

LLM推理中的可復(fù)現(xiàn)浮點(diǎn)數(shù)（Reproducible floating point numerics in LLM inference）。

圖片

LLM推理中的非確定性難題

可復(fù)現(xiàn)性是科學(xué)進(jìn)步的基石。然而，從大語言模型中獲得可復(fù)現(xiàn)的結(jié)果卻異常困難。

例如，你可能會(huì)觀察到，多次向ChatGPT提出相同的問題會(huì)得到不同的結(jié)果。

這本身并不奇怪，因?yàn)閺恼Z言模型獲得結(jié)果涉及「采樣」：

將語言模型的輸出轉(zhuǎn)換為概率分布，概率性地選擇一個(gè)token。

圖片

更令人驚訝的可能是，即使我們將temperature降至0（從而使采樣在理論上是確定性的），LLM API在實(shí)踐中仍然不是確定性的。

圖片

圖片

圖片

即使在自己的硬件上，使用像vLLM或SGLang這樣的開源推理庫(kù)運(yùn)行推理，采樣仍然不是確定性的。

圖片

但為什么LLM推理引擎不是確定性的呢？

一個(gè)常見的假設(shè)是，浮點(diǎn)數(shù)非結(jié)合性與并發(fā)執(zhí)行的某種組合，導(dǎo)致了基于哪個(gè)并發(fā)核心先完成的非確定性。

這次研究則稱之為L(zhǎng)LM推理非確定性的「并發(fā)+浮點(diǎn)數(shù)」假說。

例如，華人研究員Jiayi Yuan、Hao Li、Xinheng Ding等最近上傳了一篇arXiv預(yù)印本，其中寫道：

GPU中的浮點(diǎn)運(yùn)算表現(xiàn)出非結(jié)合性，意味著 (a+b)+c ≠ a+(b+c)，這是由于有限的精度和舍入誤差。

此屬性直接影響Transformer架構(gòu)中注意力分?jǐn)?shù)和logits的計(jì)算，其中跨多個(gè)線程的并行操作可能會(huì)根據(jù)執(zhí)行順序產(chǎn)生不同的結(jié)果。

圖片

傳送門：https://arxiv.org/abs/2506.09501

雖然這個(gè)假說有些道理，但并未揭示全貌。

例如，即使在GPU上，對(duì)相同數(shù)據(jù)重復(fù)運(yùn)行相同的矩陣乘法，也總會(huì)提供逐位相等的結(jié)果。

我們確實(shí)在使用浮點(diǎn)數(shù)，GPU確實(shí)有大量的并發(fā)計(jì)算。那為什么在這個(gè)測(cè)試中沒有看到非確定性呢？??

圖片

要理解LLM推理非確定性的元兇，我們必須更深入地探究。

不幸的是，即使是定義LLM推理的確定性也并非易事。

也許令人困惑的是，以下陳述竟能同時(shí)成立：

• GPU上的某些核（Kernel）是非確定性的。
• 然而，語言模型前向傳播中使用的所有Kernel都是確定性的。
• 此外，LLM推理服務(wù)（如vLLM）的前向傳播也可以聲稱是確定性的。
• 然而，從推理服務(wù)的用戶的角度來看，結(jié)果都是非確定性的。

這次，Thinking Machines決定要揭示LLM推理非確定性背后的真正元兇，并闡述如何克服LLM推理中的非確定性，獲得真正可復(fù)現(xiàn)的結(jié)果。

關(guān)鍵發(fā)現(xiàn)：

LLM前向傳播不需要原子加法；其非確定性真正來源是「批次大小變化」而非「原子競(jìng)爭(zhēng)」。

要想在推理服務(wù)中避免非確定性、為了使Transformer實(shí)現(xiàn)具有批處理不變性，我們必須在kernel中實(shí)現(xiàn)「批處理不變性」。

幸運(yùn)的是，我們可以假設(shè)每個(gè)逐點(diǎn)（pointwise）操作都具有批處理不變性。因此，只需要擔(dān)心涉及歸約的3個(gè)操作——RMSNorm、矩陣乘法和注意力。

它們的實(shí)現(xiàn)難度也是遞增的。每個(gè)操作都需要一些額外的考慮，才能以合理的性能實(shí)現(xiàn)批處理不變性。

批處理不變的RMSNorm：  數(shù)據(jù)并行RMSNorm

理想情況下，我們希望在并行化策略中避免核心之間的通信。

實(shí)現(xiàn)這一點(diǎn)的一種方法是為每個(gè)核心分配一個(gè)批處理元素，從而保證每個(gè)歸約都完全在單個(gè)核心內(nèi)完成。

這就是所謂的「數(shù)據(jù)并行」策略，因?yàn)槲覀冎皇茄刂粋€(gè)不需要通信的維度進(jìn)行并行化。

圖片

批處理不變的矩陣乘法：數(shù)據(jù)并行Matmul

與RMSNorm類似，矩陣乘法的標(biāo)準(zhǔn)并行策略是一種「數(shù)據(jù)并行」策略，將整個(gè)歸約保持在一個(gè)核心內(nèi)。

最直接的思考方式是將輸出張量分割成二維的分塊（tiles），并將每個(gè)分塊分配給不同的核心。然后，每個(gè)核心計(jì)算屬于該分塊的點(diǎn)積，再次在單個(gè)核心內(nèi)執(zhí)行整個(gè)歸約。

圖片

與RMSNorm不同，圍繞算術(shù)強(qiáng)度和利用張量核心（tensorcores）的額外約束，被迫分割二維分塊而不是單個(gè)輸出元素進(jìn)行，以實(shí)現(xiàn)高效的矩陣乘法kernel。

解決的核心在于，你可以將矩陣乘法看作是一個(gè)逐點(diǎn)操作后跟一個(gè)歸約。

確保矩陣乘法具有批處理不變性的最簡(jiǎn)單方法是，編譯一個(gè)kernel配置，并將其用于所有形狀。

雖然會(huì)損失一些性能，但這在大語言模型推理中通常不是災(zāi)難性的：

相比cuBLAS只損失了約20%的性能。

圖片

批處理不變的注意力機(jī)制

在為矩陣乘法獲得批處理不變性之后，注意力機(jī)制引入了兩個(gè)額外的難題——恰如其分，因?yàn)樗瑑蓚€(gè)矩陣乘法。

• 與RMSNorm和矩陣乘法僅在特征維度上進(jìn)行歸約不同，現(xiàn)在在特征維度和序列維度上進(jìn)行歸約。
• 由于上述原因，注意力機(jī)制必須處理各種影響序列處理方式的推理優(yōu)化（分塊預(yù)填充chunked prefill、前綴緩存prefix caching等）。

帶KV緩存的FlashAttention會(huì)破壞批處理不變性，根因在把「緩存KV」與「當(dāng)前KV」分開算

不同 KV 塊數(shù) → 不同掩碼/完整塊組合 → 不同規(guī)約路徑。

圖片

只要在kernel啟動(dòng)前，統(tǒng)一更新KV-cache頁(yè)表，保證任意時(shí)刻KV布局一致，就能解決這一問題。

大語言模型推理中看到的注意力形狀通常確實(shí)需要一個(gè)分裂歸約的kernel，通常稱為Split-KV或FlashDecoding。

固定數(shù)量的Split-KV策略（即FlashDecode），因?yàn)榫_的歸約策略取決于給定請(qǐng)求中處理來自序列的查詢token數(shù)量，這不幸地也破壞了批處理不變性。

圖片

如果我們的查詢長(zhǎng)度變得非常小（就像在解碼期間那樣），可能會(huì)陷入一種情況，即kernel中幾乎沒有并行性。在這些情況下，需要再次沿著歸約維度——這次是KV維度——進(jìn)行分割。分割KV維度的典型策略是計(jì)算出需要多少并行性，然后均勻地劃分KV維度。例如，如果KV長(zhǎng)度是1000，我們需要4個(gè)分割split，每個(gè)核心將處理250個(gè)元素。

此外，通常用于注意力的分裂歸約策略也對(duì)批處理不變性構(gòu)成了挑戰(zhàn)。

為了實(shí)現(xiàn)批處理不變性，不再固定分割的數(shù)量，而是固定每個(gè)分割的大小，然后得到一個(gè)可變數(shù)量的分割。

通過這種方式，可以保證無論正在處理多少個(gè)token，我們總是執(zhí)行相同的歸約順序。

圖片

這實(shí)現(xiàn)了批處理不變性，因?yàn)闅w約策略不再依賴于一次處理多少個(gè)查詢token！

用「固定塊大小」Split-KV，注意力也能像 RMSNorm/Matmul 一樣實(shí)現(xiàn)批處理不變，確定性推理。

開實(shí)現(xiàn)與實(shí)驗(yàn)

通過利用vLLM的FlexAttention后端以及torch.Library，Thinking Machines提供了一個(gè)在vLLM上進(jìn)行確定性推理的演示。

圖片

傳送門：https://github.com/thinking-machines-lab/batch_invariant_ops

補(bǔ)全結(jié)果有多大的非確定性？

使用Qwen/Qwen3-235B-A22B-Instruct-2507，在溫度為0的情況下，用提示詞「Tell me about Richard Feynman」（非思考模式）采樣1000個(gè)補(bǔ)全，每個(gè)生成1000個(gè)token。

令人驚訝的是，我們生成了80個(gè)不同的補(bǔ)全，其中最常見的出現(xiàn)了78次。

觀察補(bǔ)全結(jié)果的差異之處，我們發(fā)現(xiàn)補(bǔ)全結(jié)果實(shí)際上在前102個(gè)token上是完全相同的！第一次出現(xiàn)分歧的補(bǔ)全發(fā)生在第103個(gè)token。

所有的補(bǔ)全都生成了序列「Feynman was born on May 11, 1918, in」。

然而，992個(gè)補(bǔ)全接著生成了「Queens, New York」，而8個(gè)補(bǔ)全生成了「New York City」。

另一方面，當(dāng)啟用批處理不變kernel時(shí)，所有的1000個(gè)補(bǔ)全都是完全相同的。

性能

這次沒有投入大量精力來優(yōu)化批處理不變kernel的性能，但仍用實(shí)驗(yàn)測(cè)試了一下性能。

實(shí)驗(yàn)設(shè)置：一個(gè)帶有一塊GPU的API服務(wù)器，運(yùn)行Qwen-3-8B，并請(qǐng)求1000個(gè)序列，輸出長(zhǎng)度在90到110之間。

圖片

大部分的性能下降來自于vLLM中的FlexAttention集成尚未經(jīng)過大量?jī)?yōu)化。盡管如此，性能下降并非不可接受。

真正的同策略強(qiáng)化學(xué)習(xí)

傳送門：https://fengyao.notion.site/off-policy-rl

如果兩次相同的推理請(qǐng)求都無法做到逐位（bitwise）一致，那訓(xùn)練與推理在位級(jí)一致就更無從談起。

確定性推理讓我們可以同步改造訓(xùn)練棧，使采樣與訓(xùn)練在數(shù)值上逐位一致，從而獲得真正的同策略RL。

在Bigmath的RLVR設(shè)定下，研究人員做了實(shí)驗(yàn)：策略以Qwen 2.5-VL instruct 8B初始化，最大rollout長(zhǎng)度 4096。

無異策略校正（不做重要性加權(quán)）時(shí)，訓(xùn)練中段獎(jiǎng)勵(lì)出現(xiàn)崩塌；

加入異策略校正（importance weighting）后，訓(xùn)練可平穩(wěn)推進(jìn)；

當(dāng)讓采樣器與訓(xùn)練器逐位一致時(shí)，策略完全同源（KL=0），訓(xùn)練同樣平穩(wěn)。

同時(shí)，研究人員繪制了采樣器與訓(xùn)練器對(duì)數(shù)概率（logprobs）之間的 KL 散度：三種設(shè)定差異明顯——

• 帶重要性加權(quán)(下圖中的紅線)時(shí)，KL約在 0.001附近，偶有尖峰；
• 不加權(quán)(下圖中的綠線)時(shí)，KL 在獎(jiǎng)勵(lì)崩潰前后出現(xiàn)明顯飆升；
• 真正同策略(下圖中的藍(lán)線)時(shí)，KL始終為0，表明訓(xùn)練策略與采樣策略無分歧。

需要強(qiáng)調(diào)的是：未做重要性加權(quán)的那次運(yùn)行在Step 318左右出現(xiàn)顯著的損失峰值，同時(shí)KL散度同步陡升；而做了異策略校正或?qū)崿F(xiàn)「真正同策略」的兩種設(shè)置，RL都能持續(xù)、平滑地優(yōu)化。

圖片

萬事怕認(rèn)真

現(xiàn)代軟件系統(tǒng)層層抽象，機(jī)器學(xué)習(xí)中的非確定性與微小數(shù)值差異，往往讓人想「睜一只眼閉一只眼」：

反正系統(tǒng)本就「概率化」，多一點(diǎn)不確定也無妨？

單元測(cè)試?yán)锇補(bǔ)tol/rtol往上調(diào)一調(diào)、把訓(xùn)練與采樣間的logprob差異當(dāng)成「假陽性」，似乎也能過關(guān)。

請(qǐng)拒絕這種「算了吧」的心態(tài)。只要多做一點(diǎn)功課，我們完全可以定位并修復(fù)這些非確定性根源！

Thinking  Machines希望本文能為社區(qū)提供一套解決推理端非確定性的清晰思路，也能激勵(lì)更多人真正吃透自己的系統(tǒng)。

參考資料：

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

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

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