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