企业官网建设流程全解析

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、LayerWiseHead）；
• load_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, tokenizer=tokenizer, train_batch_size=train_batch_size) self.start_layer = start_layer

在runner.py中，CrossDecoderModel由 runner 统一装配：

base_model = get_model(self.model_args, tokenizer('Yes', add_special_tokens=False)['input_ids'][-1]) model = CrossDecoderModel( base_model, tokenizer=tokenizer, train_batch_size=self.training_args.per_device_train_batch_size, start_layer=self.model_args.start_layer )

注意两点：

1. train_batch_size直接取自per_device_train_batch_size，因此分组损失的正确性依赖 batch 内样本排列为"同一 query 的候选 passage 相邻"这一约定（详见forward一节）。
2. start_layer同时存在于模型类与RerankerModelArguments中（默认值均为 8），两者保持一致才能保证训练与推理时分层行为一致。

2.2 骨干模型：LayerWiseMiniCPMForCausalLM

base_model并非普通 MiniCPM，而是LayerWiseMiniCPMForCausalLM（modeling_minicpm_reranker.py）。它与标准MiniCPMForCausalLM的关键差异是：

• 多出口（multiple exit）：根据head_multi，lm_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（见第六节）。

LayerWiseMiniCPMConfig（configuration_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_ids=features['input_ids'], attention_mask=features['attention_mask'], position_ids=features['position_ids'] if 'position_ids' in features.keys() else None, output_hidden_states=True) all_logits = outputs.logits all_scores = [] for logits in all_logits: all_scores.append(logits[:, -1].contiguous()) return all_scores

要点解读：

• 输入字典：至少包含input_ids与attention_mask；position_ids存在时才透传（runner.py的 DataCollator 通常会构造并填充）。
• output_hidden_states=True：这是多出口的关键——只有保留各层 hidden states，LayerWiseMiniCPMForCausalLM.forward才能逐层产出 logits，最终outputs.logits是一个由多个出口 logits 组成的 tuple，每个元素形状为(batch_size, seq_len, vocab_size)或(batch_size, seq_len, 1)（取决于head_type，见第五节）。
• 取最后一位 token：logits[:, -1]取出序列末位 token 的打分，作为该 query/passage 段的相关性分数。这与 decoder-only 重排器的惯例一致——末尾 token 能"看到"前面所有 token，从而汇聚整段上下文信息。
• 返回值：all_scores是一个 Python 列表，长度等于出口层数量（num_hidden_layers - start_layer + 1，在head_multi=True时），每个元素形状为(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_states=True, cutoff_layers=cutoff_layers)中显式传入cutoff_layers，从而只计算到指定层即返回。

四、forward：训练态的多层自蒸馏损失

forward在encode之上叠加了逐层损失累加，是分层重排器训练时的关键逻辑：

def forward(self, pair=None, 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, device=grouped_logits.device, dtype=torch.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(loss=loss, scores=ranker_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_size=8）。

4.2 损失组成：排序损失 + 逐层自蒸馏

• 排序损失（每层）：target为全零，配合compute_loss（继承自AbsRerankerModel.compute_loss，见AbsModeling.py，实现为nn.CrossEntropyLoss(reduction='mean')），即把组内第 0 个候选当作正样本、其余为负样本的"第一候选在组内分数最高"的损失。每一出口层都独立施加该损失，保证每个层的打分都具备排序能力。
• 自蒸馏损失：
  • 当teacher_scores is None（官方示例默认场景）：把最后一层logits 视为教师，softmax后作为软标签，去蒸馏前面ranker_logits[:-1]各层的学生分数，损失形式为 KL 散度（-Σ p_teacher * log_softmax(p_student)）；
  • 当显式传入teacher_scores（外部教师，如knowledge_distillation=True时的教师重排器分数）：则所有层（含最后一层）都向该软标签对齐。

这一设计使得浅层在"自身排序能力"与"向深层对齐"之间取得平衡——浅层越早可用于推理，代价是分数分布与深层略有差异；深层则始终是最强排序信号。

4.3 推理态

self.training=False时loss=None，仅返回RerankerOutput(loss=None, scores=ranker_logits)。因此该模型在评估/推理阶段可直接获得多层分数，由上层选择使用全部层还是只取若干浅层。

五、头部类型与模型装配：head_type / head_multi 全解

分层重排器的"出口"如何生成，完全由load_model.py中get_model的装配逻辑决定。装配逻辑按model_type分为两条路径：

5.1 从原始模型出发（model_type="from_raw_model"）

config.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, config=config, ...) 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_type='complex'且head_multi=True：构造nn.ModuleList([LayerWiseHead(hidden_size, vocab_size) ...])，并从原lm_head复制权重初始化；
• head_type='complex'且head_multi=False：单个LayerWiseHead(hidden_size, vocab_size)，同样继承原权重；
• head_type='simple'：构造LayerWiseHead(hidden_size, 1)，并通过state_dict_back['weight'][only_for_one_logit:only_for_one_logit+1, :]从原语言模型头的词表行中截取一行（only_for_one_logit由 runner 传入，即tokenizer('Yes')最后一个 token id）来初始化单维打分头。

head_type与head_multi的组合效果汇总：

5.2 从已微调模型出发（model_type="from_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）套上 PEFT：

peft_config = LoraConfig( task_type=TaskType.CAUSAL_LM, inference_mode=False, r=model_args.lora_rank, # 默认 64 target_modules=model_args.target_modules, # 默认 ['v_proj','q_proj','k_proj','gate_proj','down_proj','o_proj','up_proj'] modules_to_save=model_args.modules_to_save, lora_alpha=model_args.lora_alpha, # 默认 16 lora_dropout=model_args.lora_dropout # 默认 0.1 ) model = get_peft_model(model, peft_config)

此外，raw_peft支持加载多个 PEFT 权重并merge_and_unload合并（模型融合场景），from_peft支持从既有 LoRA 继续训练，save_merged_lora_model=True时训练结束后由save_merged_model把 LoRA 合并回主干并保存到output_dir/merged_model。

六、分层前向的底层原理：cutoff_layers 与多出口

理解CrossDecoderModel的行为，离不开LayerWiseMiniCPMForCausalLM.forward的分层实现，其核心步骤：

1. 规范化cutoff_layers：若为None，默认[num_hidden_layers]；随后过滤掉不在[start_layer, num_hidden_layers]区间的层号，若过滤后为空则报错。
2. 前向截断：LayerWiseMiniCPMModel.forward接收cutoff_layers，在idx == max_layer时break，即只计算到最大截止层，并在每个位于cutoff_layers的层输出self.norm(hidden_states)（RMSNorm 之后的 hidden states）——这就是多出口的 hidden states 来源。
3. 逐层打分：对每个出口 hidden state，按head_type/head_multi选择对应 head：
   • simple头直接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, dim=0)取指定行，把词表维压成单维。
4. 汇总：所有出口 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(parents=True, exist_ok=True) self.trainer.train(resume_from_checkpoint=self.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_MODE=disabled train_data="\ ../example_data/prompt_based/examples.jsonl " num_train_epochs=1 per_device_train_batch_size=2 gradient_accumulation_steps=1 train_group_size=8 num_gpus=2 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 \ " cmd="torchrun --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-layerwise，checkpoint 自带分层配置，是复现官方行为的推荐路径；若要从openbmb/MiniCPM-2B-dpo-bf16原始权重起步，则改为--model_type from_raw_model。
• --start_layer 8：从第 8 层开始打分层分数，与CrossDecoderModel的默认start_layer=8保持一致；该值越小，可早停的层越多、浅层候选越多，但浅层排序质量与深层差异越大。
• --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/下的 jsonl（query/passages 对 + 可选的教师分数）。

RerankerModelArguments中还有两个与分层密切相关的参数值得注意（arguments.py）：model_type（from_raw_model/from_finetuned_model二选一）与head_multi/head_type（决定出口头的数量与形态）；save_merged_lora_model则控制是否产出可直接用于推理的合并权重。

八、与其他模块的关系与调用链小结

• 抽象基类：CrossDecoderModel->AbsRerankerModel，compute_loss由基类实现为交叉熵；RerankerOutput（loss+scores）定义于同一文件。
• 训练驱动：__main__.py->DecoderOnlyRerankerRunner->CrossDecoderModel+DecoderOnlyRerankerTrainer；trainer.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