尧图网页设计 
FlagEmbedding 分层解码器重排器 CrossDecoderModel 建模与微调全解析【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding导读本文围绕 FlagEmbedding 中 decoder-only 分层layerwise重排器reranker的模型实现展开核心对象是CrossDecoderModel。该模型以 MiniCPM 这类因果语言模型为骨干通过在多个 Transformer 层同时接出 logits 并做加权打分的方式实现逐层早停式重排可在精度与推理开销之间灵活取舍。读完本文你将掌握CrossDecoderModel.encode/forward的完整前向与损失计算逻辑、start_layer/head_multi/head_type等关键配置的作用以及如何借助 LoRA 在官方示例脚本上完成一次分层重排器的端到端微调。一、文档定位从 API 文档到源码实现本文所依托的 API 文档位于docs/source/API/finetune/reranker/decoder_only/layerwise/modeling.rst其通过 Sphinx autodoc 声明了三个核心符号FlagEmbedding.finetune.reranker.decoder_only.layerwise.CrossDecoderModel类CrossDecoderModel.encode方法CrossDecoderModel.forward方法也就是说该文档的实体内容集中在FlagEmbedding/finetune/reranker/decoder_only/layerwise/modeling.py这一个文件中。与之配套的还有同目录下的arguments.py重排器专属的模型超参LoRA、start_layer、head_type等modeling_minicpm_reranker.py分层 MiniCPM 骨干LayerWiseMiniCPMForCausalLM、LayerWiseMiniCPMConfig、LayerWiseHeadload_model.py按model_type分支加载原始模型或已微调模型并负责头部替换与 LoRA 装配runner.py组装 tokenizer / model / trainer 并执行训练与合并保存。下文将以modeling.py为主线结合这些源码逐层展开。二、CrossDecoderModel分层解码器重排器的门面类CrossDecoderModel继承自AbsRerankerModel位于FlagEmbedding/abc/finetune/reranker/AbsModeling.py后者是 FlagEmbedding 微调框架中所有 reranker 模型的抽象基类负责维护base_model、tokenizer、train_batch_size、cross_entropy等公共成员。2.1 构造参数class CrossDecoderModel(AbsRerankerModel): def __init__( self, base_model: PreTrainedModel, tokenizer: AutoTokenizer None, train_batch_size: int 4, start_layer: int 8 ): super().__init__(base_model, tokenizertokenizer, train_batch_sizetrain_batch_size) self.start_layer start_layer参数类型默认值含义base_modelPreTrainedModel必填底层预训练模型实践中为LayerWiseMiniCPMForCausalLM见 2.2tokenizerAutoTokenizerNone文本编码器为None时仍可前向但无法获得yes_loc等标记位train_batch_sizeint4训练时一个组内包含的query, passage样本对数用于把 logits 重组为分组打分start_layerint8分层打分的起始层号从该层开始逐层输出 logits在runner.py中CrossDecoderModel由 runner 统一装配base_model get_model(self.model_args, tokenizer(Yes, add_special_tokensFalse)[input_ids][-1]) model CrossDecoderModel( base_model, tokenizertokenizer, train_batch_sizeself.training_args.per_device_train_batch_size, start_layerself.model_args.start_layer )注意两点train_batch_size直接取自per_device_train_batch_size因此分组损失的正确性依赖 batch 内样本排列为同一 query 的候选 passage 相邻这一约定详见forward一节。start_layer同时存在于模型类与RerankerModelArguments中默认值均为 8两者保持一致才能保证训练与推理时分层行为一致。2.2 骨干模型LayerWiseMiniCPMForCausalLMbase_model并非普通 MiniCPM而是LayerWiseMiniCPMForCausalLMmodeling_minicpm_reranker.py。它与标准MiniCPMForCausalLM的关键差异是多出口multiple exit根据head_multilm_head可能是一个nn.ModuleList为[start_layer, num_hidden_layers]区间内的每一层各准备一个打分头LayerWiseHead即无偏置的nn.Linear三种头部类型由head_type决定见第五节raw直接复用原语言模型头的词表维权重仅用于从原始 MiniCPM 权重出发冷启动complex保留词表维度的多头可用于继续预训练/对齐场景simple把头部压缩为单维输出hidden_size - 1即纯粹的重排打分头默认微调路径。cutoff_layers控制forward接受cutoff_layers参数只运行到指定层便截断并在每个出口层输出 hidden states 后由对应 head 生成 logits见第六节。LayerWiseMiniCPMConfigconfiguration_minicpm_reranker.py新增的字段即start_layer默认 8、head_multi默认 True、head_type默认simple以及 MiniCPM 原生的scale_emb、dim_model_base、scale_depth等缩放参数。三、encode从输入特征到逐层分数列表encode是CrossDecoderModel的前向核心负责把一批 token 化后的特征送入骨干并取回每一出口层的分数def encode(self, features): if features is None: return None outputs self.model(input_idsfeatures[input_ids], attention_maskfeatures[attention_mask], position_idsfeatures[position_ids] if position_ids in features.keys() else None, output_hidden_statesTrue) all_logits outputs.logits all_scores [] for logits in all_logits: all_scores.append(logits[:, -1].contiguous()) return all_scores要点解读输入字典至少包含input_ids与attention_maskposition_ids存在时才透传runner.py的 DataCollator 通常会构造并填充。output_hidden_statesTrue这是多出口的关键——只有保留各层 hidden statesLayerWiseMiniCPMForCausalLM.forward才能逐层产出 logits最终outputs.logits是一个由多个出口 logits 组成的 tuple每个元素形状为(batch_size, seq_len, vocab_size)或(batch_size, seq_len, 1)取决于head_type见第五节。取最后一位 tokenlogits[:, -1]取出序列末位 token 的打分作为该 query/passage 段的相关性分数。这与 decoder-only 重排器的惯例一致——末尾 token 能看到前面所有 token从而汇聚整段上下文信息。返回值all_scores是一个 Python 列表长度等于出口层数量num_hidden_layers - start_layer 1在head_multiTrue时每个元素形状为(batch_size, 1)simple头或(batch_size, vocab_size)raw/complex头后续靠only_for_one_logit切行。因此encode的语义是一组分数而不是单一标量——这就是分层重排器的核心产物。从源码结构看encode不接收cutoff_layers参数默认骨干会跑满全部层再输出推理阶段的早停优化由FlagEmbedding/inference/reranker/decoder_only/layerwise.py中的compute_score_single_gpu实现它在self.model(**batch_inputs, output_hidden_statesTrue, cutoff_layerscutoff_layers)中显式传入cutoff_layers从而只计算到指定层即返回。四、forward训练态的多层自蒸馏损失forward在encode之上叠加了逐层损失累加是分层重排器训练时的关键逻辑def forward(self, pairNone, teacher_scores: Optional[Tensor] None): ranker_logits self.encode(pair) # (batch_size * num, dim) if self.training: loss 0 for logits in ranker_logits: grouped_logits logits.view(self.train_batch_size, -1) target torch.zeros(self.train_batch_size, devicegrouped_logits.device, dtypetorch.long) loss self.compute_loss(grouped_logits, target) if teacher_scores is None: teacher_scores ranker_logits[-1].view(self.train_batch_size, -1) teacher_targets torch.softmax(teacher_scores.detach(), dim-1) for logits in ranker_logits[:-1]: student_scores logits.view(self.train_batch_size, -1) loss - torch.mean(torch.sum(torch.log_softmax(student_scores, dim-1) * teacher_targets, dim-1)) else: teacher_scores torch.Tensor(teacher_scores) teacher_scores teacher_scores.view(self.train_batch_size, -1) teacher_targets torch.softmax(teacher_scores.detach(), dim-1).to(ranker_logits[-1].device) for logits in ranker_logits: student_scores logits.view(self.train_batch_size, -1) loss - torch.mean(torch.sum(torch.log_softmax(student_scores, dim-1) * teacher_targets, dim-1)) else: loss None return RerankerOutput(lossloss, scoresranker_logits)4.1 分组假设logits.view(self.train_batch_size, -1)隐含一个数据排列约定每个 batch 内同一 query 的train_group_size个候选 passage 必须相邻排列整批形状为(train_batch_size, train_group_size, ...)。这样才能 reshape 成(train_batch_size, -1)后每一行恰好是一个 query 的所有候选分数。此约定由训练数据构造与 DataCollator 保证见examples/finetune/reranker/decoder_only/layerwise.sh中的train_group_size8。4.2 损失组成排序损失 逐层自蒸馏排序损失每层target为全零配合compute_loss继承自AbsRerankerModel.compute_loss见AbsModeling.py实现为nn.CrossEntropyLoss(reductionmean)即把组内第 0 个候选当作正样本、其余为负样本的第一候选在组内分数最高的损失。每一出口层都独立施加该损失保证每个层的打分都具备排序能力。自蒸馏损失当teacher_scores is None官方示例默认场景把最后一层logits 视为教师softmax后作为软标签去蒸馏前面ranker_logits[:-1]各层的学生分数损失形式为 KL 散度-Σ p_teacher * log_softmax(p_student)当显式传入teacher_scores外部教师如knowledge_distillationTrue时的教师重排器分数则所有层含最后一层都向该软标签对齐。这一设计使得浅层在自身排序能力与向深层对齐之间取得平衡——浅层越早可用于推理代价是分数分布与深层略有差异深层则始终是最强排序信号。4.3 推理态self.trainingFalse时lossNone仅返回RerankerOutput(lossNone, scoresranker_logits)。因此该模型在评估/推理阶段可直接获得多层分数由上层选择使用全部层还是只取若干浅层。五、头部类型与模型装配head_type / head_multi 全解分层重排器的出口如何生成完全由load_model.py中get_model的装配逻辑决定。装配逻辑按model_type分为两条路径5.1 从原始模型出发model_typefrom_raw_modelconfig.use_cache False config.start_layer config.num_hidden_layers # 先置为最后一层 config.head_multi False config.head_type raw model LayerWiseMiniCPMForCausalLM.from_pretrained(model_args.model_name_or_path, configconfig, ...) config.start_layer model_args.start_layer # 再覆盖为用户配置 config.head_multi model_args.head_multi config.head_type model_args.head_type model.config config该路径先以完整语言模型形态加载原始 MiniCPM如openbmb/MiniCPM-2B-dpo-bf16随后按用户配置替换头部head_typecomplex且head_multiTrue构造nn.ModuleList([LayerWiseHead(hidden_size, vocab_size) ...])并从原lm_head复制权重初始化head_typecomplex且head_multiFalse单个LayerWiseHead(hidden_size, vocab_size)同样继承原权重head_typesimple构造LayerWiseHead(hidden_size, 1)并通过state_dict_back[weight][only_for_one_logit:only_for_one_logit1, :]从原语言模型头的词表行中截取一行only_for_one_logit由 runner 传入即tokenizer(Yes)最后一个 token id来初始化单维打分头。head_type与head_multi的组合效果汇总head_typehead_multiTruehead_multiFalseraw每层一个nn.Linear(hidden, vocab)原语言模型头单个nn.Linear(hidden, vocab)complex每层一个LayerWiseHead(hidden, vocab)单个LayerWiseHead(hidden, vocab)simple每层一个LayerWiseHead(hidden, 1)默认微调形态单个LayerWiseHead(hidden, 1)5.2 从已微调模型出发model_typefrom_finetuned_model直接LayerWiseMiniCPMForCausalLM.from_pretrained(model_args.model_name_or_path, ...)加载例如BAAI/bge-reranker-v2-minicpm-layerwise。此时 checkpoint 内已固化start_layer/head_multi/head_type由配置里的start_layer、head_multi、head_type字段携带无需手动装配头部。这是官方示例脚本layerwise.sh采用的默认路径。5.3 LoRA 装配无论哪条路径最终都会按use_lora默认True套上 PEFTpeft_config LoraConfig( task_typeTaskType.CAUSAL_LM, inference_modeFalse, rmodel_args.lora_rank, # 默认 64 target_modulesmodel_args.target_modules, # 默认 [v_proj,q_proj,k_proj,gate_proj,down_proj,o_proj,up_proj] modules_to_savemodel_args.modules_to_save, lora_alphamodel_args.lora_alpha, # 默认 16 lora_dropoutmodel_args.lora_dropout # 默认 0.1 ) model get_peft_model(model, peft_config)此外raw_peft支持加载多个 PEFT 权重并merge_and_unload合并模型融合场景from_peft支持从既有 LoRA 继续训练save_merged_lora_modelTrue时训练结束后由save_merged_model把 LoRA 合并回主干并保存到output_dir/merged_model。六、分层前向的底层原理cutoff_layers 与多出口理解CrossDecoderModel的行为离不开LayerWiseMiniCPMForCausalLM.forward的分层实现其核心步骤规范化cutoff_layers若为None默认[num_hidden_layers]随后过滤掉不在[start_layer, num_hidden_layers]区间的层号若过滤后为空则报错。前向截断LayerWiseMiniCPMModel.forward接收cutoff_layers在idx max_layer时break即只计算到最大截止层并在每个位于cutoff_layers的层输出self.norm(hidden_states)RMSNorm 之后的 hidden states——这就是多出口的 hidden states 来源。逐层打分对每个出口 hidden state按head_type/head_multi选择对应 headsimple头直接lm_head.linear_head(hidden_states)得到(batch, seq, 1)raw/complex头only_for_one_logit为None时输出词表维 logits并对hidden_size / dim_model_base做缩放raw/complex头给定only_for_one_logit时通过lm_head.weight.split(1, dim0)取指定行把词表维压成单维。汇总所有出口 logits 拼成 tuple包装进CausalLMOutputWithPast.logits返回CrossDecoderModel.encode据此逐层取末位 token 分数。这一机制同时解释了训练与推理的分层语义训练时跑满全部层并逐层算损失推理时可通过cutoff_layers只跑到第 k 层就打分在浅层信心足够时提前返回实现推理延迟与精度的权衡可进一步参考FlagEmbedding/inference/reranker/decoder_only/layerwise.py中compute_score_single_gpu的cutoff_layers处理。七、端到端微调实战以官方 layerwise 脚本为例CrossDecoderModel的完整训练流程由runner.py驱动def run(self): Path(self.training_args.output_dir).mkdir(parentsTrue, exist_okTrue) self.trainer.train(resume_from_checkpointself.training_args.resume_from_checkpoint) self.trainer.save_model() if self.model_args.save_merged_lora_model and self.training_args.process_index 0: save_merged_model(self.model_args, self.training_args.output_dir)入口为__main__.py它用HfArgumentParser解析三类参数RerankerModelArguments、AbsRerankerDataArguments、AbsRerankerTrainingArguments后交给DecoderOnlyRerankerRunner。官方示例脚本examples/finetune/reranker/decoder_only/layerwise.sh给出可直接运行的完整配置已适配当前仓库export WANDB_MODEdisabled train_data\ ../example_data/prompt_based/examples.jsonl num_train_epochs1 per_device_train_batch_size2 gradient_accumulation_steps1 train_group_size8 num_gpus2 model_args\ --model_name_or_path BAAI/bge-reranker-v2-minicpm-layerwise \ --cache_dir $HF_HUB_CACHE \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --use_flash_attn True \ --target_modules q_proj k_proj v_proj o_proj \ --save_merged_lora_model True \ --model_type from_finetuned_model \ --start_layer 8 \ --head_multi True \ --head_type simple \ --trust_remote_code True \ data_args\ --train_data $train_data \ --cache_path ~/.cache \ --train_group_size $train_group_size \ --query_max_len 512 \ --passage_max_len 512 \ --pad_to_multiple_of 8 \ --knowledge_distillation True \ --query_instruction_for_rerank A: \ --query_instruction_format {}{} \ --passage_instruction_for_rerank B: \ --passage_instruction_format {}{} \ training_args\ --output_dir ./test_decoder_only_layerwise \ --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --gradient_accumulation_steps $gradient_accumulation_steps \ --dataloader_drop_last True \ --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ --save_steps 1000 \ cmdtorchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.reranker.decoder_only.layerwise \ $model_args $data_args $training_args \ eval $cmd关键配置解读--model_type from_finetuned_model直接加载BAAI/bge-reranker-v2-minicpm-layerwisecheckpoint 自带分层配置是复现官方行为的推荐路径若要从openbmb/MiniCPM-2B-dpo-bf16原始权重起步则改为--model_type from_raw_model。--start_layer 8从第 8 层开始打分层分数与CrossDecoderModel的默认start_layer8保持一致该值越小可早停的层越多、浅层候选越多但浅层排序质量与深层差异越大。--head_multi True --head_type simple为[8, num_hidden_layers]每一层各配置一个hidden_size - 1的单维打分头LayerWiseHead对应encode返回的all_scores列表长度 num_hidden_layers - 8 1。--train_group_size 8每个 query 配 8 个候选与forward中view(train_batch_size, -1)的分组假设匹配。--knowledge_distillation True启用知识蒸馏数据流程此时forward会收到外部teacher_scores所有层的学生分数都向教师软标签对齐。--use_lora True --lora_rank 32 --lora_alpha 64 --target_modules q_proj k_proj v_proj o_proj仅对注意力投影做低秩适配减少可训练参数量target_modules默认全集为[v_proj,q_proj,k_proj,gate_proj,down_proj,o_proj,up_proj]见arguments.py的default_target_modules。--use_flash_attn True启用 FlashAttention-2 加速需环境已安装flash_attn否则应关闭LayerWiseMiniCPMConfig.__init__会在 import 失败时自动回退。--save_merged_lora_model True训练结束后由save_merged_model把 LoRA 合并回主干产物保存于output_dir/merged_model。数据侧需要注意--query_instruction_for_rerank A: 与--passage_instruction_for_rerank B: 是 decoder-only 重排器的指令前缀约定--query_max_len 512 --passage_max_len 512与--pad_to_multiple_of 8保证序列长度对齐、便于分组 reshape。训练数据格式参考examples/finetune/reranker/example_data/下的 jsonlquery/passages 对 可选的教师分数。RerankerModelArguments中还有两个与分层密切相关的参数值得注意arguments.pymodel_typefrom_raw_model/from_finetuned_model二选一与head_multi/head_type决定出口头的数量与形态save_merged_lora_model则控制是否产出可直接用于推理的合并权重。八、与其他模块的关系与调用链小结抽象基类CrossDecoderModel-AbsRerankerModelcompute_loss由基类实现为交叉熵RerankerOutputlossscores定义于同一文件。训练驱动__main__.py-DecoderOnlyRerankerRunner-CrossDecoderModelDecoderOnlyRerankerTrainertrainer.py负责 knowledge distillation 场景下的数据/教师分数组装。推理侧训练产出的合并模型由FlagEmbedding/inference/reranker/decoder_only/layerwise.py加载compute_score_single_gpu通过cutoff_layers实现分层早停打分与训练侧的多出口机制一一对应。文档侧本文件的 API 文档modeling.rst与同目录的 arguments.rst、runner.rst 共同构成 layerwise 微调模块的完整文档体系。结语CrossDecoderModel把多层同时打分 逐层自蒸馏的思想落到了 decoder-only 重排器上encode负责产出逐层分数forward负责把排序损失与软标签蒸馏损失叠加到每一层LayerWiseMiniCPMForCausalLM的多出口结构则为其提供底层支撑。配合start_layer、head_type、head_multi、cutoff_layers等配置开发者既可以在训练时让所有层共同收敛也可以在推理时按需只计算到浅层在检索场景的延迟与精度之间取得灵活平衡。建议结合上述源码路径与官方示例脚本动手运行一次微调直观感受分层分数的行为差异。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 FlagEmbedding 分层解码器重排器 CrossDecoderModel 建模与微调全解析
FlagEmbedding 分层解码器重排器 CrossDecoderModel 建模与微调全解析【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding导读本文围绕 FlagEmbedding 中 decoder-only 分层layerwise重排器reranker的模型实现展开核心对象是CrossDecoderModel。该模型以 MiniCPM 这类因果语言模型为骨干通过在多个 Transformer 层同时接出 logits 并做加权打分的方式实现逐层早停式重排可在精度与推理开销之间灵活取舍。读完本文你将掌握CrossDecoderModel.encode/forward的完整前向与损失计算逻辑、start_layer/head_multi/head_type等关键配置的作用以及如何借助 LoRA 在官方示例脚本上完成一次分层重排器的端到端微调。一、文档定位从 API 文档到源码实现本文所依托的 API 文档位于docs/source/API/finetune/reranker/decoder_only/layerwise/modeling.rst其通过 Sphinx autodoc 声明了三个核心符号FlagEmbedding.finetune.reranker.decoder_only.layerwise.CrossDecoderModel类CrossDecoderModel.encode方法CrossDecoderModel.forward方法也就是说该文档的实体内容集中在FlagEmbedding/finetune/reranker/decoder_only/layerwise/modeling.py这一个文件中。与之配套的还有同目录下的arguments.py重排器专属的模型超参LoRA、start_layer、head_type等modeling_minicpm_reranker.py分层 MiniCPM 骨干LayerWiseMiniCPMForCausalLM、LayerWiseMiniCPMConfig、LayerWiseHeadload_model.py按model_type分支加载原始模型或已微调模型并负责头部替换与 LoRA 装配runner.py组装 tokenizer / model / trainer 并执行训练与合并保存。下文将以modeling.py为主线结合这些源码逐层展开。二、CrossDecoderModel分层解码器重排器的门面类CrossDecoderModel继承自AbsRerankerModel位于FlagEmbedding/abc/finetune/reranker/AbsModeling.py后者是 FlagEmbedding 微调框架中所有 reranker 模型的抽象基类负责维护base_model、tokenizer、train_batch_size、cross_entropy等公共成员。2.1 构造参数class CrossDecoderModel(AbsRerankerModel): def __init__( self, base_model: PreTrainedModel, tokenizer: AutoTokenizer None, train_batch_size: int 4, start_layer: int 8 ): super().__init__(base_model, tokenizertokenizer, train_batch_sizetrain_batch_size) self.start_layer start_layer参数类型默认值含义base_modelPreTrainedModel必填底层预训练模型实践中为LayerWiseMiniCPMForCausalLM见 2.2tokenizerAutoTokenizerNone文本编码器为None时仍可前向但无法获得yes_loc等标记位train_batch_sizeint4训练时一个组内包含的query, passage样本对数用于把 logits 重组为分组打分start_layerint8分层打分的起始层号从该层开始逐层输出 logits在runner.py中CrossDecoderModel由 runner 统一装配base_model get_model(self.model_args, tokenizer(Yes, add_special_tokensFalse)[input_ids][-1]) model CrossDecoderModel( base_model, tokenizertokenizer, train_batch_sizeself.training_args.per_device_train_batch_size, start_layerself.model_args.start_layer )注意两点train_batch_size直接取自per_device_train_batch_size因此分组损失的正确性依赖 batch 内样本排列为同一 query 的候选 passage 相邻这一约定详见forward一节。start_layer同时存在于模型类与RerankerModelArguments中默认值均为 8两者保持一致才能保证训练与推理时分层行为一致。2.2 骨干模型LayerWiseMiniCPMForCausalLMbase_model并非普通 MiniCPM而是LayerWiseMiniCPMForCausalLMmodeling_minicpm_reranker.py。它与标准MiniCPMForCausalLM的关键差异是多出口multiple exit根据head_multilm_head可能是一个nn.ModuleList为[start_layer, num_hidden_layers]区间内的每一层各准备一个打分头LayerWiseHead即无偏置的nn.Linear三种头部类型由head_type决定见第五节raw直接复用原语言模型头的词表维权重仅用于从原始 MiniCPM 权重出发冷启动complex保留词表维度的多头可用于继续预训练/对齐场景simple把头部压缩为单维输出hidden_size - 1即纯粹的重排打分头默认微调路径。cutoff_layers控制forward接受cutoff_layers参数只运行到指定层便截断并在每个出口层输出 hidden states 后由对应 head 生成 logits见第六节。LayerWiseMiniCPMConfigconfiguration_minicpm_reranker.py新增的字段即start_layer默认 8、head_multi默认 True、head_type默认simple以及 MiniCPM 原生的scale_emb、dim_model_base、scale_depth等缩放参数。三、encode从输入特征到逐层分数列表encode是CrossDecoderModel的前向核心负责把一批 token 化后的特征送入骨干并取回每一出口层的分数def encode(self, features): if features is None: return None outputs self.model(input_idsfeatures[input_ids], attention_maskfeatures[attention_mask], position_idsfeatures[position_ids] if position_ids in features.keys() else None, output_hidden_statesTrue) all_logits outputs.logits all_scores [] for logits in all_logits: all_scores.append(logits[:, -1].contiguous()) return all_scores要点解读输入字典至少包含input_ids与attention_maskposition_ids存在时才透传runner.py的 DataCollator 通常会构造并填充。output_hidden_statesTrue这是多出口的关键——只有保留各层 hidden statesLayerWiseMiniCPMForCausalLM.forward才能逐层产出 logits最终outputs.logits是一个由多个出口 logits 组成的 tuple每个元素形状为(batch_size, seq_len, vocab_size)或(batch_size, seq_len, 1)取决于head_type见第五节。取最后一位 tokenlogits[:, -1]取出序列末位 token 的打分作为该 query/passage 段的相关性分数。这与 decoder-only 重排器的惯例一致——末尾 token 能看到前面所有 token从而汇聚整段上下文信息。返回值all_scores是一个 Python 列表长度等于出口层数量num_hidden_layers - start_layer 1在head_multiTrue时每个元素形状为(batch_size, 1)simple头或(batch_size, vocab_size)raw/complex头后续靠only_for_one_logit切行。因此encode的语义是一组分数而不是单一标量——这就是分层重排器的核心产物。从源码结构看encode不接收cutoff_layers参数默认骨干会跑满全部层再输出推理阶段的早停优化由FlagEmbedding/inference/reranker/decoder_only/layerwise.py中的compute_score_single_gpu实现它在self.model(**batch_inputs, output_hidden_statesTrue, cutoff_layerscutoff_layers)中显式传入cutoff_layers从而只计算到指定层即返回。四、forward训练态的多层自蒸馏损失forward在encode之上叠加了逐层损失累加是分层重排器训练时的关键逻辑def forward(self, pairNone, teacher_scores: Optional[Tensor] None): ranker_logits self.encode(pair) # (batch_size * num, dim) if self.training: loss 0 for logits in ranker_logits: grouped_logits logits.view(self.train_batch_size, -1) target torch.zeros(self.train_batch_size, devicegrouped_logits.device, dtypetorch.long) loss self.compute_loss(grouped_logits, target) if teacher_scores is None: teacher_scores ranker_logits[-1].view(self.train_batch_size, -1) teacher_targets torch.softmax(teacher_scores.detach(), dim-1) for logits in ranker_logits[:-1]: student_scores logits.view(self.train_batch_size, -1) loss - torch.mean(torch.sum(torch.log_softmax(student_scores, dim-1) * teacher_targets, dim-1)) else: teacher_scores torch.Tensor(teacher_scores) teacher_scores teacher_scores.view(self.train_batch_size, -1) teacher_targets torch.softmax(teacher_scores.detach(), dim-1).to(ranker_logits[-1].device) for logits in ranker_logits: student_scores logits.view(self.train_batch_size, -1) loss - torch.mean(torch.sum(torch.log_softmax(student_scores, dim-1) * teacher_targets, dim-1)) else: loss None return RerankerOutput(lossloss, scoresranker_logits)4.1 分组假设logits.view(self.train_batch_size, -1)隐含一个数据排列约定每个 batch 内同一 query 的train_group_size个候选 passage 必须相邻排列整批形状为(train_batch_size, train_group_size, ...)。这样才能 reshape 成(train_batch_size, -1)后每一行恰好是一个 query 的所有候选分数。此约定由训练数据构造与 DataCollator 保证见examples/finetune/reranker/decoder_only/layerwise.sh中的train_group_size8。4.2 损失组成排序损失 逐层自蒸馏排序损失每层target为全零配合compute_loss继承自AbsRerankerModel.compute_loss见AbsModeling.py实现为nn.CrossEntropyLoss(reductionmean)即把组内第 0 个候选当作正样本、其余为负样本的第一候选在组内分数最高的损失。每一出口层都独立施加该损失保证每个层的打分都具备排序能力。自蒸馏损失当teacher_scores is None官方示例默认场景把最后一层logits 视为教师softmax后作为软标签去蒸馏前面ranker_logits[:-1]各层的学生分数损失形式为 KL 散度-Σ p_teacher * log_softmax(p_student)当显式传入teacher_scores外部教师如knowledge_distillationTrue时的教师重排器分数则所有层含最后一层都向该软标签对齐。这一设计使得浅层在自身排序能力与向深层对齐之间取得平衡——浅层越早可用于推理代价是分数分布与深层略有差异深层则始终是最强排序信号。4.3 推理态self.trainingFalse时lossNone仅返回RerankerOutput(lossNone, scoresranker_logits)。因此该模型在评估/推理阶段可直接获得多层分数由上层选择使用全部层还是只取若干浅层。五、头部类型与模型装配head_type / head_multi 全解分层重排器的出口如何生成完全由load_model.py中get_model的装配逻辑决定。装配逻辑按model_type分为两条路径5.1 从原始模型出发model_typefrom_raw_modelconfig.use_cache False config.start_layer config.num_hidden_layers # 先置为最后一层 config.head_multi False config.head_type raw model LayerWiseMiniCPMForCausalLM.from_pretrained(model_args.model_name_or_path, configconfig, ...) config.start_layer model_args.start_layer # 再覆盖为用户配置 config.head_multi model_args.head_multi config.head_type model_args.head_type model.config config该路径先以完整语言模型形态加载原始 MiniCPM如openbmb/MiniCPM-2B-dpo-bf16随后按用户配置替换头部head_typecomplex且head_multiTrue构造nn.ModuleList([LayerWiseHead(hidden_size, vocab_size) ...])并从原lm_head复制权重初始化head_typecomplex且head_multiFalse单个LayerWiseHead(hidden_size, vocab_size)同样继承原权重head_typesimple构造LayerWiseHead(hidden_size, 1)并通过state_dict_back[weight][only_for_one_logit:only_for_one_logit1, :]从原语言模型头的词表行中截取一行only_for_one_logit由 runner 传入即tokenizer(Yes)最后一个 token id来初始化单维打分头。head_type与head_multi的组合效果汇总head_typehead_multiTruehead_multiFalseraw每层一个nn.Linear(hidden, vocab)原语言模型头单个nn.Linear(hidden, vocab)complex每层一个LayerWiseHead(hidden, vocab)单个LayerWiseHead(hidden, vocab)simple每层一个LayerWiseHead(hidden, 1)默认微调形态单个LayerWiseHead(hidden, 1)5.2 从已微调模型出发model_typefrom_finetuned_model直接LayerWiseMiniCPMForCausalLM.from_pretrained(model_args.model_name_or_path, ...)加载例如BAAI/bge-reranker-v2-minicpm-layerwise。此时 checkpoint 内已固化start_layer/head_multi/head_type由配置里的start_layer、head_multi、head_type字段携带无需手动装配头部。这是官方示例脚本layerwise.sh采用的默认路径。5.3 LoRA 装配无论哪条路径最终都会按use_lora默认True套上 PEFTpeft_config LoraConfig( task_typeTaskType.CAUSAL_LM, inference_modeFalse, rmodel_args.lora_rank, # 默认 64 target_modulesmodel_args.target_modules, # 默认 [v_proj,q_proj,k_proj,gate_proj,down_proj,o_proj,up_proj] modules_to_savemodel_args.modules_to_save, lora_alphamodel_args.lora_alpha, # 默认 16 lora_dropoutmodel_args.lora_dropout # 默认 0.1 ) model get_peft_model(model, peft_config)此外raw_peft支持加载多个 PEFT 权重并merge_and_unload合并模型融合场景from_peft支持从既有 LoRA 继续训练save_merged_lora_modelTrue时训练结束后由save_merged_model把 LoRA 合并回主干并保存到output_dir/merged_model。六、分层前向的底层原理cutoff_layers 与多出口理解CrossDecoderModel的行为离不开LayerWiseMiniCPMForCausalLM.forward的分层实现其核心步骤规范化cutoff_layers若为None默认[num_hidden_layers]随后过滤掉不在[start_layer, num_hidden_layers]区间的层号若过滤后为空则报错。前向截断LayerWiseMiniCPMModel.forward接收cutoff_layers在idx max_layer时break即只计算到最大截止层并在每个位于cutoff_layers的层输出self.norm(hidden_states)RMSNorm 之后的 hidden states——这就是多出口的 hidden states 来源。逐层打分对每个出口 hidden state按head_type/head_multi选择对应 headsimple头直接lm_head.linear_head(hidden_states)得到(batch, seq, 1)raw/complex头only_for_one_logit为None时输出词表维 logits并对hidden_size / dim_model_base做缩放raw/complex头给定only_for_one_logit时通过lm_head.weight.split(1, dim0)取指定行把词表维压成单维。汇总所有出口 logits 拼成 tuple包装进CausalLMOutputWithPast.logits返回CrossDecoderModel.encode据此逐层取末位 token 分数。这一机制同时解释了训练与推理的分层语义训练时跑满全部层并逐层算损失推理时可通过cutoff_layers只跑到第 k 层就打分在浅层信心足够时提前返回实现推理延迟与精度的权衡可进一步参考FlagEmbedding/inference/reranker/decoder_only/layerwise.py中compute_score_single_gpu的cutoff_layers处理。七、端到端微调实战以官方 layerwise 脚本为例CrossDecoderModel的完整训练流程由runner.py驱动def run(self): Path(self.training_args.output_dir).mkdir(parentsTrue, exist_okTrue) self.trainer.train(resume_from_checkpointself.training_args.resume_from_checkpoint) self.trainer.save_model() if self.model_args.save_merged_lora_model and self.training_args.process_index 0: save_merged_model(self.model_args, self.training_args.output_dir)入口为__main__.py它用HfArgumentParser解析三类参数RerankerModelArguments、AbsRerankerDataArguments、AbsRerankerTrainingArguments后交给DecoderOnlyRerankerRunner。官方示例脚本examples/finetune/reranker/decoder_only/layerwise.sh给出可直接运行的完整配置已适配当前仓库export WANDB_MODEdisabled train_data\ ../example_data/prompt_based/examples.jsonl num_train_epochs1 per_device_train_batch_size2 gradient_accumulation_steps1 train_group_size8 num_gpus2 model_args\ --model_name_or_path BAAI/bge-reranker-v2-minicpm-layerwise \ --cache_dir $HF_HUB_CACHE \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --use_flash_attn True \ --target_modules q_proj k_proj v_proj o_proj \ --save_merged_lora_model True \ --model_type from_finetuned_model \ --start_layer 8 \ --head_multi True \ --head_type simple \ --trust_remote_code True \ data_args\ --train_data $train_data \ --cache_path ~/.cache \ --train_group_size $train_group_size \ --query_max_len 512 \ --passage_max_len 512 \ --pad_to_multiple_of 8 \ --knowledge_distillation True \ --query_instruction_for_rerank A: \ --query_instruction_format {}{} \ --passage_instruction_for_rerank B: \ --passage_instruction_format {}{} \ training_args\ --output_dir ./test_decoder_only_layerwise \ --overwrite_output_dir \ --learning_rate 2e-4 \ --bf16 \ --num_train_epochs $num_train_epochs \ --per_device_train_batch_size $per_device_train_batch_size \ --gradient_accumulation_steps $gradient_accumulation_steps \ --dataloader_drop_last True \ --warmup_ratio 0.1 \ --gradient_checkpointing \ --weight_decay 0.01 \ --deepspeed ../../ds_stage0.json \ --logging_steps 1 \ --save_steps 1000 \ cmdtorchrun --nproc_per_node $num_gpus \ -m FlagEmbedding.finetune.reranker.decoder_only.layerwise \ $model_args $data_args $training_args \ eval $cmd关键配置解读--model_type from_finetuned_model直接加载BAAI/bge-reranker-v2-minicpm-layerwisecheckpoint 自带分层配置是复现官方行为的推荐路径若要从openbmb/MiniCPM-2B-dpo-bf16原始权重起步则改为--model_type from_raw_model。--start_layer 8从第 8 层开始打分层分数与CrossDecoderModel的默认start_layer8保持一致该值越小可早停的层越多、浅层候选越多但浅层排序质量与深层差异越大。--head_multi True --head_type simple为[8, num_hidden_layers]每一层各配置一个hidden_size - 1的单维打分头LayerWiseHead对应encode返回的all_scores列表长度 num_hidden_layers - 8 1。--train_group_size 8每个 query 配 8 个候选与forward中view(train_batch_size, -1)的分组假设匹配。--knowledge_distillation True启用知识蒸馏数据流程此时forward会收到外部teacher_scores所有层的学生分数都向教师软标签对齐。--use_lora True --lora_rank 32 --lora_alpha 64 --target_modules q_proj k_proj v_proj o_proj仅对注意力投影做低秩适配减少可训练参数量target_modules默认全集为[v_proj,q_proj,k_proj,gate_proj,down_proj,o_proj,up_proj]见arguments.py的default_target_modules。--use_flash_attn True启用 FlashAttention-2 加速需环境已安装flash_attn否则应关闭LayerWiseMiniCPMConfig.__init__会在 import 失败时自动回退。--save_merged_lora_model True训练结束后由save_merged_model把 LoRA 合并回主干产物保存于output_dir/merged_model。数据侧需要注意--query_instruction_for_rerank A: 与--passage_instruction_for_rerank B: 是 decoder-only 重排器的指令前缀约定--query_max_len 512 --passage_max_len 512与--pad_to_multiple_of 8保证序列长度对齐、便于分组 reshape。训练数据格式参考examples/finetune/reranker/example_data/下的 jsonlquery/passages 对 可选的教师分数。RerankerModelArguments中还有两个与分层密切相关的参数值得注意arguments.pymodel_typefrom_raw_model/from_finetuned_model二选一与head_multi/head_type决定出口头的数量与形态save_merged_lora_model则控制是否产出可直接用于推理的合并权重。八、与其他模块的关系与调用链小结抽象基类CrossDecoderModel-AbsRerankerModelcompute_loss由基类实现为交叉熵RerankerOutputlossscores定义于同一文件。训练驱动__main__.py-DecoderOnlyRerankerRunner-CrossDecoderModelDecoderOnlyRerankerTrainertrainer.py负责 knowledge distillation 场景下的数据/教师分数组装。推理侧训练产出的合并模型由FlagEmbedding/inference/reranker/decoder_only/layerwise.py加载compute_score_single_gpu通过cutoff_layers实现分层早停打分与训练侧的多出口机制一一对应。文档侧本文件的 API 文档modeling.rst与同目录的 arguments.rst、runner.rst 共同构成 layerwise 微调模块的完整文档体系。结语CrossDecoderModel把多层同时打分 逐层自蒸馏的思想落到了 decoder-only 重排器上encode负责产出逐层分数forward负责把排序损失与软标签蒸馏损失叠加到每一层LayerWiseMiniCPMForCausalLM的多出口结构则为其提供底层支撑。配合start_layer、head_type、head_multi、cutoff_layers等配置开发者既可以在训练时让所有层共同收敛也可以在推理时按需只计算到浅层在检索场景的延迟与精度之间取得灵活平衡。建议结合上述源码路径与官方示例脚本动手运行一次微调直观感受分层分数的行为差异。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考