試調(diào)優(yōu):從環(huán)境配置到性能優(yōu)化實戰(zhàn))
很多做AI訓練的朋友第一次拿到昇騰機器時第一反應(yīng)是“這不就是個NPU嘛代碼改改就能跑?!钡日嬲鲜植虐l(fā)現(xiàn)從PyTorch代碼遷移、數(shù)據(jù)管線適配到集群通信調(diào)優(yōu)、訓練穩(wěn)定性排查整套流程盤下來需要踩的坑比想象中多得多。這篇文章我就圍繞“昇騰大模型訓練調(diào)試調(diào)優(yōu)”這條主線把模型訓練全流程里的關(guān)鍵環(huán)節(jié)、實操方法和排查經(jīng)驗一次性說透希望能幫準備上手昇騰的朋友少走幾個彎路。本文既不是官方文檔的復制也不是純理論分析而是基于真實的工程實踐經(jīng)驗。內(nèi)容適合這幾類讀者一是剛拿到昇騰算力、想把手頭模型跑起來的算法工程師二是已經(jīng)在跑訓練但經(jīng)常遇到“卡住、掉卡、Loss不降”等問題的同學三是準備做集群規(guī)模訓練、正在做技術(shù)選型的架構(gòu)師。我會按實際操作的先后順序展開從環(huán)境準備、數(shù)據(jù)加載、模型遷移到訓練啟動、性能調(diào)優(yōu)和故障排查每一段都會給出可直接落地的建議。1. 環(huán)境準備與框架選型第一步就決定后續(xù)效率很多人以為昇騰訓練就是裝上驅(qū)動就能跑其實最關(guān)鍵的在于CANN版本、PyTorch版本、torch_npu版本三者之間的對齊關(guān)系。這三者就像齒輪組版本不匹配時輕則報一堆詭異的算子錯誤重則設(shè)備直接掉線。1.1 版本匹配是頭號大事昇騰的訓練軟件棧底層是CANNCompute Architecture for Neural Networks往上依次是PyTorch框架、torch_npu插件再往上才是用戶的訓練代碼。CANN版本決定了NPU驅(qū)動和運行時環(huán)境而torch_npu則負責把PyTorch的算子調(diào)用翻譯成CANN能理解的指令。這個鏈路里的任何一環(huán)版本不對齊都會讓訓練進程崩潰。我實測下來比較穩(wěn)妥的選型思路是先確定CANN版本再反推torch_npu和PyTorch的對應(yīng)關(guān)系。比如CANN 7.0.RC1通常搭配torch_npu 2.1.0.post6和PyTorch 2.1.0而更新一些的CANN 8.0.RC1版本則支持PyTorch 2.3.1或更新版本。具體組合建議直接看官方發(fā)布的版本配套表不要自己隨意組合。注意昇騰生態(tài)里很多坑都源于“默認安裝最新版”這個習慣。不要圖新鮮裝最新CANN優(yōu)先選已經(jīng)在生產(chǎn)環(huán)境被驗證過的穩(wěn)定組合。1.2 兩種開發(fā)模式的取舍現(xiàn)在昇騰上跑大模型訓練主流有兩條路線一條是基于torch_npu做PyTorch代碼最小改造保留原有代碼結(jié)構(gòu)本質(zhì)上是把昇騰當成一個“加速卡”來用另一條是用MindSpore MindFormers全家桶從框架層適配昇騰能更充分地發(fā)揮硬件特性。這兩條路線怎么選取決于你手頭的代碼現(xiàn)狀。如果你有成熟的PyTorch代碼庫團隊也熟悉PyTorch生態(tài)那走torch_npu路線成本最低。改造量通常相對可控但遇到不支持的算子時還是得手動適配如果是從零開始訓新模型或者團隊本來就打算深度定制訓練邏輯那用MindFormers自帶的并行策略、混合精度控制、斷點續(xù)訓功能會更順手底層優(yōu)化也更省心。我自己做模型訓練時前期用torch_npu快速驗證思路中期穩(wěn)定后部分場景切到MindFormers做長穩(wěn)訓練這樣兩頭的好處都能占到。2. 數(shù)據(jù)準備與加載優(yōu)化訓練瓶頸往往不在計算大模型訓練的瓶頸經(jīng)常被人忽略很多人都盯著計算卡上的算子耗時卻忘了數(shù)據(jù)加載可能正拖著整個訓練后退。昇騰訓練場景下數(shù)據(jù)管線的表現(xiàn)和GPU生態(tài)有明顯差異如果沿用GPU上的數(shù)據(jù)加載寫法很容易出現(xiàn)NPU在等數(shù)據(jù)的情況。2.1 數(shù)據(jù)管線的常見“隱形殺手”最常遇到的問題有三個。一是用普通的文件讀取配合Pillow等庫做在線解碼CPU端解碼速度跟不上NPU的消費速度二是DataLoader的num_workers設(shè)置不合理過多或過少都會影響吞吐三是缺少預取機制每一步訓練都要等數(shù)據(jù)從磁盤換上來。解決思路是把數(shù)據(jù)預處理盡量“離線化”也就是把清洗、解碼、Token化等操作在訓練前批量完成訓練時直接讀取已經(jīng)處理好的數(shù)據(jù)。對NLP模型來說可以預先將文本轉(zhuǎn)成Token ID并緩存為二進制格式對CV模型來說可以提前做Resize、歸一化等操作。這樣訓練時只需做最簡單的讀取和搬移CPU壓力會大幅下降。2.2 MindRecord與昇騰數(shù)據(jù)加速方案如果用MindSpore框架推薦用MindRecord格式來存儲訓練數(shù)據(jù)。這種格式的讀取效率比通用文件格式高不少原因在于它按樣本做了索引、支持隨機讀取和分布式分片避免了大量小文件隨機讀帶來的IO開銷。在torch_npu場景下雖然還用PyTorch的DataLoader但建議關(guān)注自定義Dataset的耗時占空比。一個很實用的排查方法訓練啟動后在日志里打印每個Step的時間如果Step耗時出現(xiàn)明顯波動先檢查Host側(cè)數(shù)據(jù)加載耗時是否過高。我遇到過一次訓練速度突然下降的問題排查到最后發(fā)現(xiàn)是數(shù)據(jù)加載時做了一次多余的復制操作。處理方法是直接把數(shù)據(jù)加載的Tensor改為非阻塞拷貝并保證數(shù)據(jù)在Host上預處理完再拷到NPU減少重復內(nèi)存搬運。這里想強調(diào)的是數(shù)據(jù)管線是否優(yōu)化到位直接影響大規(guī)模訓練的整體吞吐不能只盯著算子在算。3. 模型遷移適配從GPU到昇騰的完整轉(zhuǎn)換模型遷移是昇騰大模型訓練最核心的一步也是問題最密集的一環(huán)。PyTorch模型遷移到昇騰核心目標是用最小改動讓模型能在NPU上正常跑通訓練并保持數(shù)值精度與收斂效果。3.1 最小改造基于torch_npu的適配路徑如果你的代碼是標準的PyTorch實現(xiàn)那么遷移的第一步通常是替換設(shè)備標識。最簡單的方法是定義全局device變量然后通過model.to(device)把模型和Tensor放到昇騰設(shè)備上。import torch import torch_npu # 在代碼初始化階段指定設(shè)備 device torch.device(npu:0) model model.to(device) data data.to(device) # 原有的訓練邏輯基本不用改 output model(input_tensor) loss loss_fn(output, target) loss.backward() optimizer.step()這段代碼看起來很簡單但實際項目中的問題很少出在設(shè)備遷移本身更多出在算子兼容性上。PyTorch里某些算子尤其是一些高級索引、自定義Op在昇騰上可能沒有對應(yīng)的原生實現(xiàn)torch_npu會自動走CPU回退或報錯。遇到這種情況優(yōu)先做法是改寫模型代碼用昇騰原生支持的算子組合替代不支持的算子。3.2 用MindSpore做重寫遷移的時機如果你的模型有很多自定義結(jié)構(gòu)或者需要高性能并行策略直接遷移到MindSpore可能更劃算。MindSpore提供了模型遷移工具可以輔助轉(zhuǎn)換PyTorch模型結(jié)構(gòu)但仍需要手工校驗算子的數(shù)值一致性。重寫遷移的優(yōu)勢在于后續(xù)訓練可以直接使用MindSpore的自動并行、內(nèi)存復用、圖模式編譯等能力對大規(guī)模訓練來說性價比更高。舉個例子PyTorch的nn.Module在MindSpore里要對應(yīng)改寫成nn.Cell前向計算寫在construct方法里import mindspore import mindspore.nn as nn from mindspore import Tensor class MyModel(nn.Cell): def __init__(self): super().__init__() self.fc nn.Dense(768, 1024) self.act nn.GeLU() def construct(self, x): x self.fc(x) x self.act(x) return x這種改寫雖然是“重復造輪子”但換來的是后續(xù)訓練腳本的極大簡化。MindSpore提供了一系列高階API比如TrainOneStepCell、Model.train可以直接接管混合精度、梯度累積、分布式并行等細節(jié)。3.3 算子適配時的三個重要檢查點算子適配是整個遷移過程中最磨人的階段但核心檢查點也不復雜一是檢查模型里有沒有類似index_put_、scatter這類高級索引算子這類算子很容易出問題。建議改成masked_fill或循環(huán)加切片賦值等價的邏輯拿到昇騰上跑得更順二是檢查Loss計算中是否混用了FP32和FP16計算這會導致梯度數(shù)值異常。建議把Loss計算整體統(tǒng)一到FP32只在模型前向部分啟用混合精度三是檢查是否用到了自定義autograd.Function這類代碼往往基于CUDA實現(xiàn)昇騰無法直接運行需要改寫為昇騰兼容的算子或在CPU上做回退。我踩過一次最深的坑是模型里用了某個第三方庫做位置編碼內(nèi)部實現(xiàn)了一個自定義CUDA算子。遷移到昇騰后代碼沒有報錯但Loss始終不下降。后來用Profiling工具定位才發(fā)現(xiàn)這個自定義算子被靜默回退到了CPU執(zhí)行數(shù)據(jù)往返拷貝導致梯度計算完全錯亂。實操建議遷移完成后先跑一次固定隨機種子的短Step訓練對比GPU和NPU上的Loss曲線。Loss曲線趨勢一致且差距在可接受范圍才說明遷移正確。這一步別省。4. 訓練腳本編寫與全流程調(diào)試從單機到集群模型遷移完成只是開始真正把訓練腳本寫好并調(diào)通才是全流程調(diào)試的重頭戲。整個訓練腳本里分布式初始化、混合精度控制、梯度累積、Loss縮放這幾個環(huán)節(jié)每一個都值得認真對待。4.1 分布式訓練的初始化細節(jié)昇騰上做分布式訓練通信依賴HCCLHuawei Collective Communication Library。與NCCL類似HCCL也需要初始化一個全局通信域。如果是通過torch_npu路徑走PyTorch分布式初始化代碼基本是import torch.distributed as dist import torch_npu dist.init_process_group(backendhccl, init_methodenv://)單機多卡或集群場景下通過環(huán)境變量傳入MASTER_ADDR、MASTER_PORT、WORLD_SIZE和RANK。有一點特別容易踩坑WORLD_SIZE必須和實際參與訓練的NPU數(shù)量一致否則HCCL初始化會一直卡在等待對端上。有一次我在多機場景下忘了同步各節(jié)點的RANK編號導致HCCL建鏈直接失敗排查了很久才發(fā)現(xiàn)是rank分配錯誤。集群環(huán)境下依賴rank_table文件來管理設(shè)備信息也可以通過環(huán)境變量ASCEND_RT_VISIBLE_DEVICES來指定每臺機器上參與訓練的NPU編號。建議訓練腳本里加一段啟動自檢邏輯先遍歷所有npu:0到npu:7用簡單的AllReduce測試通信建鏈再做正式訓練。4.2 混合精度、梯度累積與Loss縮放大模型訓練幾乎都會開啟混合精度將FP32計算替換為FP16/BF16以提升計算速度和顯存利用率。昇騰對FP16和BF16的支持都很成熟但用戶需要自行管理Loss Scaling防止梯度下溢。在torch_npu路徑下可以使用torch.cuda.amp的替代實現(xiàn)昇騰提供了對應(yīng)的混合精度接口。一個可行的實現(xiàn)是from torch_npu.amp import GradScaler, autocast scaler GradScaler() with autocast(): output model(input_tensor) loss loss_fn(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()梯度累積則是把小Batch的梯度累加起來再更新參數(shù)效果上近似大Batch能在不增加顯存的情況下提升訓練穩(wěn)定性。實現(xiàn)時要特別注意累積梯度前要用optimizer.zero_grad()清零累積過程中不要執(zhí)行optimizer.step()直到累積Step數(shù)達到設(shè)定值再更新并重新清零。結(jié)合混合精度時梯度縮放涉及動態(tài)范圍變化容易踩坑。我的習慣是啟動時設(shè)置固定Scale值并開啟動態(tài)調(diào)整前幾個Step觀察溢出情況穩(wěn)定后再切換到動態(tài)模式。4.3 斷點續(xù)訓與訓練穩(wěn)定性大模型訓練動輒幾天甚至幾周斷點續(xù)訓不是“加分項”而是“必選項”。昇騰場景下斷點保存的核心是保存三樣東西模型權(quán)重、優(yōu)化器狀態(tài)、隨機數(shù)生成器狀態(tài)。如果漏掉隨機數(shù)狀態(tài)恢復訓練后數(shù)據(jù)順序會變化雖然不一定讓訓練失敗但會造成結(jié)果不可復現(xiàn)。保存優(yōu)化器狀態(tài)時要注意優(yōu)化器里可能包含指數(shù)移動平均、動態(tài)學習率等輔助狀態(tài)這些都要一并保存?;謴陀柧殨r還需要重新初始化HCCL通信域某些情況下如果斷點保存了通信域的拓撲信息恢復時也要一并恢復。實測下來最可靠的斷點保存方式是每N個Step保存一次到本地磁盤同時定期同步到共享存儲避免單點故障導致整個訓練白跑。另外恢復訓練后建議先跑幾個Step驗證Loss正常再全速跑別一恢復就直接進入長穩(wěn)階段。5. 性能調(diào)優(yōu)實戰(zhàn)從“能跑通”到“跑得快”模型能正常跑起來、Loss也在正常下降這是第一階段的勝利。但一旦進入大規(guī)模訓練階段“跑得快”就比“跑得通”更重要。昇騰性能調(diào)優(yōu)的核心思路是先定位瓶頸在計算還是通信再針對性地做優(yōu)化不盲目調(diào)參數(shù)。5.1 用Profiling工具定位瓶頸昇騰提供了Profiling工具可以采集訓練過程中的算子耗時、通信耗時、Host側(cè)耗時等信息。具體有兩種路徑一種是MindSpore Profiler適合MindSpore框架另一種是msprof工具適合整體系統(tǒng)層面的性能分析。拿到Profiling報告后我一般按照以下順序去看先看Step 耗時如果Step之間波動較大優(yōu)先排查數(shù)據(jù)加載和Host側(cè)邏輯再看通信耗時占比如果通信占比超過30%說明同步開銷太大需要調(diào)整并行策略或使用通信壓縮最后看算子耗時分布找出耗時Top 10的算子逐一判斷能否被融合或替換。有一次我優(yōu)化一個千億參數(shù)模型的訓練發(fā)現(xiàn)AllReduce占了整個Step耗時的40%多。后來把純數(shù)據(jù)并行改成張量并行加流水線并行混合模式情況才明顯好轉(zhuǎn)。通信和計算的重疊也很重要HCCL允許通信和計算并行執(zhí)行但要通過合適的Stream配置實現(xiàn)。簡單來說就是要確保數(shù)據(jù)搬運計算和通信的過程盡量重疊不要讓NPU在等通信也不要在通信時讓NPU閑著。5.2 通信優(yōu)化減少數(shù)據(jù)搬運次數(shù)通信優(yōu)化的核心原則是“減少數(shù)據(jù)搬運次數(shù)、提高單次搬運效率”。具體手段包括梯度壓縮對大梯度做量化或稀疏化處理后再通信減少通信數(shù)據(jù)量梯度分組AllReduce將梯度按層分組小梯度先通信、大梯度后通信錯開通信峰值更合理的并行策略將數(shù)據(jù)并行與模型并行結(jié)合純數(shù)據(jù)并行下每Step都要同步全量梯度通信量最大。在昇騰平臺上HCCL的效率與拓撲結(jié)構(gòu)密切相關(guān)。單機8卡的Ring AllReduce性能通常好于跨機通信所以設(shè)計模型并行時盡量把通信量大的Tensor放在同一臺機器上。跨機通信時網(wǎng)卡和交換機的帶寬也要提前確認避免因為網(wǎng)絡(luò)帶寬有限導致大規(guī)模加速比不理想。5.3 顯存優(yōu)化與計算優(yōu)化顯存是另一個瓶頸尤其在大模型訓練場景。昇騰上顯存優(yōu)化手段包括重計算Recompute把前向激活值只保存一部分反向需要時重新計算換取顯存節(jié)省混合精度把不需要高精度的Tensor切到FP16/BF16存儲顯存碎片整理調(diào)整分配策略減少碎片化提高顯存利用率流水線并行把模型切分成多個Stage放在不同設(shè)備上每臺設(shè)備只保存一部分參數(shù)和激活值。計算優(yōu)化層面算子融合是最直接的方式。昇騰提供了一個算子融合能力可以把多個連續(xù)的小算子融合成一個大的融合算子顯著減少內(nèi)核調(diào)度開銷。常見做法是把LayerNorm、Residual Add、Activation融合成一個算子或者把QKV運算合并成一個大的矩陣乘減少內(nèi)核啟動次數(shù)。同時矩陣乘的Shape對性能影響也很大建議把張量的形狀盡量對齊到昇騰的矩陣計算單元偏好。6. 常見問題與排查技巧實錄大模型訓練調(diào)試是一場持久戰(zhàn)我把自己實際操作中遇到頻率最高的四類問題整理成速查表方便你直接對照排查。問題現(xiàn)象可能原因排查思路訓練剛開始就報算子不存在模型里有昇騰不支持的算子查看報錯信息中的算子名稱改寫或替換Loss曲線異常不收斂、NaN混合精度Loss Scaling設(shè)置不當學習率過大梯度計算錯亂先關(guān)閉混合精度試跑幾個Step再逐步打開檢查Loss計算是否統(tǒng)一在FP32訓練到某一步突然卡住HCCL通信超時數(shù)據(jù)加載線程死鎖顯存溢出查看日志中的通信超時信息用Profiling看卡的Step位置多卡訓練加速比很低通信占比過高負載不均數(shù)據(jù)加載成了瓶頸用Profiling確認通信和計算耗時優(yōu)化并行策略6.1 訓練卡住的定位方法訓練“卡住”是全流程調(diào)試里最讓人頭疼的問題。表面上看進程沒有退出但Step數(shù)不再往前走。遇到這種情況我一般按這個順序排查先看NPU的利用率如果利用率很低而CPU很高大概率是數(shù)據(jù)加載卡住了檢查DataLoader和文件系統(tǒng)IO如果CPU和NPU利用率都很低大概率是卡在通信環(huán)節(jié)看看是不是有的卡掉線或通信組網(wǎng)異常如果只有某一張卡利用率異常檢查是不是模型并行不均某個Stage的計算量特別大拖慢了整體。這里要特別提一下HCCL建鏈問題。多機訓練時不同節(jié)點的設(shè)備互相通信需要走網(wǎng)卡而HCCL默認會使用某個網(wǎng)卡進行建鏈。如果這個網(wǎng)卡不通通信就會一直卡住。排查方法是在訓練啟動前用簡單的hccl_tools.py測試腳本驗證節(jié)點間的通信連通性確認沒問題再啟動正式訓練。6.2 梯度異常的定位方法梯度異常通常表現(xiàn)為兩種情況梯度爆炸導致Loss變成NaN或者梯度消失導致Loss長時間不下降。排查的第一步是逐層打印梯度數(shù)值找到異常梯度的位置。實操上可以在模型注冊一些Hook來打印各層梯度的范數(shù)值for name, param in model.named_parameters(): if param.grad is not None: grad_norm param.grad.norm().item() if grad_norm 1e4: print(fLarge grad: {name}, norm {grad_norm})這個方法雖然簡陋但能迅速縮小問題范圍。如果是某些層梯度持續(xù)異常優(yōu)先檢查這些層有沒有被混合精度影響如果是全層梯度異常則優(yōu)先檢查Loss計算和梯度累積邏輯。還有一個容易被忽略的問題權(quán)重初始化。如果某些層的初始化方差過大一開始梯度就容易爆炸尤其是在深層Transformer結(jié)構(gòu)里。所以遇到訓練初期就出NaN的情況除了檢查混合精度也要確認初始化方式是否合理。6.3 日志分析與定位調(diào)試昇騰訓練日志分析能力是基本功。建議訓練腳本統(tǒng)一用logging模塊輸出帶時間戳的日志打印每個Step的耗時和Loss。如果出現(xiàn)Step耗時突然上升需要結(jié)合日志時間線回溯當時的數(shù)據(jù)加載和通信狀態(tài)。另外昇騰的運行時日志默認帶級別可以通過環(huán)境變量調(diào)整日志級別比如把部分調(diào)試信息打開便于觀察算子執(zhí)行順序和耗時。日志量會顯著增加建議只在定位問題時臨時打開平時保持默認級別。7. 大模型訓練全流程的經(jīng)驗沉淀整套昇騰大模型訓練調(diào)試調(diào)優(yōu)走下來我的體感是昇騰已經(jīng)是一套成熟度頗高的訓練平臺不再是需要“硬啃文檔”的試驗品但它和GPU生態(tài)的差異是客觀存在的關(guān)鍵要掌握它的規(guī)律。給我留下最深的幾個經(jīng)驗第一版本對齊是“地基工程”不要在環(huán)境配置上求快一步錯后面全是連鎖反應(yīng)。拿到新機器后第一件事就是確認CANN、框架和插件的版本配套關(guān)系并保留環(huán)境配置文件方便后續(xù)復現(xiàn)。第二算子遷移是最需要耐心的環(huán)節(jié)建議把模型里所有自定義算子、第三方CUDA算子統(tǒng)一梳理出來逐個驗證昇騰兼容性。把這個工作前置到訓練啟動前能省下大量試錯時間。第三性能調(diào)優(yōu)要有數(shù)據(jù)支撐別憑感覺調(diào)參。Profiling工具一定要學會用讓數(shù)據(jù)告訴你瓶頸在哪。很多時候我們以為的計算瓶頸實際是通信或數(shù)據(jù)加載瓶頸。第四訓練穩(wěn)定性比訓練速度更重要。大模型訓練動輒數(shù)天一次掉卡造成的損失遠大于優(yōu)化帶來的收益。把斷點續(xù)訓、日志監(jiān)控、健康檢查這些“保命”功能做好優(yōu)先級高于一切花哨的優(yōu)化技巧。最后再分享一個我個人的習慣每輪全流程調(diào)試后把遇到的問題、根因、解決方案整理成一份內(nèi)部文檔形成團隊自己的“避坑手冊”。昇騰生態(tài)迭代很快這些一手經(jīng)驗往往比官方文檔更貼近實戰(zhàn)也能幫助下一次訓練啟動時少走彎路。這套“調(diào)試—沉淀—復用”的方法才是訓練全流程經(jīng)驗真正復利的地方。