Implementation:FlagOpen FlagEmbedding MiniCPM Reranker Merge

Knowledge Sources	FlagOpen_FlagEmbedding
Domains	Machine Learning, Information Retrieval, Reranking
Last Updated	2026-02-09 00:00 GMT

Overview

Merged variant of the layer-wise MiniCPM reranker model for model merging and ensemble operations.

Description

This module provides an identical implementation to the layer-wise MiniCPM reranker but is located in the merge directory, suggesting it is used for model merging operations, checkpoint consolidation, or ensemble creation. The LayerWiseMiniCPMForCausalLM class in this module supports the same layer-wise prediction capabilities as the finetune variant but may be optimized for post-training operations like parameter averaging, distillation, or checkpoint merging.

The architecture is identical to the finetune version, featuring:

Full MiniCPM transformer architecture with layer-wise outputs
Multiple attention implementations (eager, Flash Attention 2, SDPA)
Flexible head configurations (raw, complex, reranking)
Support for multi-head architectures per layer
Chat interface and generation capabilities

This separation likely facilitates different workflows where one version is used for active training while this version handles model consolidation and deployment preparation.

Usage

Use this module when merging multiple trained checkpoints, creating ensemble models, or preparing final deployment models from layer-wise trained rerankers.

Code Reference

Source Location

Repository: FlagOpen_FlagEmbedding
File: research/llm_reranker/merge/modeling_minicpm_reranker.py

Signature

class LayerWiseMiniCPMModel(MiniCPMPreTrainedModel):
    def __init__(self, config: LayerWiseMiniCPMConfig)

    def forward(self, input_ids, attention_mask=None, position_ids=None,
                past_key_values=None, inputs_embeds=None, use_cache=None,
                output_attentions=None, output_hidden_states=None,
                return_dict=None, cutoff_layers=None)

class LayerWiseMiniCPMForCausalLM(MiniCPMPreTrainedModel):
    def __init__(self, config)

    def forward(self, input_ids, attention_mask=None, position_ids=None,
                past_key_values=None, inputs_embeds=None, labels=None,
                use_cache=None, output_attentions=None,
                output_hidden_states=None, return_dict=None,
                cutoff_layers=None, only_for_one_logit=None)

    def chat(self, tokenizer, query: str, history: List[Dict] = None,
             role: str = "user", max_length: int = 4096,
             num_beams=1, do_sample=True, top_p=0.8, temperature=0.3,
             logits_processor=None, **kwargs)

Import

# Import from merge directory
from merge.modeling_minicpm_reranker import LayerWiseMiniCPMForCausalLM
from transformers import AutoTokenizer

I/O Contract

Inputs

Name	Type	Required	Description
input_ids	torch.LongTensor	Yes	Input token IDs (batch_size, seq_len)
attention_mask	torch.Tensor	No	Attention mask for padding
cutoff_layers	Union[int, List[int]]	No	Layers to extract hidden states from
labels	torch.LongTensor	No	Labels for computing loss
only_for_one_logit	int	No	Extract single logit dimension

Outputs

Name	Type	Description
loss	torch.Tensor	Training/validation loss
logits	Tuple[torch.Tensor]	Logits from each cutoff layer
hidden_states	Tuple[torch.Tensor]	Hidden states from cutoff layers
past_key_values	Tuple	Cached key-value pairs

Usage Examples

from merge.modeling_minicpm_reranker import LayerWiseMiniCPMForCausalLM
from transformers import AutoTokenizer
import torch

# Load multiple checkpoints for merging
checkpoint_paths = [
    "checkpoint-1000",
    "checkpoint-2000",
    "checkpoint-3000"
]

models = []
for path in checkpoint_paths:
    model = LayerWiseMiniCPMForCausalLM.from_pretrained(
        path,
        torch_dtype=torch.bfloat16
    )
    models.append(model)

# Parameter averaging (simple ensemble)
from collections import OrderedDict

merged_state_dict = OrderedDict()
for name, param in models[0].named_parameters():
    merged_param = torch.zeros_like(param)
    for model in models:
        merged_param += model.state_dict()[name] / len(models)
    merged_state_dict[name] = merged_param

# Load merged parameters into a new model
merged_model = LayerWiseMiniCPMForCausalLM.from_pretrained(
    checkpoint_paths[0],
    torch_dtype=torch.bfloat16
)
merged_model.load_state_dict(merged_state_dict)

# Evaluate merged model
tokenizer = AutoTokenizer.from_pretrained(checkpoint_paths[0])
query = "What is information retrieval?"
document = "Information retrieval is the process of obtaining information..."

text = f"Query: {query}\nDocument: {document}\nRelevant:"
inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)

with torch.no_grad():
    outputs = merged_model(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        cutoff_layers=[40],  # Final layer only
        only_for_one_logit=0
    )
    relevance_score = outputs.logits[0]
    print(f"Merged model score: {relevance_score}")

# Save merged model for deployment
merged_model.save_pretrained("minicpm-reranker-merged")
tokenizer.save_pretrained("minicpm-reranker-merged")

# Weighted ensemble (different weights per checkpoint)
weights = [0.5, 0.3, 0.2]  # More weight on later checkpoints
weighted_state_dict = OrderedDict()

for name, param in models[0].named_parameters():
    weighted_param = torch.zeros_like(param)
    for model, weight in zip(models, weights):
        weighted_param += model.state_dict()[name] * weight
    weighted_state_dict[name] = weighted_param

weighted_model = LayerWiseMiniCPMForCausalLM.from_pretrained(
    checkpoint_paths[0],
    torch_dtype=torch.bfloat16
)
weighted_model.load_state_dict(weighted_state_dict)
weighted_model.save_pretrained("minicpm-reranker-weighted")

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment