說實話第一次在海光 DCU 上跑微調任務我心里是有點打鼓的。以前一直在 CUDA 生態(tài)里折騰突然換到國產加速卡最怕的就是裝環(huán)境裝到懷疑人生然后好不容易裝好了loss 又是紋絲不動。但真正把「裝環(huán)境 → 數據加載 → 模型訓練 → loss 下降」這條鏈路完整跑通之后我發(fā)現 DCU 這套東西并沒有傳說中那么難搞很多坑其實都是有跡可循的。這篇就基于我自己跑通第一個微調任務的全過程把從零到 loss 下降的關鍵步驟、環(huán)境配置心得和避坑經驗一次性說清楚。這篇內容不適合誰如果你已經是 DCU 老手對 ROCm 移植、算子適配那套門兒清可以直接劃走。但如果你跟我一樣剛拿到一臺海光 DCU 機器想把手頭的模型快速跑起來又不知道從哪下手那這篇文章應該能幫你省下好幾個晚上的折騰時間。我會盡量說人話不堆術語但該給的命令、該貼的代碼、該解釋的原理一樣都不會少。1. 先把核心思路理清楚DCU 上微調到底在做什么1.1 海光 DCU 是什么跟 CUDA 有什么不一樣海光 DCUDeep Computing Unit是海光信息推出的加速計算產品專門面向 AI 訓練和推理場景。它跟英偉達 GPU 最大的區(qū)別在于軟件棧DCU 使用的是海光自研的 DTKDCU Toolkit并且深度兼容 ROCm 生態(tài)。換句話說你在 CUDA 上寫的大部分 PyTorch 代碼只要避開少數冷門算子和底層自定義 CUDA 擴展基本上可以無縫遷移過來。我第一次聽到這個概念的時候腦子里冒出來的想法是那不就是換了個torch.cuda改叫torch.dcu嗎實際上也差不多但有幾個細節(jié)必須提前知道。首先是 PyTorch 的安裝包不能從官方 PyTorch 網站直接pip install torch拉得用海光官方提供的 DTK 配套版本或者從源碼去編。其次是環(huán)境變量、設備名稱這些都要做小規(guī)模調整比如設置HCCL相關環(huán)境變量、確認torch.version.hip是否存在等。把這些前置知識理清楚后面跑起來會順很多。1.2 微調任務的整體流程拆解不管你是要做 LLM 指令微調、圖像分類模型微調還是目標檢測模型微調整體流程其實都長一個樣第一步搭環(huán)境裝驅動、裝 DTK、裝適配 PyTorch。第二步準備數據和模型把預訓練權重加載進來把你的數據集處理好封裝成 DataLoader。第三步寫訓練循環(huán)定義損失函數、優(yōu)化器、學習率調度器寫標準的 forward/backward/step 邏輯。第四步調試訓練看 loss 是否下降監(jiān)控顯存占用和日志輸出遇到問題逐步排查。這篇博文會按照我實際操作的順序來寫每一步都會給出可以直接套用的命令和代碼。我不會有任何保留凡是踩過的坑、試過沒效果的操作、最后證明有用的方案通通都拿出來講。2. 裝環(huán)境從 DTK 到 PyTorch 一次性配好2.1 先確認硬件和驅動狀態(tài)環(huán)境配置前建議先用命令確認一下 DCU 芯片類型、驅動版本和 DTK 版本。這一步很多人會跳過但后面一旦出現不明所以的報錯回頭再查會發(fā)現就是版本不匹配引起的。# 查看 DCU 設備信息在部分老版本 DTK 上命令不同 dcu-smi info # 或使用類 ROCm 的查詢命令 rocm-smi --showproductname如果dcu-smi或rocm-smi都不存在大概率是驅動沒裝或者 DTK 沒裝完整。注意DCU 驅動和 DTK 不是同一個東西驅動是內核態(tài)的DTK 是用戶態(tài)的編譯器和運行時庫。兩者版本必須匹配這個在官方文檔里有對應關系表建議嚴格按照文檔來。2.2 安裝 DTK 和配置環(huán)境變量海光 DTK 的安裝方式比較接近 ROCm 的安裝邏輯一般通過離線安裝包或者 apt 源安裝。下面給出一個常見的安裝流程# 以 root 身份執(zhí)行具體包名以實際版本為準 wget DTK 安裝包下載路徑 tar -xzf dtk版本號.tar.gz cd dtk版本號 ./install.sh安裝完成之后通常需要手動設置環(huán)境變量export PATH/opt/dtk/bin:$PATH export LD_LIBRARY_PATH/opt/dtk/lib:/opt/dtk/lib64:$LD_LIBRARY_PATH我把這個寫進/etc/profile.d/dtk.sh里這樣每次 SSH 登錄進來就不用重新設了。實測下來不寫進去的話下次開終端很容易出現libamdhip64.so: cannot open shared object file這種經典報錯非常煩人。提示建議在~/.bashrc里也加一份一樣的導出并執(zhí)行source ~/.bashrc。因為部分任務管理工具如 slurm 之類在分配節(jié)點時不會讀取/etc/profile.d下的文件提前在用戶級配置里補一份能省掉很多環(huán)境不一致的問題。2.3 安裝適配 DCU 的 PyTorch這里是最容易翻車的地方。直接pip install torch裝出來的 PyTorch 是 CUDA 版本在 DCU 上根本跑不起來因為 PyTorch 里編譯進去的底層運行時是 CUDA 的libcudart、cublas那一套。DCU 上需要的是用 HIP/ROCm 或 DTK 后端編譯的 PyTorch。海光官方通常會提供編譯好的 wheel 包在官網的社區(qū)版下載區(qū)就能找到。核心命令是pip install torch版本號dtk版本號 -f 海光官方 wheel 源地址我在實際安裝時還順手裝了torchvision和transformers如果做 NLP 微調的話會用到。這里有一個血淚教訓torchvision的版本必須跟torch嚴格匹配否則torchvision內部調用torch的算子接口會對不上輕則警告重則segmentation fault。為了驗證 PyTorch 安裝是否正常我用一個極簡代碼做了冒煙測試import torch print(torch.__version__) print(torch.cuda.is_available()) # 兼容層下這里會返回 True print(torch.version.hip) # 有輸出說明編譯了 HIP 后端 if torch.cuda.is_available(): x torch.randn(4, 4).cuda() y torch.matmul(x, x) print(y.shape)如果能看到torch.__version__帶有dtk標識、torch.version.hip有值說明環(huán)境已經通了。這一步一定要確保通過再往下走否則后面怎么調都是白費功夫。2.4 Windows 上裝深度學習環(huán)境的特別說明順便提一下有網友在「在windows系統(tǒng)裝深度學習環(huán)境」這個話題下問過我Windows 上面能不能裝 DCU 的深度學習環(huán)境就我的實測來看海光官方主推的是 Linux 環(huán)境主要是歐拉、麒麟、Ubuntu 這些發(fā)行版Windows 下雖然可以通過 WSL 曲線救國但因為驅動和 DTK 的適配問題并不推薦。如果你只是想在 Windows 上先寫寫代碼、調試邏輯可以本地裝個 CPU 版本的 PyTorch把代碼跑通再上傳到 DCU 服務器上跑真正的訓練。這個方法特別適合沒有 DCU 開發(fā)機但又想提前準備代碼的情況。3. 微調代碼的快速落地以 BERT 為例3.1 準備數據集和預訓練模型我用的是文本分類微調任務數據集是一批客服對話文本打上「咨詢」「投訴」「感謝」之類的標簽。為了快速驗證我直接用 HuggingFace 的datasets庫加載了一個小規(guī)模公開數據集免得在一開始就陷入數據清洗的泥潭。from datasets import load_dataset dataset load_dataset(imdb, splittrain[:5000]) test_dataset load_dataset(imdb, splittest[:1000])預訓練模型用了bert-base-uncased。在海光 DCU 上加載這個模型沒什么特殊之處只要transformers庫裝好了就行。但我建議在訓練前顯式把模型放到設備上import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification model_name bert-base-uncased tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForSequenceClassification.from_pretrained(model_name, num_labels2) model.to(cuda) # 在 DCU 兼容 ROCm 的軟件棧上這里就是 DCU這里必須說清楚DCU 的軟件棧在 PyTorch 里暴露出來的設備名依然是cuda所以代碼層面跟 CUDA 是統(tǒng)一的。這意味著很多現成的腳本可以原封不動地跑起來只要你沒有碰底層 CUDA API 或者自定義的 kernel。3.2 數據預處理與 DataLoader 封裝數據預處理這一塊最常見的坑是分詞器在 CPU 上處理還好但 batch 拼裝時 padding 沒處理好導致 loss 波動巨大。我通常會在DataLoader里指定一個collate_fn把 padding、attention mask 這些邏輯都封裝進去。from torch.utils.data import DataLoader from transformers import DataCollatorWithPadding data_collator DataCollatorWithPadding(tokenizertokenizer) train_loader DataLoader( dataset, batch_size32, shuffleTrue, collate_fndata_collator, )DataCollatorWithPadding會自動把同一 batch 內的句子 padding 到相同長度并生成attention_mask。這一步看似簡單但直接影響訓練穩(wěn)定性和 loss 數值扎實處理好比后面調半天學習率都管用。3.3 訓練循環(huán)的寫法與設備切換DCU 上的訓練循環(huán)跟 CUDA 上幾乎一模一樣核心代碼就是標準的三板斧。不過我要特別提醒一個問題數據放到設備上這一步驟不要只搬input_ids把attention_mask忘在 CPU 上。忘記搬 mask 的話前向計算會報錯或者產生非常詭異的 loss 值看起來很像是模型出了問題其實只是數據放置不對。from transformers import AdamW optimizer AdamW(model.parameters(), lr2e-5) model.train() for epoch in range(3): total_loss 0 for step, batch in enumerate(train_loader): batch {k: v.to(cuda) for k, v in batch.items()} outputs model(**batch) loss outputs.loss loss.backward() optimizer.step() optimizer.zero_grad() total_loss loss.item() if (step 1) % 50 0: print(fEpoch {epoch}, Step {step 1}, Loss: {total_loss / (step 1):.4f})在optimizer.zero_grad()之前注意別漏了loss.backward()這個順序要是倒了整個梯度計算就亂了。實際跑的時候如果把optimizer.step()寫在了loss.backward()之前l(fā)oss 會卡在同一個值上不去這是最常見的低水平錯誤我見過不下三次。4. 從 loss 不降到 loss 下降調試過程全記錄4.1 我第一次跑出來的 loss 曲線有多離譜說實話我第一次跑微調時loss 在 0.7 左右晃了足足兩個 epoch紋絲不動。這個時候腦子里的第一反應是是不是 DCU 的算子有問題是不是精度不對是不是驅動版本不行后來冷靜下來逐個排查才發(fā)現根本不是硬件的問題而是我的學習率沒設置對。BERT 微調的學習率常規(guī)范圍是2e-5到5e-5我之前習慣性地填了個1e-3這個學習率在從頭訓練 CNN 時可能沒啥問題但在 BERT 這種大模型上微調步子邁太大loss 極易震蕩甚至發(fā)散。改成2e-5之后再跑loss 在第一個 epoch 內就從 0.7 降到了 0.35 左右訓練曲線才算恢復正常。4.2 loss 不降的三板斧排查法如果你們也遇到 loss 掉不下去或者震蕩別急著懷疑 DCU先按這個順序自查學習率是否正確Transformer 類模型微調一般用1e-5到5e-5太大容易爆太小學不動。數據預處理是否正確檢查標簽有沒有對齊、padding 有沒有讓模型看到過多的 pad token、數據有沒有混入臟數據。模型是否真的在訓練模式調用了model.eval()忘了切回model.train()會導致 BN/Dropout 行為不對loss 和指標都會有偏差。4.3 loss 出現 NaN 怎么處理第一次跑 NLP 任務時我還碰到過一次lossnan的瞬間。排查下來發(fā)現是數據里有一條極端長的文本分詞后長度超過 512模型在截斷時遇到了一些異常 token。解決方案很簡單在DataLoader迭代時加一層try-except把壞數據打印出來然后跳過而不是讓整個訓練進程直接掛掉。for step, batch in enumerate(train_loader): try: batch {k: v.to(cuda) for k, v in batch.items()} outputs model(**batch) loss outputs.loss except Exception as e: print(fSkipping bad batch at step {step}: {e}) continue這種做法在調試期非常實用但生產級別的訓練里我更推薦直接做數據清洗從源頭掐斷臟數據。5. 常見問題速查與獨家心得5.1 DCU 環(huán)境問題排查速查表癥狀可能原因解決方案torch.cuda.is_available()為 FalsePyTorch 裝錯版本重新安裝帶 dtk 標識的 PyTorch啟動訓練時Segmentation faulttorch 與 torchvision 版本不匹配嚴格配對版本或卸載 torchvision 重裝DCU out of memorybatch_size 過大調小 batch_size或使用梯度累積第一次迭代特別慢算子編譯中 / 緩存未生效多跑幾次或者設置HIP_CACHE_PATH指向高速磁盤訓練時 CPU 占用 100%數據加載瓶頸增加num_workers或啟用pin_memoryTrue5.2 DCU 上的顯存優(yōu)化心得DCU 的顯存管理與 CUDA 類似一旦 OOM最直接的辦法是減 batch size。但有些時候 batch size 不能無限小太小會導致 BN 統(tǒng)計不穩(wěn)定、訓練收斂慢。我慣用的方案是「梯度累積」本質上就是用時間換空間把一次大 batch 的梯度拆成幾次小 batch 累加起來再統(tǒng)一做參數更新。accumulation_steps 4 optimizer.zero_grad() for step, batch in enumerate(train_loader): batch {k: v.to(cuda) for k, v in batch.items()} outputs model(**batch) loss outputs.loss / accumulation_steps # 歸一化 loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()這里要特別注意的是除以accumulation_steps那一步。如果不除累積之后的總梯度相當于放大了accumulation_steps倍會導致學習率失控loss 直接飛掉。5.3 power loss 自動關閉的教訓網上有個熱詞叫power loss 自動關閉我最初沒在意直到有次訓練跑到一半機器突然斷電關機所有進度丟失才意識到問題有多嚴重。在電力和機房環(huán)境不穩(wěn)定的時候訓練前務必看一下電源策略有條件的話建議配 UPS而不是一味地盯著代碼層面做斷點續(xù)訓。斷點續(xù)訓當然也要做給代碼加個模型和優(yōu)化器的定期保存能讓你在斷電之后從最近一次 checkpoint 恢復而不是從頭再來if step % 500 0: torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, step: step, }, checkpoint.pt)恢復的時候也很簡單checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])還有一個值得留意的點torch.save和torch.load在 DCU 上默認保存到 CPU 顯存、加載時也先落到 CPU如果直接加載到cuda上偶爾會碰到設備不匹配的報錯。穩(wěn)妥的做法是加載時加一個map_location參數傳到 CPU 再手動to(cuda)不要嫌多這一步。6. 寫在最后的幾條經驗整個第一輪微調任務跑通之后我有幾個比較深的體會。第一DCU 沒有想象中那么「特殊」它在軟件生態(tài)上已經盡量向 CUDA 兼容很多坑其實是通用的不是你換了加速卡才會遇到。第二裝環(huán)境階段遇到的問題九成都是版本不匹配引起的盯緊 torch、torchvision、DTK 三者的對應關系能避開大部分雷區(qū)。第三loss 不降的時候先別懷疑硬件先把學習率、數據、模型模式這三板斧檢查完再往深處查。最后再分享一個小技巧給代碼加日志的時候除了記錄 loss一定要記錄當前學習率和顯存占用。DCU 上可以用torch.cuda.memory_allocated()實時打印顯存這樣一旦發(fā)生 OOM 或者 loss 突變你翻日志就能快速定位是哪一步出的問題不用滿頭大汗地重新設斷點去那么一條條跑。如果你也正在海光 DCU 上跑自己的第一個微調任務希望這篇總結能幫你少走點彎路。有問題可以留言交流我看到了會回復。