化學(xué)習(xí)訓(xùn)練)
最近在昇騰NPU上折騰 vLLM-ascend 配合 GRPO 強(qiáng)化學(xué)習(xí)訓(xùn)練也就是很多人常掛在嘴邊的“GPRO”訓(xùn)練。實(shí)際上 TRL 框架里對應(yīng)的是 GRPO群體相對策略優(yōu)化。這段鏈路跑通之后采樣速度比直接用原生生成方法快了五六倍整個訓(xùn)練迭代節(jié)奏完全不一樣了。先說明這套東西是干什么用的一個 LLM 在 SFT 之后通常還要用強(qiáng)化學(xué)習(xí)做一輪偏好對齊。GRPO 這類輕量算法不需要像 PPO 那樣額外維護(hù)價值網(wǎng)絡(luò)而是通過對同一個 prompt 采樣一組回答用組內(nèi)相對獎勵來更新策略。TRL 已經(jīng)實(shí)現(xiàn)了 GRPO 的訓(xùn)練流程但默認(rèn)采樣還是走 transformers 的 generate。對于大規(guī)模 RL 任務(wù)采樣本身是最大的開銷所以這里引入 vLLM-ascend 作為生成引擎利用它的連續(xù)批處理和 PagedAttention 來加速采樣。這篇內(nèi)容主要面向兩類人一是想在昇騰 NPU 上跑 LLM 對齊訓(xùn)練的同學(xué)二是在其他國產(chǎn)芯片上搞推理優(yōu)化、需要理解 vLLM 與訓(xùn)練框架如何銜接的朋友。我會把環(huán)境搭建、版本匹配、訓(xùn)練代碼、模型部署、性能調(diào)優(yōu)和踩坑記錄都講清楚盡量給出能直接照抄的配置。1. 整體方案設(shè)計與技術(shù)選型1.1 這條鏈路到底要解決什么問題在昇騰 NPU 上做 LLM 強(qiáng)化學(xué)習(xí)常規(guī)思路是直接用 PyTorch transformers 訓(xùn)練采樣也靠 transformers 的 generate 方法。但 GRPO 有個特點(diǎn)每個 prompt 要采樣多組輸出如果組大小 n 設(shè)為 4意味著模型離線生成的數(shù)量是訓(xùn)練更新數(shù)量的 4 倍。如果采樣過程太慢訓(xùn)練會一直等生成NPU 利用率上不去。同時transformers 原生的 generate 在 NPU 上對動態(tài) batch 和長序列的支持沒有專門優(yōu)化很容易出現(xiàn)顯存碎片和低吞吐。vLLM-ascend 是 vLLM 在昇騰上的適配版本它把連續(xù)批處理和 paged KV cache 帶到了 NPU。這個方案要解決的就是在單卡 NPU 上把“采樣”和“更新”兩個環(huán)節(jié)串起來互不拖后腿。實(shí)話說一開始我把問題想簡單了。我以為只要把 TRL 的 device 參數(shù)改成 npu 就能跑結(jié)果發(fā)現(xiàn) TRL 的生成后端還是走 transformers 的接口采樣速度根本沒上去。直到我把use_vllmTrue打開并把vllm_device指到npu:0才真正把 vLLM-ascend 的能力接進(jìn)訓(xùn)練流程。1.2 為什么選 TRL 的 GRPO 而不是其他方案RLHF 類方案里PPO、DPO、GRPO 是最常被提到的三個。PPO 需要額外的 critic 模型除了策略模型還要在每一步估計價值函數(shù)顯存占用高在昇騰上的算子適配也更復(fù)雜。DPO 雖然簡單但它是離線偏好優(yōu)化沒有在線探索通常放在 RLHF 之前做輕量對齊。GRPO 介于兩者之間它在線采樣但沒有 critic 模型顯存占用顯著降低訓(xùn)練穩(wěn)定性也不錯。我最終選 TRL 的 GRPOTrainer還有一個現(xiàn)實(shí)考慮它把 vLLM 生成后端集成得比較干凈。在GRPOConfig里開啟use_vllmTruetrainer 內(nèi)部就會創(chuàng)建一個 vLLM 引擎來采樣。這個設(shè)計讓我不需要魔改訓(xùn)練主循環(huán)只要把 vLLM 的 device 指到 NPU并把模型路徑傳給 vLLM 就行。相比自己寫一個采樣進(jìn)程再和訓(xùn)練進(jìn)程通信這個方案更省事。當(dāng)然TRL 的 GRPOTrainer 也不是沒有缺點(diǎn)。它內(nèi)部封裝了很多細(xì)節(jié)報錯信息比較抽象一旦 vLLM 和 trainer 之間的接口版本不匹配問題很難定位。所以選擇它之前必須接受“版本鎖死”這件事。1.3 vLLM-ascend 在訓(xùn)練鏈路里的定位一個完整 GRPO step 包含采樣、打分、更新三步。vLLM-ascend 只負(fù)責(zé)第一步采樣。它會常駐在訓(xùn)練進(jìn)程內(nèi)接收 prompt 列表返回采樣完的文本列表。TRL 拿到這些文本后再交給獎勵函數(shù)打分。vLLM-ascend 帶來的提升主要來自三個方面連續(xù)批處理能動態(tài)拼接長短不同的請求減少等待浪費(fèi)PagedAttention 把 KV cache 拆成固定大小的 page顯存利用率比原生 generate 高很多昇騰上有針對 flash attention 和全連接算子的融合優(yōu)化長序列生成時開銷更穩(wěn)定。在訓(xùn)練鏈路里vLLM-ascend 不是額外起一個服務(wù)進(jìn)程而是以庫的形式被 import 進(jìn)來。這么做的好處是數(shù)據(jù)和參數(shù)都在同一個 Python 進(jìn)程內(nèi)交互不需要把文本序列化成 JSON 走網(wǎng)絡(luò)省掉通信開銷。訓(xùn)練進(jìn)程和 vLLM 引擎共享同一塊 NPU 顯存這是它最有價值也最容易翻車的地方。1.4 版本匹配的教訓(xùn)這條鏈路最大的坑就是版本組合。vLLM-ascend 是緊接著某個 vLLM 版本發(fā)布適配包的TRL 內(nèi)部對 vLLM 的調(diào)用接口又會隨版本變化。我一開始裝的是比較新的 TRL 開發(fā)版結(jié)果它對 vLLM 的初始化參數(shù)變了vLLM-ascend 還沒同步導(dǎo)致啟動即報錯。這個問題的本質(zhì)是三方項(xiàng)目版本聯(lián)動的滯后。要等 TRL 先適配新參數(shù)然后 vLLM-ascend 再跟上來中間往往有窗口期。所以版本不是越新越好穩(wěn)定才是第一優(yōu)先級。我建議固定一個“經(jīng)過驗(yàn)證的組合”。比如Python 3.10 torch 2.1 對應(yīng)版本的 torch_npu CANN 7.x vLLM-ascend 某穩(wěn)定版 TRL 0.10 到 0.14 之間。在跑真實(shí)訓(xùn)練前先用一個小模型跑 20 步小實(shí)驗(yàn)確認(rèn)生成和訓(xùn)練都沒問題再上大規(guī)模任務(wù)。版本沖撞的時候優(yōu)先看 vLLM-ascend 的 release notes它支持的 vLLM 版本范圍是硬邊界。2. 昇騰NPU上的環(huán)境搭建與驗(yàn)證2.1 基礎(chǔ)環(huán)境準(zhǔn)備CANN、torch_npu和Python環(huán)境昇騰 NPU 上的 PyTorch 應(yīng)用依賴兩層底層組件CANN 提供算子庫和 runtimetorch_npu 是 PyTorch 的昇騰擴(kuò)展插件。安裝順序不能亂先裝 CANN再裝 torch_npu。如果你用的是 Docker 鏡像從官方鏡像啟動會比較快但版本也要和宿主機(jī)驅(qū)動匹配。我習(xí)慣用 conda 創(chuàng)建一個專門的 Python 3.10 環(huán)境不污染系統(tǒng)環(huán)境。依賴裝完后先做一個最小驗(yàn)證python -c import torch, torch_npu; print(torch.npu.device_count())輸出不為 0說明驅(qū)動和插件基本可用。然后手動把一個隨機(jī)模型放到npu:0上跑一次 forward 和 backward確認(rèn)梯度能正?；貍?。這一步雖然基礎(chǔ)但能避免后續(xù)把環(huán)境問題誤判成代碼問題。我見過不少人在同一個腳本里既跑 transformers 又跑 vLLM最后報錯是 torch_npu 和 PyTorch 版本沖突白白排查一整天。2.2 安裝 vLLM-ascend 與 TRLvLLM-ascend 的安裝有兩種方式預(yù)編譯 wheel 包和源碼編譯。如果你的 CANN 版本和 Python 版本都在它支持的列表里優(yōu)先用預(yù)編譯包省時間。源碼編譯需要 cmake、gcc 等工具鏈時間較長但遇到 wheel 沒有對應(yīng)架構(gòu)時就得走這條路。pip install vllm-ascendTRL 的安裝也很簡單但版本一定要提前鎖定。我在實(shí)驗(yàn)中用穩(wěn)定版然后強(qiáng)制把 transformers 退回到一個和 torch_npu 兼容的版本pip install trl0.10,0.15 pip install transformers4.40.0為什么這么在意 transformers因?yàn)?TRL 依賴 transformers而 transformers 新版本可能會調(diào)用 torch_npu 還沒實(shí)現(xiàn)的新接口。我遇到過幾次莫名其妙的 SEGV最后定位都是 transformers 版本太新。所以鎖版本不是保守是在給后續(xù)調(diào)試減少變量。2.3 用一個小模型驗(yàn)證鏈路環(huán)境裝好后先別急著寫 GRPO 訓(xùn)練代碼而是先用一個小模型單獨(dú)驗(yàn)證 vLLM-ascend 能否正常加載和生成。我用的是某個開源 2B instruct 模型tokenizer 和權(quán)重都從本地目錄讀取from vllm import LLM, SamplingParams llm LLM( modelpath/to/2b-model, devicenpu, gpu_memory_utilization0.3, max_model_len2048, trust_remote_codeTrue, ) params SamplingParams(n4, max_tokens256, temperature0.6) outputs llm.generate([請用一兩句話解釋什么是相對策略優(yōu)化], params) for o in outputs: for out in o.outputs: print(out.text)這里n4是 GRPO 里最常用的做法對同一個 prompt 采樣 4 條回答。如果這段代碼能在 NPU 上順利跑完說明 vLLM-ascend 的生成引擎沒問題。這一步最容易翻車的地方是device參數(shù)。vLLM-ascend 早期版本要求寫devicenpu新版又改成了通過環(huán)境變量VLLM_TARGET_DEVICEnpu來聲明。不同版本行為不一樣一定以你本地版本的文檔為準(zhǔn)。如果兩個都寫上通常會報重復(fù)配置。2.4 環(huán)境變量與運(yùn)行時配置昇騰環(huán)境下有幾個環(huán)境變量會影響整條鏈路。ASCEND_RT_VISIBLE_DEVICES用于指定當(dāng)前 Python 進(jìn)程可見的 NPU 卡號類似 CUDA_VISIBLE_DEVICES。PYTORCH_NPU_ALLOC_CONF可以控制顯存分配策略例如expandable_segments:True能減少碎片化。在啟動訓(xùn)練腳本之前我通常會在 shell 里做一套統(tǒng)一設(shè)置export ASCEND_RT_VISIBLE_DEVICES0 export PYTORCH_NPU_ALLOC_CONFexpandable_segments:True export VLLM_TARGET_DEVICEnpuexpandable_segments:True這個參數(shù)我用下來比較關(guān)鍵。GRPO 訓(xùn)練時張量頻繁申請和釋放顯存會逐漸碎片化最終導(dǎo)致明明還有富余卻申請失敗。開啟后頁面碎片減少但也會帶來一點(diǎn)額外內(nèi)存開銷需要自己權(quán)衡。3. GRPO訓(xùn)練與vLLM-ascend的聯(lián)合實(shí)現(xiàn)3.1 準(zhǔn)備訓(xùn)練數(shù)據(jù)與獎勵函數(shù)GRPO 訓(xùn)練首先需要一組訓(xùn)練數(shù)據(jù)格式可以很簡單每條數(shù)據(jù)包含一個 prompt 和對應(yīng)的參考答案。獎勵函數(shù)我一般先用規(guī)則版因?yàn)樵谝粋€開源 2B 模型上驗(yàn)證鏈路不需要加載額外的獎勵模型跑通了再換真實(shí) RM。示例獎勵函數(shù)判斷回答長度是否在合理區(qū)間以及是否包含關(guān)鍵標(biāo)志詞。獎勵分?jǐn)?shù)取 0 到 1 之間的浮點(diǎn)數(shù)然后 GRPO 會在組內(nèi)做歸一化。from datasets import Dataset train_data [ {prompt: 一列火車每小時行駛120公里行駛360公里需要多少小時, answer: 3小時}, {prompt: 如果一個矩形的長是6米寬是4米面積是多少, answer: 24平方米}, ] dataset Dataset.from_list(train_data) def reward_func(prompts, completions, **kwargs): rewards [] golden kwargs.get(answer) for completion in completions: score 0.0 if 10 len(completion) 200: score 0.5 if golden and golden in completion: score 0.5 rewards.append(score) return rewards有幾個細(xì)節(jié)容易被忽略獎勵函數(shù)接收的prompts和completions都是 list必須返回同樣長度的 list。如果用了kwargs.get(answer)訓(xùn)練數(shù)據(jù)里必須有answer字段并且獎勵函數(shù)要在初始化時允許接收額外參數(shù)。3.2 初始化策略模型、參考模型和獎勵模型GRPO 里用于更新的策略模型和用于約束結(jié)果的參考模型初始權(quán)重是同一個模型所以代碼里要加載兩份權(quán)重。昇騰上加載 transformers 模型時我把它們放到 NPU 上from transformers import AutoModelForCausalLM, AutoTokenizer model_path path/to/2b-model tokenizer AutoTokenizer.from_pretrained(model_path, trust_remote_codeTrue) policy_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue ) ref_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue )這里有個容易忽略的點(diǎn)TRL 內(nèi)部更新策略模型時會保留一份“凍結(jié)”的參考模型所以即使初始權(quán)重一樣兩個對象必須獨(dú)立不能用同一個對象反復(fù)賦值否則參考模型會被意外更新。如果你的 NPU 顯存比較緊張參考模型可以包在torch.inference_mode(modeTrue)里運(yùn)行減少梯度相關(guān)開銷。不過 TRL 的 GRPOTrainer 內(nèi)部也會對 ref_model 做凍結(jié)處理所以大部分情況直接用就行。3.3 配置 GRPOTrainer 和 vLLM 生成后端接下來是核心配置。TRL 的GRPOConfig支持很多參數(shù)這里只列重點(diǎn)from trl import GRPOConfig, GRPOTrainer training_args GRPOConfig( output_dir./grpo_ascend_logs, per_device_train_batch_size4, gradient_accumulation_steps2, max_steps80, learning_rate5e-6, adam_beta10.9, adam_beta20.99, warmup_steps5, use_vllmTrue, vllm_devicenpu:0, vllm_gpu_memory_utilization0.35, vllm_max_model_len2048, vllm_dtypeauto, logging_steps1, ) trainer GRPOTrainer( modelpolicy_model, ref_modelref_model, reward_funcsreward_func, argstraining_args, train_datasetdataset, processing_classtokenizer, ) trainer.train()use_vllmTrue這一步就是讓 TRL 不再走原生 generate而是創(chuàng)建一個 vLLM-ascend 引擎。vllm_device必須寫成npu:0如果寫成 cpu 或 cuda會直接報錯。vllm_gpu_memory_utilization是關(guān)鍵設(shè)太大會導(dǎo)致后面訓(xùn)練進(jìn)程 OOM設(shè)太小采樣會很慢。如果你的 TRL 版本里沒有use_vllm這個參數(shù)說明版本偏舊需要升級或者改用vllm_config的方式。升級前先看 vLLM-ascend 支持的版本范圍別盲目升到最新。3.4 訓(xùn)練循環(huán)里發(fā)生了什么采樣、獎勵歸一化、策略更新跑起來之后訓(xùn)練器每步大致做這些事從訓(xùn)練集取一個 batch 的 prompt送到 vLLM-ascend 的引擎對每個 prompt 采樣 n 條回答拿到 completions 后獎勵函數(shù)批量計算分?jǐn)?shù)然后在組內(nèi)做零均值化得到每個 token 位置的優(yōu)勢值最后用策略梯度更新策略模型。組內(nèi)歸一化是 GRPO 的核心。假設(shè)一個 prompt 采樣了 4 條回答獎勵分?jǐn)?shù)分別是 0.6、0.5、0.5、0.2歸一化后就是相對獎勵。這樣訓(xùn)練信號來自“這條回答比同組的其他回答好多少”不依賴絕對分?jǐn)?shù)也不需要額外的基線模型。這個過程里vLLM 生成的 tokens 數(shù)量會直接影響每一步的時間。我自己實(shí)測在 2B 模型、max_new_tokens256、單卡昇騰上vLLM-ascend 采樣 128 個 prompt 平均耗時約 3 秒而原生 transformers generate 大概要 16 秒。換用 vLLM 后訓(xùn)練瓶頸從生成變成了反向傳播這更符合預(yù)期。3.5 踩坑記錄vLLM與Trainer之間的設(shè)備沖突聯(lián)調(diào)階段最典型的問題是 vLLM 引擎和訓(xùn)練進(jìn)程出現(xiàn)設(shè)備爭搶。vLLM-ascend 啟動時會申請一部分 NPU 顯存作為 KV cache這部分顯存不會立刻釋放。如果訓(xùn)練時再加載優(yōu)化器顯存很容易超。我第一次跑的時候把vllm_gpu_memory_utilization設(shè)成 0.6結(jié)果訓(xùn)練到第 3 步直接 OOM。后來調(diào)成 0.35并把訓(xùn)練 batch 從 8 改成 4才穩(wěn)定下來。另一個坑是 tokenizer 不一致。TRL trainer 用自己的 tokenizer 對 prompt 做 encodevLLM 生成時也會有自己的 tokenizer 配置。如果兩個 tokenizer 并不是同一個 checkpoint 的同一個版本生成結(jié)果里會出現(xiàn)亂 token。解決辦法是顯式給 GRPOTrainer 傳入和 vLLM 同樣的 tokenizer 路徑并確保 trust_remote_code 一致。4. 訓(xùn)練產(chǎn)物部署從checkpoint到推理服務(wù)4.1 checkpoint的保存與導(dǎo)出訓(xùn)練結(jié)束后TRL 會在輸出目錄保存 checkpoint。如果想用 vLLM-ascend 做線上推理最好把權(quán)重導(dǎo)出成 vLLM 能直接加載的格式比如 safetensors并確保 tokenizer 文件在同一目錄。通常trainer.save_model()會把模型和 tokenizer 都保存到指定目錄。保存完之后檢查一下目錄結(jié)構(gòu)確認(rèn)有config.json、tokenizer.json和*.safetensors。如果沒有 tokenizer 文件vLLM 啟動時會嘗試從 model 目錄找找不到就報錯。我一般還會單獨(dú)做一次本地加載測試用 transformers 加載保存出來的模型跑一次生成確認(rèn)權(quán)重沒有損壞。這一步雖然多花幾分鐘但能避免部署階段才發(fā)現(xiàn) checkpoint 不完整。4.2 用 vLLM-ascend 啟動推理服務(wù)線上推理和訓(xùn)練時的 vLLM 配置不一樣線上沒有訓(xùn)練進(jìn)程搶顯存可以把gpu_memory_utilization調(diào)高給 KV cache 留更多空間。我用下面的命令啟動vllm serve /path/to/saved-model \ --device npu \ --max-model-len 2048 \ --gpu-memory-utilization 0.85 \ --enable-prefix-caching \ --trust-remote-code如果希望用 Python 代碼控制服務(wù)生命周期也可以直接寫from vllm import LLM, SamplingParams llm LLM( model/path/to/saved-model, devicenpu, gpu_memory_utilization0.85, max_model_len2048, enable_prefix_cachingTrue, trust_remote_codeTrue, )啟動后可以用 curl 測試接口確認(rèn)模型能正常響應(yīng)curl http://localhost:8000/v1/completions \ -H Content-Type: application/json \ -d {model: saved-model, prompt: 你好介紹一下你自己, max_tokens: 64}4.3 部署時的性能配置與注意事項(xiàng)線上推理的顯存分配邏輯和訓(xùn)練完全不同但有幾個點(diǎn)值得專門提醒。max_model_len不能拍腦袋。它決定了 KV cache 的上限設(shè)太大會浪費(fèi)顯存設(shè)太小會把長 prompt 截斷。我會先統(tǒng)計線上請求的序列長度分布取 P99 值加上生成的最大 tokens 數(shù)再留一點(diǎn)余量。max_num_seqs控制同一時刻能并發(fā)幾條序列。這個值太大會導(dǎo)致 GPU 算不過來反而增加排隊延遲太小又浪費(fèi)顯存。建議從小往大調(diào)觀察請求的 end-to-end 時延。enable_prefix_caching對多輪對話場景很有效。當(dāng)用戶發(fā)新消息時歷史消息的前綴 KV cache 可以復(fù)用減少重復(fù)計算。但在純單輪生成場景中收益不明顯。5. 性能優(yōu)化與問題排查實(shí)操5.1 先量化瓶頸生成吞吐和訓(xùn)練吞吐怎么測性能調(diào)優(yōu)不能靠感覺建議先把兩個指標(biāo)拆開生成吞吐tokens/s和訓(xùn)練吞吐samples/s。生成吞吐可以通過 vLLM 的 log 直接看或者用llm.generate跑一批固定長度的生成長度然后用總 token 數(shù)除以耗時。訓(xùn)練吞吐可以用trainer.state里的 global_step 和運(yùn)行時間來算或者直接盯日志里的step打印。我一般會單獨(dú)測兩步先測 vLLM 生成 200 條 completion 的耗時再測訓(xùn)練端一個 step 更新權(quán)重的耗時。誰的耗時高瓶頸就在誰身上。舉個例子如果生成 200 條 completion 花了 8 秒而一個 step 更新權(quán)重只花 2 秒那生成端就是瓶頸優(yōu)先調(diào) vLLM 參數(shù)。反過來如果生成只要 1 秒但 step 更新要 5 秒那就要從訓(xùn)練側(cè)優(yōu)化比如梯度累積和混合精度。5.2 生成端調(diào)優(yōu)max_model_len、gpu_memory_utilization、前綴緩存生成端最重要的參數(shù)是max_model_len。它決定了 KV cache 能覆蓋的最大序列長度。我的做法是先統(tǒng)計訓(xùn)練數(shù)據(jù)里最長的 prompt 長度加上 max_tokens取 1.1 倍余量再向上取整。gpu_memory_utilization則是給 vLLM 劃多少顯存。在訓(xùn)練鏈路里它不是越高越好。我最終的取值是 0.35相比默認(rèn)的 0.9 會損失一部分并發(fā)能力但換來了訓(xùn)練側(cè)的穩(wěn)定。如果你的訓(xùn)練 batch 很小可以適當(dāng)提高到 0.5。vLLM 的 prefix caching 也很值得開。GRPO 里同一個 prompt 會反復(fù)采樣多個 completion前面前綴的 KV cache 可以復(fù)用。開啟之后在固定 prompt 集合的反復(fù)采樣場景下生成延遲能再降 20% 左右。對應(yīng) vLLM 啟動參數(shù)是--enable-prefix-caching在 GRPOConfig 里可以傳一個 dict 到vllm_extra_args。5.3 訓(xùn)練端調(diào)優(yōu)批量大小、梯度累積、混合精度訓(xùn)練端一個常見誤區(qū)是把per_device_train_batch_size調(diào)得很大來提升吞吐。但在 GRPO 訓(xùn)練里batch 大會導(dǎo)致 vLLM 一次性要生成大量 completions顯存壓力反而更大。建議保持 batch 在 4 到 8 之間再用gradient_accumulation_steps把有效 batch 撐大?；旌暇壬蠒N騰 NPU 對 FP16 和 BF16 支持比較完整。策略更新階段用 BF16 比 FP16 更穩(wěn)避免 loss 突然 NaN。設(shè)置 transformers 的torch_dtypebfloat16同時把GRPOConfig.bf16True打開。如果某些算子不支持 BF16可以給關(guān)鍵模塊單獨(dú)保持 FP16但這會增加代碼復(fù)雜度。還有一個容易忽略的點(diǎn)學(xué)習(xí)率的設(shè)定。GRPO 訓(xùn)練中策略模型很容易被獎勵信號推偏學(xué)習(xí)率不要太大。我在 2B 模型上用 5e-6配合 warmup效果比較穩(wěn)。如果獎勵函數(shù)本身噪聲大可以再降到 1e-6。5.4 顯存分配策略vLLM與訓(xùn)練進(jìn)程如何共存一張 NPU 里跑兩個進(jìn)程不現(xiàn)實(shí)因?yàn)?TRL 的訓(xùn)練進(jìn)程和 vLLM 引擎在同一個進(jìn)程內(nèi)它們共享同一塊顯存。顯存分配的核心是給 vLLM 的 KV cache 留一個合理上限同時給訓(xùn)練權(quán)重、梯度、優(yōu)化器狀態(tài)預(yù)留出明確空間。以 2B 模型、BF16 單卡為例模型權(quán)重約 4GBAdamW 優(yōu)化器狀態(tài)約 8GB梯度約 4GB一共 16GB如果 NPU 總顯存是 32GB那 vLLM 最多用 10GB大約 0.3 到 0.4 的 util。這個比例不是死的要按模型大小重新算。我建議寫一個簡單的腳本統(tǒng)計每部分峰值?？梢韵汝P(guān)掉 vLLM只跑訓(xùn)練看訓(xùn)練側(cè)的顯存峰值再單獨(dú)跑 vLLM看生成側(cè)的峰值。兩邊相加小于總顯存才算留出了安全余量。很多 OOM 都是因?yàn)榈凸懒擞?xùn)練側(cè)的 AdamW 狀態(tài)這部分往往比模型權(quán)重還大。5.5 常見問題速查表現(xiàn)象可能原因解決辦法vLLM 啟動報 device 不匹配未設(shè)置devicenpu或環(huán)境變量檢查 vLLM-ascend 文檔統(tǒng)一用devicenpu訓(xùn)練到第幾步就 OOMvLLM 顯存分配太大調(diào)低vllm_gpu_memory_utilization減小訓(xùn)練 batch生成結(jié)果亂碼tokenizer 不一致給 vLLM 和 Trainer 傳同一個 tokenizer 路徑算子不支持報錯transformers 或 torch_npu 版本過新回退到 vLLM-ascend release 鎖定的版本loss 出現(xiàn) NaNFP16 溢出切換 BF16檢查 reward 是否出現(xiàn) infvLLM 和訓(xùn)練速度差距過大prefix caching 未啟用開啟緩存并復(fù)用同 prompt 的采樣顯存碎片嚴(yán)重反復(fù)調(diào)整 util 導(dǎo)致設(shè)置PYTORCH_NPU_ALLOC_CONFexpandable_segments:True部署接口響應(yīng)慢max_num_seqs 或 gpu_memory_utilization 不合理調(diào)大gpu-memory-utilization控制并發(fā)序列數(shù)這些是我實(shí)際調(diào)試中高頻出現(xiàn)的問題。如果你也在昇騰 NPU 上跑這條鏈路大概率會碰到其中幾個。6. 實(shí)測效果與優(yōu)化心得6.1 性能數(shù)據(jù)對比用某開源 2B 模型單一昇騰 NPUprompt 長度為 128max_new_tokens256GRPO 組大小 n4我這邊實(shí)測的數(shù)據(jù)大致如下配置采樣耗時128 prompt一個訓(xùn)練 step 總耗時原生 generate約 16s約 22svLLM-ascend 默認(rèn)約 5.2s約 8.1svLLM-ascend 開 prefix cache約 3.9s約 6.7s數(shù)據(jù)不算驚艷但對一個 2B 規(guī)模的實(shí)驗(yàn)來說采樣部分提速約 4 倍整個訓(xùn)練 step 提速約 3 倍。如果你的 prompt 更長、采樣數(shù)量更多提速會更明顯。6.2 可復(fù)用的優(yōu)化經(jīng)驗(yàn)第一版本鎖死比追新重要。vLLM-ascend 和 TRL 的組合每換一個版本都可能帶來接口變化我建議把整套依賴版本寫進(jìn) requirements 或 docker 鏡像不要散落在個人環(huán)境里。第二顯存分配要有預(yù)算表。先算訓(xùn)練側(cè)峰值再給 vLLM 定 util別讓 vLLM 一上來吃掉大塊顯存。你會發(fā)現(xiàn)顯存預(yù)算表比調(diào)參日志更有用。第三獎勵函數(shù)盡量先用規(guī)則版跑通鏈路再換獎勵模型。規(guī)則函數(shù)簡單、可控能快速暴露鏈路問題一上來就用大 RM遇到問題會分不清是采樣問題還是打分問題。6.3 后續(xù)擴(kuò)展與個人體會如果你打算把這個方案往更大模型和更長序列上遷移還有幾個方向可以嘗試用托管內(nèi)存做更大 KV cache把 vLLM 引擎放到獨(dú)立進(jìn)程通過跨進(jìn)程通信傳 completion這樣訓(xùn)練進(jìn)程和設(shè)備分配更干凈或者考慮多卡版本用數(shù)據(jù)并行把采樣任務(wù)切分到多個 NPU 上。我個人接下來準(zhǔn)備把獎勵模型替換成可學(xué)習(xí)的 RM并對比不同采樣組大小對訓(xùn)練穩(wěn)定性的影響。最后再分享一個小技巧無論你用什么版本組合第一步永遠(yuǎn)先跑通 100 步小實(shí)驗(yàn)不要上來就全量訓(xùn)練。小步實(shí)驗(yàn)?zāi)軒湍惆扬@存、梯度、采樣三塊都看到峰值再放大規(guī)模心里就有底了。 在昇騰NPU上用vLLM-ascend加速GRPO強(qiáng)化學(xué)習(xí)訓(xùn)練
最近在昇騰NPU上折騰 vLLM-ascend 配合 GRPO 強(qiáng)化學(xué)習(xí)訓(xùn)練也就是很多人常掛在嘴邊的“GPRO”訓(xùn)練。實(shí)際上 TRL 框架里對應(yīng)的是 GRPO群體相對策略優(yōu)化。這段鏈路跑通之后采樣速度比直接用原生生成方法快了五六倍整個訓(xùn)練迭代節(jié)奏完全不一樣了。先說明這套東西是干什么用的一個 LLM 在 SFT 之后通常還要用強(qiáng)化學(xué)習(xí)做一輪偏好對齊。GRPO 這類輕量算法不需要像 PPO 那樣額外維護(hù)價值網(wǎng)絡(luò)而是通過對同一個 prompt 采樣一組回答用組內(nèi)相對獎勵來更新策略。TRL 已經(jīng)實(shí)現(xiàn)了 GRPO 的訓(xùn)練流程但默認(rèn)采樣還是走 transformers 的 generate。對于大規(guī)模 RL 任務(wù)采樣本身是最大的開銷所以這里引入 vLLM-ascend 作為生成引擎利用它的連續(xù)批處理和 PagedAttention 來加速采樣。這篇內(nèi)容主要面向兩類人一是想在昇騰 NPU 上跑 LLM 對齊訓(xùn)練的同學(xué)二是在其他國產(chǎn)芯片上搞推理優(yōu)化、需要理解 vLLM 與訓(xùn)練框架如何銜接的朋友。我會把環(huán)境搭建、版本匹配、訓(xùn)練代碼、模型部署、性能調(diào)優(yōu)和踩坑記錄都講清楚盡量給出能直接照抄的配置。1. 整體方案設(shè)計與技術(shù)選型1.1 這條鏈路到底要解決什么問題在昇騰 NPU 上做 LLM 強(qiáng)化學(xué)習(xí)常規(guī)思路是直接用 PyTorch transformers 訓(xùn)練采樣也靠 transformers 的 generate 方法。但 GRPO 有個特點(diǎn)每個 prompt 要采樣多組輸出如果組大小 n 設(shè)為 4意味著模型離線生成的數(shù)量是訓(xùn)練更新數(shù)量的 4 倍。如果采樣過程太慢訓(xùn)練會一直等生成NPU 利用率上不去。同時transformers 原生的 generate 在 NPU 上對動態(tài) batch 和長序列的支持沒有專門優(yōu)化很容易出現(xiàn)顯存碎片和低吞吐。vLLM-ascend 是 vLLM 在昇騰上的適配版本它把連續(xù)批處理和 paged KV cache 帶到了 NPU。這個方案要解決的就是在單卡 NPU 上把“采樣”和“更新”兩個環(huán)節(jié)串起來互不拖后腿。實(shí)話說一開始我把問題想簡單了。我以為只要把 TRL 的 device 參數(shù)改成 npu 就能跑結(jié)果發(fā)現(xiàn) TRL 的生成后端還是走 transformers 的接口采樣速度根本沒上去。直到我把use_vllmTrue打開并把vllm_device指到npu:0才真正把 vLLM-ascend 的能力接進(jìn)訓(xùn)練流程。1.2 為什么選 TRL 的 GRPO 而不是其他方案RLHF 類方案里PPO、DPO、GRPO 是最常被提到的三個。PPO 需要額外的 critic 模型除了策略模型還要在每一步估計價值函數(shù)顯存占用高在昇騰上的算子適配也更復(fù)雜。DPO 雖然簡單但它是離線偏好優(yōu)化沒有在線探索通常放在 RLHF 之前做輕量對齊。GRPO 介于兩者之間它在線采樣但沒有 critic 模型顯存占用顯著降低訓(xùn)練穩(wěn)定性也不錯。我最終選 TRL 的 GRPOTrainer還有一個現(xiàn)實(shí)考慮它把 vLLM 生成后端集成得比較干凈。在GRPOConfig里開啟use_vllmTruetrainer 內(nèi)部就會創(chuàng)建一個 vLLM 引擎來采樣。這個設(shè)計讓我不需要魔改訓(xùn)練主循環(huán)只要把 vLLM 的 device 指到 NPU并把模型路徑傳給 vLLM 就行。相比自己寫一個采樣進(jìn)程再和訓(xùn)練進(jìn)程通信這個方案更省事。當(dāng)然TRL 的 GRPOTrainer 也不是沒有缺點(diǎn)。它內(nèi)部封裝了很多細(xì)節(jié)報錯信息比較抽象一旦 vLLM 和 trainer 之間的接口版本不匹配問題很難定位。所以選擇它之前必須接受“版本鎖死”這件事。1.3 vLLM-ascend 在訓(xùn)練鏈路里的定位一個完整 GRPO step 包含采樣、打分、更新三步。vLLM-ascend 只負(fù)責(zé)第一步采樣。它會常駐在訓(xùn)練進(jìn)程內(nèi)接收 prompt 列表返回采樣完的文本列表。TRL 拿到這些文本后再交給獎勵函數(shù)打分。vLLM-ascend 帶來的提升主要來自三個方面連續(xù)批處理能動態(tài)拼接長短不同的請求減少等待浪費(fèi)PagedAttention 把 KV cache 拆成固定大小的 page顯存利用率比原生 generate 高很多昇騰上有針對 flash attention 和全連接算子的融合優(yōu)化長序列生成時開銷更穩(wěn)定。在訓(xùn)練鏈路里vLLM-ascend 不是額外起一個服務(wù)進(jìn)程而是以庫的形式被 import 進(jìn)來。這么做的好處是數(shù)據(jù)和參數(shù)都在同一個 Python 進(jìn)程內(nèi)交互不需要把文本序列化成 JSON 走網(wǎng)絡(luò)省掉通信開銷。訓(xùn)練進(jìn)程和 vLLM 引擎共享同一塊 NPU 顯存這是它最有價值也最容易翻車的地方。1.4 版本匹配的教訓(xùn)這條鏈路最大的坑就是版本組合。vLLM-ascend 是緊接著某個 vLLM 版本發(fā)布適配包的TRL 內(nèi)部對 vLLM 的調(diào)用接口又會隨版本變化。我一開始裝的是比較新的 TRL 開發(fā)版結(jié)果它對 vLLM 的初始化參數(shù)變了vLLM-ascend 還沒同步導(dǎo)致啟動即報錯。這個問題的本質(zhì)是三方項(xiàng)目版本聯(lián)動的滯后。要等 TRL 先適配新參數(shù)然后 vLLM-ascend 再跟上來中間往往有窗口期。所以版本不是越新越好穩(wěn)定才是第一優(yōu)先級。我建議固定一個“經(jīng)過驗(yàn)證的組合”。比如Python 3.10 torch 2.1 對應(yīng)版本的 torch_npu CANN 7.x vLLM-ascend 某穩(wěn)定版 TRL 0.10 到 0.14 之間。在跑真實(shí)訓(xùn)練前先用一個小模型跑 20 步小實(shí)驗(yàn)確認(rèn)生成和訓(xùn)練都沒問題再上大規(guī)模任務(wù)。版本沖撞的時候優(yōu)先看 vLLM-ascend 的 release notes它支持的 vLLM 版本范圍是硬邊界。2. 昇騰NPU上的環(huán)境搭建與驗(yàn)證2.1 基礎(chǔ)環(huán)境準(zhǔn)備CANN、torch_npu和Python環(huán)境昇騰 NPU 上的 PyTorch 應(yīng)用依賴兩層底層組件CANN 提供算子庫和 runtimetorch_npu 是 PyTorch 的昇騰擴(kuò)展插件。安裝順序不能亂先裝 CANN再裝 torch_npu。如果你用的是 Docker 鏡像從官方鏡像啟動會比較快但版本也要和宿主機(jī)驅(qū)動匹配。我習(xí)慣用 conda 創(chuàng)建一個專門的 Python 3.10 環(huán)境不污染系統(tǒng)環(huán)境。依賴裝完后先做一個最小驗(yàn)證python -c import torch, torch_npu; print(torch.npu.device_count())輸出不為 0說明驅(qū)動和插件基本可用。然后手動把一個隨機(jī)模型放到npu:0上跑一次 forward 和 backward確認(rèn)梯度能正?；貍?。這一步雖然基礎(chǔ)但能避免后續(xù)把環(huán)境問題誤判成代碼問題。我見過不少人在同一個腳本里既跑 transformers 又跑 vLLM最后報錯是 torch_npu 和 PyTorch 版本沖突白白排查一整天。2.2 安裝 vLLM-ascend 與 TRLvLLM-ascend 的安裝有兩種方式預(yù)編譯 wheel 包和源碼編譯。如果你的 CANN 版本和 Python 版本都在它支持的列表里優(yōu)先用預(yù)編譯包省時間。源碼編譯需要 cmake、gcc 等工具鏈時間較長但遇到 wheel 沒有對應(yīng)架構(gòu)時就得走這條路。pip install vllm-ascendTRL 的安裝也很簡單但版本一定要提前鎖定。我在實(shí)驗(yàn)中用穩(wěn)定版然后強(qiáng)制把 transformers 退回到一個和 torch_npu 兼容的版本pip install trl0.10,0.15 pip install transformers4.40.0為什么這么在意 transformers因?yàn)?TRL 依賴 transformers而 transformers 新版本可能會調(diào)用 torch_npu 還沒實(shí)現(xiàn)的新接口。我遇到過幾次莫名其妙的 SEGV最后定位都是 transformers 版本太新。所以鎖版本不是保守是在給后續(xù)調(diào)試減少變量。2.3 用一個小模型驗(yàn)證鏈路環(huán)境裝好后先別急著寫 GRPO 訓(xùn)練代碼而是先用一個小模型單獨(dú)驗(yàn)證 vLLM-ascend 能否正常加載和生成。我用的是某個開源 2B instruct 模型tokenizer 和權(quán)重都從本地目錄讀取from vllm import LLM, SamplingParams llm LLM( modelpath/to/2b-model, devicenpu, gpu_memory_utilization0.3, max_model_len2048, trust_remote_codeTrue, ) params SamplingParams(n4, max_tokens256, temperature0.6) outputs llm.generate([請用一兩句話解釋什么是相對策略優(yōu)化], params) for o in outputs: for out in o.outputs: print(out.text)這里n4是 GRPO 里最常用的做法對同一個 prompt 采樣 4 條回答。如果這段代碼能在 NPU 上順利跑完說明 vLLM-ascend 的生成引擎沒問題。這一步最容易翻車的地方是device參數(shù)。vLLM-ascend 早期版本要求寫devicenpu新版又改成了通過環(huán)境變量VLLM_TARGET_DEVICEnpu來聲明。不同版本行為不一樣一定以你本地版本的文檔為準(zhǔn)。如果兩個都寫上通常會報重復(fù)配置。2.4 環(huán)境變量與運(yùn)行時配置昇騰環(huán)境下有幾個環(huán)境變量會影響整條鏈路。ASCEND_RT_VISIBLE_DEVICES用于指定當(dāng)前 Python 進(jìn)程可見的 NPU 卡號類似 CUDA_VISIBLE_DEVICES。PYTORCH_NPU_ALLOC_CONF可以控制顯存分配策略例如expandable_segments:True能減少碎片化。在啟動訓(xùn)練腳本之前我通常會在 shell 里做一套統(tǒng)一設(shè)置export ASCEND_RT_VISIBLE_DEVICES0 export PYTORCH_NPU_ALLOC_CONFexpandable_segments:True export VLLM_TARGET_DEVICEnpuexpandable_segments:True這個參數(shù)我用下來比較關(guān)鍵。GRPO 訓(xùn)練時張量頻繁申請和釋放顯存會逐漸碎片化最終導(dǎo)致明明還有富余卻申請失敗。開啟后頁面碎片減少但也會帶來一點(diǎn)額外內(nèi)存開銷需要自己權(quán)衡。3. GRPO訓(xùn)練與vLLM-ascend的聯(lián)合實(shí)現(xiàn)3.1 準(zhǔn)備訓(xùn)練數(shù)據(jù)與獎勵函數(shù)GRPO 訓(xùn)練首先需要一組訓(xùn)練數(shù)據(jù)格式可以很簡單每條數(shù)據(jù)包含一個 prompt 和對應(yīng)的參考答案。獎勵函數(shù)我一般先用規(guī)則版因?yàn)樵谝粋€開源 2B 模型上驗(yàn)證鏈路不需要加載額外的獎勵模型跑通了再換真實(shí) RM。示例獎勵函數(shù)判斷回答長度是否在合理區(qū)間以及是否包含關(guān)鍵標(biāo)志詞。獎勵分?jǐn)?shù)取 0 到 1 之間的浮點(diǎn)數(shù)然后 GRPO 會在組內(nèi)做歸一化。from datasets import Dataset train_data [ {prompt: 一列火車每小時行駛120公里行駛360公里需要多少小時, answer: 3小時}, {prompt: 如果一個矩形的長是6米寬是4米面積是多少, answer: 24平方米}, ] dataset Dataset.from_list(train_data) def reward_func(prompts, completions, **kwargs): rewards [] golden kwargs.get(answer) for completion in completions: score 0.0 if 10 len(completion) 200: score 0.5 if golden and golden in completion: score 0.5 rewards.append(score) return rewards有幾個細(xì)節(jié)容易被忽略獎勵函數(shù)接收的prompts和completions都是 list必須返回同樣長度的 list。如果用了kwargs.get(answer)訓(xùn)練數(shù)據(jù)里必須有answer字段并且獎勵函數(shù)要在初始化時允許接收額外參數(shù)。3.2 初始化策略模型、參考模型和獎勵模型GRPO 里用于更新的策略模型和用于約束結(jié)果的參考模型初始權(quán)重是同一個模型所以代碼里要加載兩份權(quán)重。昇騰上加載 transformers 模型時我把它們放到 NPU 上from transformers import AutoModelForCausalLM, AutoTokenizer model_path path/to/2b-model tokenizer AutoTokenizer.from_pretrained(model_path, trust_remote_codeTrue) policy_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue ) ref_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue )這里有個容易忽略的點(diǎn)TRL 內(nèi)部更新策略模型時會保留一份“凍結(jié)”的參考模型所以即使初始權(quán)重一樣兩個對象必須獨(dú)立不能用同一個對象反復(fù)賦值否則參考模型會被意外更新。如果你的 NPU 顯存比較緊張參考模型可以包在torch.inference_mode(modeTrue)里運(yùn)行減少梯度相關(guān)開銷。不過 TRL 的 GRPOTrainer 內(nèi)部也會對 ref_model 做凍結(jié)處理所以大部分情況直接用就行。3.3 配置 GRPOTrainer 和 vLLM 生成后端接下來是核心配置。TRL 的GRPOConfig支持很多參數(shù)這里只列重點(diǎn)from trl import GRPOConfig, GRPOTrainer training_args GRPOConfig( output_dir./grpo_ascend_logs, per_device_train_batch_size4, gradient_accumulation_steps2, max_steps80, learning_rate5e-6, adam_beta10.9, adam_beta20.99, warmup_steps5, use_vllmTrue, vllm_devicenpu:0, vllm_gpu_memory_utilization0.35, vllm_max_model_len2048, vllm_dtypeauto, logging_steps1, ) trainer GRPOTrainer( modelpolicy_model, ref_modelref_model, reward_funcsreward_func, argstraining_args, train_datasetdataset, processing_classtokenizer, ) trainer.train()use_vllmTrue這一步就是讓 TRL 不再走原生 generate而是創(chuàng)建一個 vLLM-ascend 引擎。vllm_device必須寫成npu:0如果寫成 cpu 或 cuda會直接報錯。vllm_gpu_memory_utilization是關(guān)鍵設(shè)太大會導(dǎo)致后面訓(xùn)練進(jìn)程 OOM設(shè)太小采樣會很慢。如果你的 TRL 版本里沒有use_vllm這個參數(shù)說明版本偏舊需要升級或者改用vllm_config的方式。升級前先看 vLLM-ascend 支持的版本范圍別盲目升到最新。3.4 訓(xùn)練循環(huán)里發(fā)生了什么采樣、獎勵歸一化、策略更新跑起來之后訓(xùn)練器每步大致做這些事從訓(xùn)練集取一個 batch 的 prompt送到 vLLM-ascend 的引擎對每個 prompt 采樣 n 條回答拿到 completions 后獎勵函數(shù)批量計算分?jǐn)?shù)然后在組內(nèi)做零均值化得到每個 token 位置的優(yōu)勢值最后用策略梯度更新策略模型。組內(nèi)歸一化是 GRPO 的核心。假設(shè)一個 prompt 采樣了 4 條回答獎勵分?jǐn)?shù)分別是 0.6、0.5、0.5、0.2歸一化后就是相對獎勵。這樣訓(xùn)練信號來自“這條回答比同組的其他回答好多少”不依賴絕對分?jǐn)?shù)也不需要額外的基線模型。這個過程里vLLM 生成的 tokens 數(shù)量會直接影響每一步的時間。我自己實(shí)測在 2B 模型、max_new_tokens256、單卡昇騰上vLLM-ascend 采樣 128 個 prompt 平均耗時約 3 秒而原生 transformers generate 大概要 16 秒。換用 vLLM 后訓(xùn)練瓶頸從生成變成了反向傳播這更符合預(yù)期。3.5 踩坑記錄vLLM與Trainer之間的設(shè)備沖突聯(lián)調(diào)階段最典型的問題是 vLLM 引擎和訓(xùn)練進(jìn)程出現(xiàn)設(shè)備爭搶。vLLM-ascend 啟動時會申請一部分 NPU 顯存作為 KV cache這部分顯存不會立刻釋放。如果訓(xùn)練時再加載優(yōu)化器顯存很容易超。我第一次跑的時候把vllm_gpu_memory_utilization設(shè)成 0.6結(jié)果訓(xùn)練到第 3 步直接 OOM。后來調(diào)成 0.35并把訓(xùn)練 batch 從 8 改成 4才穩(wěn)定下來。另一個坑是 tokenizer 不一致。TRL trainer 用自己的 tokenizer 對 prompt 做 encodevLLM 生成時也會有自己的 tokenizer 配置。如果兩個 tokenizer 并不是同一個 checkpoint 的同一個版本生成結(jié)果里會出現(xiàn)亂 token。解決辦法是顯式給 GRPOTrainer 傳入和 vLLM 同樣的 tokenizer 路徑并確保 trust_remote_code 一致。4. 訓(xùn)練產(chǎn)物部署從checkpoint到推理服務(wù)4.1 checkpoint的保存與導(dǎo)出訓(xùn)練結(jié)束后TRL 會在輸出目錄保存 checkpoint。如果想用 vLLM-ascend 做線上推理最好把權(quán)重導(dǎo)出成 vLLM 能直接加載的格式比如 safetensors并確保 tokenizer 文件在同一目錄。通常trainer.save_model()會把模型和 tokenizer 都保存到指定目錄。保存完之后檢查一下目錄結(jié)構(gòu)確認(rèn)有config.json、tokenizer.json和*.safetensors。如果沒有 tokenizer 文件vLLM 啟動時會嘗試從 model 目錄找找不到就報錯。我一般還會單獨(dú)做一次本地加載測試用 transformers 加載保存出來的模型跑一次生成確認(rèn)權(quán)重沒有損壞。這一步雖然多花幾分鐘但能避免部署階段才發(fā)現(xiàn) checkpoint 不完整。4.2 用 vLLM-ascend 啟動推理服務(wù)線上推理和訓(xùn)練時的 vLLM 配置不一樣線上沒有訓(xùn)練進(jìn)程搶顯存可以把gpu_memory_utilization調(diào)高給 KV cache 留更多空間。我用下面的命令啟動vllm serve /path/to/saved-model \ --device npu \ --max-model-len 2048 \ --gpu-memory-utilization 0.85 \ --enable-prefix-caching \ --trust-remote-code如果希望用 Python 代碼控制服務(wù)生命周期也可以直接寫from vllm import LLM, SamplingParams llm LLM( model/path/to/saved-model, devicenpu, gpu_memory_utilization0.85, max_model_len2048, enable_prefix_cachingTrue, trust_remote_codeTrue, )啟動后可以用 curl 測試接口確認(rèn)模型能正常響應(yīng)curl http://localhost:8000/v1/completions \ -H Content-Type: application/json \ -d {model: saved-model, prompt: 你好介紹一下你自己, max_tokens: 64}4.3 部署時的性能配置與注意事項(xiàng)線上推理的顯存分配邏輯和訓(xùn)練完全不同但有幾個點(diǎn)值得專門提醒。max_model_len不能拍腦袋。它決定了 KV cache 的上限設(shè)太大會浪費(fèi)顯存設(shè)太小會把長 prompt 截斷。我會先統(tǒng)計線上請求的序列長度分布取 P99 值加上生成的最大 tokens 數(shù)再留一點(diǎn)余量。max_num_seqs控制同一時刻能并發(fā)幾條序列。這個值太大會導(dǎo)致 GPU 算不過來反而增加排隊延遲太小又浪費(fèi)顯存。建議從小往大調(diào)觀察請求的 end-to-end 時延。enable_prefix_caching對多輪對話場景很有效。當(dāng)用戶發(fā)新消息時歷史消息的前綴 KV cache 可以復(fù)用減少重復(fù)計算。但在純單輪生成場景中收益不明顯。5. 性能優(yōu)化與問題排查實(shí)操5.1 先量化瓶頸生成吞吐和訓(xùn)練吞吐怎么測性能調(diào)優(yōu)不能靠感覺建議先把兩個指標(biāo)拆開生成吞吐tokens/s和訓(xùn)練吞吐samples/s。生成吞吐可以通過 vLLM 的 log 直接看或者用llm.generate跑一批固定長度的生成長度然后用總 token 數(shù)除以耗時。訓(xùn)練吞吐可以用trainer.state里的 global_step 和運(yùn)行時間來算或者直接盯日志里的step打印。我一般會單獨(dú)測兩步先測 vLLM 生成 200 條 completion 的耗時再測訓(xùn)練端一個 step 更新權(quán)重的耗時。誰的耗時高瓶頸就在誰身上。舉個例子如果生成 200 條 completion 花了 8 秒而一個 step 更新權(quán)重只花 2 秒那生成端就是瓶頸優(yōu)先調(diào) vLLM 參數(shù)。反過來如果生成只要 1 秒但 step 更新要 5 秒那就要從訓(xùn)練側(cè)優(yōu)化比如梯度累積和混合精度。5.2 生成端調(diào)優(yōu)max_model_len、gpu_memory_utilization、前綴緩存生成端最重要的參數(shù)是max_model_len。它決定了 KV cache 能覆蓋的最大序列長度。我的做法是先統(tǒng)計訓(xùn)練數(shù)據(jù)里最長的 prompt 長度加上 max_tokens取 1.1 倍余量再向上取整。gpu_memory_utilization則是給 vLLM 劃多少顯存。在訓(xùn)練鏈路里它不是越高越好。我最終的取值是 0.35相比默認(rèn)的 0.9 會損失一部分并發(fā)能力但換來了訓(xùn)練側(cè)的穩(wěn)定。如果你的訓(xùn)練 batch 很小可以適當(dāng)提高到 0.5。vLLM 的 prefix caching 也很值得開。GRPO 里同一個 prompt 會反復(fù)采樣多個 completion前面前綴的 KV cache 可以復(fù)用。開啟之后在固定 prompt 集合的反復(fù)采樣場景下生成延遲能再降 20% 左右。對應(yīng) vLLM 啟動參數(shù)是--enable-prefix-caching在 GRPOConfig 里可以傳一個 dict 到vllm_extra_args。5.3 訓(xùn)練端調(diào)優(yōu)批量大小、梯度累積、混合精度訓(xùn)練端一個常見誤區(qū)是把per_device_train_batch_size調(diào)得很大來提升吞吐。但在 GRPO 訓(xùn)練里batch 大會導(dǎo)致 vLLM 一次性要生成大量 completions顯存壓力反而更大。建議保持 batch 在 4 到 8 之間再用gradient_accumulation_steps把有效 batch 撐大?；旌暇壬蠒N騰 NPU 對 FP16 和 BF16 支持比較完整。策略更新階段用 BF16 比 FP16 更穩(wěn)避免 loss 突然 NaN。設(shè)置 transformers 的torch_dtypebfloat16同時把GRPOConfig.bf16True打開。如果某些算子不支持 BF16可以給關(guān)鍵模塊單獨(dú)保持 FP16但這會增加代碼復(fù)雜度。還有一個容易忽略的點(diǎn)學(xué)習(xí)率的設(shè)定。GRPO 訓(xùn)練中策略模型很容易被獎勵信號推偏學(xué)習(xí)率不要太大。我在 2B 模型上用 5e-6配合 warmup效果比較穩(wěn)。如果獎勵函數(shù)本身噪聲大可以再降到 1e-6。5.4 顯存分配策略vLLM與訓(xùn)練進(jìn)程如何共存一張 NPU 里跑兩個進(jìn)程不現(xiàn)實(shí)因?yàn)?TRL 的訓(xùn)練進(jìn)程和 vLLM 引擎在同一個進(jìn)程內(nèi)它們共享同一塊顯存。顯存分配的核心是給 vLLM 的 KV cache 留一個合理上限同時給訓(xùn)練權(quán)重、梯度、優(yōu)化器狀態(tài)預(yù)留出明確空間。以 2B 模型、BF16 單卡為例模型權(quán)重約 4GBAdamW 優(yōu)化器狀態(tài)約 8GB梯度約 4GB一共 16GB如果 NPU 總顯存是 32GB那 vLLM 最多用 10GB大約 0.3 到 0.4 的 util。這個比例不是死的要按模型大小重新算。我建議寫一個簡單的腳本統(tǒng)計每部分峰值?？梢韵汝P(guān)掉 vLLM只跑訓(xùn)練看訓(xùn)練側(cè)的顯存峰值再單獨(dú)跑 vLLM看生成側(cè)的峰值。兩邊相加小于總顯存才算留出了安全余量。很多 OOM 都是因?yàn)榈凸懒擞?xùn)練側(cè)的 AdamW 狀態(tài)這部分往往比模型權(quán)重還大。5.5 常見問題速查表現(xiàn)象可能原因解決辦法vLLM 啟動報 device 不匹配未設(shè)置devicenpu或環(huán)境變量檢查 vLLM-ascend 文檔統(tǒng)一用devicenpu訓(xùn)練到第幾步就 OOMvLLM 顯存分配太大調(diào)低vllm_gpu_memory_utilization減小訓(xùn)練 batch生成結(jié)果亂碼tokenizer 不一致給 vLLM 和 Trainer 傳同一個 tokenizer 路徑算子不支持報錯transformers 或 torch_npu 版本過新回退到 vLLM-ascend release 鎖定的版本loss 出現(xiàn) NaNFP16 溢出切換 BF16檢查 reward 是否出現(xiàn) infvLLM 和訓(xùn)練速度差距過大prefix caching 未啟用開啟緩存并復(fù)用同 prompt 的采樣顯存碎片嚴(yán)重反復(fù)調(diào)整 util 導(dǎo)致設(shè)置PYTORCH_NPU_ALLOC_CONFexpandable_segments:True部署接口響應(yīng)慢max_num_seqs 或 gpu_memory_utilization 不合理調(diào)大gpu-memory-utilization控制并發(fā)序列數(shù)這些是我實(shí)際調(diào)試中高頻出現(xiàn)的問題。如果你也在昇騰 NPU 上跑這條鏈路大概率會碰到其中幾個。6. 實(shí)測效果與優(yōu)化心得6.1 性能數(shù)據(jù)對比用某開源 2B 模型單一昇騰 NPUprompt 長度為 128max_new_tokens256GRPO 組大小 n4我這邊實(shí)測的數(shù)據(jù)大致如下配置采樣耗時128 prompt一個訓(xùn)練 step 總耗時原生 generate約 16s約 22svLLM-ascend 默認(rèn)約 5.2s約 8.1svLLM-ascend 開 prefix cache約 3.9s約 6.7s數(shù)據(jù)不算驚艷但對一個 2B 規(guī)模的實(shí)驗(yàn)來說采樣部分提速約 4 倍整個訓(xùn)練 step 提速約 3 倍。如果你的 prompt 更長、采樣數(shù)量更多提速會更明顯。6.2 可復(fù)用的優(yōu)化經(jīng)驗(yàn)第一版本鎖死比追新重要。vLLM-ascend 和 TRL 的組合每換一個版本都可能帶來接口變化我建議把整套依賴版本寫進(jìn) requirements 或 docker 鏡像不要散落在個人環(huán)境里。第二顯存分配要有預(yù)算表。先算訓(xùn)練側(cè)峰值再給 vLLM 定 util別讓 vLLM 一上來吃掉大塊顯存。你會發(fā)現(xiàn)顯存預(yù)算表比調(diào)參日志更有用。第三獎勵函數(shù)盡量先用規(guī)則版跑通鏈路再換獎勵模型。規(guī)則函數(shù)簡單、可控能快速暴露鏈路問題一上來就用大 RM遇到問題會分不清是采樣問題還是打分問題。6.3 后續(xù)擴(kuò)展與個人體會如果你打算把這個方案往更大模型和更長序列上遷移還有幾個方向可以嘗試用托管內(nèi)存做更大 KV cache把 vLLM 引擎放到獨(dú)立進(jìn)程通過跨進(jìn)程通信傳 completion這樣訓(xùn)練進(jìn)程和設(shè)備分配更干凈或者考慮多卡版本用數(shù)據(jù)并行把采樣任務(wù)切分到多個 NPU 上。我個人接下來準(zhǔn)備把獎勵模型替換成可學(xué)習(xí)的 RM并對比不同采樣組大小對訓(xùn)練穩(wěn)定性的影響。最后再分享一個小技巧無論你用什么版本組合第一步永遠(yuǎn)先跑通 100 步小實(shí)驗(yàn)不要上來就全量訓(xùn)練。小步實(shí)驗(yàn)?zāi)軒湍惆扬@存、梯度、采樣三塊都看到峰值再放大規(guī)模心里就有底了。