基于 ms-swift 框架微调 Embedding 模型（Ubuntu22.04）：从原理到实践

作者：吴业亮
博客：wuyeliang.blog.csdn.net

一、核心原理

1. Embedding模型微调的本质

Embedding模型的核心是将文本转化为低维稠密向量，使语义相似的文本向量距离更近、语义相异的更远。微调的目标是让预训练Embedding模型适配特定领域/任务（如金融、医疗文本相似度匹配），核心逻辑是：

冻结模型主干网络（减少计算量），仅训练轻量化适配器（如LoRA）；
通过对比学习损失（InfoNCE） 优化向量空间：构造"查询句-正样本-负样本"三元组，让查询句与正样本向量余弦相似度最大化，与负样本最小化。

2. ms-swift框架核心

ms-swift（MindSpore Swift）是华为基于MindSpore打造的轻量化大模型微调框架，适配Embedding模型的核心优势：

原生支持LoRA/QLoRA等高效微调方法，显存占用降低90%+；
内置Embedding任务专用损失函数（对比损失、余弦损失）；
一键式配置化训练，无需大量底层代码开发；
兼容主流Embedding模型（如BERT、RoBERTa、m3e等）。

3. 关键技术：LoRA（低秩适配）

Embedding模型通常参数量大（如BERT-base有1.1亿参数），全量微调显存成本高。LoRA的核心是：

冻结主干模型所有参数；
在注意力层的权重矩阵旁插入低秩矩阵（如rank=8）；
仅训练这些低秩矩阵（参数量仅数万~数十万）；
推理时将LoRA权重合并回主干模型，无额外延迟。

二、环境准备（Ubuntu22.04）

1. 系统依赖安装

首先安装基础依赖，确保系统环境完整：

bash 复制代码

# 更新系统源
sudo apt update && sudo apt upgrade -y

# 安装基础编译工具、Python依赖
sudo apt install -y python3-pip python3-dev gcc g++ git libopenmpi-dev

# 安装CUDA（如需GPU训练，推荐11.6/11.8，需匹配显卡驱动）
# 注：若用CPU训练，可跳过CUDA安装
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
sudo sh cuda_11.8.0_520.61.05_linux.run --silent --toolkit
echo "export PATH=/usr/local/cuda-11.8/bin:\$PATH" >> ~/.bashrc
echo "export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:\$LD_LIBRARY_PATH" >> ~/.bashrc
source ~/.bashrc

# 验证CUDA（GPU环境）
nvcc -V

2. Python环境配置

推荐使用虚拟环境隔离依赖：

bash 复制代码

# 安装虚拟环境工具
pip3 install virtualenv

# 创建并激活虚拟环境
virtualenv -p python3.9 ms-swift-env  # 推荐Python3.8~3.10
source ms-swift-env/bin/activate

# 升级pip
pip install --upgrade pip setuptools wheel

3. 安装MindSpore

根据硬件环境选择CPU/GPU版本（需匹配CUDA版本）：

bash 复制代码

# GPU版本（CUDA11.8）
pip install mindspore-gpu==2.2.14 -i https://pypi.tuna.tsinghua.edu.cn/simple

# CPU版本（无GPU时）
# pip install mindspore==2.2.14 -i https://pypi.tuna.tsinghua.edu.cn/simple

# 验证MindSpore安装
python -c "import mindspore; print(mindspore.__version__)"

4. 安装ms-swift框架

从源码安装最新版本（功能更全）：

bash 复制代码

# 克隆源码
git clone https://gitee.com/mindspore/swift.git
cd swift

# 安装ms-swift核心依赖
pip install -e .[all] -i https://pypi.tuna.tsinghua.edu.cn/simple

# 验证安装
python -c "import swift; print(swift.__version__)"

三、数据准备

1. 数据集格式要求

ms-swift适配Embedding微调的标准格式为JSONL（每行一个样本），支持两种格式：

格式1：三元组（推荐，对比学习）

json 复制代码

{"query": "如何办理信用卡", "positive": "信用卡办理流程", "negative": "手机充值方法"}
{"query": "Ubuntu22.04安装CUDA", "positive": "Ubuntu22.04 CUDA11.8安装教程", "negative": "Windows11安装Python"}

query：查询句；
positive：语义相似的正样本；
negative：语义相异的负样本。

格式2：成对标注（二分类）

json 复制代码

{"sentence1": "苹果手机电池更换", "sentence2": "iPhone换电池教程", "label": 1.0}
{"sentence1": "苹果手机电池更换", "sentence2": "华为手机充电技巧", "label": 0.0}

label：0（不相似）~1（相似）的相似度标注。

2. 数据集划分

将数据集分为训练集（train.jsonl）、验证集（dev.jsonl），比例建议9:1，示例目录结构：

复制代码

data/
├── train.jsonl
└── dev.jsonl

3. 示例数据集（快速测试）

创建测试用的小数据集（data/train.jsonl）：

bash 复制代码

mkdir -p data
cat > data/train.jsonl << EOF
{"query": "机器学习入门", "positive": "零基础学机器学习", "negative": "烹饪技巧大全"}
{"query": "Python读取Excel", "positive": "Python使用pandas读取Excel文件", "negative": "Java实现多线程"}
{"query": "ms-swift微调Embedding", "positive": "基于ms-swift的Embedding模型微调教程", "negative": "王者荣耀出装攻略"}
EOF

# 创建验证集（data/dev.jsonl）
cat > data/dev.jsonl << EOF
{"query": "MindSpore安装", "positive": "Ubuntu22.04安装MindSpore GPU版本", "negative": "抖音短视频制作"}
EOF

四、微调配置与训练

1. 配置文件修改

ms-swift通过YAML配置文件定义训练参数，复制官方Embedding微调配置模板并修改：

bash 复制代码

# 复制配置模板（以m3e-base模型为例，主流中文Embedding模型）
cp examples/embedding/run_embedding.yaml ./embedding_finetune.yaml

编辑embedding_finetune.yaml核心参数（关键部分）：

yaml 复制代码

# 基础配置
task: embedding  # 任务类型
model_type: m3e_base  # 模型类型（支持bert_base_chinese、roberta_wwm_ext等）
model_name_or_path: m3e-base  # 预训练模型路径/名称（自动下载）
dataset:
  train_file: ./data/train.jsonl  # 训练集路径
  dev_file: ./data/dev.jsonl      # 验证集路径
  max_seq_length: 128            # 文本最大长度
  type: triplet                  # 数据集类型：triplet（三元组）/pair（成对标注）

# 微调策略（LoRA）
finetune_type: lora  # 微调类型：lora/qlora/full（全量）
lora:
  rank: 8            # LoRA秩（越小参数量越少）
  alpha: 32          # LoRA缩放系数
  dropout: 0.1       # Dropout概率
  target_modules: ['q', 'v']  # 针对注意力层的q/v矩阵训练

# 训练参数
train:
  batch_size: 2      # 批次大小（GPU显存不足时减小）
  eval_batch_size: 2
  learning_rate: 1e-4  # 学习率（LoRA建议1e-4~5e-4）
  epochs: 10           # 训练轮数
  loss_type: info_nce  # 损失函数：info_nce（对比损失）/cosine（余弦损失）
  optimizer: adamw     # 优化器
  weight_decay: 0.01   # 权重衰减
  lr_scheduler_type: cosine  # 学习率调度器

# 输出配置
output_dir: ./output/embedding_finetune  # 模型输出路径
save_strategy: epoch                     # 按轮数保存模型
save_total_limit: 3                      # 最多保存3个模型
logging_steps: 10                        # 日志打印步数
eval_strategy: epoch                     # 按轮数验证
device: 0                                # GPU卡号（CPU设为cpu）

2. 启动微调训练

执行训练命令（基于配置文件）：

bash 复制代码

# GPU训练（推荐）
swift run --config ./embedding_finetune.yaml

# CPU训练（无GPU时）
swift run --config ./embedding_finetune.yaml --device cpu

3. 训练过程说明

训练日志示例：

复制代码

2025-12-05 10:00:00 - INFO - Loading m3e-base model...
2025-12-05 10:00:10 - INFO - Training epoch 1/10, step 1/2, loss: 2.3026
2025-12-05 10:00:15 - INFO - Training epoch 1/10, step 2/2, loss: 1.8975
2025-12-05 10:00:20 - INFO - Evaluating dev set...
2025-12-05 10:00:25 - INFO - Dev set cosine similarity: 0.852
2025-12-05 10:00:25 - INFO - Saving model to ./output/embedding_finetune/checkpoint-1

训练完成后，output/embedding_finetune目录下会生成：

checkpoint-*：模型权重文件；
config.json：模型配置；
training_log.txt：训练日志。

五、模型验证与推理

1. 模型验证（余弦相似度计算）

编写验证脚本（eval_embedding.py），对比微调前后的向量相似度：

python 复制代码

import numpy as np
from mindspore import context
from swift import Swift
from swift.llm import EmbeddingModel

# 配置环境（GPU/CPU）
context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=0)
# context.set_context(mode=context.GRAPH_MODE, device_target="CPU")  # CPU环境

# 加载微调后的模型
model = EmbeddingModel(
    model_type="m3e_base",
    model_name_or_path="./output/embedding_finetune",  # 微调后模型路径
    finetune_type="lora"
)

# 测试文本对
text_pairs = [
    ("机器学习入门", "零基础学机器学习"),
    ("机器学习入门", "烹饪技巧大全")
]

# 计算余弦相似度
def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))

for text1, text2 in text_pairs:
    vec1 = model.encode(text1)
    vec2 = model.encode(text2)
    sim = cosine_similarity(vec1, vec2)
    print(f"文本1: {text1}\n文本2: {text2}\n相似度: {sim:.4f}\n")

运行验证脚本：

bash 复制代码

python eval_embedding.py

预期输出（微调后正样本相似度更高，负样本更低）：

复制代码

文本1: 机器学习入门
文本2: 零基础学机器学习
相似度: 0.9258

文本1: 机器学习入门
文本2: 烹饪技巧大全
相似度: 0.1023

2. 推理部署（文本向量生成）

编写推理脚本（infer_embedding.py），生成任意文本的Embedding向量：

python 复制代码

from mindspore import context
from swift import Swift
from swift.llm import EmbeddingModel

# 环境配置
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU", device_id=0)

# 加载模型
model = EmbeddingModel(
    model_type="m3e_base",
    model_name_or_path="./output/embedding_finetune",
    finetune_type="lora",
    normalize=True  # 向量归一化（推荐）
)

# 生成文本向量
text = "基于ms-swift微调Embedding模型"
vec = model.encode(text)

print(f"文本: {text}")
print(f"向量维度: {vec.shape}")
print(f"向量前10维: {vec[:10]}")

运行推理脚本：

bash 复制代码

python infer_embedding.py

输出示例：

复制代码

文本: 基于ms-swift微调Embedding模型
向量维度: (768,)
向量前10维: [0.0235 0.0189 -0.0456 0.0321 0.0098 -0.0123 0.0567 -0.0289 0.0105 0.0342]

六、常见问题与优化

1. 显存不足

减小batch_size（如设为1）；
使用QLoRA（配置finetune_type: qlora），量化模型权重至4/8位；
缩短max_seq_length（如从128改为64）；
改用CPU训练（速度慢，但无显存限制）。

2. 模型效果差

增加训练数据量（至少1000+样本）；
调整LoRA参数（rank改为16，alpha改为64）；
增大学习率（如5e-4）或训练轮数；
优化数据集质量（确保正负样本区分度高）。

3. 模型下载失败

手动下载预训练模型（如m3e-base）到本地，修改model_name_or_path为本地路径；
配置国内镜像源（如清华镜像）。

七、总结

本文从原理（Embedding微调、LoRA、ms-swift核心）到实践（Ubuntu22.04环境搭建、数据准备、配置训练、验证推理）完整覆盖了ms-swift微调Embedding模型的流程。核心要点：

ms-swift通过配置化实现低资源高效微调，无需大量代码；
LoRA是Embedding模型微调的首选策略，平衡显存与效果；
对比学习损失（InfoNCE）是Embedding微调的核心损失函数；
数据集质量直接决定微调效果，需保证正负样本区分度。

后续可扩展方向：

适配更大的Embedding模型（如m3e-large、BERT-large）；
结合增量预训练进一步优化领域适配；
将微调后的模型部署为API服务（如FastAPI）。