Matroyshka Re-Ranker基于 Mistral-7B 的可配置深度与宽度 LLM 重排序器FlagEmbedding 实战指南【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding导读本篇技术指南围绕 FlagEmbedding 仓库中 research/Matroyshka_reranker 目录下的Matroyshka Re-Ranker展开介绍一种能够按照用户配置在运行时自定义模型层数与各层序列长度的轻量级 LLM 重排序器。重排序Re-ranking不同于向量检索它以「查询 文档」对为输入、直接输出相关性分数。读完本文你将掌握Matroyshka Re-Ranker 的推理调用方式compress_ratio/compress_layers/cutoff_layers三组核心配置、基于「级联自蒸馏Cascaded Self-distillation 因子化补偿Factorized Compensation」的两阶段微调流程以及隐藏在源码中的 Token 压缩与逐层打分原理。背景为什么需要「可变形」的 LLM Re-Ranker与嵌入模型不同重排序模型将查询query与文档passage同时作为输入直接输出二者的相关性分数而非生成向量表示。Matroyshka Re-Ranker 的输出分数可以通过 sigmoid 函数映射到 [0, 1] 区间便于与其他模块统一解读对应实现见 inference/rank_model.py 中从 FlagEmbedding 复用的sigmoid。其核心设计目标是运行时的灵活配置runtime customization深度维度可以选择从模型的哪些层cutoff_layers输出分数宽度维度可以在指定的层compress_layers上按比例compress_ratio压缩 token 序列长度减少后续层的计算量。这样的设计使得同一个模型在离线高精度重排不压缩、用深层输出与在线低延迟重排按需压缩层数与序列长度之间按需切换支持轻量化的灵活配置。环境安装Matroyshka Re-Ranker 基于 Mistral-7B 架构开发推荐使用 conda 创建独立 Python 3.10 环境conda create -n reranker python3.10 conda activate reranker pip install -r requirements.txt依赖清单见 requirements.txt关键依赖包括transformers4.41.1、peft0.10.0、deepspeed、flash-attn、datasets、modelscope、bitsandbytes等。其中flash-attn为可选加速项训练脚本可通过--use_flash_attn False关闭bitsandbytes用于量化场景。预训练模型列表仓库提供了三个在公开数据上微调好的 Matroyshka Re-Ranker 模型可直接用于推理模型说明BAAI/Matroyshka-ReRanker-passage在 MS MARCO passage 上微调BAAI/Matroyshka-ReRanker-document在 MS MARCO document 上微调BAAI/Matroyshka-ReRanker-beir面向通用检索场景微调三个模型均以mistralai/Mistral-7B-v0.1为底座差异在于训练数据与适配场景。使用者可以按检索任务类型短文、长文、通用选择对应的 checkpoint。推理使用三组核心配置参数进入推理目录并启动交互式 Pythoncd ./inference python使用预训练模型进行推理from rank_model import MatroyshkaReranker compress_ratio 2 # config your compress ratio compress_layers [8, 16] # cofig your layers to compress cutoff_layers [20, 24] # config your layers to output reranker MatroyshkaReranker( model_name_or_pathBAAI/Matroyshka-ReRanker-passage, peft_path[ ./models/Matroyshka-ReRanker-passage/compensate/layer/full ] use_fp16True, cache_dir./model_cache, compress_ratiocompress_ratio, compress_layerscompress_layers, cutoff_layerscutoff_layers ) score reranker.compute_score([query, passage]) print(score) scores reranker.compute_score([[what is panda?, hi], [what is panda?, The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.]]) print(scores)注意peft_path列表项之间在原文档示例中缺少逗号实际使用时每个路径之间需要用英文逗号分隔如上例所示。三个核心参数的含义参数类型作用compress_ratioint指定层的 token 序列压缩比例支持[1, 2, 4, 8]1表示不压缩compress_layersList[int]需要执行 token 压缩的 Transformer 层索引列表cutoff_layersList[int]从中截取隐藏状态并输出分数的层索引列表对应「深度」配置从 inference/rank_model.py 的构造函数签名可以看到更完整的推理参数体系use_fp16/use_bf16控制半精度query_instruction_for_rerankA: 与passage_instruction_for_rerankB: 分别给查询与文档添加前缀指令max_length默认 512、query_max_length默认为max_length * 3 // 4batch_size默认 128normalize为 True 时用 sigmoid 把分数归一化到 [0, 1]。此外from_rawTrue表示从原始 Mistral 权重初始化逐层打分头CostWiseHeadstart_layer4对应打分层起始索引。输入与打分流程源码级compute_score的核心实现在compute_score_single_gpurank_model.py关键流程如下预处理query 与 passage 分别按query_max_length与max_length截断分词再按长度降序排序以减少 padding构造输入拼接[BOS] query \n passage \n prompt默认 prompt 为Predict whether passage B contains an answer to query A.并记录query_lengths与prompt_lengths供压缩阶段切分自动显存适配先用一个试探 batch 前向若触发RuntimeError或torch.cuda.OutOfMemoryError则自动将 batch_size 缩小为原来的 3/4batch_size batch_size * 3 // 4保证在有限显存下也能运行多输出头打分模型以output_hidden_statesTrue前向对每个cutoff_layers层的 logits 用last_logit_pool_lightweight取最后一个有效 token 的 logit 作为分数因此一次前向可以同时产出多个深度配置下的分数归一化若normalizeTrue对每个分数执行 sigmoid 映射到 [0, 1]。层压缩与 Token 压缩的底层实现推理端定制了基于 Mistral 的模型族位于 inference/mistral_model.py 与 inference/mistral_config.pyCostWiseMistralConfigmistral_config.py继承自MistralConfig新增start_layer打分层起始默认 18、layer_sep层间隔默认 18、layer_wise是否逐层输出三个配置字段CostWiseMistralModel.forwardmistral_model.py逐层前向时若idx in compress_layer且compress_ratio不为 1则调用token_compress对 hidden states 与 attention mask 执行序列压缩并重建位置编码当layer_wise开启时会在每个cutoff_layers层的输出处保存经 RMSNorm 后的 hidden states 与对应的 attention masktoken_compressmistral_model.py先根据query_lengths/prompt_lengths从序列中分离出 passage 区域将 passage 的 hidden states 按compress_ratio分组后加权求和权重来自注意力分数即all_self_attns layer_outputs[1][:, :, -1, :]同时保留 query 与 prompt 的原始 hidden states得到压缩后的新序列与新 attention mask。这一机制正是「在指定层按比例缩短 passage 序列」的宽度压缩实现CostWiseMistralForCausalLMmistral_model.py当layer_wiseTrue时lm_head被替换为nn.ModuleList每个元素是一个CostWiseHead单层线性层将 hidden_size 映射到 1 维分数对应一个可输出的层前向时对每个保存的 hidden state 独立计算分数得到一组 logits。从源码结构可以推断cutoff_layers中任何层只要落在[start_layer, num_hidden_layers]区间内且与layer_sep对齐都会被挂上一个独立的打分头越深的层语义越丰富但也意味着更大的计算开销这构成了「深度方向」的弹性。两阶段微调流程Matroyshka Re-Ranker 的训练方法包含两大特色级联自蒸馏Cascaded Self-distillation先让模型学会在每个可输出层、多种压缩比例下都能给出可靠分数因子化补偿Factorized Compensation在自蒸馏基础上对层压缩与 token 压缩带来的信息损失进行针对性补偿微调。阶段一Cascaded Self-distillation进入自蒸馏训练目录cd self_distillation以 8 卡 torchrun 为例train_data_path... your_huggingface_token... torchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_self_distillation \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --train_data ${train_data_path} \ --learning_rate 2e-4 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 100 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed /share/chaofan/code/stage/stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios 1 2 4 8 \ --compress_layers 4 8 12 16 20 24 28 \ --train_method distill_fix_layer_teacher本仓库已内置一份可参考的 DeepSpeed 配置 finetune/self_distillation/stage1.json。关键参数说明参数说明--model_name_or_path底座模型官方为mistralai/Mistral-7B-v0.1--query_max_len/--passage_max_len查询与文档的最大 token 长度默认 32 / 192--train_group_size每个 query 对应的候选文档数1 个正例 若干负例--start_layer/--layer_sep/--layer_wise逐层打分头配置从第 4 层开始、间隔 1 层、逐层输出--compress_ratios 1 2 4 8训练时随机采样使用的压缩比例集合--compress_layers训练时可能执行压缩的层集合--train_method distill_fix_layer_teacher蒸馏策略固定层作为教师逐步蒸馏--loss_type only logits仅使用 logits 计算损失自蒸馏的源码实现训练入口为 finetune/self_distillation/run.py模型封装在BiEncoderModelfinetune/self_distillation/modeling.py逐层打分头初始化load_model.py 对应的自蒸馏版本当layer_wiseTrue时把lm_head替换为从start_layer到num_hidden_layers的CostWiseHead列表并用原始lm_head中Yestoken 对应的权重行初始化每个打分头——这解释了推理 prompt 中「Predict whether ... Yes/No」的设计模型以「Yes」token 的 logit 作为相关性分数训练时的随机弹性modeling.pyencode中compress_ratiorandom.choice(self.compress_ratios)、compress_layer[random.choice([0, 1]) * i for i in self.compress_layers]即每个 batch 随机决定压缩比例与是否在某层压缩让模型学会在不同配置下都能打分级联蒸馏损失modeling.py对所有输出层 logits 按正样本位置计算交叉熵损失同时按distill_fix_layer_teacher策略用固定间隔的教师层 softmax 分数作为软标签对学生层施加 KL 型蒸馏损失从而把「深层教师」的能力逐级传递到「浅层学生」数据组织finetune/self_distillation/data.pyTrainDatasetForReranker读取 JSON 格式的query / pos / neg数据可选pos_scores/neg_scores/prompt字段每个样本构造train_group_size条「查询 文档 prompt」序列并记录 query 长度与 prompt 长度供训练时 token 压缩切分使用RerankCollator负责统一 padding 与组装pair / query_lengths / prompt_lengths / teacher_scores。阶段二Factorized Compensation层补偿与 Token 压缩补偿自蒸馏得到的./result_self_distillation会在推理或微调中作为raw_peft基底使用。补偿阶段将「压缩带来的能力损失」拆分为层压缩补偿与token 压缩补偿两部分分别训练。训练前先定义公共变量cd finetune/compensation train_data_path... your_huggingface_token... raw_peft_path../../self_distillation/result_self_distillationToken 压缩补偿--finetune_type tokencompress_ratio2 torchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_compensation_token_compress_ratio_${compress_ratio} \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --raw_peft ${raw_peft_path} \ --train_data ${train_data_path} \ --learning_rate 2e-5 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 500 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios ${compress_ratio} \ --compress_layers 4 8 12 16 20 24 28 \ --train_method normal \ --finetune_type token层补偿--finetune_type layertorchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_compensation_layer \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --raw_peft ${raw_peft_path} \ --train_data ${train_data_path} \ --learning_rate 2e-5 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 500 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios 1 \ --compress_layers 4 8 12 16 20 24 28 \ --train_method normal \ --finetune_type layer两个脚本的差异集中在三点--finetune_typetoken或layer、--compress_ratiostoken 压缩补偿按目标比例设置层补偿固定为 1 即不压缩、--output_dir。层补偿与 token 压缩补偿可以在推理时叠加使用见下文「使用自训练模型」示例。补偿的源码实现参数解析补偿脚本新增finetune_typelayer/token与raw_peft参数见 finetune/compensation/arguments.py基底模型加载finetune/compensation/load_model.py通过--raw_peft将自蒸馏阶段的 LoRA 权重合并进模型后再套上新的 LoRA即补偿模块是叠加在自蒸馏基底之上的增量参数训练损失finetune/compensation/modeling.pytrain_method normal下仅对逐层 logits 计算正样本位置的交叉熵损失不再引入教师蒸馏专注修复压缩带来的分数偏差训练时同样随机化压缩层与压缩比例以覆盖多种配置组合。使用自训练模型进行推理自训练完成后把自蒸馏 LoRA 与补偿 LoRA 按序传给peft_pathcd ./inference pythonfrom rank_model import MatroyshkaReranker compress_ratio 2 # config your compress ratio compress_layers [8, 16] # cofig your layers to compress cutoff_layers [20, 24] # config your layers to output reranker MatroyshkaReranker( model_name_or_pathmistralai/Mistral-7B-v0.1, peft_path[ ./finetune/self_distillation/result_self_distillation, ./finetune/compensation/result_compensation_token_compress_ratio_2, ], use_fp16True, cache_dir./model_cache, compress_ratiocompress_ratio, compress_layerscompress_layers, cutoff_layerscutoff_layers ) score reranker.compute_score([query, passage]) print(score) scores reranker.compute_score([[what is panda?, hi], [what is panda?, The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.]]) print(scores)MatroyshkaReranker在加载时会对peft_path中的每个 LoRA 依次执行PeftModel.from_pretrained与merge_and_unloadrank_model.py因此peft_path的顺序应与训练顺序一致先是自蒸馏基底再是补偿模块。调参与效果权衡建议压缩比例选择compress_ratio越大passage 序列越短、推理越快但信息损失越多。建议线上服务从2起步依据评测指标如 nDCG、Recall与延迟预算逐步上调到4/8压缩层位置compress_layers应尽量靠后设置——前期层负责基础语义提取过早压缩会损失细节官方训练采用[4, 8, 12, 16, 20, 24, 28]覆盖整个网络但以间隔采样输出层深度cutoff_layers的层越深、分数越可靠但每多一个输出层就多一次打分头的计算可在推理时一次性传入多个候选层利用多输出头机制对比不同深度的分数质量再为线上选定一个深度显存与批次推理端遇到 OOM 会自动将 batch_size 缩减至 3/4训练端建议开启--gradient_checkpointing并配合 DeepSpeed stage1示例见 finetune/self_distillation/stage1.json精度与归一化半精度use_fp16/use_bf16是官方推荐默认源码中若两者均为 False 会自动回退到 fp16 并告警若下游需要概率化解读请开启normalizeTrue。小结Matroyshka Re-Ranker 的核心贡献在于把 LLM 重排序器的「深度」与「宽度」都变成运行时可配置的维度通过cutoff_layers控制打分层数通过compress_layerscompress_ratio控制序列压缩配套「级联自蒸馏 因子化补偿」的两阶段训练让模型在所有配置组合下都保持可靠。读者可以依据 research/Matroyshka_reranker/README.md 的脚本与本文的源码分析从预训练模型直接推理起步再按需复现完整的两阶段微调流程最终将其接入自己的检索与 RAG 管线中作为精排模块。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Matroyshka Re-Ranker：基于 Mistral-7B 的可配置深度与宽度 LLM 重排序器（FlagEmbedding 实战指南）
Matroyshka Re-Ranker基于 Mistral-7B 的可配置深度与宽度 LLM 重排序器FlagEmbedding 实战指南【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding导读本篇技术指南围绕 FlagEmbedding 仓库中 research/Matroyshka_reranker 目录下的Matroyshka Re-Ranker展开介绍一种能够按照用户配置在运行时自定义模型层数与各层序列长度的轻量级 LLM 重排序器。重排序Re-ranking不同于向量检索它以「查询 文档」对为输入、直接输出相关性分数。读完本文你将掌握Matroyshka Re-Ranker 的推理调用方式compress_ratio/compress_layers/cutoff_layers三组核心配置、基于「级联自蒸馏Cascaded Self-distillation 因子化补偿Factorized Compensation」的两阶段微调流程以及隐藏在源码中的 Token 压缩与逐层打分原理。背景为什么需要「可变形」的 LLM Re-Ranker与嵌入模型不同重排序模型将查询query与文档passage同时作为输入直接输出二者的相关性分数而非生成向量表示。Matroyshka Re-Ranker 的输出分数可以通过 sigmoid 函数映射到 [0, 1] 区间便于与其他模块统一解读对应实现见 inference/rank_model.py 中从 FlagEmbedding 复用的sigmoid。其核心设计目标是运行时的灵活配置runtime customization深度维度可以选择从模型的哪些层cutoff_layers输出分数宽度维度可以在指定的层compress_layers上按比例compress_ratio压缩 token 序列长度减少后续层的计算量。这样的设计使得同一个模型在离线高精度重排不压缩、用深层输出与在线低延迟重排按需压缩层数与序列长度之间按需切换支持轻量化的灵活配置。环境安装Matroyshka Re-Ranker 基于 Mistral-7B 架构开发推荐使用 conda 创建独立 Python 3.10 环境conda create -n reranker python3.10 conda activate reranker pip install -r requirements.txt依赖清单见 requirements.txt关键依赖包括transformers4.41.1、peft0.10.0、deepspeed、flash-attn、datasets、modelscope、bitsandbytes等。其中flash-attn为可选加速项训练脚本可通过--use_flash_attn False关闭bitsandbytes用于量化场景。预训练模型列表仓库提供了三个在公开数据上微调好的 Matroyshka Re-Ranker 模型可直接用于推理模型说明BAAI/Matroyshka-ReRanker-passage在 MS MARCO passage 上微调BAAI/Matroyshka-ReRanker-document在 MS MARCO document 上微调BAAI/Matroyshka-ReRanker-beir面向通用检索场景微调三个模型均以mistralai/Mistral-7B-v0.1为底座差异在于训练数据与适配场景。使用者可以按检索任务类型短文、长文、通用选择对应的 checkpoint。推理使用三组核心配置参数进入推理目录并启动交互式 Pythoncd ./inference python使用预训练模型进行推理from rank_model import MatroyshkaReranker compress_ratio 2 # config your compress ratio compress_layers [8, 16] # cofig your layers to compress cutoff_layers [20, 24] # config your layers to output reranker MatroyshkaReranker( model_name_or_pathBAAI/Matroyshka-ReRanker-passage, peft_path[ ./models/Matroyshka-ReRanker-passage/compensate/layer/full ] use_fp16True, cache_dir./model_cache, compress_ratiocompress_ratio, compress_layerscompress_layers, cutoff_layerscutoff_layers ) score reranker.compute_score([query, passage]) print(score) scores reranker.compute_score([[what is panda?, hi], [what is panda?, The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.]]) print(scores)注意peft_path列表项之间在原文档示例中缺少逗号实际使用时每个路径之间需要用英文逗号分隔如上例所示。三个核心参数的含义参数类型作用compress_ratioint指定层的 token 序列压缩比例支持[1, 2, 4, 8]1表示不压缩compress_layersList[int]需要执行 token 压缩的 Transformer 层索引列表cutoff_layersList[int]从中截取隐藏状态并输出分数的层索引列表对应「深度」配置从 inference/rank_model.py 的构造函数签名可以看到更完整的推理参数体系use_fp16/use_bf16控制半精度query_instruction_for_rerankA: 与passage_instruction_for_rerankB: 分别给查询与文档添加前缀指令max_length默认 512、query_max_length默认为max_length * 3 // 4batch_size默认 128normalize为 True 时用 sigmoid 把分数归一化到 [0, 1]。此外from_rawTrue表示从原始 Mistral 权重初始化逐层打分头CostWiseHeadstart_layer4对应打分层起始索引。输入与打分流程源码级compute_score的核心实现在compute_score_single_gpurank_model.py关键流程如下预处理query 与 passage 分别按query_max_length与max_length截断分词再按长度降序排序以减少 padding构造输入拼接[BOS] query \n passage \n prompt默认 prompt 为Predict whether passage B contains an answer to query A.并记录query_lengths与prompt_lengths供压缩阶段切分自动显存适配先用一个试探 batch 前向若触发RuntimeError或torch.cuda.OutOfMemoryError则自动将 batch_size 缩小为原来的 3/4batch_size batch_size * 3 // 4保证在有限显存下也能运行多输出头打分模型以output_hidden_statesTrue前向对每个cutoff_layers层的 logits 用last_logit_pool_lightweight取最后一个有效 token 的 logit 作为分数因此一次前向可以同时产出多个深度配置下的分数归一化若normalizeTrue对每个分数执行 sigmoid 映射到 [0, 1]。层压缩与 Token 压缩的底层实现推理端定制了基于 Mistral 的模型族位于 inference/mistral_model.py 与 inference/mistral_config.pyCostWiseMistralConfigmistral_config.py继承自MistralConfig新增start_layer打分层起始默认 18、layer_sep层间隔默认 18、layer_wise是否逐层输出三个配置字段CostWiseMistralModel.forwardmistral_model.py逐层前向时若idx in compress_layer且compress_ratio不为 1则调用token_compress对 hidden states 与 attention mask 执行序列压缩并重建位置编码当layer_wise开启时会在每个cutoff_layers层的输出处保存经 RMSNorm 后的 hidden states 与对应的 attention masktoken_compressmistral_model.py先根据query_lengths/prompt_lengths从序列中分离出 passage 区域将 passage 的 hidden states 按compress_ratio分组后加权求和权重来自注意力分数即all_self_attns layer_outputs[1][:, :, -1, :]同时保留 query 与 prompt 的原始 hidden states得到压缩后的新序列与新 attention mask。这一机制正是「在指定层按比例缩短 passage 序列」的宽度压缩实现CostWiseMistralForCausalLMmistral_model.py当layer_wiseTrue时lm_head被替换为nn.ModuleList每个元素是一个CostWiseHead单层线性层将 hidden_size 映射到 1 维分数对应一个可输出的层前向时对每个保存的 hidden state 独立计算分数得到一组 logits。从源码结构可以推断cutoff_layers中任何层只要落在[start_layer, num_hidden_layers]区间内且与layer_sep对齐都会被挂上一个独立的打分头越深的层语义越丰富但也意味着更大的计算开销这构成了「深度方向」的弹性。两阶段微调流程Matroyshka Re-Ranker 的训练方法包含两大特色级联自蒸馏Cascaded Self-distillation先让模型学会在每个可输出层、多种压缩比例下都能给出可靠分数因子化补偿Factorized Compensation在自蒸馏基础上对层压缩与 token 压缩带来的信息损失进行针对性补偿微调。阶段一Cascaded Self-distillation进入自蒸馏训练目录cd self_distillation以 8 卡 torchrun 为例train_data_path... your_huggingface_token... torchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_self_distillation \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --train_data ${train_data_path} \ --learning_rate 2e-4 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 100 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed /share/chaofan/code/stage/stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios 1 2 4 8 \ --compress_layers 4 8 12 16 20 24 28 \ --train_method distill_fix_layer_teacher本仓库已内置一份可参考的 DeepSpeed 配置 finetune/self_distillation/stage1.json。关键参数说明参数说明--model_name_or_path底座模型官方为mistralai/Mistral-7B-v0.1--query_max_len/--passage_max_len查询与文档的最大 token 长度默认 32 / 192--train_group_size每个 query 对应的候选文档数1 个正例 若干负例--start_layer/--layer_sep/--layer_wise逐层打分头配置从第 4 层开始、间隔 1 层、逐层输出--compress_ratios 1 2 4 8训练时随机采样使用的压缩比例集合--compress_layers训练时可能执行压缩的层集合--train_method distill_fix_layer_teacher蒸馏策略固定层作为教师逐步蒸馏--loss_type only logits仅使用 logits 计算损失自蒸馏的源码实现训练入口为 finetune/self_distillation/run.py模型封装在BiEncoderModelfinetune/self_distillation/modeling.py逐层打分头初始化load_model.py 对应的自蒸馏版本当layer_wiseTrue时把lm_head替换为从start_layer到num_hidden_layers的CostWiseHead列表并用原始lm_head中Yestoken 对应的权重行初始化每个打分头——这解释了推理 prompt 中「Predict whether ... Yes/No」的设计模型以「Yes」token 的 logit 作为相关性分数训练时的随机弹性modeling.pyencode中compress_ratiorandom.choice(self.compress_ratios)、compress_layer[random.choice([0, 1]) * i for i in self.compress_layers]即每个 batch 随机决定压缩比例与是否在某层压缩让模型学会在不同配置下都能打分级联蒸馏损失modeling.py对所有输出层 logits 按正样本位置计算交叉熵损失同时按distill_fix_layer_teacher策略用固定间隔的教师层 softmax 分数作为软标签对学生层施加 KL 型蒸馏损失从而把「深层教师」的能力逐级传递到「浅层学生」数据组织finetune/self_distillation/data.pyTrainDatasetForReranker读取 JSON 格式的query / pos / neg数据可选pos_scores/neg_scores/prompt字段每个样本构造train_group_size条「查询 文档 prompt」序列并记录 query 长度与 prompt 长度供训练时 token 压缩切分使用RerankCollator负责统一 padding 与组装pair / query_lengths / prompt_lengths / teacher_scores。阶段二Factorized Compensation层补偿与 Token 压缩补偿自蒸馏得到的./result_self_distillation会在推理或微调中作为raw_peft基底使用。补偿阶段将「压缩带来的能力损失」拆分为层压缩补偿与token 压缩补偿两部分分别训练。训练前先定义公共变量cd finetune/compensation train_data_path... your_huggingface_token... raw_peft_path../../self_distillation/result_self_distillationToken 压缩补偿--finetune_type tokencompress_ratio2 torchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_compensation_token_compress_ratio_${compress_ratio} \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --raw_peft ${raw_peft_path} \ --train_data ${train_data_path} \ --learning_rate 2e-5 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 500 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios ${compress_ratio} \ --compress_layers 4 8 12 16 20 24 28 \ --train_method normal \ --finetune_type token层补偿--finetune_type layertorchrun --nproc_per_node 8 \ run.py \ --output_dir ./result_compensation_layer \ --model_name_or_path mistralai/Mistral-7B-v0.1 \ --raw_peft ${raw_peft_path} \ --train_data ${train_data_path} \ --learning_rate 2e-5 \ --num_train_epochs 1 \ --per_device_train_batch_size 4 \ --gradient_accumulation_steps 4 \ --dataloader_drop_last True \ --query_max_len 32 \ --passage_max_len 192 \ --train_group_size 16 \ --logging_steps 1 \ --save_steps 500 \ --save_total_limit 50 \ --ddp_find_unused_parameters False \ --gradient_checkpointing \ --deepspeed stage1.json \ --warmup_ratio 0.1 \ --bf16 \ --use_lora True \ --lora_rank 32 \ --lora_alpha 64 \ --loss_type only logits \ --use_flash_attn False \ --target_modules q_proj k_proj v_proj o_proj down_proj up_proj gate_proj linear_head \ --token ${your_huggingface_token} \ --cache_dir ../../model_cache \ --cache_path ../../data_cache \ --padding_side right \ --start_layer 4 \ --layer_sep 1 \ --layer_wise True \ --compress_ratios 1 \ --compress_layers 4 8 12 16 20 24 28 \ --train_method normal \ --finetune_type layer两个脚本的差异集中在三点--finetune_typetoken或layer、--compress_ratiostoken 压缩补偿按目标比例设置层补偿固定为 1 即不压缩、--output_dir。层补偿与 token 压缩补偿可以在推理时叠加使用见下文「使用自训练模型」示例。补偿的源码实现参数解析补偿脚本新增finetune_typelayer/token与raw_peft参数见 finetune/compensation/arguments.py基底模型加载finetune/compensation/load_model.py通过--raw_peft将自蒸馏阶段的 LoRA 权重合并进模型后再套上新的 LoRA即补偿模块是叠加在自蒸馏基底之上的增量参数训练损失finetune/compensation/modeling.pytrain_method normal下仅对逐层 logits 计算正样本位置的交叉熵损失不再引入教师蒸馏专注修复压缩带来的分数偏差训练时同样随机化压缩层与压缩比例以覆盖多种配置组合。使用自训练模型进行推理自训练完成后把自蒸馏 LoRA 与补偿 LoRA 按序传给peft_pathcd ./inference pythonfrom rank_model import MatroyshkaReranker compress_ratio 2 # config your compress ratio compress_layers [8, 16] # cofig your layers to compress cutoff_layers [20, 24] # config your layers to output reranker MatroyshkaReranker( model_name_or_pathmistralai/Mistral-7B-v0.1, peft_path[ ./finetune/self_distillation/result_self_distillation, ./finetune/compensation/result_compensation_token_compress_ratio_2, ], use_fp16True, cache_dir./model_cache, compress_ratiocompress_ratio, compress_layerscompress_layers, cutoff_layerscutoff_layers ) score reranker.compute_score([query, passage]) print(score) scores reranker.compute_score([[what is panda?, hi], [what is panda?, The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.]]) print(scores)MatroyshkaReranker在加载时会对peft_path中的每个 LoRA 依次执行PeftModel.from_pretrained与merge_and_unloadrank_model.py因此peft_path的顺序应与训练顺序一致先是自蒸馏基底再是补偿模块。调参与效果权衡建议压缩比例选择compress_ratio越大passage 序列越短、推理越快但信息损失越多。建议线上服务从2起步依据评测指标如 nDCG、Recall与延迟预算逐步上调到4/8压缩层位置compress_layers应尽量靠后设置——前期层负责基础语义提取过早压缩会损失细节官方训练采用[4, 8, 12, 16, 20, 24, 28]覆盖整个网络但以间隔采样输出层深度cutoff_layers的层越深、分数越可靠但每多一个输出层就多一次打分头的计算可在推理时一次性传入多个候选层利用多输出头机制对比不同深度的分数质量再为线上选定一个深度显存与批次推理端遇到 OOM 会自动将 batch_size 缩减至 3/4训练端建议开启--gradient_checkpointing并配合 DeepSpeed stage1示例见 finetune/self_distillation/stage1.json精度与归一化半精度use_fp16/use_bf16是官方推荐默认源码中若两者均为 False 会自动回退到 fp16 并告警若下游需要概率化解读请开启normalizeTrue。小结Matroyshka Re-Ranker 的核心贡献在于把 LLM 重排序器的「深度」与「宽度」都变成运行时可配置的维度通过cutoff_layers控制打分层数通过compress_layerscompress_ratio控制序列压缩配套「级联自蒸馏 因子化补偿」的两阶段训练让模型在所有配置组合下都保持可靠。读者可以依据 research/Matroyshka_reranker/README.md 的脚本与本文的源码分析从预训练模型直接推理起步再按需复现完整的两阶段微调流程最终将其接入自己的检索与 RAG 管线中作为精排模块。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考