Just-in-time compilationJAX JIT 編譯原理與實戰(zhàn)指南【免費下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項目地址: https://gitcode.com/GitHub_Trending/ja/jax導(dǎo)讀jax.jit是 JAX 中最核心的變換之一它把一段 Python 函數(shù)在追蹤tracing階段規(guī)約成中間表示 jaxpr再交給 XLA 編譯器做即時Just-In-Time編譯最終生成針對 CPU/GPU/TPU 高度優(yōu)化的可執(zhí)行代碼。本篇文章圍繞 docs/jit-compilation.md 展開先講透 JAX 變換與 jaxpr 的原理再手把手演示 SELU 算子從逐 op 執(zhí)行到 JIT 加速的完整過程并深入剖析「為什么不能無腦 JIT 一切」、靜態(tài)參數(shù)static_argnums/static_argnames的使用時機以及 JIT 緩存與重編譯的行為邊界。讀完你將掌握 JAX JIT 的正確打開方式并能用jax.make_jaxpr觀察函數(shù)內(nèi)部到底發(fā)生了什么。JAX 變換是如何工作的JAX 允許我們變換 Python 函數(shù)其秘密在于JAX 會把每個函數(shù)規(guī)約成一系列 primitive原語操作每個 primitive 代表一個最基本的計算單元。而jax.make_jaxpr就是觀察這一過程的窗口——它返回函數(shù)的 jaxprJAX 的中間表示讓我們直觀看到追蹤結(jié)果。import jax import jax.numpy as jnp global_list [] def log2(x): global_list.append(x) ln_x jnp.log(x) ln_2 jnp.log(2.0) return ln_x / ln_2 print(jax.make_jaxpr(log2)(3.0))輸出會是一段類似下面這樣的 jaxpr{ lambda ; a:f32[]. let b:f32[] log a c:f32[] log 1.0 d:f32[] div b c in (d,) }關(guān)于 jaxpr 各字段lambda、類型、let 綁定、in 返回值的語義可以參考 docs/601/jaxpr.md 的詳細說明。關(guān)鍵點副作用不會進入 jaxpr注意上面log2函數(shù)體里有g(shù)lobal_list.append(x)但 jaxpr 里完全沒有與它對應(yīng)的內(nèi)容。這并非 bug而是特性JAX 變換被設(shè)計為只理解**無副作用功能純**的代碼。如果對「純函數(shù)」「副作用」這些術(shù)語不熟悉可以在倉庫的 docs/notebooks/Common_Gotchas_in_JAX.md 的 Pure Functions 一節(jié)找到通俗解釋。不純函數(shù)在 JAX 變換下是危險的它們可能靜默失敗或者產(chǎn)生像Tracer 泄漏tracer leak那樣令人困惑的下游錯誤而且 JAX 往往無法自動檢測到副作用的存在。文檔給出的官方替代方案是想要調(diào)試打印用jax.debug.print其實現(xiàn)位于 jax/_src/debug.py想表達通用副作用、接受性能損失用jax.experimental.io_callback實現(xiàn)在 jax/_src/callback.py 的io_callback想檢查 tracer 泄漏、接受性能損失用jax.check_tracer_leaks。追蹤tracing的微觀機制追蹤時JAX 會用tracer 對象包裝每個參數(shù)。這些 tracer 會記錄函數(shù)調(diào)用期間發(fā)生在普通 Python 層面對它們執(zhí)行的所有 JAX 操作然后 JAX 用這些記錄重建整個函數(shù)重建的產(chǎn)物就是 jaxpr。由于 tracer 不記錄 Python 副作用副作用自然不會出現(xiàn)在 jaxpr 中——但注意副作用在追蹤過程中仍然真實發(fā)生了。一個典型例子是 Python 的print()def log2_with_print(x): print(printed x:, x) ln_x jnp.log(x) ln_2 jnp.log(2.0) return ln_x / ln_2 print(jax.make_jaxpr(log2_with_print)(3.))你會發(fā)現(xiàn)打印出來的x是一個Traced對象——這正是 JAX 內(nèi)部機制在起作用?！窹ython 代碼至少會被執(zhí)行一次」嚴格來說是實現(xiàn)細節(jié)不應(yīng)依賴它但在調(diào)試時可以利用這一點打印中間值。jaxpr 只反映「參數(shù)給定」的那一次執(zhí)行jaxpr 捕獲的是函數(shù)在給定參數(shù)上的那次執(zhí)行路徑。如果函數(shù)里有 Python 條件分支jaxpr 只會包含實際走到的那個分支def log2_if_rank_2(x): if x.ndim 2: ln_x jnp.log(x) ln_2 jnp.log(2.0) return ln_x / ln_2 else: return x print(jax.make_jaxpr(log2_if_rank_2)(jax.numpy.array([1, 2, 3])))傳入的是 1 維數(shù)組x.ndim 2為假所以 jaxpr 只會包含else分支的return x。這暗示了一個重要的推論jaxpr 的形狀/類型是由追蹤時的輸入決定的一旦輸入 shape 或 dtype 改變JAX 就不得不重新追蹤和重新編譯。從源碼看jax.make_jaxpr的實現(xiàn)位于 jax/_src/api.py 的make_jaxpr函數(shù)約 L2136 起。它本質(zhì)上是對jit(fun, static_argnumsstatic_argnums).trace(...)的一次封裝先做一次純追蹤再把常量consts重新合并回 jaxpr 返回給調(diào)用者。其 docstring 也明確說明jaxpr 是「基于帶 let 綁定的簡單類型化一階 lambda 演算」的追蹤中間表示make_jaxpr返回的是抽象到ShapedArray級別的追蹤結(jié)果。JIT 編譯一個函數(shù)JAX 允許同一份代碼跑在 CPU/GPU/TPU 上但默認逐 op 把操作發(fā)給加速器這限制了 XLA 編譯器的優(yōu)化空間。下面以深度學(xué)習(xí)常用的SELUScaled Exponential Linear Unit算子為例import jax import jax.numpy as jnp def selu(x, alpha1.67, lambda_1.05): return lambda_ * jnp.where(x 0, x, alpha * jnp.exp(x) - alpha) x jnp.arange(1000000) %timeit selu(x).block_until_ready()這段代碼的問題在于每次只把單個 op 發(fā)給加速器XLA 無法看到全局做融合優(yōu)化。我們的目標自然是把盡可能多的代碼交給 XLA 編譯器讓它做整體優(yōu)化。jax.jit就是為此設(shè)計的selu_jit jax.jit(selu) # 先預(yù)編譯一次再計時 selu_jit(x).block_until_ready() %timeit selu_jit(x).block_until_ready()剛才發(fā)生了什么selu_jit jax.jit(selu)得到selu的編譯版本一個被包裝的函數(shù)。調(diào)用一次selu_jit(x)JAX 在這里做追蹤——它必須有真實輸入才能用 tracer 包裝。得到的 jaxpr 再由 XLA 編譯成針對 GPU/TPU 優(yōu)化過的代碼隨后立即執(zhí)行以滿足這次調(diào)用。之后的每次調(diào)用都直接復(fù)用已編譯代碼完全跳過 Python 實現(xiàn)。如果不單獨做 warm-up基準測試會把編譯時間也計進去——雖然因為循環(huán)很多次整體仍會更快但那就不是公平對比了。計時對編譯版本測速。注意這里用了block_until_ready()這是因為 JAX 采用異步派發(fā)async dispatch需要顯式等待結(jié)果就緒再計時。關(guān)于異步派發(fā)機制可以參閱 docs/async_dispatch.rstblock_until_ready的定義見 jax/_src/api.py 與 jax/_src/array.py。值得說明的是jax.jit在 jax/_src/api.py 中的簽名還支持更多參數(shù)完整的默認值如下節(jié)選in_shardings/out_shardings輸入輸出的分片sharding規(guī)格配合jax.sharding使用static_argnums/static_argnames標記靜態(tài)編譯期常量參數(shù)donate_argnums/donate_argnames標記可捐贈的緩沖區(qū)幫助 XLA 復(fù)用輸入內(nèi)存、降低峰值內(nèi)存keep_unused默認False未被函數(shù)使用的參數(shù)會從編譯產(chǎn)物中剔除、不傳上設(shè)備device/backend顯式指定運行設(shè)備或后端cpu/gpu/tpuinline嵌套 jitted 函數(shù)的內(nèi)聯(lián)策略默認False即jax.Inline.AUTOcompiler_options傳遞給 XLA 編譯器的選項字典。其中donate_argnums是文檔未展開但實踐中很重要的能力捐贈后不能再復(fù)用這些緩沖區(qū)JAX 會在你嘗試復(fù)用時報錯。更多細節(jié)可以參考倉庫中的 docs/buffer_donation.md。為什么不能無腦 JIT 一切看完上面的加速效果你可能會想干脆給所有函數(shù)都套上jax.jit得了。先來看兩個 JIT 會失敗的例子# 對 x 的「值」做條件判斷 def f(x): if x 0: return x else: return 2 * x jax.jit(f)(10) # 報錯 # 循環(huán)條件依賴 x 和 n 的值 def g(x, n): i 0 while i n: i 1 return x i jax.jit(g)(10, 20) # 報錯根因用運行時值控制追蹤期流程兩個例子的共同問題是試圖用運行時runtime值來控制追蹤期trace-time的程序流程。在 JIT 內(nèi)部被追蹤的值如這里的x、n只能通過它們的靜態(tài)屬性——例如shape或dtype——來影響控制流而不能通過它們的具體數(shù)值。if x 0這樣的判斷發(fā)生在追蹤期此時x是 tracer對它做布爾比較會直接拋出 tracer 錯誤。關(guān)于 Python 控制流與 JAX 的交互細節(jié)請參閱 docs/control-flow.md。解法一改寫代碼或使用 lax 控制流應(yīng)對這個問題一種方式是改寫代碼、避免對值做條件判斷另一種是使用 docs/201/control-flow.md 中介紹的特殊控制流原語比如jax.lax.cond其實現(xiàn)位于 jax/_src/lax/control_flow/conditionals.py。解法二只 JIT 編譯函數(shù)的一部分有時候改寫不現(xiàn)實那就考慮只 JIT 函數(shù)中計算最昂貴的部分。比如循環(huán)體是整個函數(shù)的熱點就只 JIT 循環(huán)體但要小心下一節(jié)講的緩存問題避免適得其反# 循環(huán)條件依賴 x 和 n但循環(huán)體是 JIT 的 jax.jit def loop_body(prev_i): return prev_i 1 def g_inner_jitted(x, n): i 0 while i n: i loop_body(i) return x i g_inner_jitted(10, 20)外層while i n仍是普通 Python 循環(huán)可以按值判斷內(nèi)層loop_body被 JIT 編譯熱點計算獲得了加速。這是「部分 JIT」的典型模式。把參數(shù)標記為 static靜態(tài)參數(shù)如果確實需要 JIT 一個「對輸入值做條件判斷」的函數(shù)可以告訴 JAX對某個輸入使用抽象程度更低更具體的 tracer。方法是指定static_argnums按位置索引或static_argnames按參數(shù)名。代價是顯著的靜態(tài)參數(shù)的每個不同取值都會產(chǎn)生不同的 jaxpr 和編譯產(chǎn)物JAX 不得不為每個新值重新編譯。所以只有在該函數(shù)只會遇到有限的靜態(tài)取值集合時這才是個好策略。f_jit_correct jax.jit(f, static_argnums0) print(f_jit_correct(10))g_jit_correct jax.jit(g, static_argnames[n]) print(g_jit_correct(10, 20))以裝飾器形式使用時用裝飾器工廠模式j(luò)ax.jit(static_argnames[n]) def g_jit_decorated(x, n): i 0 while i n: i 1 return x i print(g_jit_decorated(10, 20))源碼層面的約定與限制結(jié)合 jax/_src/api.py 中jit的 docstring靜態(tài)參數(shù)還有一些容易踩坑的約定靜態(tài)參數(shù)必須是可哈希的實現(xiàn)__hash__和__eq__且不可變因為它們的值會參與編譯緩存鍵compilation cache key的計算。文檔特別強調(diào)非數(shù)組類型或數(shù)組容器之外的參數(shù)必須標記為 static否則無法被正確追蹤。如果只給了static_argnums而沒給static_argnames或反之JAX 會用inspect.signature(fun)自動推斷對應(yīng)的參數(shù)名/位置如果兩者都給了則只把顯式列出的參數(shù)當(dāng)作靜態(tài)不再推斷。從 JAX v0.8.1 起jit支持省略函數(shù)參數(shù)的裝飾器工廠寫法即jax.jit(static_argnames[n])而非partial(jax.jit, ...)上面的示例正是官方推薦的現(xiàn)代寫法舊版本則常用functools.partial實現(xiàn)同樣效果。JAX 對fun持有弱引用作為緩存鍵的一部分因此fun必須可被弱引用weakly-referenceable。JIT 與緩存第一次 JIT 調(diào)用有編譯開銷所以理解jax.jit何時、如何緩存編譯結(jié)果是用好它的關(guān)鍵。緩存的基本規(guī)則假設(shè)f jax.jit(g)首次調(diào)用f時完成追蹤 編譯XLA 代碼被緩存后續(xù)調(diào)用f直接復(fù)用緩存代碼不再重復(fù)編譯——這就是jax.jit攤平編譯前期成本的方式。如果指定了static_argnums那么只有靜態(tài)參數(shù)取值與緩存一致時才復(fù)用任何一個靜態(tài)值變化都會觸發(fā)重編譯。如果靜態(tài)參數(shù)取值范圍很大程序花在編譯上的時間可能比逐 op 執(zhí)行還多——這是常見的性能陷阱。不要在循環(huán)里對臨時函數(shù)調(diào)用 jit文檔明確警告避免在循環(huán)或其他 Python 作用域內(nèi)對臨時函數(shù)調(diào)用jax.jit。原因在于緩存依賴函數(shù)的哈希當(dāng)?shù)葍r函數(shù)被反復(fù)重新定義哈希不同時緩存就失效了導(dǎo)致每次循環(huán)迭代都重新編譯from functools import partial def unjitted_loop_body(prev_i): return prev_i 1 def g_inner_jitted_partial(x, n): i 0 while i n: # 別這么做每次 partial 返回的函數(shù)哈希都不同 i jax.jit(partial(unjitted_loop_body))(i) return x i def g_inner_jitted_lambda(x, n): i 0 while i n: # 別這么做lambda 每次也返回哈希不同的函數(shù) i jax.jit(lambda x: unjitted_loop_body(x))(i) return x i def g_inner_jitted_normal(x, n): i 0 while i n: # 這樣沒問題JAX 能找到緩存的編譯函數(shù) i jax.jit(unjitted_loop_body)(i) return x i print(jit called in a loop with partials:) %timeit g_inner_jitted_partial(10, 20).block_until_ready() print(jit called in a loop with lambdas:) %timeit g_inner_jitted_lambda(10, 20).block_until_ready() print(jit called in a loop with caching:) %timeit g_inner_jitted_normal(10, 20).block_until_ready()結(jié)論很直觀partial和lambda每次都會產(chǎn)生新的函數(shù)對象、新的哈希緩存形同虛設(shè)而直接傳入同一個模塊級函數(shù)時JAX 能穩(wěn)定命中緩存。緩存機制的底層佐證從源碼角度看JAX 在派發(fā)層大量使用帶緩存的裝飾器來復(fù)用編譯產(chǎn)物例如 jax/_src/dispatch.py 中的xla_primitive_callableL98 附近使用util.cache()緩存 primitive 的 callable并有util.test_event(xla_primitive_callable_cache_miss)這樣的測試探針標記緩存未命中事件同一文件還有多處util.weakref_lru_cache/util.cache(max_size2048, ...)用于緩存各類派發(fā)中間結(jié)果。這正是「同一函數(shù)對象反復(fù)jit能命中緩存、哈希不同的臨時函數(shù)會反復(fù)編譯」這一文檔結(jié)論在實現(xiàn)層面的體現(xiàn)。小結(jié)JIT 的正確打開方式盡量 JIT 大片代碼把盡可能多的計算交給 XLA讓它做算子融合與設(shè)備級優(yōu)化用block_until_ready()配合%timeit得到不含異步派發(fā)誤差的公平基準。保持函數(shù)純追蹤期只記錄 primitive 操作副作用既不進 jaxpr 也會帶來 tracer 泄漏等隱患調(diào)試打印用jax.debug.print通用副作用用jax.experimental.io_callback?？刂屏靼础胳o態(tài)屬性」而非「值」走需要按值分支/循環(huán)時改寫代碼、改用jax.lax.cond等原語或只 JIT 熱點內(nèi)層部分。靜態(tài)參數(shù)要克制static_argnums/static_argnames只適用于取值集合有限、可哈希的場景否則會陷入重編譯泥潭。把jit用在穩(wěn)定、可緩存的函數(shù)對象上避免在循環(huán)內(nèi)用partial/lambda現(xiàn)造函數(shù)再 JIT。如果需要更系統(tǒng)地學(xué)習(xí)可以繼續(xù)閱讀倉庫中 docs/201/jit.md 關(guān)于jit進階語義、docs/601/jaxpr.md 關(guān)于 jaxpr 語言以及 docs/async_dispatch.rst 關(guān)于異步派發(fā)的說明。【免費下載鏈接】jaxComposable transformations of PythonNumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more項目地址: https://gitcode.com/GitHub_Trending/ja/jax創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考