FlagEmbedding 编码器架构 Embedder 微调全解析：Base 与 M3 双路径 API 指南

FlagEmbedding 编码器架构 Embedder 微调全解析：Base 与 M3 双路径 API 指南

【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding

本文以 FlagEmbedding 官方 API 文档 docs/source/API/finetune/embedder/encoder_only.rst 为骨架，系统讲解FlagEmbedding.finetune.embedder.encoder_only模块下两类编码器（Encoder-Only）Embedder 微调路径：面向普通双塔模型的Base路径（BiEncoderOnlyEmbedderModel）与面向 BGE-M3 的多向量M3路径（EncoderOnlyEmbedderM3Model）。读完本文，你将掌握两大路径的类层次、核心方法与训练参数、损失计算原理，以及可直接运行的微调启动命令，能够在 FlagEmbedding 框架内独立配置并启动一次编码器架构 Embedder 的微调任务。

一、模块定位：Encoder-Only 微调在 FlagEmbedding 中的角色

FlagEmbedding 的finetune.embedder目录下分为encoder_only与decoder_only两大分支：前者面向 BERT/RoBERTa 这类双向编码器架构，将句子编码为固定维度的向量；后者则面向 LLM 类 decoder 架构。本文聚焦的 encoder_only 目录结构如下：

FlagEmbedding/finetune/embedder/encoder_only/ ├── __init__.py ├── base/ # Base 路径：普通双塔编码器 │ ├── modeling.py # BiEncoderOnlyEmbedderModel │ ├── runner.py # EncoderOnlyEmbedderRunner │ └── trainer.py # EncoderOnlyEmbedderTrainer └── m3/ # M3 路径：BGE-M3 多向量统一微调 ├── arguments.py # M3 专用 Model/Training 参数 ├── modeling.py # EncoderOnlyEmbedderM3Model 及推理类 ├── runner.py # EncoderOnlyEmbedderM3Runner └── trainer.py # EncoderOnlyEmbedderM3Trainer

该模块的 API 文档（base.rst 与 m3.rst）通过 Sphinx autodoc 生成，分别索引了Modeling、Runner、Trainer（Base）以及Arguments、Modeling、Runner、Trainer（M3）等子章节，所有类均继承自抽象基类目录 FlagEmbedding/abc/finetune/embedder 中的AbsEmbedderModel、AbsEmbedderRunner、AbsEmbedderTrainer。

二、Base 路径：BiEncoderOnlyEmbedderModel双塔编码器模型

2.1 类的定位与构造参数

BiEncoderOnlyEmbedderModel定义于 modeling.py，继承AbsEmbedderModel，是"查询-文档"双塔结构的核心封装。其构造参数与含义如下表：

参数	默认值	说明
base_model	必填	用于训练的基座模型（PreTrainedModel），通常为 BERT 系编码器
tokenizer	None	分词器实例
negatives_cross_device	False	是否跨设备共享负样本计算损失
temperature	1.0	控制分数缩放的温度系数
sub_batch_size	-1	编码时的子批次大小，为负则不切分
kd_loss_type	'kl_div'	知识蒸馏损失类型
use_mrl	False	是否使用 Matryoshka Representation Learning 训练
mrl_dims	[]	MRL 各层输出维度列表
sentence_pooling_method	'cls'	句向量池化方式，可选cls/mean/last_token
normalize_embeddings	False	是否对向量做 L2 归一化

值得注意：temperature与normalize_embeddings的默认值在 AbsArguments.py 中分别被命令行默认覆盖为0.02与True，即实际微调脚本默认使用 0.02 的温度并归一化向量。

2.2 核心方法逐一解析

encode(features)是模型的前向编码入口。从源码实现看（modeling.py），它有三条分支路径：

1. 子批次切分：当sub_batch_size > 0时，按 attention mask 长度将 batch 切成多个子批次分别前向，再torch.cat拼接，用于显存受限场景；
2. 整批编码：直接对整批输入做一次前向；
3. 列表输入：features为 list（每组特征长度不同）时逐组编码后拼接。

若启用use_mrl，则按mrl_dims依次截取前dim维（超出原始维度时告警并截断到原始维度），可选归一化后返回多个维度的向量列表；否则返回单条向量并（可选）归一化。

_sentence_embedding(last_hidden_state, attention_mask)实现三种池化（modeling.py）：

• cls：取last_hidden_state[:, 0]；
• mean：按 attention mask 加权求平均；
• last_token：先判断是否为左侧 padding（left_padding），左 padding 直接取末位，否则取每个序列最后一个有效 token（attention_mask.sum(dim=1) - 1定位）。

compute_score(q_reps, p_reps)与_compute_similarity(q_reps, p_reps)：相似度采用内积（torch.matmul），再除以temperature得到分数矩阵；compute_loss(scores, target)直接使用交叉熵（torch.nn.CrossEntropyLoss(reduction='mean')）。gradient_checkpointing_enable与enable_input_require_grads分别透传给底层 HuggingFace 模型，用于显存优化与梯度检查点兼容；save(output_dir)将权重克隆到 CPU 后调用save_pretrained落盘。

2.3 Runner 与 Trainer

EncoderOnlyEmbedderRunner（runner.py）负责组装训练管线：

• load_tokenizer_and_model()：用AutoTokenizer/AutoModel/AutoConfig从model_args.model_name_or_path加载，构建BiEncoderOnlyEmbedderModel，将训练参数逐一注入；开启gradient_checkpointing时调用enable_input_require_grads()；设置fix_position_embedding时遍历参数将含position_embeddings的权重requires_grad置为False；
• load_trainer()：构造EncoderOnlyEmbedderTrainer，并在same_dataset_within_batch为真时注册EmbedderTrainerCallbackForDataRefresh回调，保证同一 batch 内的样本来自同一数据集。

EncoderOnlyEmbedderTrainer（trainer.py）覆写_save：先调用模型的save(output_dir)保存权重，再由主进程保存 tokenizer，并额外把training_args以training_args.bin存入输出目录，便于后续恢复训练配置。

三、M3 路径：EncoderOnlyEmbedderM3Model多向量统一微调

3.1 与 Base 的本质区别

BGE-M3 同时产出稠密（Dense）、**稀疏（Sparse）**与 **ColBERT（多向量）**三类表示。因此EncoderOnlyEmbedderM3Model（modeling.py）在构造时接收的不再是单个base_model，而是一个 dict：{'model': ..., 'colbert_linear': ..., 'sparse_linear': ...}。其中：

• colbert_linear：Linear(hidden_size, hidden_size 或 colbert_dim)，将 token 隐层投影为 ColBERT 向量；
• sparse_linear：Linear(hidden_size, 1)，为每个 token 产出稀疏词权重。

unified_finetuning=True时三者联合训练；为False时只保留model，colbert_linear/sparse_linear置空，等价于纯稠密微调。此外该模型明确禁止 MRL：构造函数中if use_mrl is True: raise NotImplementedError。

3.2 三类表示的编码实现

• 稠密：_dense_embedding复用与 Base 相同的三种池化（cls/mean/last_token）；
• 稀疏：_sparse_embedding（modeling.py）先经sparse_linear+ ReLU 得到 token 权重，再scatter到 vocab 维度的稀疏向量上。训练态用torch.scatter，推理态用scatter_reduce(..., reduce="amax")（避免原地操作破坏梯度，详见代码内注释引用的 issue #1364）；随后将cls/eos/pad/unk特殊 token 的权重清零；
• ColBERT：_colbert_embedding取last_hidden_state[:, 1:]（跳过 [CLS]）经colbert_linear投影，并与 mask 相乘屏蔽 padding。

3.3 打分与损失：三种分数 + 集成分数

compute_dense_score、compute_sparse_score均为内积除以温度；compute_colbert_score（modeling.py）用torch.einsum('qin,pjn->qipj', q_reps, p_reps)计算 token 级相似度矩阵，对 passage 维度取 max（晚期交互），再对 query token 求和并除以 query mask 中有效 token 数。三者加权组合为最终分数：

dense_score * dense_weight + sparse_score * sparse_weight + colbert_score * colbert_weight

默认权重为 dense=1.0、sparse=0.3、colbert=1.0（compute_score签名默认值）；ensemble_score亦按dense + 0.3 * sparse + colbert合成集成分数。

forward()（modeling.py）的损失逻辑最值得关注：训练态下，dense / sparse / colbert 三路分别用compute_loss_func计算损失（negatives_cross_device或no_in_batch_neg_flag会切换为跨设备/无 in-batch 负样本损失），再计算 ensemble 损失，最终：

loss = (loss + ensemble_loss + 0.1 * sparse_loss + colbert_loss) / 4

若开启use_self_distill且self.step > self_distill_start_step，则以 ensemble 分数（detach 后 softmax）作为软标签，对三路分数各算一次 KL 散度自蒸馏损失，叠加后整体减半。teacher_scores非空时则走知识蒸馏：以教师分数 softmax 作为teacher_targets。

3.4 M3 专用参数类

EncoderOnlyEmbedderM3ModelArguments与EncoderOnlyEmbedderM3TrainingArguments（arguments.py）在抽象参数之上扩展了：

参数	默认值	说明
colbert_dim	-1	ColBERT 线性层输出维度，≤0 时沿用hidden_size
unified_finetuning	False	是否统一微调三路表示
use_self_distill	False	统一微调时是否使用自蒸馏
fix_encoder	False	冻结编码器，仅训练 colbert/sparse 线性层
self_distill_start_step	-1	自蒸馏启动的步数阈值

其中fix_encoder在EncoderOnlyEmbedderM3Runner.load_tokenizer_and_model（m3/runner.py）中实现：遍历参数时仅放行名字含colbert_linear或sparse_linear的权重。同时该 Runner 的静态方法get_model会从本地路径或 HuggingFace Hub 拉取模型，新建两个线性层；若模型目录下已存在colbert_linear.pt与sparse_linear.pt（由save方法保存），则自动加载续训，否则视为全新初始化并打印提示。

3.5 推理封装：EncoderOnlyEmbedderM3ModelForInference

该子类（modeling.py）重写forward，通过return_dense/return_sparse/return_colbert_vecs三个开关按需输出，且断言三者至少一个为真。truncate_dim可对 dense/colbert 向量做维度截断（兼容 Matryoshka 场景）；return_sparse_embedding控制稀疏输出是完整 embedding 还是仅 token 权重。进入该推理分支时会强制self.training = False，保证稀疏计算走非原地路径。

四、M3 Trainer 与 Base Trainer 的保存差异

EncoderOnlyEmbedderM3Trainer的_save与 Base 版行为一致（调用model.save(output_dir)、保存 tokenizer 与training_args.bin），差异集中在EncoderOnlyEmbedderM3Model.save（modeling.py）：除主干权重外，unified_finetuning模式下还会额外保存colbert_linear.pt与sparse_linear.pt两个独立文件，这正是下一轮训练时get_model能加载续训的前提。

五、实战：从命令行启动 Encoder-Only 微调

仓库提供了开箱即用的脚本 base.sh 与 m3.sh。两者公共的数据与训练配置如下（测试用途，正式训练请调大 epoch 与 batch）：

export WANDB_MODE=disabled train_data="\ ../example_data/retrieval \ ../example_data/sts/sts.jsonl \ ../example_data/classification-no_in_batch_neg \ ../example_data/clustering-no_in_batch_neg " num_train_epochs=4 per_device_train_batch_size=2 num_gpus=2 data_args="\ --train_data $train_data \ --cache_path ~/.cache \ --train_group_size 8 \ --query_max_len 512 \ --passage_max_len 512 \ --pad_to_multiple_of 8 \ "

Base 路径（对应文档encoder_only/base）启动命令：

torchrun --nproc_per_node 2 \ -m FlagEmbedding.finetune.embedder.encoder_only.base \ --model_name_or_path BAAI/bge-large-en-v1.5 \ --query_instruction_for_retrieval 'Represent this sentence for searching relevant passages: ' \ --output_dir ./test_encoder_only_base_bge-large-en-v1.5 \ --overwrite_output_dir --learning_rate 1e-5 --fp16 \ --num_train_epochs 4 --per_device_train_batch_size 2 \ --dataloader_drop_last True --warmup_ratio 0.1 \ --gradient_checkpointing --deepspeed ../../ds_stage0.json \ --logging_steps 1 --save_steps 1000 \ --negatives_cross_device --temperature 0.02 \ --sentence_pooling_method cls --normalize_embeddings True \ --kd_loss_type kl_div

M3 路径（对应文档encoder_only/m3）在 Base 基础上追加多向量微调参数：

torchrun --nproc_per_node 2 \ -m FlagEmbedding.finetune.embedder.encoder_only.m3 \ --model_name_or_path BAAI/bge-m3 \ --output_dir ./test_encoder_only_m3_bge-m3 \ --learning_rate 1e-5 --fp16 --num_train_epochs 4 \ --per_device_train_batch_size 2 --dataloader_drop_last True \ --warmup_ratio 0.1 --gradient_checkpointing \ --deepspeed ../../ds_stage0.json --logging_steps 1 --save_steps 1000 \ --negatives_cross_device --temperature 0.02 \ --sentence_pooling_method cls --normalize_embeddings True \ --kd_loss_type m3_kd_loss \ --unified_finetuning True --use_self_distill True \ --fix_encoder False --self_distill_start_step 0

5.1 核心训练参数速查表

以下参数定义于 AbsArguments.py，适用于两条路径：

参数	默认值	说明
negatives_cross_device	False	跨设备共享负样本（多卡时等价于扩大 batch 的负样本数）
temperature	0.02	相似度分数缩放温度
fix_position_embedding	False	冻结 position embeddings 参数
sentence_pooling_method	cls	池化方式，可选cls/mean/last_token
normalize_embeddings	True	是否归一化输出向量
sub_batch_size	None	训练编码子批次大小
kd_loss_type	kl_div	蒸馏损失，可选kl_div/m3_kd_loss
use_mrl/mrl_dims	False/[]	Matryoshka 表示学习开关与维度列表（M3 模型不支持）
train_data	必填	训练数据路径，要求每条含query、pos: List[str]、neg: List[str]字段
train_group_size	8	每组样本数（含正负样本）
query_max_len/passage_max_len	32/128	查询/文档最大长度
query_instruction_for_retrieval	None	查询侧指令前缀
knowledge_distillation	False	数据含pos_scores/neg_scores时启用蒸馏
same_dataset_within_batch	False	同一 batch 样本来自同一数据集（多数据集训练时防止跨集互相充当负样本）

六、总结与延伸阅读

encoder_only模块为 BERT 系编码器提供了两条成熟微调路径：Base 路径以BiEncoderOnlyEmbedderModel的"内积相似度 + 交叉熵"双塔范式适配各类稠密检索任务，支持 MRL、跨设备负样本、知识蒸馏与三种池化；M3 路径则在EncoderOnlyEmbedderM3Model中实现 dense/sparse/ColBERT 三路联合训练、加权集成打分、自蒸馏与统一的fix_encoder冻结策略，是 BGE-M3 能力向微调场景的完整开放接口。

若希望进一步深入，可依次阅读：

• 抽象基类定义：FlagEmbedding/abc/finetune/embedder（AbsEmbedderModel、AbsEmbedderRunner、AbsEmbedderTrainer、AbsArguments）；
• 解码器（LLM）架构微调对照：decoder_only；
• 推理侧封装：FlagEmbedding/inference/embedder/encoder_only；
• 完整 API 文档索引：docs/source/API/finetune/embedder.rst。

【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding