態(tài)形狀（Dynamic Shapes）核心概念詳解：SymInt、Guards 與編譯期符號(hào)推理)
PyTorch 動(dòng)態(tài)形狀Dynamic Shapes核心概念詳解SymInt、Guards 與編譯期符號(hào)推理【免費(fèi)下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文基于 PyTorch 官方文檔 dynamic_shapes_core_concepts.md 整理并擴(kuò)充面向torch.compile編譯器棧的工程者與進(jìn)階用戶系統(tǒng)講解動(dòng)態(tài)形狀的兩大核心原語(yǔ)——符號(hào)整數(shù)SymInt與守衛(wèi)Guards——以及 Runtime Asserts、Hint 值、動(dòng)態(tài)行為診斷與整體架構(gòu)。讀完后你將能夠理解一次torch.compile編譯中符號(hào)形狀是如何從分配符號(hào) → 算子傳播 → 守衛(wèi)約束 → 簡(jiǎn)化與安裝這條鏈路流動(dòng)的并能結(jié)合倉(cāng)庫(kù)源碼定位mark_dynamic、ShapeEnv、SymNode等關(guān)鍵實(shí)現(xiàn)。1. 符號(hào)整數(shù)SymInt用代數(shù)表達(dá)未定的尺寸符號(hào)整數(shù)Symbolic Integers簡(jiǎn)稱 Symints用于表示可以跨取一個(gè)范圍的變量是動(dòng)態(tài)形狀體系的基礎(chǔ)數(shù)據(jù)單元。文檔給出的經(jīng)典例子是x torch.randn(5, 5) # shape: [5, 5] torch._dynamo.decorators.mark_dynamic(x, 0) x torch.randn(5, 5) # shape: [s0, 5]第 0 維被符號(hào)化 y torch.cat([x, x], dim0) # shape: [2*s0, 5]形狀以 Sympy 表達(dá)式呈現(xiàn)可以看到cat之后張量的形狀不再是具體整數(shù)而是一個(gè)符號(hào)表達(dá)式2*s0——形狀計(jì)算在編譯期以符號(hào)形式完成而不是等到運(yùn)行時(shí)才確定。文檔同時(shí)強(qiáng)調(diào)了一個(gè)容易被忽略的邊界z x * y會(huì)直接報(bào)錯(cuò)因?yàn)槲覀冎乐瘘c(diǎn)乘法要求兩個(gè)張量形狀相同而編譯期可以靜態(tài)證明s0 ! 2 * s0。細(xì)心的讀者會(huì)指出當(dāng)s0 0時(shí)兩者其實(shí)相等。文檔給出的解釋是——這種特例之所以被安全忽略源于 PyTorch 的zero-one specialization0/1 特化機(jī)制詳見(jiàn)姊妹篇 動(dòng)態(tài)形狀的 0/1 特化。1.1 mark_dynamic 的源碼行為mark_dynamic定義在 torch/_dynamo/decorators.py簽名為def mark_dynamic( t: Any, index: int | list[Any] | tuple[Any], *, hint_override: int | None None, min: int | None None, max: int | None None, specialize_on: list[Any] | None None, ) - None:從源碼注釋和實(shí)現(xiàn)可以確認(rèn)文檔中如果已知min/max可以指定的說(shuō)法min/max會(huì)被寫(xiě)入張量的_dynamo_unbacked_bounds屬性直接收窄該符號(hào)的取值范圍見(jiàn) decorators.py 中 bounds 的寫(xiě)入邏輯index可以是單個(gè)維度、維度列表或元組支持一次標(biāo)記多個(gè)維該 API 被標(biāo)注forbid_in_graph——即所有mark_dynamic調(diào)用必須在torch.compile之前完成嘗試在可 trace 的函數(shù)內(nèi)調(diào)用會(huì)顯式拋錯(cuò)行為受torch._dynamo.config.dynamic_shapes控制配置為False時(shí)與mark_dynamic并用會(huì)拋異常docstring 注明該支持將最終實(shí)現(xiàn)進(jìn)階參數(shù)hint_override會(huì)替換首次示例輸入推導(dǎo)出的 size hint且會(huì)影響 Inductor 代碼生成決策autotuning、歸約策略因此會(huì)進(jìn)入FxGraphCache的緩存鍵通過(guò)ShapeEnv.var_to_hint_override進(jìn)階參數(shù)specialize_on支持一次泛型 trace 多份特化編譯例如specialize_on[lambda x: x 8, lambda x: x 16]會(huì)生成一份 Dynamo trace 和兩次后端編譯運(yùn)行時(shí)輸入命中條件即派發(fā)到特化版本。此外mark_dynamic與弱標(biāo)記maybe_mark_dynamic存在優(yōu)先級(jí)規(guī)則對(duì)同一維度后者的標(biāo)記會(huì)被mark_dynamic接管并強(qiáng)制執(zhí)行更強(qiáng)的動(dòng)態(tài)語(yǔ)義。2. Guards保證已編譯代碼圖仍然有效在torch.compile中g(shù)uard守衛(wèi)是確保編譯代碼圖依然有效的機(jī)制。把某個(gè)變量動(dòng)態(tài)化后其默認(rèn)取值范圍是[-inf, inf]此時(shí)任何輸入尺寸都復(fù)用同一份編譯產(chǎn)物但一旦函數(shù)內(nèi)部出現(xiàn)了依賴具體數(shù)值的分支就必須把分支條件記錄成 guard。文檔中的例子def foo(x): if x 5: return x / 2 return x / 3第一次調(diào)用foo(6)走x / 2分支編譯器隨之添加 guardx 5之后再調(diào)用foo(4)時(shí)該 guard 求值為假觸發(fā)重新編譯。這就是守衛(wèi)破壞 → recompile的基本模型。2.1 守衛(wèi)在架構(gòu)中的位置文檔Overall Architecture一節(jié)給出的六步流程中g(shù)uard 出現(xiàn)在第 4、5、6 步Dynamo 編譯一個(gè) frame 時(shí)分配一個(gè)ShapeEnv掛在FakeTensorMode上來(lái)跟蹤符號(hào)形狀根據(jù)策略決定在入口處為張量分配符號(hào)尺寸符號(hào)尺寸隨算子傳播同時(shí)維護(hù)兩套表示用于符號(hào)計(jì)算導(dǎo)出的FX IR與用于推理的Sympy 表達(dá)式Dynamo trace 期間或 Inductor 優(yōu)化期間產(chǎn)生的條件會(huì)轉(zhuǎn)成 guard來(lái)源既有 Python 也有 Cguard 可以反過(guò)來(lái)簡(jiǎn)化符號(hào)變量——例如斷言s0 4之后所有s0的出現(xiàn)都可以被4替換trace 與優(yōu)化完成后所有 guard 隨編譯產(chǎn)物一起安裝只有當(dāng)全部 guard 求值為真時(shí)才允許復(fù)用。值得注意的是第 5 步guard 不只是失效檢測(cè)器它同時(shí)是符號(hào)簡(jiǎn)化的信息源這解釋了為什么守衛(wèi)系統(tǒng)必須保持精確——一條錯(cuò)誤安裝的守衛(wèi)會(huì)讓后續(xù)推理建立在錯(cuò)誤的前提上。guard 的實(shí)際求值與安裝代碼集中在 torch/_dynamo/guards.py而符號(hào)層面的守衛(wèi)記錄則落在ShapeEnv見(jiàn) torch/fx/experimental/symbolic_shapes.py。3. Runtime Asserts向編譯器提供已知事實(shí)當(dāng)你明確知道某些事實(shí)例如 batch size 一定小于 100時(shí)可以用 runtime assert 主動(dòng)告知編譯器def foo(batch_size): torch._check(batch_size 100) if batch_size 100: return do_something return do_something_else()torch._check引入的斷言會(huì)讓系統(tǒng)做三件事文檔Value Ranges and Constraints一節(jié)的描述嘗試用等價(jià)表達(dá)式替換未回溯unbacked符號(hào)依據(jù)斷言收窄符號(hào)的取值范圍value range refinement記住那些恒為真的布爾表達(dá)式供后續(xù)守衛(wèi)與簡(jiǎn)化使用。這與 guard 的區(qū)別在于guard 是從條件分支被動(dòng)誘發(fā)出來(lái)的約束而 runtime assert 是用戶主動(dòng)注入的先驗(yàn)知識(shí)能在沒(méi)有 Python 分支的情況下同樣驅(qū)動(dòng)符號(hào)推理。4. Hint 值編譯期的具體樣例Hint value 指編譯過(guò)程中實(shí)際已知的具體數(shù)值。符號(hào)形狀雖然以 Sympy 表達(dá)式參與計(jì)算但 JIT 編譯器在做表達(dá)式?jīng)Q策時(shí)例如選擇歸約策略、判斷循環(huán)上界的可行性常常需要一個(gè)具體數(shù)字——hint 就提供了這個(gè)編譯時(shí)樣例值使得不同維度的調(diào)用可以復(fù)用同一份編譯產(chǎn)物而無(wú)需反復(fù)重編譯。這與mark_dynamic的hint_override參數(shù)直接對(duì)應(yīng)不傳時(shí) hint 默認(rèn)取首次示例輸入在該維度上的實(shí)際值顯式傳入hint_override則替換掉這個(gè)默認(rèn) hint并因此進(jìn)入 FxGraphCache 的緩存鍵因?yàn)?hint 會(huì)影響 Inductor 的 autotuning 與歸約策略選擇。5. 動(dòng)態(tài)行為的整體工作方式Dynamic Behavior Overview文檔對(duì)動(dòng)態(tài)形狀到底何時(shí)發(fā)生、如何診斷、如何控制給出了完整的操作層面總結(jié)這里逐條展開(kāi)5.1 默認(rèn)靜態(tài)、尺寸變化才動(dòng)態(tài)PyTorch默認(rèn)假設(shè)靜態(tài)形狀。當(dāng)檢測(cè)到尺寸變化時(shí)Dynamo 會(huì)嘗試以動(dòng)態(tài)輸入重新編譯但若存在條件分支或缺少對(duì)動(dòng)態(tài)形狀的支持這次重編譯可能失敗。要診斷過(guò)度特化overspecialization可以設(shè)置TORCH_LOGSdynamic觀察日志中的 eval 條目——它們指示守衛(wèi)是何時(shí)、因?yàn)槭裁幢惶砑拥?。日志格式本身的說(shuō)明見(jiàn) 動(dòng)態(tài)形狀調(diào)試tlparse 與 TORCH_LOGS。5.2 提前標(biāo)記 vs 兩種 dynamic 開(kāi)關(guān)預(yù)期某維會(huì)是動(dòng)態(tài)時(shí)用torch._dynamo.mark_dynamic(tensor, dim)提前標(biāo)記已知上下限時(shí)同時(shí)給出min/maxtorch.compile(dynamicFalse)關(guān)閉自動(dòng)動(dòng)態(tài)形狀每個(gè)新尺寸都會(huì)觸發(fā)一次重編譯——簡(jiǎn)單、可預(yù)測(cè)但編譯次數(shù)隨尺寸增長(zhǎng)torch.compile(dynamicTrue)盡可能多地使用動(dòng)態(tài)形狀最適合小型模型文檔明確提醒它對(duì)大型模型未必合適可能帶來(lái)崩潰或性能問(wèn)題。5.3 按來(lái)源白名單dynamic_sources 與 static_sources對(duì)含圖間斷點(diǎn)graph breaks的大模型有時(shí)很難找到該動(dòng)態(tài)標(biāo)記哪些輸入。此時(shí)可以按來(lái)源source名做白名單且由于 source 名在圖間斷前后保持穩(wěn)定動(dòng)態(tài)性可以跨斷點(diǎn)保持。文檔提到的兩個(gè)變量與配置對(duì)在 torch/compiler/config.py 中定義動(dòng)態(tài)白名單環(huán)境變量TORCH_COMPILE_DYNAMIC_SOURCES或torch.compiler.config.dynamic_sources。取值為逗號(hào)分隔的 source 名列表例如L[x], L[y]也支持正則例如L\[x.*\], L\[y.*\]它甚至可以把普通整數(shù)標(biāo)記為動(dòng)態(tài)。這個(gè)白名單優(yōu)先級(jí)高于dynamicFalse、force_nn_module_property_static_shapes與force_parameter_static_shapes等其它開(kāi)關(guān)。靜態(tài)鏡像TORCH_COMPILE_STATIC_SOURCES/torch.compiler.config.static_sources把列出的來(lái)源釘在靜態(tài)上接受同樣的 source 名、正則與:N逐維語(yǔ)法且優(yōu)先于自動(dòng)動(dòng)態(tài)形狀、PGO 與dynamicTrue。當(dāng) PGO 或dynamicTrue誤判了某來(lái)源、而你希望它保持靜態(tài)時(shí)這是對(duì)應(yīng)的反向閥門(mén)。5.4 eager_then_compile stance讓框架替你推導(dǎo)動(dòng)態(tài)性文檔指出的另一個(gè)務(wù)實(shí)選項(xiàng)是eager_then_compilestance——愿意接受首個(gè) batch 的性能代價(jià)換取框架自動(dòng)推導(dǎo)出哪些輸入該動(dòng)態(tài)。入口是torch.compiler.set_stance實(shí)現(xiàn)見(jiàn) torch/compiler/init.py。從 docstring 可確認(rèn)可選 stance 包括stance行為default正常編譯force_eager忽略所有torch.compile指令eager_on_recompile需要重編譯時(shí)退化為 eager命中緩存的編譯產(chǎn)物仍照常使用fail_on_recompile觸發(fā)重編譯即拋錯(cuò)eager_then_compile首次調(diào)用 eager、后續(xù)調(diào)用編譯從前兩次調(diào)用的差異中推斷動(dòng)態(tài)性避免第一次調(diào)用浪費(fèi)在靜態(tài)編譯上aot_eager_then_compile首次調(diào)用走 AOT eager可獲得 activation checkpointing 的內(nèi)存收益后續(xù)編譯set_stance可作函數(shù)、上下文管理器或裝飾器使用但不能在torch.compile區(qū)域內(nèi)調(diào)用。6. 整體架構(gòu)符號(hào)形狀的五步工作流文檔Overall Architecture一節(jié)給出了符號(hào)形狀的端到端工作流這是理解整套機(jī)制的骨架ShapeEnv 分配Dynamo 編譯一個(gè) frame 時(shí)分配ShapeEnv掛在FakeTensorMode上負(fù)責(zé)跟蹤符號(hào)形狀入口符號(hào)化根據(jù)策略決定在入口處為張量分配哪些符號(hào)尺寸符號(hào)傳播符號(hào)尺寸穿過(guò)算子傳播同時(shí)維護(hù) FX IR供符號(hào)計(jì)算導(dǎo)出與 Sympy 表達(dá)式供推理兩套表示守衛(wèi)歸納Dynamo trace 或 Inductor 優(yōu)化期間產(chǎn)生的條件轉(zhuǎn)成 guard來(lái)源橫跨 Python 與 C守衛(wèi)簡(jiǎn)化符號(hào)斷言s0 4之后所有s0都可以替換為4守衛(wèi)安裝trace 與優(yōu)化結(jié)束后全部 guard 隨編譯代碼安裝全部為真才允許復(fù)用。配套文檔可以從兩個(gè)方向繼續(xù)深入動(dòng)態(tài)形狀入門(mén)backed/unbacked 的區(qū)別與動(dòng)態(tài)形狀故障排查。7. 內(nèi)部 API 類層次文檔給出了 Python 與 C 兩側(cè)對(duì)照的類層次這是閱讀動(dòng)態(tài)形狀源碼的導(dǎo)航圖7.1 Python 側(cè)SymInt/SymFloat/SymBool用戶可見(jiàn)類模擬int/float/bool的行為。兩個(gè)SymInt相加產(chǎn)生一個(gè)新的SymInt符號(hào)化地跟蹤這次整數(shù)加法SymNode內(nèi)部結(jié)構(gòu)可通過(guò)symint.node訪問(wèn)保存實(shí)際的符號(hào)跟蹤信息。SymNode是類型擦除的因此方便表達(dá)混合類型的運(yùn)算ShapeEnv每次編譯一份的上下文狀態(tài)跟蹤所有自由符號(hào)與迄今累積的全部守衛(wèi)。每個(gè)SymNode都記錄它所屬的ShapeEnv但反向不成立——SymNode只有在參與某個(gè) guard 時(shí)才會(huì)被保留使用。7.2 C 側(cè)c10::SymInt/SymFloat/SymBool與 Python 對(duì)應(yīng)模擬int/float/boolc10::SymNode/SymNodeImpl對(duì)應(yīng) Python 的SymNode沒(méi)有 C 版 ShapeEnv為便于調(diào)試整套符號(hào)推理設(shè)施留在 Python 一側(cè)。由此可推斷一條明確的開(kāi)發(fā)約束任何希望被make_fx等工具 trace 的代碼都必須能處理流經(jīng)它的SymInt/SymFloat/SymBool——例如不能對(duì)符號(hào)尺寸做int(x)之外的假設(shè)性斷言也不能對(duì) Sympy 表達(dá)式使用只對(duì)普通整數(shù)成立的分支。8. 取值范圍與約束Value Ranges and Constraints符號(hào)變量維護(hù)取值范圍描述其可能的取值集合。文檔給出的默認(rèn)值是尺寸類的 unbackedSymInt取值范圍[0, Inf]普通 unbackedSymInt取值范圍[-Inf, Inf]。當(dāng)斷言發(fā)生如torch._check(x y)時(shí)系統(tǒng)依次執(zhí)行嘗試用等價(jià)表達(dá)式替換 unbacked 符號(hào)依據(jù)斷言收窄取值范圍記住恒為真的布爾表達(dá)式。min/max參數(shù)mark_dynamic、torch._check斷言、guard 三者都是對(duì)取值范圍這一中心數(shù)據(jù)結(jié)構(gòu)的寫(xiě)入路徑這也是第 5 步guard 簡(jiǎn)化符號(hào)能夠成立的基礎(chǔ)范圍收窄到單點(diǎn)即可執(zhí)行常數(shù)替換。8.1 關(guān)鍵文件索引文檔末尾給出的Important files清單結(jié)合倉(cāng)庫(kù)實(shí)際路徑整理如下關(guān)注點(diǎn)文件C SymInt APIc10/core/SymInt.h同目錄SymFloat.h、SymBool.hPython SymInt APItorch/init.py查找SymInt/SymFloat/SymBoolC 銜接層plumbingc10/core/SymNodeImpl.h、torch/csrc/utils/python_symnode.h、torch/csrc/jit/python/init.cppPython 基礎(chǔ)設(shè)施ShapeEnv/SymNode 核心torch/fx/experimental/symbolic_shapes.pyShapeEnv定義于 該文件第 3957 行 附近其它重要文件torch/_subclasses/fake_tensor.py、torch/_meta_registrations.py、各算子的 decomps 與 PrimTorch refs9. 小結(jié)從文檔脈絡(luò)到源碼入口把整份核心概念文檔串起來(lái)動(dòng)態(tài)形狀的機(jī)制可以濃縮為一句ShapeEnv在每次編譯中建立符號(hào)世界算子傳播產(chǎn)生 Sympy 表達(dá)式Python/C 兩側(cè)的條件沉淀為 guardsguards 再反過(guò)來(lái)收窄符號(hào)范圍、簡(jiǎn)化表達(dá)式最終 guard 集合作為復(fù)用條件隨編譯產(chǎn)物安裝。用戶側(cè)的三個(gè)杠桿——mark_dynamic含min/max、torch._checkruntime assert、dynamic_sources/static_sources白名單——本質(zhì)上都是在向這套符號(hào)推理系統(tǒng)注入先驗(yàn)。進(jìn)一步閱讀建議按此路徑展開(kāi)先讀 動(dòng)態(tài)形狀基礎(chǔ)backed/unbacked 與 進(jìn)階用法 鞏固 API 層遇到守衛(wèi)破壞/重編譯問(wèn)題時(shí)查 Troubleshooting 與 GuardON 錯(cuò)誤想理解 0/1 特化為何讓s0 0的邊界情況無(wú)關(guān)緊要讀 0/1 特化專章動(dòng)手排查時(shí)用TORCH_LOGSdynamic配合 tlparse/TORCH_LOGS 調(diào)試指南 查看 eval 條目最后以 symbolic_shapes.py 與 torch/_dynamo/guards.py 作為源碼級(jí)入口精讀實(shí)現(xiàn)?！久赓M(fèi)下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 PyTorch 動(dòng)態(tài)形狀（Dynamic Shapes）核心概念詳解：SymInt、Guards 與編譯期符號(hào)推理
PyTorch 動(dòng)態(tài)形狀Dynamic Shapes核心概念詳解SymInt、Guards 與編譯期符號(hào)推理【免費(fèi)下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch本文基于 PyTorch 官方文檔 dynamic_shapes_core_concepts.md 整理并擴(kuò)充面向torch.compile編譯器棧的工程者與進(jìn)階用戶系統(tǒng)講解動(dòng)態(tài)形狀的兩大核心原語(yǔ)——符號(hào)整數(shù)SymInt與守衛(wèi)Guards——以及 Runtime Asserts、Hint 值、動(dòng)態(tài)行為診斷與整體架構(gòu)。讀完后你將能夠理解一次torch.compile編譯中符號(hào)形狀是如何從分配符號(hào) → 算子傳播 → 守衛(wèi)約束 → 簡(jiǎn)化與安裝這條鏈路流動(dòng)的并能結(jié)合倉(cāng)庫(kù)源碼定位mark_dynamic、ShapeEnv、SymNode等關(guān)鍵實(shí)現(xiàn)。1. 符號(hào)整數(shù)SymInt用代數(shù)表達(dá)未定的尺寸符號(hào)整數(shù)Symbolic Integers簡(jiǎn)稱 Symints用于表示可以跨取一個(gè)范圍的變量是動(dòng)態(tài)形狀體系的基礎(chǔ)數(shù)據(jù)單元。文檔給出的經(jīng)典例子是x torch.randn(5, 5) # shape: [5, 5] torch._dynamo.decorators.mark_dynamic(x, 0) x torch.randn(5, 5) # shape: [s0, 5]第 0 維被符號(hào)化 y torch.cat([x, x], dim0) # shape: [2*s0, 5]形狀以 Sympy 表達(dá)式呈現(xiàn)可以看到cat之后張量的形狀不再是具體整數(shù)而是一個(gè)符號(hào)表達(dá)式2*s0——形狀計(jì)算在編譯期以符號(hào)形式完成而不是等到運(yùn)行時(shí)才確定。文檔同時(shí)強(qiáng)調(diào)了一個(gè)容易被忽略的邊界z x * y會(huì)直接報(bào)錯(cuò)因?yàn)槲覀冎乐瘘c(diǎn)乘法要求兩個(gè)張量形狀相同而編譯期可以靜態(tài)證明s0 ! 2 * s0。細(xì)心的讀者會(huì)指出當(dāng)s0 0時(shí)兩者其實(shí)相等。文檔給出的解釋是——這種特例之所以被安全忽略源于 PyTorch 的zero-one specialization0/1 特化機(jī)制詳見(jiàn)姊妹篇 動(dòng)態(tài)形狀的 0/1 特化。1.1 mark_dynamic 的源碼行為mark_dynamic定義在 torch/_dynamo/decorators.py簽名為def mark_dynamic( t: Any, index: int | list[Any] | tuple[Any], *, hint_override: int | None None, min: int | None None, max: int | None None, specialize_on: list[Any] | None None, ) - None:從源碼注釋和實(shí)現(xiàn)可以確認(rèn)文檔中如果已知min/max可以指定的說(shuō)法min/max會(huì)被寫(xiě)入張量的_dynamo_unbacked_bounds屬性直接收窄該符號(hào)的取值范圍見(jiàn) decorators.py 中 bounds 的寫(xiě)入邏輯index可以是單個(gè)維度、維度列表或元組支持一次標(biāo)記多個(gè)維該 API 被標(biāo)注forbid_in_graph——即所有mark_dynamic調(diào)用必須在torch.compile之前完成嘗試在可 trace 的函數(shù)內(nèi)調(diào)用會(huì)顯式拋錯(cuò)行為受torch._dynamo.config.dynamic_shapes控制配置為False時(shí)與mark_dynamic并用會(huì)拋異常docstring 注明該支持將最終實(shí)現(xiàn)進(jìn)階參數(shù)hint_override會(huì)替換首次示例輸入推導(dǎo)出的 size hint且會(huì)影響 Inductor 代碼生成決策autotuning、歸約策略因此會(huì)進(jìn)入FxGraphCache的緩存鍵通過(guò)ShapeEnv.var_to_hint_override進(jìn)階參數(shù)specialize_on支持一次泛型 trace 多份特化編譯例如specialize_on[lambda x: x 8, lambda x: x 16]會(huì)生成一份 Dynamo trace 和兩次后端編譯運(yùn)行時(shí)輸入命中條件即派發(fā)到特化版本。此外mark_dynamic與弱標(biāo)記maybe_mark_dynamic存在優(yōu)先級(jí)規(guī)則對(duì)同一維度后者的標(biāo)記會(huì)被mark_dynamic接管并強(qiáng)制執(zhí)行更強(qiáng)的動(dòng)態(tài)語(yǔ)義。2. Guards保證已編譯代碼圖仍然有效在torch.compile中g(shù)uard守衛(wèi)是確保編譯代碼圖依然有效的機(jī)制。把某個(gè)變量動(dòng)態(tài)化后其默認(rèn)取值范圍是[-inf, inf]此時(shí)任何輸入尺寸都復(fù)用同一份編譯產(chǎn)物但一旦函數(shù)內(nèi)部出現(xiàn)了依賴具體數(shù)值的分支就必須把分支條件記錄成 guard。文檔中的例子def foo(x): if x 5: return x / 2 return x / 3第一次調(diào)用foo(6)走x / 2分支編譯器隨之添加 guardx 5之后再調(diào)用foo(4)時(shí)該 guard 求值為假觸發(fā)重新編譯。這就是守衛(wèi)破壞 → recompile的基本模型。2.1 守衛(wèi)在架構(gòu)中的位置文檔Overall Architecture一節(jié)給出的六步流程中g(shù)uard 出現(xiàn)在第 4、5、6 步Dynamo 編譯一個(gè) frame 時(shí)分配一個(gè)ShapeEnv掛在FakeTensorMode上來(lái)跟蹤符號(hào)形狀根據(jù)策略決定在入口處為張量分配符號(hào)尺寸符號(hào)尺寸隨算子傳播同時(shí)維護(hù)兩套表示用于符號(hào)計(jì)算導(dǎo)出的FX IR與用于推理的Sympy 表達(dá)式Dynamo trace 期間或 Inductor 優(yōu)化期間產(chǎn)生的條件會(huì)轉(zhuǎn)成 guard來(lái)源既有 Python 也有 Cguard 可以反過(guò)來(lái)簡(jiǎn)化符號(hào)變量——例如斷言s0 4之后所有s0的出現(xiàn)都可以被4替換trace 與優(yōu)化完成后所有 guard 隨編譯產(chǎn)物一起安裝只有當(dāng)全部 guard 求值為真時(shí)才允許復(fù)用。值得注意的是第 5 步guard 不只是失效檢測(cè)器它同時(shí)是符號(hào)簡(jiǎn)化的信息源這解釋了為什么守衛(wèi)系統(tǒng)必須保持精確——一條錯(cuò)誤安裝的守衛(wèi)會(huì)讓后續(xù)推理建立在錯(cuò)誤的前提上。guard 的實(shí)際求值與安裝代碼集中在 torch/_dynamo/guards.py而符號(hào)層面的守衛(wèi)記錄則落在ShapeEnv見(jiàn) torch/fx/experimental/symbolic_shapes.py。3. Runtime Asserts向編譯器提供已知事實(shí)當(dāng)你明確知道某些事實(shí)例如 batch size 一定小于 100時(shí)可以用 runtime assert 主動(dòng)告知編譯器def foo(batch_size): torch._check(batch_size 100) if batch_size 100: return do_something return do_something_else()torch._check引入的斷言會(huì)讓系統(tǒng)做三件事文檔Value Ranges and Constraints一節(jié)的描述嘗試用等價(jià)表達(dá)式替換未回溯unbacked符號(hào)依據(jù)斷言收窄符號(hào)的取值范圍value range refinement記住那些恒為真的布爾表達(dá)式供后續(xù)守衛(wèi)與簡(jiǎn)化使用。這與 guard 的區(qū)別在于guard 是從條件分支被動(dòng)誘發(fā)出來(lái)的約束而 runtime assert 是用戶主動(dòng)注入的先驗(yàn)知識(shí)能在沒(méi)有 Python 分支的情況下同樣驅(qū)動(dòng)符號(hào)推理。4. Hint 值編譯期的具體樣例Hint value 指編譯過(guò)程中實(shí)際已知的具體數(shù)值。符號(hào)形狀雖然以 Sympy 表達(dá)式參與計(jì)算但 JIT 編譯器在做表達(dá)式?jīng)Q策時(shí)例如選擇歸約策略、判斷循環(huán)上界的可行性常常需要一個(gè)具體數(shù)字——hint 就提供了這個(gè)編譯時(shí)樣例值使得不同維度的調(diào)用可以復(fù)用同一份編譯產(chǎn)物而無(wú)需反復(fù)重編譯。這與mark_dynamic的hint_override參數(shù)直接對(duì)應(yīng)不傳時(shí) hint 默認(rèn)取首次示例輸入在該維度上的實(shí)際值顯式傳入hint_override則替換掉這個(gè)默認(rèn) hint并因此進(jìn)入 FxGraphCache 的緩存鍵因?yàn)?hint 會(huì)影響 Inductor 的 autotuning 與歸約策略選擇。5. 動(dòng)態(tài)行為的整體工作方式Dynamic Behavior Overview文檔對(duì)動(dòng)態(tài)形狀到底何時(shí)發(fā)生、如何診斷、如何控制給出了完整的操作層面總結(jié)這里逐條展開(kāi)5.1 默認(rèn)靜態(tài)、尺寸變化才動(dòng)態(tài)PyTorch默認(rèn)假設(shè)靜態(tài)形狀。當(dāng)檢測(cè)到尺寸變化時(shí)Dynamo 會(huì)嘗試以動(dòng)態(tài)輸入重新編譯但若存在條件分支或缺少對(duì)動(dòng)態(tài)形狀的支持這次重編譯可能失敗。要診斷過(guò)度特化overspecialization可以設(shè)置TORCH_LOGSdynamic觀察日志中的 eval 條目——它們指示守衛(wèi)是何時(shí)、因?yàn)槭裁幢惶砑拥?。日志格式本身的說(shuō)明見(jiàn) 動(dòng)態(tài)形狀調(diào)試tlparse 與 TORCH_LOGS。5.2 提前標(biāo)記 vs 兩種 dynamic 開(kāi)關(guān)預(yù)期某維會(huì)是動(dòng)態(tài)時(shí)用torch._dynamo.mark_dynamic(tensor, dim)提前標(biāo)記已知上下限時(shí)同時(shí)給出min/maxtorch.compile(dynamicFalse)關(guān)閉自動(dòng)動(dòng)態(tài)形狀每個(gè)新尺寸都會(huì)觸發(fā)一次重編譯——簡(jiǎn)單、可預(yù)測(cè)但編譯次數(shù)隨尺寸增長(zhǎng)torch.compile(dynamicTrue)盡可能多地使用動(dòng)態(tài)形狀最適合小型模型文檔明確提醒它對(duì)大型模型未必合適可能帶來(lái)崩潰或性能問(wèn)題。5.3 按來(lái)源白名單dynamic_sources 與 static_sources對(duì)含圖間斷點(diǎn)graph breaks的大模型有時(shí)很難找到該動(dòng)態(tài)標(biāo)記哪些輸入。此時(shí)可以按來(lái)源source名做白名單且由于 source 名在圖間斷前后保持穩(wěn)定動(dòng)態(tài)性可以跨斷點(diǎn)保持。文檔提到的兩個(gè)變量與配置對(duì)在 torch/compiler/config.py 中定義動(dòng)態(tài)白名單環(huán)境變量TORCH_COMPILE_DYNAMIC_SOURCES或torch.compiler.config.dynamic_sources。取值為逗號(hào)分隔的 source 名列表例如L[x], L[y]也支持正則例如L\[x.*\], L\[y.*\]它甚至可以把普通整數(shù)標(biāo)記為動(dòng)態(tài)。這個(gè)白名單優(yōu)先級(jí)高于dynamicFalse、force_nn_module_property_static_shapes與force_parameter_static_shapes等其它開(kāi)關(guān)。靜態(tài)鏡像TORCH_COMPILE_STATIC_SOURCES/torch.compiler.config.static_sources把列出的來(lái)源釘在靜態(tài)上接受同樣的 source 名、正則與:N逐維語(yǔ)法且優(yōu)先于自動(dòng)動(dòng)態(tài)形狀、PGO 與dynamicTrue。當(dāng) PGO 或dynamicTrue誤判了某來(lái)源、而你希望它保持靜態(tài)時(shí)這是對(duì)應(yīng)的反向閥門(mén)。5.4 eager_then_compile stance讓框架替你推導(dǎo)動(dòng)態(tài)性文檔指出的另一個(gè)務(wù)實(shí)選項(xiàng)是eager_then_compilestance——愿意接受首個(gè) batch 的性能代價(jià)換取框架自動(dòng)推導(dǎo)出哪些輸入該動(dòng)態(tài)。入口是torch.compiler.set_stance實(shí)現(xiàn)見(jiàn) torch/compiler/init.py。從 docstring 可確認(rèn)可選 stance 包括stance行為default正常編譯force_eager忽略所有torch.compile指令eager_on_recompile需要重編譯時(shí)退化為 eager命中緩存的編譯產(chǎn)物仍照常使用fail_on_recompile觸發(fā)重編譯即拋錯(cuò)eager_then_compile首次調(diào)用 eager、后續(xù)調(diào)用編譯從前兩次調(diào)用的差異中推斷動(dòng)態(tài)性避免第一次調(diào)用浪費(fèi)在靜態(tài)編譯上aot_eager_then_compile首次調(diào)用走 AOT eager可獲得 activation checkpointing 的內(nèi)存收益后續(xù)編譯set_stance可作函數(shù)、上下文管理器或裝飾器使用但不能在torch.compile區(qū)域內(nèi)調(diào)用。6. 整體架構(gòu)符號(hào)形狀的五步工作流文檔Overall Architecture一節(jié)給出了符號(hào)形狀的端到端工作流這是理解整套機(jī)制的骨架ShapeEnv 分配Dynamo 編譯一個(gè) frame 時(shí)分配ShapeEnv掛在FakeTensorMode上負(fù)責(zé)跟蹤符號(hào)形狀入口符號(hào)化根據(jù)策略決定在入口處為張量分配哪些符號(hào)尺寸符號(hào)傳播符號(hào)尺寸穿過(guò)算子傳播同時(shí)維護(hù) FX IR供符號(hào)計(jì)算導(dǎo)出與 Sympy 表達(dá)式供推理兩套表示守衛(wèi)歸納Dynamo trace 或 Inductor 優(yōu)化期間產(chǎn)生的條件轉(zhuǎn)成 guard來(lái)源橫跨 Python 與 C守衛(wèi)簡(jiǎn)化符號(hào)斷言s0 4之后所有s0都可以替換為4守衛(wèi)安裝trace 與優(yōu)化結(jié)束后全部 guard 隨編譯代碼安裝全部為真才允許復(fù)用。配套文檔可以從兩個(gè)方向繼續(xù)深入動(dòng)態(tài)形狀入門(mén)backed/unbacked 的區(qū)別與動(dòng)態(tài)形狀故障排查。7. 內(nèi)部 API 類層次文檔給出了 Python 與 C 兩側(cè)對(duì)照的類層次這是閱讀動(dòng)態(tài)形狀源碼的導(dǎo)航圖7.1 Python 側(cè)SymInt/SymFloat/SymBool用戶可見(jiàn)類模擬int/float/bool的行為。兩個(gè)SymInt相加產(chǎn)生一個(gè)新的SymInt符號(hào)化地跟蹤這次整數(shù)加法SymNode內(nèi)部結(jié)構(gòu)可通過(guò)symint.node訪問(wèn)保存實(shí)際的符號(hào)跟蹤信息。SymNode是類型擦除的因此方便表達(dá)混合類型的運(yùn)算ShapeEnv每次編譯一份的上下文狀態(tài)跟蹤所有自由符號(hào)與迄今累積的全部守衛(wèi)。每個(gè)SymNode都記錄它所屬的ShapeEnv但反向不成立——SymNode只有在參與某個(gè) guard 時(shí)才會(huì)被保留使用。7.2 C 側(cè)c10::SymInt/SymFloat/SymBool與 Python 對(duì)應(yīng)模擬int/float/boolc10::SymNode/SymNodeImpl對(duì)應(yīng) Python 的SymNode沒(méi)有 C 版 ShapeEnv為便于調(diào)試整套符號(hào)推理設(shè)施留在 Python 一側(cè)。由此可推斷一條明確的開(kāi)發(fā)約束任何希望被make_fx等工具 trace 的代碼都必須能處理流經(jīng)它的SymInt/SymFloat/SymBool——例如不能對(duì)符號(hào)尺寸做int(x)之外的假設(shè)性斷言也不能對(duì) Sympy 表達(dá)式使用只對(duì)普通整數(shù)成立的分支。8. 取值范圍與約束Value Ranges and Constraints符號(hào)變量維護(hù)取值范圍描述其可能的取值集合。文檔給出的默認(rèn)值是尺寸類的 unbackedSymInt取值范圍[0, Inf]普通 unbackedSymInt取值范圍[-Inf, Inf]。當(dāng)斷言發(fā)生如torch._check(x y)時(shí)系統(tǒng)依次執(zhí)行嘗試用等價(jià)表達(dá)式替換 unbacked 符號(hào)依據(jù)斷言收窄取值范圍記住恒為真的布爾表達(dá)式。min/max參數(shù)mark_dynamic、torch._check斷言、guard 三者都是對(duì)取值范圍這一中心數(shù)據(jù)結(jié)構(gòu)的寫(xiě)入路徑這也是第 5 步guard 簡(jiǎn)化符號(hào)能夠成立的基礎(chǔ)范圍收窄到單點(diǎn)即可執(zhí)行常數(shù)替換。8.1 關(guān)鍵文件索引文檔末尾給出的Important files清單結(jié)合倉(cāng)庫(kù)實(shí)際路徑整理如下關(guān)注點(diǎn)文件C SymInt APIc10/core/SymInt.h同目錄SymFloat.h、SymBool.hPython SymInt APItorch/init.py查找SymInt/SymFloat/SymBoolC 銜接層plumbingc10/core/SymNodeImpl.h、torch/csrc/utils/python_symnode.h、torch/csrc/jit/python/init.cppPython 基礎(chǔ)設(shè)施ShapeEnv/SymNode 核心torch/fx/experimental/symbolic_shapes.pyShapeEnv定義于 該文件第 3957 行 附近其它重要文件torch/_subclasses/fake_tensor.py、torch/_meta_registrations.py、各算子的 decomps 與 PrimTorch refs9. 小結(jié)從文檔脈絡(luò)到源碼入口把整份核心概念文檔串起來(lái)動(dòng)態(tài)形狀的機(jī)制可以濃縮為一句ShapeEnv在每次編譯中建立符號(hào)世界算子傳播產(chǎn)生 Sympy 表達(dá)式Python/C 兩側(cè)的條件沉淀為 guardsguards 再反過(guò)來(lái)收窄符號(hào)范圍、簡(jiǎn)化表達(dá)式最終 guard 集合作為復(fù)用條件隨編譯產(chǎn)物安裝。用戶側(cè)的三個(gè)杠桿——mark_dynamic含min/max、torch._checkruntime assert、dynamic_sources/static_sources白名單——本質(zhì)上都是在向這套符號(hào)推理系統(tǒng)注入先驗(yàn)。進(jìn)一步閱讀建議按此路徑展開(kāi)先讀 動(dòng)態(tài)形狀基礎(chǔ)backed/unbacked 與 進(jìn)階用法 鞏固 API 層遇到守衛(wèi)破壞/重編譯問(wèn)題時(shí)查 Troubleshooting 與 GuardON 錯(cuò)誤想理解 0/1 特化為何讓s0 0的邊界情況無(wú)關(guān)緊要讀 0/1 特化專章動(dòng)手排查時(shí)用TORCH_LOGSdynamic配合 tlparse/TORCH_LOGS 調(diào)試指南 查看 eval 條目最后以 symbolic_shapes.py 與 torch/_dynamo/guards.py 作為源碼級(jí)入口精讀實(shí)現(xiàn)?！久赓M(fèi)下載鏈接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration項(xiàng)目地址: https://gitcode.com/GitHub_Trending/py/pytorch創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考