領域特定嵌入微調與 NVIDIA Nemotron – 一天內完成

TL;DR: NVIDIA 與 Hugging Face 發布了一個六步驟、單 GPU 的管線，能在一天內對 10 億參數的 Llama‑Nemotron‑Embed‑1B‑v2 模型進行合成領域資料的微調，提升 Recall@10/NDCG@10 超過 10%，在真實企業資料集上最高提升 26%。

概覽 – 為何領域特定嵌入微調很重要

一般用途的嵌入模型在網路規模語意上表現優異，但缺乏合約、製造日誌或內部分類等領域所需的細緻區分。透過領域特定資料的微調可彌補此差距，提升 Retrieval‑Augmented Generation (RAG) 流程的檢索品質。全新配方自動化資料產生、硬負樣本挖掘、多跳處理、訓練、評估與部署，全部僅需一張 NVIDIA Ampere 級別的 GPU。

快速存取連結

整合開源元件

組件 角色
NeMo Data Designer 產生來自原始文件的合成問答對
NeMo Automodel 訓練雙編碼器嵌入模型
BEIR 提供標準化的評估基準
NeMo Export‑Deploy 將檢查點轉換為 ONNX/TensorRT
NVIDIA NIM 透過相容 OpenAI 的嵌入端點提供模型服務

前置條件

  • 領域文件目錄（.txt、.md 等）
  • 免費的 NVIDIA API 金鑰，取得自 https://build.nvidia.com/
  • NVIDIA Ampere 或更新的 GPU，具備 ≥ 80 GB VRAM（已在 A100‑80GB 與 H100‑80GB 測試）

步驟式管線

1️⃣ 產生合成訓練資料

此管線使用 LLM nvidia/nemotron-3-nano-30b-a3b 讀取每份文件，產出高品質的 QA 配對，無需任何人工標記。

nemotron embed sdg -c default corpus_dir=./data/my_domain_docs

四階段 SDG 管線（由 NeMo Data Designer 實作）會產生 1‑3 跳的不同複雜度問題，並給予品質分數。僅保留高於可設定門檻的配對供訓練使用。

2️⃣ 挖掘硬負樣本

硬負樣本是基礎模型將其排在正向答案附近的非相關段落。挖掘流程如下：

  1. 使用基礎模型對所有查詢與語料段落進行嵌入。
  2. 計算相似度分數。
  3. 遮蔽真陽性。
  4. 套用 95% 邊際過濾，以避免偽陰性。
  5. 選取前 k（預設 5）最高分的非陽性作為硬負樣本。
nemotron embed prep -c default

此指令同時將資料切分（80% 訓練、20% 測試）並將多跳問題展開為獨立的 (query, positive) 範例，對每個正例保留相同的硬負樣本。

3️⃣ 了解多跳問題

多跳查詢（2‑3 跳）需要模型檢索多個相關段落。將每一跳展開為獨立訓練實例，可讓模型將 所有 相關文件視為正例，提升在複雜使用者查詢中呈現完整答案集合的能力。

4️⃣ 微調雙編碼器

訓練使用對比損失，溫度係數為 0.02，迫使模型將正例與硬負樣本明顯分離。

nemotron embed finetune -c default

關鍵超參數（預設值已針對示例資料集調校）:

參數 預設值
Epochs 3（對於較大語料建議 1–2 epoch）
Learning rate 1e‑5
Warmup steps 5（≈ 5‑10 % 總步數）
Global batch size 128
Passages per query 5（1 正例 + 4 硬負樣本）

5️⃣ 評估檢索提升

評估使用 BEIR 框架於保留測試集上執行，報告 nDCG、Recall、Precision 與 MAP，k 分別為 1、5、10、100。

nemotron embed eval -c default

合成 NVDocs 結果（基礎 → 微調後）:

  • NDCG@10: 0.555 → 0.616（+10.9 %）
  • Recall@10: 0.630 → 0.693（+10.0 %）
  • 所有 k 值皆觀察到類似提升。

真實 Atlassian 案例：在公開的 JIRA 資料集上微調後，Recall@60 從 0.751 提升至 0.951，提升 26.7%，僅使用單張 A100‑80GB GPU。

疑難排解

  • 合成品質低 → 改善文件格式或使用更強大的 LLM。
  • 資料不足 → 增加來源文件並重新執行 SDG。
  • 過度擬合 → 將 epoch 降至 1‑2，或提升品質門檻。
  • 學習率不佳 → 嘗試 0.5 倍或 2 倍的預設值。

6️⃣ 匯出與部署

將微調後的檢查點轉換為 ONNX（opset 17），並可選擇編譯 TensorRT 引擎以獲得最高吞吐量。

nemotron embed export -c default            # 只匯出 ONNX
nemotron embed export -c default export_to_trt=false
nemotron embed export -c default quant_cfg=fp8   # FP8 量化

使用 NVIDIA NIM 部署模型，NIM 會公開相容 OpenAI 的 /v1/embeddings 端點。

nemotron embed deploy -c default

範例請求:

curl -X POST http://localhost:8000/v1/embeddings \
  -H "Content-Type: application/json" \
  -d '{"input": ["What cooling is needed for 8 H100 GPUs in a 2U chassis?"], "model": "custom", "input_type": "query"}'

內建的 NIM 精度檢查會再次以 BEIR 評估已部署服務，確保轉換未造成效能退化。

端到端指令摘要

# 1. 合成資料生成
nemotron embed sdg -c default corpus_dir=./data/my_docs

# 2. 資料準備（切分、硬負樣本挖掘、展開）
nemotron embed prep -c default

# 3. 微調嵌入模型
nemotron embed finetune -c default

# 4. 評估基礎與微調後檢查點
nemotron embed eval -c default

# 5. 匯出至 ONNX/TensorRT
nemotron embed export -c default

# 6. 使用 NVIDIA NIM 部署
nemotron embed deploy -c default

資源與時間估算

階段 需要 GPU 嗎？ 大約時間（單卡 A100‑80GB）
合成資料生成 否（API） 約 1 小時（視語料大小而定）
資料準備（硬負樣本挖掘） 是（≈ 40 GB VRAM） 約 5 分鐘
微調 是（80 GB VRAM） 約 1 小時
評估 是（≈ 40 GB VRAM） 約 5 分鐘
匯出 是（≈ 40 GB VRAM） 約 5 分鐘
部署 是（≈ 40 GB VRAM） 約 5 分鐘
總計: < 24 小時；對於約 500 份文件的中等規模語料，整個流程可在 2–3 小時內完成。

實務要點

此配方證明領域適應的嵌入不再是需要多週、多 GPU 的工程。藉由合成資料生成、硬負樣本挖掘與 NVIDIA 最佳化工具，實務者可在單張 GPU、不到一天的運算時間內，於標準指標上提升超過 10%，在企業工作負載上提升超過 25%。

自己動手試試

將倉庫克隆下來，取得 NVIDIA API 金鑰，並將管線指向自己的文件集合。即用的 nvidia/Retrieval-Synthetic-NVDocs-v1 資料集讓您立即上手實驗。歡迎對 Nemotron、NeMo Data Designer 與 NeMo Automodel 倉庫提交貢獻與加星。

Sources