訓(xùn)練腳本剛跑起來(lái)終端就飄過(guò)一行WARNING:tensorflow: Not using XLA:CPU for cluster很多人的第一反應(yīng)是“我的分布式集群是不是壞了”“CPU 是不是不支持 XLA”。其實(shí)這是 TensorFlow 在編譯優(yōu)化階段打出來(lái)的一條典型 warning意思是你圖中有一部分計(jì)算子圖cluster本來(lái)可以在 XLA:CPU 設(shè)備上編譯執(zhí)行但最終沒(méi)有被 XLA 接管退回到了默認(rèn)的普通算子執(zhí)行路徑。這不是 error程序大概率還會(huì)正常跑。但如果你跑的是訓(xùn)練任務(wù)尤其是大規(guī)模卷積、Transformer、圖像生成這類(lèi)計(jì)算密集模型這行 warning 很值得較真一下它往往意味著你白等了。這篇內(nèi)容我會(huì)從 XLA 的 cluster 到底是什么講起給出我實(shí)際排查用的命令和配置再聊幾個(gè)“日志消失但速度反而變差”的隱藏坑希望能幫你從“消警告”走向“真提速”。1. 先拆開(kāi)這句話XLA、CPU 與 cluster 分別指什么1.1 cluster 不是分布式集群而是計(jì)算子圖第一次看到“cluster”這個(gè)詞很容易聯(lián)想到分布式訓(xùn)練里那套節(jié)點(diǎn)集群。但在這里它指的是計(jì)算圖里的一批連續(xù)算子。TensorFlow 訓(xùn)練時(shí)會(huì)把模型表達(dá)成一張有向計(jì)算圖。XLA 這個(gè)編譯器不會(huì)一個(gè)一個(gè)算子去優(yōu)化而是把圖中可以合并的算子切成一個(gè)一個(gè)小組每個(gè)小組就是一個(gè) cluster。它就像把好幾個(gè)菜打包成一份“套餐”交給同一個(gè)廚師一次做完而不是每個(gè)菜都重新開(kāi)一次火。所以Not using XLA:CPU for cluster的意思是這份“套餐”沒(méi)有交給 XLA 大廚而是回到了默認(rèn)流程按老辦法一個(gè)算子一個(gè)算子運(yùn)行。1.2 為什么訓(xùn)練過(guò)程中會(huì)冒出來(lái)這個(gè) warning 最常出現(xiàn)在兩種場(chǎng)景你打開(kāi)了 XLA 自動(dòng)聚類(lèi)autoclusteringGrappler 優(yōu)化器掃描計(jì)算圖后發(fā)現(xiàn)某個(gè) cluster 可以做 XLA 編譯但運(yùn)行時(shí)環(huán)境不滿足條件你用了tf.function(jit_compileTrue)或者tf.xla.experimental.jit_scope()這種顯式開(kāi)啟方式但某個(gè)片段沒(méi)有成功走 XLA 執(zhí)行路徑。TensorFlow 在決定“一個(gè) cluster 是否要用 XLA:CPU”時(shí)需要同時(shí)檢查很多東西XLA 設(shè)備是否注冊(cè)、cluster 里的算子是否都有對(duì)應(yīng) XLA kernel、shape 是否靜態(tài)可推斷、是否有循環(huán)或控制流依賴等。任何一個(gè)條件不滿足它就會(huì)放棄這個(gè) cluster并打出這么一行 warning。1.3 它和不報(bào)錯(cuò)是兩件事很多人會(huì)誤以為“只是 warning不用管”。確實(shí)如果你只是做個(gè)幾十步的小 demo那無(wú)所謂。但如果訓(xùn)練一個(gè)要跑十幾個(gè)小時(shí)的模型XLA 的融合優(yōu)化有時(shí)候能帶來(lái) 20%~50% 的端到端收益這時(shí)候每一行 warning 都可能是性能損失點(diǎn)。更關(guān)鍵的是這個(gè) warning 經(jīng)常把真正的問(wèn)題藏起來(lái)。真正導(dǎo)致 cluster 無(wú)法編譯的原因并不會(huì)默認(rèn)出現(xiàn)在終端而是被 TF 自己的日志級(jí)別過(guò)濾掉了。所以看到它第一件事不是找一堆環(huán)境變量亂試而是搞清楚到底是設(shè)備沒(méi)注冊(cè)、聚類(lèi)沒(méi)開(kāi)、還是某個(gè) cluster 本身不滿足編譯條件。2. 定位根因日志背后的三層開(kāi)關(guān)我在處理這類(lèi)問(wèn)題時(shí)習(xí)慣按三層來(lái)排查設(shè)備注冊(cè)、自動(dòng)聚類(lèi)、單個(gè) cluster 的可編譯性。2.1 第一層X(jué)LA 設(shè)備有沒(méi)有注冊(cè)XLA:CPU 雖然名字里帶 CPU但它并不是一個(gè)真實(shí)物理設(shè)備而是邏輯上的編譯后端。TF 里有一類(lèi)虛擬設(shè)備叫XLA_CPU、XLA_GPU用來(lái)承接 XLA 編譯后的 executable。如果沒(méi)有注冊(cè)這類(lèi)設(shè)備那訓(xùn)練時(shí)根本沒(méi)地方“運(yùn)行”編譯好的 XLA kernel自然就會(huì)出現(xiàn) warning。先跑一段代碼確認(rèn)一下import tensorflow as tf print(physical devices:, tf.config.list_physical_devices()) print(logical devices:, tf.config.list_logical_devices())正常情況下你會(huì)看到類(lèi)似輸出Physical devices: - CPU - GPU (如果有機(jī)型卡) Logical devices: - CPU:0 - GPU:0 - XLA_CPU:0 - XLA_GPU:0如果Logical devices里沒(méi)有XLA_CPU:0那基本可以斷定設(shè)備注冊(cè)這一步就卡住了。老版本 TF 2.4~2.7 里XLA 設(shè)備默認(rèn)是隱藏的需要通過(guò)環(huán)境變量把它暴露出來(lái)這個(gè)我后面會(huì)給出具體寫(xiě)法。2.2 第二層Grappler 自動(dòng)聚類(lèi)有沒(méi)有開(kāi)即使設(shè)備注冊(cè)了XLA 也不會(huì)默認(rèn)接管所有計(jì)算。TensorFlow 里有一個(gè)叫 Grappler 的優(yōu)化器它負(fù)責(zé)掃描tf.function對(duì)應(yīng)的圖并把能編譯的算子合并成 cluster。合并動(dòng)作就叫“自動(dòng)聚類(lèi)”。打開(kāi)自動(dòng)聚類(lèi)的方式很簡(jiǎn)單import tensorflow as tf tf.config.optimizer.set_jit(True) # 或 tf.config.optimizer.set_jit(autoclustering)這行代碼要放在模型構(gòu)建和訓(xùn)練步驟定義之前。你可以把它理解為對(duì) Grappler 說(shuō)“看到合適的 cluster就去嘗試 XLA 編譯?！比绻阋呀?jīng)開(kāi)了自動(dòng)聚類(lèi)設(shè)備列表里也有 XLA_CPU但還是看到 warning那問(wèn)題通常出在第三層某個(gè) cluster 本身編譯不了。2.3 第三層具體哪個(gè) cluster 被拒默認(rèn)情況下 TensorFlow 的日志級(jí)別比較高很多關(guān)鍵信息被壓掉了。我會(huì)先把日志級(jí)別調(diào)到最安靜但保留 debug 信息export TF_CPP_MIN_LOG_LEVEL0 export TF_CPP_MIN_VLOG_LEVEL1 python train.py 21 | grep -E XLA|cluster|Not usingTF_CPP_MIN_VLOG_LEVEL1會(huì)讓底層輸出大量調(diào)試日志跑起來(lái)會(huì)刷屏所以最好配合 grep 看。這時(shí)候你可能會(huì)看到更具體的描述比如cluster requires static shapescluster has unsupported opEnable XLA to get more accurate …這些才是根因。我遇到過(guò)不少情況是tf.shape使用了動(dòng)態(tài)維度或者某個(gè)自定義 op 沒(méi)有注冊(cè) XLA kernel導(dǎo)致整個(gè) cluster 被整體放棄。如果你一味地去加環(huán)境變量而不去看這一層的日志永遠(yuǎn)只是在表面打轉(zhuǎn)。3. 我的實(shí)測(cè)路徑從警告到正常執(zhí)行的四種做法下面我按從傳統(tǒng)到現(xiàn)代的路徑給出四種實(shí)際可用的配置方式。它們不是互斥的但適用的 TensorFlow 版本和場(chǎng)景不同。3.1 老版本方案用環(huán)境變量注冊(cè) XLA 設(shè)備在 TF 2.4 到 2.7 附近XLA 設(shè)備默認(rèn)是隱藏的。常見(jiàn)做法是在啟動(dòng)訓(xùn)練前加一個(gè)環(huán)境變量TF_XLA_FLAGS--tf_xla_enable_xla_devices python train.py這個(gè) flag 的作用就是讓XLA_CPU、XLA_GPU這類(lèi)邏輯設(shè)備出現(xiàn)在運(yùn)行時(shí)環(huán)境中讓后面的 XLA cluster 編譯結(jié)果有地方可放。如果你用的是 Docker記得在容器啟動(dòng)時(shí)傳進(jìn)去docker run -e TF_XLA_FLAGS--tf_xla_enable_xla_devices ...在代碼里提前設(shè)置也可以但必須在 import TensorFlow 之前import os os.environ[TF_XLA_FLAGS] --tf_xla_enable_xla_devices import tensorflow as tf注意新版 TensorFlow2.10 之后已經(jīng)默認(rèn)注冊(cè) XLA 設(shè)備這個(gè) flag 不一定再起作用。所以當(dāng)你在新版本里發(fā)現(xiàn)這個(gè)參數(shù)沒(méi)效果不要繼續(xù)糾結(jié)直接看下一種做法。3.2 現(xiàn)代主流做法打開(kāi)自動(dòng)聚類(lèi)從 TF 2.8 開(kāi)始我首選tf.config.optimizer.set_jit。它可以接受布爾值和字符串import tensorflow as tf tf.config.optimizer.set_jit(True) # 或 tf.config.optimizer.set_jit(autoclustering) # 之后正常定義模型 model tf.keras.Sequential([...]) model.compile(...)這樣設(shè)置后Grappler 會(huì)對(duì)tf.function里的連續(xù)算子做聚類(lèi)嘗試。如果某個(gè) cluster 成功你會(huì)看到執(zhí)行日志里出現(xiàn)XlaLaunch算子如果失敗控制臺(tái)會(huì)繼續(xù)出現(xiàn)Not using XLA:CPU for cluster但此時(shí)你可以參考上一節(jié)的方法進(jìn)一步定位具體 cluster 失敗原因。這種方式的優(yōu)點(diǎn)是省心適合已有訓(xùn)練腳本、不想大改代碼的情況。缺點(diǎn)是你無(wú)法精確控制哪些 cluster 編譯、哪些不編譯有時(shí)候小算子也會(huì)被打包反而帶來(lái)額外開(kāi)銷(xiāo)。3.3 顯式強(qiáng)約束jit_compileTrue如果你希望“這段計(jì)算必須走 XLA”那就不要再用隱式 warning 來(lái)猜了直接給tf.function加約束tf.function(jit_compileTrue) def train_step(x, y): with tf.GradientTape() as tape: pred model(x, trainingTrue) loss loss_fn(y, pred) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables)) return lossjit_compileTrue會(huì)強(qiáng)制 TensorFlow 把這個(gè)函數(shù)中的可編譯部分交給 XLA 編譯。如果編譯失敗它通常不會(huì)靜默退到普通路徑而是直接拋錯(cuò)或者打出更加明確的編譯錯(cuò)誤。這非常適合用來(lái)“逼問(wèn)”到底哪里不支持 XLA。還有一個(gè)更老的 APIwith tf.xla.experimental.jit_scope(): result my_compute(x, y)它適合局部控制某一段代碼但優(yōu)先級(jí)和表達(dá)能力都不如jit_compileTrue老項(xiàng)目中遇到可以認(rèn)識(shí)一下新項(xiàng)目我不太推薦。提示jit_compileTrue不等于所有代碼都能被 XLA 編譯。遇到while_loop、動(dòng)態(tài) shape、某些自定義 TF op 時(shí)仍然可能失敗。它最大的價(jià)值是讓你更快看到失敗原因而不是保證一定成功。3.4 組合拳老腳本怎么低成本接入如果是維護(hù)一個(gè)半年沒(méi)動(dòng)過(guò)的訓(xùn)練腳本我會(huì)建議這樣一步步來(lái)不要一下全改先加環(huán)境變量TF_XLA_FLAGS--tf_xla_enable_xla_devices保證設(shè)備注冊(cè)再在訓(xùn)練腳本頂部加tf.config.optimizer.set_jit(True)如果還看到 warning挑最耗時(shí)的訓(xùn)練函數(shù)臨時(shí)加jit_compileTrue根據(jù)編譯錯(cuò)誤逐個(gè)處理比如把動(dòng)態(tài) sequence 做 padding把不支持的 op 拆到j(luò)it_compile函數(shù)外面。這套組合我實(shí)測(cè)下來(lái)大部分“Not using XLA:CPU for cluster”都能在半小時(shí)內(nèi)定位到原因。日志確實(shí)會(huì)消失但真正的檢查才剛剛開(kāi)始。4. 日志消失不等于提速三個(gè)容易誤判的細(xì)節(jié)很多時(shí)候我們把 warning 搞沒(méi)了就收工結(jié)果整體訓(xùn)練時(shí)間一點(diǎn)沒(méi)少。這里有幾個(gè)非常隱蔽的坑值得單獨(dú)說(shuō)。4.1 動(dòng)態(tài) shape 導(dǎo)致反復(fù)重編譯XLA 是一個(gè)靜態(tài)編譯編譯器它在編譯時(shí)會(huì)固定輸入 tensor 的 shape。你的模型如果 batch size 固定、輸入圖片尺寸固定那編譯一次可以反復(fù)用。但如果你一邊訓(xùn)練一邊動(dòng)態(tài)改變輸入長(zhǎng)度比如 NLP 里常見(jiàn)的 variable sequence lengthXLA 每遇到一個(gè)新 shape 就可能重新編譯一次。重新編譯的時(shí)間往往比你省下的時(shí)間還要多最后你會(huì)看到 “日志沒(méi)了但每步 step 反而慢了不少”。我甚至見(jiàn)過(guò) CPU 上因?yàn)轭l繁 recompile訓(xùn)練時(shí)間翻倍的案例。解決辦法有兩個(gè)方向在數(shù)據(jù) pipeline 里對(duì)序列做統(tǒng)一 padding把 shape 固定下來(lái)把變長(zhǎng)部分拆到 XLA 編譯范圍之外只讓固定 shape 的核心矩陣運(yùn)算走 XLA。固定 shape 對(duì)于 XLA 來(lái)說(shuō)幾乎是第一原則性能影響比編譯器版本大得多。4.2 不支持的算子會(huì)讓整個(gè) cluster 全回退XLA:CPU 不是萬(wàn)能的有些 TensorFlow 算子它沒(méi)有實(shí)現(xiàn) kernel比如部分字符串處理、某些高級(jí)索引操作、稀疏相關(guān)算子。如果 Grappler 把一個(gè) cluster 里的幾十個(gè)算子全部包在一起只要里面有一個(gè)算子 XLA 不支持整個(gè) cluster 就會(huì)回退之前所有合并努力全部白費(fèi)。這就是為什么有時(shí)候你只改了一小段代碼warning 突然消失又突然回來(lái)。問(wèn)題很可能不是 XLA 本身而是你圖里的一個(gè)不起眼的小算子。排查時(shí)建議把模型拆成更小的tf.function并配合jit_compileTrue分段測(cè)試。哪一段報(bào)錯(cuò)就說(shuō)明問(wèn)題在哪一段。然后要么把不支持的算子挪出該函數(shù)要么用 XLA 支持的等價(jià)操作重寫(xiě)。4.3 XLA:CPU 和 oneDNN 不是一回事別混著用CPU 訓(xùn)練還有一個(gè)常見(jiàn)的加速通道叫 oneDNN以前叫 MKL-DNN。很多人看到兩個(gè)都是 CPU 方向就以為相關(guān)。但 oneDNN 優(yōu)化的是單個(gè)算子內(nèi)部的數(shù)學(xué)實(shí)現(xiàn)比如卷積、矩陣乘XLA 優(yōu)化的是整個(gè) cluster 的算子融合和內(nèi)存規(guī)劃。兩者的關(guān)系是并行的不是替代關(guān)系。你單獨(dú)設(shè)置TF_ENABLE_ONEDNN_OPTS也不會(huì)消除Not using XLA:CPU for cluster因?yàn)檫@句話說(shuō)的是 XLA 編譯器有沒(méi)有接管 cluster和 oneDNN 無(wú)關(guān)。如果你追求性價(jià)比可以兩個(gè)都打開(kāi)。但注意同時(shí)開(kāi)啟時(shí) profile 結(jié)果可能比較難讀得區(qū)分XlaLaunch和 oneDNN 優(yōu)化算子分別占了多少時(shí)間。4.4 用 Profiler 驗(yàn)證而不是看 warning判斷 XLA 到底有沒(méi)有生效最可靠的方法不是看終端有沒(méi)有 warning而是看執(zhí)行期算子列表里有沒(méi)有XlaLaunch。我自己會(huì)這樣跑一輪import tensorflow as tf tf.profiler.experimental.start(/tmp/xla_profile) for step in range(step_num): train_step(data) tf.profiler.experimental.stop()然后打開(kāi) TensorBoardtensorboard --logdir/tmp/xla_profile在 Trace Viewer 里搜索XlaLaunch。如果你看到多個(gè)XlaLaunch算子出現(xiàn)說(shuō)明 cluster 確實(shí)被 XLA 接管了如果完全沒(méi)有那即使 warning 被壓掉了也只是被“隱藏”不是被“解決”。這個(gè)檢查方式我越來(lái)越依賴因?yàn)樗苤苯臃从秤?xùn)練時(shí)間到底花在哪。只看 warning 的話很容易被表面現(xiàn)象誤導(dǎo)。5. 給不同版本 TF 的最終建議5.1 快速對(duì)照表不同 TensorFlow 版本對(duì) XLA:CPU 的處理方式差異很大。我整理了一張表方便你按自己的版本快速對(duì)號(hào)入座場(chǎng)景推薦做法注意點(diǎn)TF 2.4 ~ 2.7設(shè)置TF_XLA_FLAGS--tf_xla_enable_xla_devices必須在 import TensorFlow 前設(shè)置TF 2.8 ~ 2.10使用tf.config.optimizer.set_jit(True)可以配合舊 flag 過(guò)渡TF 2.11優(yōu)先tf.config.optimizer.set_jit(autoclustering)XLA 設(shè)備默認(rèn)注冊(cè)特定計(jì)算塊使用tf.function(jit_compileTrue)編譯失敗時(shí)會(huì)快速暴露原因動(dòng)態(tài) shape 模型padding 固定 shape動(dòng)態(tài) shape 是 XLA 最大天敵訓(xùn)練沒(méi)提速使用 profiler 查XlaLaunch不要只看 warning5.2 我自己現(xiàn)在的“不折騰”組合我現(xiàn)在維護(hù)訓(xùn)練腳本時(shí)默認(rèn)會(huì)這樣做訓(xùn)練步整體包進(jìn)tf.function全局開(kāi)tf.config.optimizer.set_jit(autoclustering)對(duì)最耗時(shí)的矩陣大塊單獨(dú)加一個(gè)jit_compileTrue的小函數(shù)所有可變長(zhǎng)的輸入全部在數(shù)據(jù)管線里 padding 到固定 shape每改一次模型結(jié)構(gòu)用 profiler 看一遍XlaLaunch覆蓋范圍。這套組合犧牲了一點(diǎn)靈活性但勝在穩(wěn)定、可復(fù)現(xiàn)。至少我最近跑的幾個(gè) CPU 訓(xùn)練任務(wù)里XLA 的收益能穩(wěn)定落在 10%~30% 之間而且不會(huì)再被“Not using XLA:CPU for cluster”這種提示反復(fù)干擾。老實(shí)說(shuō)這一行 warning 本身不復(fù)雜復(fù)雜的是它背后牽扯的編譯機(jī)制、算子支持和 shape 策略。只要你不把“消 warning”當(dāng)最終目標(biāo)而是把“讓 XlaLaunch 真正跑起來(lái)、并且讓訓(xùn)練總時(shí)間降下去”當(dāng)目標(biāo)問(wèn)題解決起來(lái)反而會(huì)順利得多。