ms-swift嵌入模型训练：Embedding任务实战

1. 引言：为什么需要高质量的Embedding模型？

在当前的大模型应用生态中，我们常常关注生成式任务，比如对话、写作、代码生成等。但有一类看似低调却至关重要的任务——Embedding（嵌入）建模，正在支撑着搜索、推荐、语义匹配、RAG（检索增强生成）等核心场景。

你有没有遇到过这些问题：

用户输入“我想买一台轻薄本”，系统却返回了游戏本？
检索文档时，关键词匹配上了，但内容完全不相关？
RAG系统召回的内容驴唇不对马嘴，导致大模型输出错误答案？

这些问题的背后，往往不是大模型本身的问题，而是向量检索环节的语义理解能力不足。而这一切，都取决于你的Embedding模型是否足够智能。

今天，我们就用ms-swift这个强大的微调框架，手把手带你完成一次高质量Embedding模型的训练实战。不需要复杂的理论推导，也不需要从零写代码，全程基于命令行和配置化操作，适合工程落地。

2. ms-swift是什么？它为何适合Embedding任务？

2.1 一个全链路支持的微调框架

ms-swift是魔搭社区推出的一站式大模型微调与部署框架，支持超过600个纯文本模型和300个多模态模型的训练、推理、评测、量化与部署。

它的优势不仅在于“支持多”，更在于“做得深”：

✅ 支持LoRA、QLoRA、DoRA等多种轻量微调方式
✅ 内置GaLore、Q-Galore等显存优化技术，降低长序列训练成本
✅ 提供vLLM、SGLang、LMDeploy三大推理引擎加速
✅ 覆盖预训练、SFT、DPO、KTO、RM、CPO、SimPO、ORPO等主流训练范式
✅原生支持Embedding和Reranker任务

这意味着，你不需要自己搭建训练流程、处理数据格式、实现损失函数，只需要准备好数据集和模型，剩下的交给ms-swift。

2.2 Embedding任务的独特挑战

传统的分类或生成任务关注的是“输出结果”，而Embedding任务关注的是“中间表示”。它的目标是让语义相近的文本在向量空间中距离更近，语义不同的远离。

这就带来几个关键要求：

使用合适的损失函数（如对比学习Contrastive Loss）
构造正负样本对（Positive/Negative Pairs）
训练过程中进行归一化处理
输出固定维度的向量表示

幸运的是，ms-swift已经内置了这些能力，只需通过参数配置即可启用。

3. 实战准备：环境与数据

3.1 环境安装

# 创建虚拟环境 conda create -n swift python=3.10 conda activate swift # 安装ms-swift（推荐清华源加速） pip install 'ms-swift[all]' -U -i https://pypi.tuna.tsinghua.edu.cn/simple

如果你希望使用最新功能，也可以从源码安装：

git clone https://github.com/modelscope/ms-swift.git cd swift pip install -e '.[llm]' -i https://pypi.tuna.tsinghua.edu.cn/simple

验证安装成功：

swift --help

你应该能看到sft,pt,rlhf,infer,export等子命令。

3.2 数据集选择与准备

Embedding任务通常使用成对数据（Pairwise Data），常见格式如下：

sentence1	sentence2	label
如何重装Windows系统？	怎么重新安装操作系统？	1
苹果是一种水果	iPhone是苹果公司产品	0

其中label=1表示语义相似，label=0表示不相似。

ms-swift支持多种公开数据集，例如：

bge-m3：北航发布的多语言检索数据集
mteb：大规模文本嵌入基准数据集
nli-zh：中文自然语言推断数据集

我们以nli-zh为例，构造一个中文语义相似度任务：

swift sft \ --model Qwen/Qwen2-7B-Instruct \ --train_type lora \ --dataset AI-ModelScope/nli-zh \ --task_name embedding \ --embedding_norm true \ --embedding_normalize_dim -1 \ --output_dir output_embedding \ --num_train_epochs 3 \ --per_device_train_batch_size 4 \ --learning_rate 2e-5 \ --lora_rank 8 \ --lora_alpha 32 \ --max_length 512 \ --save_steps 100 \ --eval_steps 100 \ --logging_steps 10

4. Embedding训练全流程详解

4.1 核心参数说明

以下是Embedding任务的关键参数解析：

参数	含义	推荐值
`--task_name embedding`	明确指定任务类型为embedding	必填
`--embedding_norm`	是否对输出向量做L2归一化	`true`
`--embedding_normalize_dim`	归一化维度	`-1`（最后一维）
`--pooling_strategy`	向量池化策略	`last_attention`或`cls`
`--loss_type`	损失函数类型	`contrastive`/`cosent`
`--margin`	对比学习边距	`0.5`
`--scale`	相似度缩放系数	`20.0`

💡 小贴士：cosent损失函数在实践中表现优于传统contrastive，建议优先使用。

4.2 完整训练命令示例

CUDA_VISIBLE_DEVICES=0 swift sft \ --model Qwen/Qwen2-7B-Instruct \ --train_type lora \ --dataset AI-ModelScope/nli-zh#1000 \ --task_name embedding \ --loss_type cosent \ --pooling_strategy last_attention \ --embedding_norm true \ --embedding_normalize_dim -1 \ --output_dir output_embedding \ --num_train_epochs 3 \ --per_device_train_batch_size 4 \ --per_device_eval_batch_size 4 \ --gradient_accumulation_steps 8 \ --learning_rate 2e-5 \ --lora_rank 8 \ --lora_alpha 32 \ --target_modules all-linear \ --max_length 512 \ --save_steps 100 \ --eval_steps 100 \ --logging_steps 10 \ --warmup_ratio 0.1 \ --dataloader_num_workers 4 \ --torch_dtype bfloat16 \ --fp16 false \ --bf16 true

解释几个关键点：

#1000表示只取前1000条数据用于快速实验
bfloat16提升训练稳定性，尤其适合A100/H100
gradient_accumulation_steps=8在单卡上模拟更大batch size
target_modules all-linear对所有线性层添加LoRA适配器

4.3 自定义数据集格式

如果你想用自己的数据，只需将数据组织为JSONL格式：

{"sentence1": "今天天气真好", "sentence2": "外面阳光明媚", "label": 1} {"sentence1": "我要去吃饭", "sentence2": "电脑坏了怎么办", "label": 0}

然后通过本地路径加载：

--dataset /path/to/your/dataset.jsonl

并确保字段名正确映射。如果字段不同，可通过自定义dataset注册机制处理（参考文档自定义数据集）。

5. 训练过程监控与效果评估

5.1 日志输出解读

训练启动后，你会看到类似以下日志：

Train: 100%|██████████| 250/250 [12:34<00:00, 0.33it/s] {'train_loss': 0.123, 'epoch': 1.0, 'global_step': 250} Eval: 100%|██████████| 50/50 [01:15<00:00, 0.66it/s] {'eval_loss': 0.118, 'eval_accuracy': 0.892, 'eval_pearson': 0.765}

重点关注：

eval_loss：越低越好
eval_accuracy：分类准确率（适用于二分类相似判断）
eval_pearson：预测相似度与真实标签的相关性，越高越好

5.2 使用MTEB进行专业评测

ms-swift集成了EvalScope作为评测后端，支持MTEB（Massive Text Embedding Benchmark）标准评测。

训练完成后，可运行：

swift eval \ --model output_embedding/checkpoint-250 \ --eval_backend EvalScope \ --eval_dataset mteb-zh \ --infer_backend vllm \ --batch_size 32

输出将包含多个子任务得分，如：

STS（语义文本相似度）
Clustering（聚类）
Retrieval（检索）

最终得到一个综合排名，便于横向比较不同模型版本。

6. 模型推理与服务部署

6.1 本地推理测试

训练好的Embedding模型可以用于生成句向量：

swift infer \ --adapters output_embedding/checkpoint-250 \ --stream false \ --infer_backend pt \ --max_new_tokens 1

注意：Embedding模型不需要生成新token，所以max_new_tokens=1即可，实际输出是隐藏状态的池化结果。

你可以传入一段文本，获取其768维（或其他维度）的向量表示，用于后续的相似度计算。

6.2 部署为API服务

使用vLLM加速部署：

swift deploy \ --model output_embedding/checkpoint-250 \ --infer_backend vllm \ --port 8080 \ --host 0.0.0.0

部署成功后，可通过HTTP请求获取嵌入向量：

curl -X POST http://localhost:8080/embeddings \ -H "Content-Type: application/json" \ -d '{"input": "如何提高工作效率？"}'

响应示例：

{ "object": "list", "data": [ { "object": "embedding", "embedding": [0.12, -0.45, ..., 0.67], "index": 0 } ], "model": "qwen2-7b-instruct-embedding" }

这就可以直接接入你的搜索引擎、推荐系统或RAG pipeline。