尧图网页设计 
FlagEmbedding BGE-M3 微调建模源码解析EncoderOnlyEmbedderM3Model 的三路表征与统一微调实现【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding本篇技术指南以 FlagEmbedding 官方 API 文档 modeling.rst 为骨架结合仓库内 M3 微调模块 的完整源码深入剖析 EncoderOnlyEmbedderM3Model 及其推理子类 EncoderOnlyEmbedderM3ModelForInference 的架构设计、三路表征Dense / Sparse / ColBERT的生成与打分机制、统一微调Unified Fine-tuning与自蒸馏损失的计算细节并给出基于 examples/finetune/embedder/encoder_only/m3.sh 的可复现实战配置。读者读完即可掌握 M3 模型在 FlagEmbedding 框架中的建模原理、关键参数含义及二次开发切入点。一、模块定位M3 微调建模在 FlagEmbedding 中的角色FlagEmbedding 将 BGE-M3 的多向量混合检索能力封装在FlagEmbedding.finetune.embedder.encoder_only.m3包中该包仅包含 5 个文件modeling.py本文主角定义EncoderOnlyEmbedderM3Model与EncoderOnlyEmbedderM3ModelForInferencearguments.py模型参数与训练参数的数据类runner.py负责模型加载、tokenizer 加载与 Trainer 装配trainer.py自定义 Trainer 的保存逻辑main.py命令行入口支持torchrun -m FlagEmbedding.finetune.embedder.encoder_only.m3直接启动训练。从继承关系看EncoderOnlyEmbedderM3Model继承自抽象基类AbsEmbedderModelAbsModeling.py因此它天然具备基类提供的能力in-batch negatives 损失、跨设备 negatives 损失negatives_cross_device、知识蒸馏损失分发kd_loss_type、MRL 支持等。但 M3 是一个多向量、三路表征的模型其核心差异体现在除了常规 dense 向量外还通过两个额外线性层产出 sparse 词权重与 ColBERT token 向量并据此设计了专用的m3_kd_loss蒸馏方式与四路损失加权方案。这正是本模块区别于 base embedder 建模的精华所在。二、模型构造与关键超参EncoderOnlyEmbedderM3Model.__init__的完整签名modeling.py参数默认值含义base_model必填由 Runner 装配好的 dict含model、colbert_linear、sparse_linear三个组件tokenizerNone训练所用 tokenizernegatives_cross_deviceFalse是否启用跨设备负样本需要先初始化分布式环境temperature1.0打分时的温度系数控制分数缩放sub_batch_size-1编码时的子批次大小为负则不拆分子批次省显存技巧kd_loss_typem3_kd_loss蒸馏损失类型M3 默认专用m3_kd_loss也可用kl_divuse_mrlFalse是否启用 MRL 训练——注意 M3 构造函数中若置 True 会直接raise NotImplementedErrormrl_dims[]MRL 各层维度M3 当前不支持sentence_pooling_methodclsdense 向量池化方式normalize_embeddingsFalse是否对 dense / colbert 向量做 L2 归一化unified_finetuningTrue是否启用统一微调训练 sparse 与 colbert 头use_self_distillFalse是否启用自蒸馏self_distill_start_step-1自蒸馏起始步数关键实现细节组件装配unified_finetuningTrue时self.model / self.colbert_linear / self.sparse_linear全部取自base_model为False时只保留 backbone两个线性头置为None模型退化为纯 dense 微调。sparse 头是Linear(hidden_size, 1)输出经relu得到非负 token 权重colbert 头是Linear(hidden_size, hidden_size)维度可通过colbert_dim覆盖见 runner.py。构造时即持有self.cross_entropy torch.nn.CrossEntropyLoss(reductionmean)用于最终的对比学习损失同时记录self.vocab_size供 sparse 向量 one-hot 展开使用。配套参数类arguments.pyEncoderOnlyEmbedderM3ModelArguments继承AbsEmbedderModelArguments额外只有colbert_dim: int -1colbert 线性层输出维度≤0 时用 hidden_sizeEncoderOnlyEmbedderM3TrainingArguments继承AbsEmbedderTrainingArguments新增unified_finetuning、use_self_distill、fix_encoder冻结 backbone、只训练两个头、self_distill_start_step四个开关。三、三路表征的生成_dense_embedding/_sparse_embedding/_colbert_embedding三个下划线私有方法分别实现 Dense、Sparse、ColBERT 三种表征统一在_encode与encode中串联调用。3.1 Dense池化得到句向量_dense_embedding从 backbone 的last_hidden_state提取句向量支持三种池化cls直接取第 0 个 token 的隐状态last_hidden_state[:, 0]mean按 attention_mask 加权求和后除以有效 token 数last_token根据 padding 方向取最后一个有效 token——先判断是否为左侧 paddingattention_mask[:, -1].sum() attention_mask.shape[0]左侧 padding 取[:, -1]否则用attention_mask.sum(dim1) - 1定位每个样本的末尾位置索引后 gather。其余池化方式会抛出NotImplementedError这是推理端与训练端保持一致性的基础。3.2 Sparse词级权重展开为词表向量_sparse_embedding是理解 M3 稀疏检索的关键将 last_hidden_state 过sparse_linear后接relu得到每个 token 的非负权重token_weightsreturn_embeddingFalse时直接返回 token 权重推理时可配合 BM25 风格加权训练与推理分支使用不同实现训练态初始化(batch, seq_len, vocab_size)零张量用torch.scatter把 token 权重按input_ids写入对应词表位置再沿序列维度取max汇聚推理态采用scatter_reduce(..., reduceamax)直接在(batch, vocab_size)上取每个词的最大权重。源码注释明确指出这一优化源自 issue #1364并强调训练态不能使用该路径否则会触发RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation最后将cls / eos / pad / unk四个特殊 token 对应的词表维度权重清零sparse_embedding[:, unused_tokens] * 0.。得到的 sparse 向量可直接与 dense 向量拼接或独立建稀疏索引实现词级精确匹配能力。3.3 ColBERTtoken 级向量_colbert_embedding将 last_hidden_state 去掉[CLS]位[:, 1:]后过colbert_linear投影再按 attention_mask 去掉 padding tokencolbert_vecs * mask[:, 1:][:, :, None]得到(batch, seq_len, dim)的 token 向量序列用于后期交互late interaction式的细粒度匹配。3.4 统一入口_encode与encode_encode单次前向last_hidden_state self.model(**features, return_dictTrue).last_hidden_state dense_vecs self._dense_embedding(last_hidden_state, features[attention_mask]) sparse_vecs self._sparse_embedding(last_hidden_state, features[input_ids]) # unified_finetuning 时 colbert_vecs self._colbert_embedding(last_hidden_state, features[attention_mask]) # unified_finetuning 时 if self.normalize_embeddings: dense_vecs F.normalize(dense_vecs, dim-1) colbert_vecs F.normalize(colbert_vecs, dim-1)encode则处理两种输入形态dict时按sub_batch_size切分避免 OOM或整体编码list[dict]时逐个编码后torch.cat。统一微调开启时返回(dense_vecs, sparse_vecs, colbert_vecs)三元组关闭时返回(dense_vecs, None, None)并全部调用.contiguous()保证内存布局。四、三种打分与集成打分compute_*_score/ensemble_score4.1 Dense / Sparse 打分compute_dense_score与compute_sparse_score逻辑一致先由_compute_similarity计算内积相似度二维输入走matmul(q, p.T)更高维走matmul(q, p.transpose(-2,-1))再除以temperature并view(q_reps.size(0), -1)还原为(batch, batch*group_size)的得分矩阵。4.2 ColBERT 打分后期交互compute_colbert_score实现 ColBERT 的 MaxSim 后期交互token_scores torch.einsum(qin,pjn-qipj, q_reps, p_reps) # 全量 token 对点积 scores, _ token_scores.max(-1) # 每个 query token 取最相似 passage token scores scores.sum(1) / q_mask[:, 1:].sum(-1, keepdimTrue) # 求和并按 query 有效 token 数归一 scores scores / self.temperature其中q_mask来自_get_queries_attention_mask当 queries 是list[dict]时会按 padding 方向把各子批的 attention_mask 右侧或左侧 pad 到同一长度再拼接保证归一化分母正确。4.3 三路集成打分compute_score提供带权重的即时集成dense_score * dense_weight sparse_score * sparse_weight colbert_score * colbert_weight默认权重为dense1.0, sparse0.3, colbert1.0。而ensemble_score则要求三路分数作为参数传入硬编码dense 0.3*sparse colbert三个分数缺一即抛ValueError。这组权重与官方 BGE-M3 推理侧的0.3稀疏权重设计保持一致在训练中也被复用为集成损失。五、前向传播与统一微调损失forward/compute_loss5.1 forward 主流程forward是训练核心先编码 query 与 passagepassage 形状为(batch*group_size, dim)随后仅在self.training分支计算损失teacher_targets 构造若传入teacher_scoresreshape 为(batch, group_size)后detach()并softmax化为分布否则为None走纯对比学习。损失函数选择no_in_batch_neg_flag为 True 时用_compute_no_in_batch_neg_loss否则按negatives_cross_device选择跨设备或批内负样本实现这些实现在基类 AbsModeling.py 中。三路损失 集成损失unified_finetuning开启时dense loss由compute_dense_score计算sparse loss由compute_sparse_score计算加权系数 0.1colbert loss由compute_colbert_score计算并传入q_maskself._get_queries_attention_mask(queries)ensemble loss用ensemble_score融合三路分数后再算对比损失。最终loss (loss ensemble_loss 0.1 * sparse_loss colbert_loss) / 4。跨设备截取开启negatives_cross_device时按process_rank从全局得分矩阵中截取本进程的 dense 得分参与集成源码注释引用了 issue #1410 的 bug 修复no_in_batch_neg_flag下需先按group_size取对角线 passage。自蒸馏use_self_distill且self.step self_distill_start_step时以集成得分softmax作为教师分布分别对 dense / sparse / colbert 三路计算kl_div蒸馏损失并累加最后整体再除 2。5.2 compute_loss 与 m3_kd_losscompute_loss即对得分矩阵与目标索引求交叉熵。而 M3 专用的蒸馏损失m3_kd_loss实现在基类 AbsModeling.py它以group_size为步长构造目标位置逐组对得分加掩码后计算逐样本交叉熵并用教师分布teacher_targets[:, i]对每组损失加权求和——这使得「同一 query 的多个正负样本」各自获得与教师置信度匹配的梯度贡献比单纯 KL 散度更贴合 BGE-M3 的蒸馏训练。5.3 内存优化钩子gradient_checkpointing_enable与enable_input_require_grads分别透传 backbone 的梯度检查点与输入梯度要求。Runner 在加载模型后会根据训练参数调用它们见 runner.py并支持通过fix_position_embedding/fix_encoder冻结特定参数。六、模型保存save与 Trainer 落盘save先以 CPU 克隆方式保存 backbone 权重save_pretrained随后在unified_finetuning开启时将两个头的状态字典分别落盘为colbert_linear.pt与sparse_linear.pt。这正好与加载逻辑对称——Runner 的get_model会检查模型目录下是否同时存在这两个文件存在则恢复权重否则提示「参数为新初始化确保是训练而非推理加载」。Trainer 的_save则调用self.model.save(output_dir)同时保存 tokenizer 与training_args.bin保证 checkpoint 可完整恢复训练。七、推理子类EncoderOnlyEmbedderM3ModelForInferenceEncoderOnlyEmbedderM3ModelForInference继承训练类并覆写 forward将训练语义切换为按需产出表征的推理语义参数默认值说明return_denseTrue返回 dense 句向量return_sparseFalse返回 sparse 向量词表维度return_colbert_vecsFalse返回 colbert token 向量return_sparse_embeddingFalse透传给_sparse_embeddingTrue 返回展开的词表向量False 只返回 token 权重truncate_dimNone截断输出维度兼容 Matryoshka 表征学习类模型实现要点入口断言return_dense or return_sparse or return_colbert_vecs至少一项为 True强制self.training False从而让 sparse 走推理态的高效scatter_reduce(amax)路径再次呼应 issue #1364 的优化只做一次 backbone 前向按需组合输出字典{dense_vecs, sparse_vecs, colbert_vecs}normalize_embeddings开启时对 dense 与 colbert 做 L2 归一化。该子类与推理侧FlagEmbedding.inference.embedder.decoder_only/encoder_only中的封装共同构成 M3 从训练到部署的完整链路。八、实战统一微调完整命令与参数说明仓库提供了可直接运行的示例脚本 examples/finetune/embedder/encoder_only/m3.sh其核心配置如下路径均相对该脚本所在目录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_epochs4 per_device_train_batch_size2 num_gpus2 model_args--model_name_or_path BAAI/bge-m3 --cache_dir $HF_HUB_CACHE 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 \ --knowledge_distillation False \ training_args\ --output_dir ./test_encoder_only_m3_bge-m3 \ --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --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 \ cmdtorchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.m3 \ $model_args $data_args $training_args关键参数逐条解读对应前文源码--unified_finetuning True开启 sparse / colbert 头训练此时 forward 才会计算四路损失--use_self_distill True --self_distill_start_step 0从第 0 步起启用集成分数自蒸馏--kd_loss_type m3_kd_loss使用上文 5.2 的 M3 专用蒸馏损失--negatives_cross_device跨设备负样本需配合torchrun多卡--temperature 0.02打分温度直接进入compute_*_score的分母--sentence_pooling_method cls --normalize_embeddings True与模型构造参数一一对应--train_group_size 8每个 query 对应 1 正 7 负与基类get_local_score中的group_size计算一致--gradient_checkpointing训练时由 Runner 调用model.enable_input_require_grads()配合梯度检查点。数据格式要求每个训练样本需包含query: str、pos: List[str]、neg: List[str]字段见 AbsArguments.py混合 retrieval / sts / classification / clustering 四种任务时classification-no_in_batch_neg这类数据集对应no_in_batch_neg_flag路径源码在 forward 中会据此选择_compute_no_in_batch_neg_loss。若设置same_dataset_within_batchRunner 还会注册数据刷新回调runner.py。九、总结与延伸阅读EncoderOnlyEmbedderM3Model的设计精髓可概括为三点三路表征解耦dense 池化 sparse 词权重 colbert 后期交互、统一微调联合训练四路损失加权 可选自蒸馏、训练/推理实现分离sparse 的高效scatter_reduce推理路径与显存友好的子批次编码。理解这一建模层是进一步阅读推理封装、评估脚本或二次开发的基础。推理侧封装见 FlagEmbedding/inference/embedderM3 对应 encoder_only/m3.py基类损失与蒸馏见 AbsModeling.py训练入口与参数见main.py、arguments.py端到端示例见 examples/finetune/embedder/encoder_only/m3.sh 与 m3_same_dataset.sh官方 API 文档见 modeling.rst本文章节一一对应其中的方法清单。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 FlagEmbedding BGE-M3 微调建模源码解析：EncoderOnlyEmbedderM3Model 的三路表征与统一微调实现
FlagEmbedding BGE-M3 微调建模源码解析EncoderOnlyEmbedderM3Model 的三路表征与统一微调实现【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding本篇技术指南以 FlagEmbedding 官方 API 文档 modeling.rst 为骨架结合仓库内 M3 微调模块 的完整源码深入剖析 EncoderOnlyEmbedderM3Model 及其推理子类 EncoderOnlyEmbedderM3ModelForInference 的架构设计、三路表征Dense / Sparse / ColBERT的生成与打分机制、统一微调Unified Fine-tuning与自蒸馏损失的计算细节并给出基于 examples/finetune/embedder/encoder_only/m3.sh 的可复现实战配置。读者读完即可掌握 M3 模型在 FlagEmbedding 框架中的建模原理、关键参数含义及二次开发切入点。一、模块定位M3 微调建模在 FlagEmbedding 中的角色FlagEmbedding 将 BGE-M3 的多向量混合检索能力封装在FlagEmbedding.finetune.embedder.encoder_only.m3包中该包仅包含 5 个文件modeling.py本文主角定义EncoderOnlyEmbedderM3Model与EncoderOnlyEmbedderM3ModelForInferencearguments.py模型参数与训练参数的数据类runner.py负责模型加载、tokenizer 加载与 Trainer 装配trainer.py自定义 Trainer 的保存逻辑main.py命令行入口支持torchrun -m FlagEmbedding.finetune.embedder.encoder_only.m3直接启动训练。从继承关系看EncoderOnlyEmbedderM3Model继承自抽象基类AbsEmbedderModelAbsModeling.py因此它天然具备基类提供的能力in-batch negatives 损失、跨设备 negatives 损失negatives_cross_device、知识蒸馏损失分发kd_loss_type、MRL 支持等。但 M3 是一个多向量、三路表征的模型其核心差异体现在除了常规 dense 向量外还通过两个额外线性层产出 sparse 词权重与 ColBERT token 向量并据此设计了专用的m3_kd_loss蒸馏方式与四路损失加权方案。这正是本模块区别于 base embedder 建模的精华所在。二、模型构造与关键超参EncoderOnlyEmbedderM3Model.__init__的完整签名modeling.py参数默认值含义base_model必填由 Runner 装配好的 dict含model、colbert_linear、sparse_linear三个组件tokenizerNone训练所用 tokenizernegatives_cross_deviceFalse是否启用跨设备负样本需要先初始化分布式环境temperature1.0打分时的温度系数控制分数缩放sub_batch_size-1编码时的子批次大小为负则不拆分子批次省显存技巧kd_loss_typem3_kd_loss蒸馏损失类型M3 默认专用m3_kd_loss也可用kl_divuse_mrlFalse是否启用 MRL 训练——注意 M3 构造函数中若置 True 会直接raise NotImplementedErrormrl_dims[]MRL 各层维度M3 当前不支持sentence_pooling_methodclsdense 向量池化方式normalize_embeddingsFalse是否对 dense / colbert 向量做 L2 归一化unified_finetuningTrue是否启用统一微调训练 sparse 与 colbert 头use_self_distillFalse是否启用自蒸馏self_distill_start_step-1自蒸馏起始步数关键实现细节组件装配unified_finetuningTrue时self.model / self.colbert_linear / self.sparse_linear全部取自base_model为False时只保留 backbone两个线性头置为None模型退化为纯 dense 微调。sparse 头是Linear(hidden_size, 1)输出经relu得到非负 token 权重colbert 头是Linear(hidden_size, hidden_size)维度可通过colbert_dim覆盖见 runner.py。构造时即持有self.cross_entropy torch.nn.CrossEntropyLoss(reductionmean)用于最终的对比学习损失同时记录self.vocab_size供 sparse 向量 one-hot 展开使用。配套参数类arguments.pyEncoderOnlyEmbedderM3ModelArguments继承AbsEmbedderModelArguments额外只有colbert_dim: int -1colbert 线性层输出维度≤0 时用 hidden_sizeEncoderOnlyEmbedderM3TrainingArguments继承AbsEmbedderTrainingArguments新增unified_finetuning、use_self_distill、fix_encoder冻结 backbone、只训练两个头、self_distill_start_step四个开关。三、三路表征的生成_dense_embedding/_sparse_embedding/_colbert_embedding三个下划线私有方法分别实现 Dense、Sparse、ColBERT 三种表征统一在_encode与encode中串联调用。3.1 Dense池化得到句向量_dense_embedding从 backbone 的last_hidden_state提取句向量支持三种池化cls直接取第 0 个 token 的隐状态last_hidden_state[:, 0]mean按 attention_mask 加权求和后除以有效 token 数last_token根据 padding 方向取最后一个有效 token——先判断是否为左侧 paddingattention_mask[:, -1].sum() attention_mask.shape[0]左侧 padding 取[:, -1]否则用attention_mask.sum(dim1) - 1定位每个样本的末尾位置索引后 gather。其余池化方式会抛出NotImplementedError这是推理端与训练端保持一致性的基础。3.2 Sparse词级权重展开为词表向量_sparse_embedding是理解 M3 稀疏检索的关键将 last_hidden_state 过sparse_linear后接relu得到每个 token 的非负权重token_weightsreturn_embeddingFalse时直接返回 token 权重推理时可配合 BM25 风格加权训练与推理分支使用不同实现训练态初始化(batch, seq_len, vocab_size)零张量用torch.scatter把 token 权重按input_ids写入对应词表位置再沿序列维度取max汇聚推理态采用scatter_reduce(..., reduceamax)直接在(batch, vocab_size)上取每个词的最大权重。源码注释明确指出这一优化源自 issue #1364并强调训练态不能使用该路径否则会触发RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation最后将cls / eos / pad / unk四个特殊 token 对应的词表维度权重清零sparse_embedding[:, unused_tokens] * 0.。得到的 sparse 向量可直接与 dense 向量拼接或独立建稀疏索引实现词级精确匹配能力。3.3 ColBERTtoken 级向量_colbert_embedding将 last_hidden_state 去掉[CLS]位[:, 1:]后过colbert_linear投影再按 attention_mask 去掉 padding tokencolbert_vecs * mask[:, 1:][:, :, None]得到(batch, seq_len, dim)的 token 向量序列用于后期交互late interaction式的细粒度匹配。3.4 统一入口_encode与encode_encode单次前向last_hidden_state self.model(**features, return_dictTrue).last_hidden_state dense_vecs self._dense_embedding(last_hidden_state, features[attention_mask]) sparse_vecs self._sparse_embedding(last_hidden_state, features[input_ids]) # unified_finetuning 时 colbert_vecs self._colbert_embedding(last_hidden_state, features[attention_mask]) # unified_finetuning 时 if self.normalize_embeddings: dense_vecs F.normalize(dense_vecs, dim-1) colbert_vecs F.normalize(colbert_vecs, dim-1)encode则处理两种输入形态dict时按sub_batch_size切分避免 OOM或整体编码list[dict]时逐个编码后torch.cat。统一微调开启时返回(dense_vecs, sparse_vecs, colbert_vecs)三元组关闭时返回(dense_vecs, None, None)并全部调用.contiguous()保证内存布局。四、三种打分与集成打分compute_*_score/ensemble_score4.1 Dense / Sparse 打分compute_dense_score与compute_sparse_score逻辑一致先由_compute_similarity计算内积相似度二维输入走matmul(q, p.T)更高维走matmul(q, p.transpose(-2,-1))再除以temperature并view(q_reps.size(0), -1)还原为(batch, batch*group_size)的得分矩阵。4.2 ColBERT 打分后期交互compute_colbert_score实现 ColBERT 的 MaxSim 后期交互token_scores torch.einsum(qin,pjn-qipj, q_reps, p_reps) # 全量 token 对点积 scores, _ token_scores.max(-1) # 每个 query token 取最相似 passage token scores scores.sum(1) / q_mask[:, 1:].sum(-1, keepdimTrue) # 求和并按 query 有效 token 数归一 scores scores / self.temperature其中q_mask来自_get_queries_attention_mask当 queries 是list[dict]时会按 padding 方向把各子批的 attention_mask 右侧或左侧 pad 到同一长度再拼接保证归一化分母正确。4.3 三路集成打分compute_score提供带权重的即时集成dense_score * dense_weight sparse_score * sparse_weight colbert_score * colbert_weight默认权重为dense1.0, sparse0.3, colbert1.0。而ensemble_score则要求三路分数作为参数传入硬编码dense 0.3*sparse colbert三个分数缺一即抛ValueError。这组权重与官方 BGE-M3 推理侧的0.3稀疏权重设计保持一致在训练中也被复用为集成损失。五、前向传播与统一微调损失forward/compute_loss5.1 forward 主流程forward是训练核心先编码 query 与 passagepassage 形状为(batch*group_size, dim)随后仅在self.training分支计算损失teacher_targets 构造若传入teacher_scoresreshape 为(batch, group_size)后detach()并softmax化为分布否则为None走纯对比学习。损失函数选择no_in_batch_neg_flag为 True 时用_compute_no_in_batch_neg_loss否则按negatives_cross_device选择跨设备或批内负样本实现这些实现在基类 AbsModeling.py 中。三路损失 集成损失unified_finetuning开启时dense loss由compute_dense_score计算sparse loss由compute_sparse_score计算加权系数 0.1colbert loss由compute_colbert_score计算并传入q_maskself._get_queries_attention_mask(queries)ensemble loss用ensemble_score融合三路分数后再算对比损失。最终loss (loss ensemble_loss 0.1 * sparse_loss colbert_loss) / 4。跨设备截取开启negatives_cross_device时按process_rank从全局得分矩阵中截取本进程的 dense 得分参与集成源码注释引用了 issue #1410 的 bug 修复no_in_batch_neg_flag下需先按group_size取对角线 passage。自蒸馏use_self_distill且self.step self_distill_start_step时以集成得分softmax作为教师分布分别对 dense / sparse / colbert 三路计算kl_div蒸馏损失并累加最后整体再除 2。5.2 compute_loss 与 m3_kd_losscompute_loss即对得分矩阵与目标索引求交叉熵。而 M3 专用的蒸馏损失m3_kd_loss实现在基类 AbsModeling.py它以group_size为步长构造目标位置逐组对得分加掩码后计算逐样本交叉熵并用教师分布teacher_targets[:, i]对每组损失加权求和——这使得「同一 query 的多个正负样本」各自获得与教师置信度匹配的梯度贡献比单纯 KL 散度更贴合 BGE-M3 的蒸馏训练。5.3 内存优化钩子gradient_checkpointing_enable与enable_input_require_grads分别透传 backbone 的梯度检查点与输入梯度要求。Runner 在加载模型后会根据训练参数调用它们见 runner.py并支持通过fix_position_embedding/fix_encoder冻结特定参数。六、模型保存save与 Trainer 落盘save先以 CPU 克隆方式保存 backbone 权重save_pretrained随后在unified_finetuning开启时将两个头的状态字典分别落盘为colbert_linear.pt与sparse_linear.pt。这正好与加载逻辑对称——Runner 的get_model会检查模型目录下是否同时存在这两个文件存在则恢复权重否则提示「参数为新初始化确保是训练而非推理加载」。Trainer 的_save则调用self.model.save(output_dir)同时保存 tokenizer 与training_args.bin保证 checkpoint 可完整恢复训练。七、推理子类EncoderOnlyEmbedderM3ModelForInferenceEncoderOnlyEmbedderM3ModelForInference继承训练类并覆写 forward将训练语义切换为按需产出表征的推理语义参数默认值说明return_denseTrue返回 dense 句向量return_sparseFalse返回 sparse 向量词表维度return_colbert_vecsFalse返回 colbert token 向量return_sparse_embeddingFalse透传给_sparse_embeddingTrue 返回展开的词表向量False 只返回 token 权重truncate_dimNone截断输出维度兼容 Matryoshka 表征学习类模型实现要点入口断言return_dense or return_sparse or return_colbert_vecs至少一项为 True强制self.training False从而让 sparse 走推理态的高效scatter_reduce(amax)路径再次呼应 issue #1364 的优化只做一次 backbone 前向按需组合输出字典{dense_vecs, sparse_vecs, colbert_vecs}normalize_embeddings开启时对 dense 与 colbert 做 L2 归一化。该子类与推理侧FlagEmbedding.inference.embedder.decoder_only/encoder_only中的封装共同构成 M3 从训练到部署的完整链路。八、实战统一微调完整命令与参数说明仓库提供了可直接运行的示例脚本 examples/finetune/embedder/encoder_only/m3.sh其核心配置如下路径均相对该脚本所在目录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_epochs4 per_device_train_batch_size2 num_gpus2 model_args--model_name_or_path BAAI/bge-m3 --cache_dir $HF_HUB_CACHE 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 \ --knowledge_distillation False \ training_args\ --output_dir ./test_encoder_only_m3_bge-m3 \ --overwrite_output_dir \ --learning_rate 1e-5 \ --fp16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --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 \ cmdtorchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.embedder.encoder_only.m3 \ $model_args $data_args $training_args关键参数逐条解读对应前文源码--unified_finetuning True开启 sparse / colbert 头训练此时 forward 才会计算四路损失--use_self_distill True --self_distill_start_step 0从第 0 步起启用集成分数自蒸馏--kd_loss_type m3_kd_loss使用上文 5.2 的 M3 专用蒸馏损失--negatives_cross_device跨设备负样本需配合torchrun多卡--temperature 0.02打分温度直接进入compute_*_score的分母--sentence_pooling_method cls --normalize_embeddings True与模型构造参数一一对应--train_group_size 8每个 query 对应 1 正 7 负与基类get_local_score中的group_size计算一致--gradient_checkpointing训练时由 Runner 调用model.enable_input_require_grads()配合梯度检查点。数据格式要求每个训练样本需包含query: str、pos: List[str]、neg: List[str]字段见 AbsArguments.py混合 retrieval / sts / classification / clustering 四种任务时classification-no_in_batch_neg这类数据集对应no_in_batch_neg_flag路径源码在 forward 中会据此选择_compute_no_in_batch_neg_loss。若设置same_dataset_within_batchRunner 还会注册数据刷新回调runner.py。九、总结与延伸阅读EncoderOnlyEmbedderM3Model的设计精髓可概括为三点三路表征解耦dense 池化 sparse 词权重 colbert 后期交互、统一微调联合训练四路损失加权 可选自蒸馏、训练/推理实现分离sparse 的高效scatter_reduce推理路径与显存友好的子批次编码。理解这一建模层是进一步阅读推理封装、评估脚本或二次开发的基础。推理侧封装见 FlagEmbedding/inference/embedderM3 对应 encoder_only/m3.py基类损失与蒸馏见 AbsModeling.py训练入口与参数见main.py、arguments.py端到端示例见 examples/finetune/embedder/encoder_only/m3.sh 与 m3_same_dataset.sh官方 API 文档见 modeling.rst本文章节一一对应其中的方法清单。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考