如何在一天内构建领域专用嵌入模型
Build a Domain-Specific Embedding Model in Under a Day
NVIDIA 发布一套在单张 GPU 上不到一天完成领域嵌入模型微调的流程,基于 Llama-Nemotron-Embed-1B-v2,无需人工标注。流程用 nemotron-3-nano-30b-a3b 从领域文档生成合成问答数据、挖掘难负样本并做多跳展开,再用对比学习微调,实测 nDCG@10 与 Recall@10 均提升约 10%。
NVIDIA 给出了从合成数据、难负样本挖掘到 ONNX/TensorRT 部署的完整嵌入模型微调配方,附真实评测数字。
如果你正在构建一个 RAG(检索增强生成)系统,你很可能已经撞上了这堵墙:一切都能跑通……直到它不行。通用嵌入模型被训练来理解互联网,而不是你的合同、制造日志、专有化学配方或内部分类体系。它们能捕捉宽泛的语义相似度,却无法理解在你所在领域中至关重要的细粒度差异。当现成的模型无法有效捕捉领域特定细微差别时,微调嵌入模型可以提升检索管线的性能。尽管嵌入对 RAG 性能如此关键,整个流程却仍然出奇地碎片化,所需技能相当专业,时间投入也令人生畏。
只需一块 GPU 和不到一天的训练时间,你就能把一个通用嵌入模型转变为真正理解你所在领域的模型,而且无需人工标注。为了帮你快速上手,我们还发布了一个开箱即用的合成训练数据集,它是使用这套完全相同的流程从 NVIDIA 公开文档生成的。使用这份数据和这套方法,我们在 Recall@10 和 NDCG@10 上都看到了超过 10% 的提升。Atlassian 将这套方法应用于他们的JIRA 数据集进行微调,将 Recall@60 从 0.751 提升到 0.951,提升了 26%——而且只用了单块 GPU。
🔗数据集和代码快速链接:
🧑💻开源项目方法集成了:
- NeMo Data Designer 用于合成数据生成
- NeMo Automodel 用于嵌入模型训练
- BEIR 用于信息检索评估
- NeMo Export-Deploy 用于 ONNX/TensorRT 转换
- NVIDIA NIM 用于生产推理服务
📋前提条件:
- 一个包含领域文档的目录(文本文件 - .txt、.md 或类似格式)
- 一个有效的 NVIDIA API 密钥(可在 build.nvidia.com 免费获取)
- NVIDIA Ampere GPU or newer with at least 80GB memory (with Compute Capability >= 8.0)
- 本教程已在 1xA100(80GB)和 1xH100(80GB)上测试通过
读完本文后,你将学会如何:
📄 在没有标注数据的情况下从领域文档生成训练数据
🎯 使用难负例挖掘进行有效的对比训练
🔗 通过多跳查询提升嵌入质量
⚙️ 微调双编码器嵌入模型
📊 评估微调是否改善了检索
🚀 在你的管线中部署微调后的模型
⚙️环境设置
在本教程中,我们将微调基础模型 Llama-Nemotron-Embed-1B-v2 —— 一个在质量和推理成本之间取得平衡的 10 亿参数嵌入模型。要开始使用,请遵循这份设置指南。
📚 第 1 步:从文档生成训练数据
微调嵌入模型需要数千个(查询,相关文档)对。大多数用例并没有现成可用的这类数据。手动创建这类数据成本高昂、速度缓慢,而且常常会受到标注者个人对“相关”理解的影响而带有偏差。
与其手工标注数据,你可以使用 LLM(nvidia/nemotron-3-nano-30b-a3b)来阅读你的文档,并自动生成高质量的合成问答对。
nemotron embed sdg -c default corpus_dir=./data/my_domain_docs
它是如何工作的?
在后台,这运行的是一个由 NeMo Data Designer 驱动的四阶段合成数据生成(SDG)流水线:
输出是什么样的?
源文档片段:
H100 GPU 的热设计功耗(TDP)在 SXM 形态下为 700W。散热方案必须确保在持续负载下结温保持在 83°C 以下。对于每节点超过 4 块 GPU 的高密度部署,建议采用液冷,因为在标准 2U 机箱配置中,风冷无法消散足够的热量。
生成的问答对:
{
"question": "What cooling approach is recommended when deploying more than 4 H100 GPUs per server node?",
"answer": "Liquid cooling is recommended for dense deployments exceeding 4 GPUs per node, as air cooling cannot dissipate sufficient heat in standard 2U chassis configurations.",
"query_type": "contextual",
"reasoning_type": "factual",
"question_complexity": 3,
"segment_ids": [1],
"quality_score": 8.5
}
{
"question": "How does the 700W TDP of the H100 SXM constrain the choice between air and liquid cooling in multi-GPU configurations?",
"answer": "The 700W TDP generates substantial heat that must be dissipated to keep junction temperatures below 83°C. In dense configurations exceeding 4 GPUs per node, air cooling in standard 2U chassis cannot handle this thermal load, making liquid cooling necessary.",
"query_type": "multi_hop",
"reasoning_type": "causal",
"question_complexity": 4,
"segment_ids": [1, 2],
"hop_count": 2,
"quality_score": 9.0
}
注意其中的区别:第一个问题是简单的事实查找。第二个问题则需要多跳的因果推理。该流水线会同时生成这两种类型,并支持可配置的复杂度级别(2–5)和跳数(1–3)。随后每个问答对都会经过质量评估,获得相关性、准确性、上下文支持度和清晰度等子分数,以及一个总体分数。只有达到阈值的问答对才会被纳入训练。
⛏️ 第 2 步:挖掘难负样本(以及它们为何重要)
如果只用正样本对(查询 + 正确文档)来训练嵌入模型,它学会的是区分明显不同的文档,但在困难情况下会失败——那些看起来相关但并非正确答案的段落。在真实的检索系统中,这些近似但不正确的文档正是导致糟糕答案的根源。难负样本挖掘会找出这些容易混淆的段落,让模型学会区分它们。
nemotron embed prep -c default
上述命令会自动运行三个子步骤:
2a. 训练 / 验证 / 测试划分
生成的问答对被划分为训练集(80%)和测试集(20%)。测试集被格式化为与 BEIR 兼容的基准,以便在第 5 步中进行标准化评估。
2b. 难负样本挖掘
使用基础嵌入模型,该流水线会:
- 对语料库中的每个查询和每个段落进行嵌入。
- 计算每个查询与所有段落之间的相似度。
- 屏蔽掉每个查询已标注的正样本文档。
- 应用边界过滤:任何非正样本文档,如果其得分高于最小正样本得分的 95%,都会被剔除。这个排除区可以防止假负样本——即那些与正样本非常接近、可能实际上相关的未标注段落。
- 从剩余的候选项中,选取得分最高的 top-k 文档作为难负样本(默认每个查询 5 个)。
结果是:难负样本是最相似的非正样本段落,同时仍安全地低于正样本得分上限。它们是当前模型认为高度相关、但并非标注答案的段落。
为什么这有效:用简单负样本(完全不相关的段落)训练无法让模型学到新东西。用难负样本训练则迫使模型学习在你的领域中至关重要的细微差别。例如,在医学语料库中,关于“二甲双胍治疗 2 型糖尿病的剂量”的问题,其难负样本可能是关于“二甲双胍副作用”或“胰岛素治疗 1 型糖尿病的剂量”的内容——相近但存在关键差异。95% 的边界上限可以防止挖掘器选择与正样本过于接近的段落,因为这些段落实际上可能是正确答案,只是在 SDG 过程中未被标注。
2c. 多跳展开
多跳问题会引用多个正样本文档。例如,像“第3.2节中的热管理系统与第5.1节中描述的功率限制有何关联?”这样的问题就有两个正样本段落。
展开会为每个(查询,正样本文档)对生成一个训练示例,因此对比损失会独立地看待每个正样本。一个带有2个正样本文档的问题会变成2个训练示例,每个示例拥有相同的难负样本但不同的正样本。
最终输出是一个可直接用于训练的JSON文件:
{
"question_id": "q42_0",
"question": "How does the thermal management system in Section 3.2 relate to the power constraints described in Section 5.1?",
"pos_doc": [{"id": "Section 3.2"}],
"neg_doc": [{"id": "d_x7y8z9"}, {"id": "d_m4n5o6"}, {"id": "d_p1q2r3"}, {"id": "d_s4t5u6"}, {"id": "d_v7w8x9"}]
},
{
"question_id": "q42_1",
"question": "How does the thermal management system in Section 3.2 relate to the power constraints described in Section 5.1?",
"pos_doc": [{"id": "Section 5.1"}],
"neg_doc": [{"id": "d_x7y8z9"}, {"id": "d_m4n5o6"}, {"id": "d_p1q2r3"}, {"id": "d_s4t5u6"}, {"id": "d_v7w8x9"}]
}
🔍 第3步:理解多跳问题及其为何能提升检索效果
标准的嵌入微调会为每个段落生成一个问题,并训练模型将它们匹配起来。这对于简单的事实查找是有效的,但真实用户会提出跨越多个文档或章节的复杂问题。如果模型只见过单跳训练数据,它在为这些复杂查询检索所有相关段落时会很吃力。
SDG流水线默认生成1到3跳的问题:
- 1跳:“H100 SXM的TDP是多少?”——由单个段落即可回答。
- 2跳:“H100的TDP与高密度部署中的冷却需求有何关联?”——需要连接两个段落中的信息。
- 3跳:“考虑到TDP、冷却约束和机架密度限制,标准数据中心单排最多可部署多少块H100 GPU?”——综合三个段落的信息。
每一跳都通过其各自的上下文摘要和段落ID进行追踪,因此训练数据保留了完整的推理链。在展开(第2c步)之后,每个(问题,相关段落)对都会成为一个独立的训练信号,教会模型这些段落全部都与该多跳查询相关。
微调后的模型学会检索上下文相关的文档,而不仅仅是词汇上相似的文档。
🧠 第4步:微调嵌入模型
nemotron embed finetune -c default
对比学习的工作原理
训练使用带有对比损失的双编码器架构。
0.02的温度值是刻意设置得比较激进的,它会产生非常锐利的概率分布。这样效果很好,因为第2步得到的难负样本质量很高:它们确实是容易混淆的段落,模型需要较强的梯度才能学会区分它们。
关键超参数
| 参数 | 默认值 | 说明 |
|---|---|---|
| 轮数 | 3 | 默认值针对小型示例数据集进行了调优;对于真实世界的数据,1–2轮通常就足够了,以避免过拟合。 |
| 学习率 | 1e-5 | 调优:可以尝试默认值的两倍和一半 |
| 学习率预热步数 | 5 | 设置为微调总步数的5-10%,以获得更好的早期训练稳定性 |
| 全局批量大小 | 128 | |
| 每个查询的段落数 | 5 | 1个正样本 + 4个难负样本 |
检查点保存频率
如果配置中省略了ckpt_every_steps,检查点保存频率会自动设置:
- 映射式数据集(已知长度):默认每个轮次保存一次。
- 可迭代数据集(未知长度):默认训练期间保存两次。
这意味着你可以先用一个小型语料库(50–100个文档)快速进行概念验证,之后再扩大规模,而无需手动调整检查点设置。
📈 第5步:衡量改进效果
微调真的有帮助吗?让我们在留出测试集上运行标准化评估,对比基础模型与微调后的检查点,来一探究竟:
nemotron embed eval -c default
评估使用 BEIR 框架,并在 k = 1、5、10 和 100 时计算四个标准信息检索指标:
- nDCG@k:排序质量——最佳文档是否排在最前面?
- Recall@k:覆盖率——相关文档中有多大比例出现在前 k 个结果中?
- Precision@k:准确率——前 k 个结果中实际相关的比例是多少?
- MAP@k:所有查询的平均精度
成功的微调通常能在 <1 天内使 nDCG@10 和 Recall@10 提升 10%。
使用 Retrieval Synthetic NVDocs 的结果:
📊 Comparison (Base -> Fine-tuned)
============================================================
NDCG:
NDCG@1: 0.55178 → 0.60796 (+0.05618, +10.2%)
NDCG@5: 0.51894 → 0.57689 (+0.05795, +11.2%)
NDCG@10: 0.55506 → 0.61559 (+0.06053, +10.9%)
NDCG@100: 0.60617 → 0.66567 (+0.05950, +9.8%)
Recall:
Recall@1: 0.28478 → 0.31547 (+0.03069, +10.8%)
Recall@5: 0.54486 → 0.60288 (+0.05802, +10.6%)
Recall@10: 0.62979 → 0.69296 (+0.06317, +10.0%)
Recall@100: 0.81421 → 0.87020 (+0.05599, +6.9%)
如果数字没有提升怎么办?
该流程让迭代变得简单:
- SDG 中质量分数低? 检查你的文档质量——干净、格式良好的文本能生成更好的合成数据。尝试更大、更强大的 LLM。
- 训练数据不够? 向语料库添加更多文档,并重新运行 Stage 0。
- 过拟合? 默认的 3 个 epoch 是针对小型示例数据集校准的;对于大多数真实世界数据,1–2 个 epoch 就足够了。也可以考虑提高质量阈值,只保留最好的训练样本。
- 学习率不对? 试试将默认学习率翻倍或减半。
🏆 真实世界成果:Atlassian
该方案已由 Atlassian 在真实企业数据上验证。他们按照上述相同阶段,使用单块 NVIDIA A100 80GB GPU,将本流程应用于在 公开 Jira 数据集上微调 Llama-Nemotron-Embed-1B-v2
Recall@60 从 0.751 跃升至 0.951——提升了 26.7%。
微调后的模型对 95.1% 的查询能在前 60 个结果中检索到正确文档,而基础模型仅为 75.1%。对于支撑 Jira 搜索的检索系统而言,这直接为数百万用户带来了更相关的结果。更多细节请见他们的博客文章Advancing semantic search for millions of Rovo users。
🚀 第 6 步:导出与部署
PyTorch 检查点很适合评估,但对于生产环境来说太慢。最后两个阶段会转换模型,并将其部署在 API 之后。
导出为 ONNX / TensorRT
nemotron embed export -c default
这会将微调后的检查点导出为 ONNX(opset 17)。可选地,它会编译一个 TensorRT 引擎以实现最大推理吞吐量,并可为批量大小(1–64)和序列长度(3–256)配置优化配置文件:
# ONNX only (runs anywhere)
nemotron embed export -c default export_to_trt=false
# FP8 quantization for further speedup
nemotron embed export -c default quant_cfg=fp8
使用 NVIDIA NIM 部署
导出的模型部署在 NVIDIA 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 提供 OpenAI 兼容的 API,你可以将其直接接入任何使用该 embeddings API 格式的现有 RAG 流程——无需更改代码。
验证部署准确率
该流程包含一个 NIM 准确率验证步骤,会对已部署的端点运行相同的 BEIR 评估:
nemotron embed eval -c default eval_nim=true eval_base=false
这能捕捉 ONNX/TensorRT 转换带来的任何准确率损失。在容差范围内匹配的指标(@1 为 0.03,@5+ 为 0.01)会标记为通过;超出转换噪声的偏差会被标记出来。
整合起来
完整的嵌入微调流水线可以通过六条命令运行,从原始文档到部署模型。
# 1. Generate synthetic training data from your documents
nemotron embed sdg -c default corpus_dir=./data/my_docs
# 2. Prepare the training data (split data, mine hard negatives, unroll)
nemotron embed prep -c default
# 3. Fine-tune the embedding model
nemotron embed finetune -c default
# 4. Evaluate the base vs. fine-tuned model
nemotron embed eval -c default
# 5. Export the optimized model
nemotron embed export -c default
# 6. Deploy the model
nemotron embed deploy -c default
预计时间和资源
| 阶段 | 需要GPU? | 预计时间 | 说明 |
|---|---|---|---|
| SDG | 否(使用API) | 约1小时 | 因语料库大小和API速率限制而异 |
| 数据准备 | 是(40 GB显存) | 约5分钟 | 在GPU上进行难负例挖掘 |
| 微调 | 是(80 GB显存) | 约1小时 | 因数据集大小和训练轮数而异 |
| 评估 | 是(40 GB显存) | 约5分钟 | |
| 导出 | 是(40 GB显存) | 约5分钟 | TensorRT需要NGC容器 |
| 部署 | 是(40 GB显存) | 约5分钟 | NIM容器启动 |
总计:不到一天,其中大部分时间是无须人工干预的训练。对于小型语料库(约500份文档),整个流水线大约在2–3小时内完成。
该流水线可以端到端运行,但每个阶段也可以根据你的起点独立执行。例如,如果你有原始文档,可以从合成数据生成(SDG)开始;而已经包含难负例的数据集可以跳过前面的步骤,直接进行微调。由于每个阶段都使用JSON、BEIR和ONNX等标准格式,因此很容易集成自定义组件,或在其他工作流中复用中间输出。该配方在运行方式上也很灵活,支持在本地机器、Docker容器内或基于Slurm的集群上执行。
自己动手试试
如果你有领域文档并且手头有些时间,今天就可以生成你的第一批合成训练数据!完整流水线——从文档到已部署的领域适配嵌入模型——在单块GPU上不到一天即可运行。你可以从我们现成的nvidia/Retrieval-Synthetic-NVDocs-v1数据集开始,立即试用该流水线。告诉我们你构建了什么。
如果你觉得有用,请为Nemotron、NeMo Data Designer和NeMo Automodel点亮仓库星标。
来源:Hugging Face Blog · hf-mirror.com


