:讓多元 AI 芯片即插即用 PyTorch)
1. 多元芯片跑 PyTorch 的真實困境搞過深度學習部署的人大概都有這種體會手里攢了一堆不同品牌的加速卡想在同一套 PyTorch 訓練腳本里把它們都用起來結(jié)果發(fā)現(xiàn)每換一種芯片就得改一遍代碼、重裝一遍環(huán)境、重新調(diào)一遍算子。這事兒說起來簡單做起來能把人逼瘋。我最早接觸這個痛點是在一個多卡異構(gòu)的小集群項目上。當時機器里插了不同廠商的加速卡本意是想讓訓練任務(wù)能靈活調(diào)度到任意空閑設(shè)備上。理想很豐滿現(xiàn)實是每塊卡的驅(qū)動棧、運行時庫、算子實現(xiàn)都不一樣PyTorch 官方對非主流芯片的支持又參差不齊。最后的結(jié)果就是同一份模型代碼在 A 卡上跑得通換到 B 卡就報算子不支持好不容易把 B 卡調(diào)通了C 卡又因為內(nèi)存對齊方式不同直接崩掉。這個問題的本質(zhì)是PyTorch 的碎片化。PyTorch 本身是一個上層框架它依賴底層的設(shè)備后端來完成張量計算、內(nèi)存管理和算子調(diào)度。當?shù)讓有酒寤ò碎T時框架和芯片之間就出現(xiàn)了一道鴻溝——每款芯片都需要一套專門的適配層而適配層的質(zhì)量、覆蓋的算子范圍、支持的 PyTorch 版本又各不相同。開發(fā)者被迫在“框架版本”“芯片型號”“算子兼容性”這三個維度里反復橫跳。FlagOS 的 Torch-FL 就是沖著這個痛點來的。它的核心思路是在 PyTorch 和多元芯片之間插入一層統(tǒng)一的抽象讓上層框架看到的永遠是同一套接口底層芯片的差異被這層抽象吃掉。用一句話概括就是——讓多元 AI 芯片對 PyTorch 實現(xiàn)“即插即用”。這篇文章我會從設(shè)計思路、核心機制、實操落地、踩坑排查幾個角度把 Torch-FL 這套方案拆開講透。不管你是剛接觸異構(gòu)計算的新手還是已經(jīng)被多芯片適配折磨過的老手應(yīng)該都能從中找到能直接抄作業(yè)的東西。2. Torch-FL 的整體設(shè)計思路拆解2.1 為什么要在框架和芯片之間加一層要理解 Torch-FL 的價值得先搞清楚 PyTorch 和芯片之間到底發(fā)生了什么。PyTorch 在執(zhí)行一個算子時大致經(jīng)歷這幾個階段Python 層調(diào)用 → 框架調(diào)度層 → 設(shè)備后端 → 驅(qū)動 → 硬件。其中“設(shè)備后端”這一層就是碎片化的重災(zāi)區(qū)。以主流的加速卡為例PyTorch 通過一套私有的設(shè)備接口來對接這套接口包含了內(nèi)存分配器、流管理、算子內(nèi)核注冊等一整套機制。換一款芯片就意味著要重新實現(xiàn)這一整套東西。更麻煩的是PyTorch 的版本迭代很快每次大版本升級都可能改動設(shè)備接口的簽名適配層就得跟著改。Torch-FL 的做法是在 PyTorch 的設(shè)備后端之上再抽象一層。它定義了一套統(tǒng)一的設(shè)備描述和算子注冊規(guī)范各芯片廠商只需要按照這套規(guī)范提供自己的實現(xiàn)PyTorch 側(cè)則通過 Torch-FL 的統(tǒng)一入口來調(diào)用。這樣一來框架層不需要關(guān)心底層是哪款芯片芯片層也不需要關(guān)心上層跑的是什么模型。這里有個關(guān)鍵點Torch-FL 不是去改 PyTorch 的源碼而是通過插件化的方式掛載進去。這意味著你不需要維護一個 fork 版本的 PyTorch升級框架版本時也不會因為適配層的改動而卡住。2.2 虛擬設(shè)備抽象到底抽象了什么標題里提到的“虛擬設(shè)備”是 Torch-FL 的核心概念之一。很多人第一次聽到這個詞會以為是某種模擬器其實不是。虛擬設(shè)備在這里指的是對物理芯片的一層邏輯封裝它向上暴露統(tǒng)一的設(shè)備接口向下屏蔽具體的硬件差異。具體來說虛擬設(shè)備抽象了這幾樣東西設(shè)備發(fā)現(xiàn)與枚舉不管底層插了幾種卡、每種幾張Torch-FL 都會把它們統(tǒng)一注冊成一組邏輯設(shè)備PyTorch 側(cè)看到的是一份扁平的設(shè)備列表。內(nèi)存管理不同芯片的內(nèi)存分配策略、對齊要求、顯存回收機制都不一樣。虛擬設(shè)備層統(tǒng)一了內(nèi)存申請和釋放的接口把對齊、池化這些細節(jié)封裝在內(nèi)部。算子分發(fā)同一個算子在不同芯片上可能有不同的實現(xiàn)。虛擬設(shè)備層維護一張算子映射表根據(jù)當前設(shè)備類型把調(diào)用路由到對應(yīng)的內(nèi)核實現(xiàn)。流與同步多芯片場景下的流管理和同步是個大坑。虛擬設(shè)備層提供了統(tǒng)一的流抽象讓上層可以用同一套代碼管理不同芯片上的異步執(zhí)行。這套抽象帶來的直接好處是PyTorch 側(cè)的設(shè)備相關(guān)代碼可以做到零修改。你原來寫的.to(device)、torch.cuda.stream()這類調(diào)用在 Torch-FL 環(huán)境下依然能用只是背后的設(shè)備可能是任意一款被支持的芯片。2.3 方案選型背后的取舍做這種統(tǒng)一抽象層業(yè)界其實有幾條不同的技術(shù)路線。一條是編譯期方案通過代碼生成把不同芯片的算子編譯成統(tǒng)一中間表示另一條是運行期方案在運行時動態(tài)分發(fā)算子調(diào)用。Torch-FL 走的是運行期路線這個選擇有它的道理。編譯期方案的優(yōu)勢是性能上限高因為算子可以被充分優(yōu)化。但它的缺點是靈活性差——每支持一款新芯片就要重新編譯一遍而且對動態(tài)形狀、動態(tài)控制流的支持很有限。運行期方案雖然有一定的分發(fā)開銷但勝在靈活新芯片接入只需要提供運行時的算子實現(xiàn)不需要動編譯工具鏈。Torch-FL 在運行期方案的基礎(chǔ)上做了一些優(yōu)化來降低分發(fā)開銷。比如算子映射表在初始化階段就構(gòu)建好運行時只做一次查表再比如對高頻算子做了緩存避免重復的路徑解析。實測下來這套機制帶來的額外開銷在大多數(shù)場景下可以控制在個位數(shù)百分比。另一個取舍是關(guān)于算子覆蓋策略。Torch-FL 沒有追求一次性覆蓋 PyTorch 的全部算子而是優(yōu)先覆蓋高頻核心算子邊緣算子通過回退機制處理。這個策略很務(wù)實——PyTorch 有上千個算子但實際模型里常用的也就那么一兩百個。先把這些搞定就能覆蓋絕大多數(shù)場景。3. 核心機制與關(guān)鍵實現(xiàn)細節(jié)3.1 設(shè)備注冊與發(fā)現(xiàn)流程Torch-FL 啟動時會做一次設(shè)備掃描把系統(tǒng)里所有可用的加速設(shè)備枚舉出來。這個過程分幾步走驅(qū)動探測通過各芯片廠商提供的運行時接口檢測設(shè)備是否存在、驅(qū)動版本是否滿足要求。能力查詢讀取每款設(shè)備的計算能力、內(nèi)存容量、支持的算子集等信息。邏輯設(shè)備注冊把物理設(shè)備映射成 Torch-FL 的邏輯設(shè)備分配統(tǒng)一的設(shè)備 ID。算子表構(gòu)建根據(jù)設(shè)備能力為每款設(shè)備構(gòu)建可用的算子映射表。這個流程里有個細節(jié)值得注意設(shè)備 ID 的分配是穩(wěn)定的。也就是說同一臺機器上多次啟動同一塊物理卡拿到的邏輯 ID 是一樣的。這個特性對需要固定設(shè)備編號的訓練腳本很重要否則每次重啟都要改配置。設(shè)備注冊完成后PyTorch 側(cè)通過torch.device就能訪問到這些邏輯設(shè)備。比如torch.device(fl:0)就代表第一個 Torch-FL 邏輯設(shè)備至于它背后是哪個品牌的卡上層不需要關(guān)心。3.2 算子分發(fā)的實現(xiàn)原理算子分發(fā)是 Torch-FL 最核心的機制。當一個 PyTorch 算子被調(diào)用時Torch-FL 需要決定用哪個底層實現(xiàn)來執(zhí)行它。這個決策過程大致是這樣的首先框架層會把算子調(diào)用轉(zhuǎn)換成 Torch-FL 的內(nèi)部表示包含算子名稱、輸入張量的設(shè)備信息、數(shù)據(jù)類型、形狀等元數(shù)據(jù)。然后分發(fā)器根據(jù)設(shè)備類型去查算子映射表找到對應(yīng)的內(nèi)核實現(xiàn)。如果找到了就直接調(diào)用如果沒找到就走回退路徑?;赝寺窂接袔追N策略CPU 回退把張量搬到 CPU 上執(zhí)行執(zhí)行完再搬回來。這個策略簡單但性能差只適合極少數(shù)邊緣算子。通用內(nèi)核回退用一套跨平臺的通用內(nèi)核實現(xiàn)來執(zhí)行。性能介于原生實現(xiàn)和 CPU 回退之間。報錯提示如果連通用內(nèi)核都沒有就明確報錯告訴用戶哪個算子不支持。實操心得在接入新芯片時建議先用一個覆蓋常用算子的測試模型跑一遍看看哪些算子走了回退路徑。如果回退比例過高說明這款芯片的算子覆蓋還不夠需要補充實現(xiàn)。算子映射表的構(gòu)建是有優(yōu)先級的。同一款芯片可能提供多個版本的算子實現(xiàn)比如一個高精度版本和一個高性能版本。Torch-FL 會根據(jù)當前的精度要求和性能配置來選擇。這個機制在混合精度訓練場景下特別有用。3.3 內(nèi)存管理的統(tǒng)一抽象內(nèi)存管理是異構(gòu)計算里最容易出問題的地方。不同芯片的顯存分配粒度、對齊要求、是否支持統(tǒng)一內(nèi)存尋址這些差異如果處理不好輕則性能下降重則直接崩潰。Torch-FL 的內(nèi)存抽象層做了這幾件事統(tǒng)一分配接口上層通過fl_malloc這類接口申請內(nèi)存具體怎么分配由底層決定。對齊處理根據(jù)設(shè)備要求自動做內(nèi)存對齊避免因為對齊問題導致的性能損失或錯誤。內(nèi)存池維護一個跨設(shè)備的內(nèi)存池減少頻繁分配釋放帶來的開銷。生命周期管理跟蹤每塊內(nèi)存的歸屬設(shè)備在跨設(shè)備拷貝時做正確的同步。這里有個容易踩的坑跨設(shè)備拷貝的同步問題。當你把一個張量從設(shè)備 A 拷到設(shè)備 B 時如果 A 上還有未完成的異步操作在寫這塊內(nèi)存直接拷貝會讀到臟數(shù)據(jù)。Torch-FL 在拷貝接口里內(nèi)置了同步邏輯會等待源設(shè)備上的相關(guān)操作完成后再執(zhí)行拷貝。但這個同步是有開銷的如果拷貝頻繁性能會受影響。3.4 流與同步的統(tǒng)一處理多芯片場景下的流管理比單芯片復雜得多。每款芯片有自己的流模型有的支持多流并發(fā)有的流之間還有隱式依賴。Torch-FL 提供了一套統(tǒng)一的流抽象讓上層可以用同一套 API 管理不同芯片上的異步執(zhí)行。統(tǒng)一流抽象的核心是一個流注冊表。每個邏輯設(shè)備可以注冊多條流流與流之間的依賴關(guān)系通過事件來管理。當上層發(fā)起一個異步操作時Torch-FL 會把它分配到對應(yīng)的流上并記錄操作之間的依賴。同步方面Torch-FL 提供了設(shè)備內(nèi)同步和跨設(shè)備同步兩種機制。設(shè)備內(nèi)同步就是常規(guī)的流同步跨設(shè)備同步則需要通過事件來協(xié)調(diào)??缭O(shè)備同步的開銷通常比較大所以在設(shè)計并行策略時要盡量減少跨設(shè)備的數(shù)據(jù)依賴。4. 從零搭建 Torch-FL 實操環(huán)境4.1 環(huán)境準備與依賴檢查動手之前先把基礎(chǔ)環(huán)境理清楚。Torch-FL 對系統(tǒng)環(huán)境有一些基本要求我按實際踩坑經(jīng)驗整理了一份檢查清單檢查項要求檢查命令操作系統(tǒng)主流 Linux 發(fā)行版uname -aPython 版本3.8 及以上python --versionPyTorch 版本與 Torch-FL 版本匹配python -c import torch; print(torch.__version__)芯片驅(qū)動廠商提供的最新穩(wěn)定版各廠商查詢命令編譯工具鏈GCC 9 及以上gcc --version依賴檢查里最容易出問題的是PyTorch 版本匹配。Torch-FL 的不同版本對 PyTorch 有明確的版本要求裝錯了版本會在導入時直接報錯。建議先查清楚 Torch-FL 的版本說明再決定裝哪個版本的 PyTorch。另一個常見問題是驅(qū)動版本。有些芯片的驅(qū)動更新比較頻繁新驅(qū)動可能改了運行時接口導致 Torch-FL 的適配層不兼容。穩(wěn)妥的做法是使用 Torch-FL 官方驗證過的驅(qū)動版本不要盲目追新。4.2 安裝步驟與配置要點環(huán)境檢查通過后就可以開始安裝了。整個安裝過程分三步第一步安裝 PyTorch 基礎(chǔ)環(huán)境如果你用的是 conda建議單獨建一個環(huán)境避免和系統(tǒng)里的其他 Python 包沖突conda create -n torchfl python3.10 conda activate torchfl然后安裝對應(yīng)版本的 PyTorch。注意這里要裝的是 CPU 版本或者與你的主設(shè)備匹配的版本Torch-FL 會在運行時接管設(shè)備管理pip install torch2.1.0 torchvision0.16.0第二步安裝 Torch-FLTorch-FL 通常以 wheel 包的形式提供直接 pip 安裝即可pip install torch-fl安裝完成后驗證一下是否裝好了python -c import torch_fl; print(torch_fl.__version__)第三步配置設(shè)備Torch-FL 的配置文件通常放在~/.torch_fl/config.yaml。一個典型的配置長這樣devices: - type: vendor_a count: 2 memory_fraction: 0.9 - type: vendor_b count: 1 memory_fraction: 0.8 fallback: enable_cpu_fallback: true enable_generic_kernel: true logging: level: info path: /var/log/torch_fl.log配置里幾個關(guān)鍵參數(shù)說明一下memory_fraction控制每塊設(shè)備上 Torch-FL 能使用的顯存比例。留一點余量給系統(tǒng)和其他進程避免 OOM。enable_cpu_fallback是否允許算子回退到 CPU。調(diào)試階段建議開著生產(chǎn)環(huán)境可以關(guān)掉以便及早發(fā)現(xiàn)問題。enable_generic_kernel是否啟用通用內(nèi)核回退。這個比 CPU 回退性能好建議開啟。4.3 驗證安裝與基礎(chǔ)測試裝完之后別急著跑大模型先用一個小測試確認環(huán)境是通的。我一般用這段代碼做冒煙測試import torch import torch_fl # 查看 Torch-FL 識別到的設(shè)備 print(可用設(shè)備:, torch_fl.list_devices()) # 創(chuàng)建一個張量并放到 Torch-FL 設(shè)備上 device torch.device(fl:0) x torch.randn(1000, 1000, devicedevice) y torch.randn(1000, 1000, devicedevice) # 做一個矩陣乘法 z torch.matmul(x, y) print(計算結(jié)果形狀:, z.shape) print(計算所在設(shè)備:, z.device) # 測試跨設(shè)備拷貝 if len(torch_fl.list_devices()) 1: device2 torch.device(fl:1) z2 z.to(device2) print(拷貝后設(shè)備:, z2.device)這段代碼覆蓋了設(shè)備發(fā)現(xiàn)、張量創(chuàng)建、算子執(zhí)行、跨設(shè)備拷貝幾個核心路徑。如果都能跑通說明基礎(chǔ)環(huán)境沒問題。注意如果list_devices()返回空列表大概率是驅(qū)動沒裝好或者配置文件有問題。先檢查驅(qū)動再看配置文件里的設(shè)備類型名稱是否和實際硬件匹配。4.4 跑通第一個訓練任務(wù)冒煙測試通過后可以拿一個真實的模型來試。建議從簡單的 CNN 或者小型的 Transformer 開始不要一上來就上大模型。下面是一個在 Torch-FL 上跑 MNIST 訓練的簡化示例import torch import torch.nn as nn import torch.optim as optim import torch_fl device torch.device(fl:0) class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, 3, padding1) self.conv2 nn.Conv2d(32, 64, 3, padding1) self.fc nn.Linear(64 * 7 * 7, 10) self.pool nn.MaxPool2d(2) def forward(self, x): x self.pool(torch.relu(self.conv1(x))) x self.pool(torch.relu(self.conv2(x))) x x.view(x.size(0), -1) return self.fc(x) model SimpleCNN().to(device) optimizer optim.Adam(model.parameters(), lr1e-3) criterion nn.CrossEntropyLoss() # 假設(shè) train_loader 已經(jīng)準備好 for epoch in range(5): for data, target in train_loader: data, target data.to(device), target.to(device) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() print(fEpoch {epoch}, Loss: {loss.item()})這段代碼和標準的 PyTorch 訓練代碼幾乎一模一樣唯一的區(qū)別是設(shè)備指定成了fl:0。這就是 Torch-FL 想要達到的效果——上層代碼零修改。5. 常見問題與排查技巧實錄5.1 設(shè)備識別不到怎么辦這是最常見的問題表現(xiàn)是list_devices()返回空或者缺少預(yù)期的設(shè)備。排查思路按這個順序走先查驅(qū)動。用廠商提供的工具確認設(shè)備在系統(tǒng)層面是可見的。如果系統(tǒng)都看不到設(shè)備Torch-FL 自然也無能為力。再查權(quán)限。有些設(shè)備需要特定的用戶組權(quán)限才能訪問。檢查當前用戶是否在對應(yīng)的組里設(shè)備文件的權(quán)限是否正確。然后查配置。確認config.yaml里的設(shè)備類型名稱和實際硬件匹配。不同廠商的設(shè)備類型標識不一樣寫錯了就識別不到。最后看日志。Torch-FL 的日志里會記錄設(shè)備探測的詳細過程包括每一步的成功和失敗原因。日志級別調(diào)到 debug 能看到更多信息。5.2 算子不支持的回退與替代方案遇到算子不支持時Torch-FL 會打印一條警告說明哪個算子走了回退路徑。如果只是偶爾幾個邊緣算子回退影響不大。但如果核心算子頻繁回退性能會明顯下降。處理策略分幾種情況如果是通用算子檢查是不是 Torch-FL 版本太舊升級到最新版可能就支持了。如果是自定義算子需要自己實現(xiàn)對應(yīng)的內(nèi)核或者用一組基礎(chǔ)算子來組合替代。如果是邊緣算子可以接受回退但要注意回退帶來的性能損失和精度差異。我整理了一份常見回退算子及其替代方案的對照表回退算子替代方案注意事項某些特殊激活函數(shù)用基礎(chǔ)激活函數(shù)組合注意數(shù)值穩(wěn)定性非標準卷積變體拆解為標準卷積可能增加計算量特殊歸一化層手動實現(xiàn)歸一化注意訓練和推理的差異自定義損失函數(shù)用基礎(chǔ)運算組合檢查梯度是否正確5.3 性能不達預(yù)期的調(diào)優(yōu)思路Torch-FL 跑起來之后如果性能不如預(yù)期可以從這幾個方向排查先看回退比例。如果大量算子走了回退路徑性能肯定上不去。用 Torch-FL 提供的 profiling 工具統(tǒng)計一下回退比例超過 10% 就要重視了。再看內(nèi)存拷貝??缭O(shè)備拷貝是性能殺手。檢查模型里有沒有不必要的設(shè)備間數(shù)據(jù)搬運能合并的合并能避免的避免。然后看流配置。如果設(shè)備支持多流并發(fā)但代碼里只用了單流就浪費了并行能力。合理配置多流可以提升吞吐。最后看批大小。批大小太小會導致設(shè)備利用率不足太大又可能 OOM。找到一個平衡點很重要。5.4 多芯片混合訓練的注意事項多芯片混合訓練是 Torch-FL 的強項但也是坑最多的地方。幾個關(guān)鍵注意點設(shè)備間的性能差異不同芯片的算力可能差很多。如果簡單地做數(shù)據(jù)并行慢的那塊卡會成為瓶頸。可以考慮按算力比例分配數(shù)據(jù)量。通信開銷跨設(shè)備通信的開銷通常比設(shè)備內(nèi)通信大得多。梯度同步、參數(shù)廣播這些操作要盡量減少頻率。精度一致性不同芯片的浮點運算實現(xiàn)可能有細微差異混合訓練時要注意梯度累積的精度問題。故障恢復多芯片場景下任何一塊卡出問題都可能影響整個訓練。要做好 checkpoint 和故障恢復機制。實操心得混合訓練初期建議先用小規(guī)模數(shù)據(jù)跑通流程確認各設(shè)備都能正常工作、通信正常、精度正常再逐步擴大規(guī)模。不要一上來就上全量數(shù)據(jù)出了問題很難定位。6. 工具選型與生態(tài)適配經(jīng)驗6.1 Torch-FL 與其他適配方案的對比市面上做 PyTorch 多芯片適配的方案不止 Torch-FL 一家選型時可以從這幾個維度對比維度Torch-FL編譯期方案廠商私有方案接入成本低插件化高需重新編譯中依賴廠商支持靈活性高運行期分發(fā)低靜態(tài)編譯中性能上限中高高高算子覆蓋逐步完善取決于編譯范圍取決于廠商框架升級兼容性好需重新適配可能滯后選型的核心考量是你的場景更看重靈活性還是極致性能。如果是快速驗證、多芯片混跑Torch-FL 這類運行期方案更合適。如果是單一芯片、追求極致性能編譯期方案可能更好。6.2 與現(xiàn)有訓練框架的集成Torch-FL 設(shè)計上是可以和現(xiàn)有訓練框架共存的。如果你用的是 HuggingFace Trainer、PyTorch Lightning 這類高層框架通常只需要改設(shè)備指定那一行代碼。以 PyTorch Lightning 為例在 Trainer 里指定加速器from pytorch_lightning import Trainer trainer Trainer( acceleratorfl, devices2, strategyddp )Lightning 會通過 Torch-FL 的插件接口來管理設(shè)備。前提是 Torch-FL 提供了對應(yīng)的 Lightning 插件這個需要確認版本兼容性。對于自定義的訓練循環(huán)集成就更簡單了基本上就是把torch.device(cuda)換成torch.device(fl:0)其他代碼不用動。6.3 版本升級與兼容性維護Torch-FL 和 PyTorch 都在快速迭代版本兼容性是個持續(xù)要關(guān)注的問題。我的經(jīng)驗是鎖定版本組合生產(chǎn)環(huán)境不要用 latest用經(jīng)過驗證的版本組合。把 PyTorch 版本、Torch-FL 版本、驅(qū)動版本都固定下來。關(guān)注發(fā)布說明每次升級前仔細看 release notes特別是 breaking changes 部分?;叶壬壪仍跍y試環(huán)境驗證確認沒問題再推到生產(chǎn)。保留回滾方案升級前做好環(huán)境備份出問題能快速回滾。這套流程看起來繁瑣但比起生產(chǎn)環(huán)境出故障的代價這點準備工作完全值得。7. 實際項目中的經(jīng)驗體會我在幾個實際項目里用過 Torch-FL有一些體會是文檔里不會寫的。第一不要指望一次配置就完美。異構(gòu)環(huán)境的變數(shù)太多驅(qū)動、固件、系統(tǒng)庫的版本組合千差萬別。做好反復調(diào)試的心理準備把每次成功的配置記錄下來形成自己的配置庫。第二性能調(diào)優(yōu)要有耐心。Torch-FL 的開箱性能可能不是最優(yōu)的需要根據(jù)具體場景調(diào)。內(nèi)存池大小、流數(shù)量、批大小這些參數(shù)都要試。我一般會做一個參數(shù)掃描找到當前場景下的最優(yōu)組合。第三日志是你的朋友。Torch-FL 的日志信息很豐富遇到問題先看日志。把日志級別調(diào)到 debug很多問題的原因一目了然。第四社區(qū)和文檔要結(jié)合起來看。官方文檔覆蓋了主要功能但一些邊緣場景和最新問題往往在社區(qū)里才有答案。遇到卡住的問題搜一下社區(qū)討論大概率有人踩過同樣的坑。最后分享一個小技巧在接入新芯片時先跑一個算子覆蓋測試把模型里用到的所有算子列出來逐個確認是否支持。這個測試可以在正式訓練前就發(fā)現(xiàn)問題避免訓練到一半才報錯。測試腳本可以基于 Torch-FL 的算子查詢接口來寫把模型的計算圖導出后遍歷所有節(jié)點逐個查詢支持情況。這個習慣幫我省了很多返工的時間。