Loading a LoRA CrossEncoder checkpoint does not restore a trained `modules_to_save` classification head

[!NOTE]

Generated by AI and edited by human

Summary

Sentence Transformers can train a CrossEncoder with a PEFT LoRA adapter and then save the resulting adapter checkpoint. This checkpoint may include:

1. LoRA weights for selected transformer layers.
2. A separately trained sequence-classification head, configured through PEFT's modules_to_save.

For a sequence-classification model, the classification head is commonly named score. In this report, score.weight is the linear layer that converts the model's final hidden state into the relevance/classification score returned by the CrossEncoder.

When a checkpoint trained with modules_to_save=["score"] is loaded through:

CrossEncoder("path/to/adapter-checkpoint")

the trained score weight is not restored into the module that PEFT actually uses during inference. Instead, that active classification head remains newly initialized. This means predictions use an untrained random head, even though the checkpoint contains the trained head weight.

The checkpoint loads with the following Transformers load report:

score.weight                         | UNEXPECTED |
score.modules_to_save.default.weight | MISSING    |

Here:

• UNEXPECTED means a weight exists in the checkpoint but does not match a parameter expected by the model currently being loaded.
• MISSING means the model expects a parameter but did not receive a saved value for it, so Transformers initializes it from scratch.
• score.modules_to_save.default.weight is the active PEFT-managed copy of the classification head for the default adapter.

In this scenario, the UNEXPECTED entry is not safe to ignore, because it corresponds to the trained classification head and the MISSING parameter is the head actually used in forward passes.

Context

Environment

sentence-transformers: 1d9647561b931252f868b9495172216161d26c82
peft: 0.20.0
transformers: 5.15.1
torch: 2.13.0+cu130
Python: 3.14

How the checkpoint was created

The base model is a Hugging Face sequence-classification model used through Sentence Transformers' CrossEncoder.

The adapter was configured as follows:

from peft import LoraConfig

lora_config = LoraConfig(
    r=8,
    lora_alpha=8,
    lora_dropout=0,
    target_modules="all-linear",
    modules_to_save=["score"],
    bias="none",
    task_type="SEQ_CLS",
)

modules_to_save=["score"] tells PEFT to train and save the complete score classification layer in addition to the LoRA weights.

The resulting adapter_config.json contains:

{
  "modules_to_save": ["score"],
  "task_type": "SEQ_CLS"
}

The adapter checkpoint stores the trained classification-head tensor under this key:

base_model.model.score.weight

This is the standard PEFT checkpoint key for a module listed in modules_to_save; PEFT maps it to the active adapter-specific copy when loading normally.

Reproduction

from pathlib import Path

import torch
from safetensors import safe_open
from sentence_transformers import CrossEncoder

checkpoint_path = Path("path/to/adapter-checkpoint")

# Read the trained classification head directly from the saved adapter.
with safe_open(
    checkpoint_path / "adapter_model.safetensors",
    framework="pt",
    device="cpu",
) as source:
    saved_score_weight = source.get_tensor(
        "base_model.model.score.weight"
    ).float()

# Load through Sentence Transformers' automatic checkpoint loading path.
encoder = CrossEncoder(
    checkpoint_path,
    model_kwargs={"torch_dtype": "float32"},
)

model = encoder.model
assert model is not None

# PEFT wraps modules listed in `modules_to_save`.
score = model.score

print(type(score))
print(score.active_adapters)
print(score.modules_to_save.keys())

# This should be True: the active inference head should equal the saved,
# trained classification head.
active_score_weight = (
    score.modules_to_save["default"].weight.detach().cpu().float()
)
print(torch.equal(active_score_weight, saved_score_weight))

Actual behavior

During CrossEncoder(checkpoint_path) loading, Transformers reports:

score.weight                         | UNEXPECTED |
score.modules_to_save.default.weight | MISSING    |

The model's score module is loaded as:

peft.utils.other.ModulesToSaveWrapper

Its active adapter is default, and it contains:

score.modules_to_save["default"].weight

However, this active weight does not equal the trained weight stored in the checkpoint:

torch.equal(active_score_weight, saved_score_weight)
# False

The tensors differ substantially, rather than only by floating-point serialization noise. For example, the maximum absolute element-wise difference in one reproduction was approximately 0.095; for normally restored float32 checkpoint weights, this value should be exactly 0.0.

Why this changes inference results

PEFT's ModulesToSaveWrapper uses the active adapter-specific copy in its forward path:

def _forward_wrapped(self, *args, **kwargs):
    if not self.active_adapters:
        return self._forward_wrapped_passthrough(*args, **kwargs)
    return self.modules_to_save[self.active_adapters[0]](*args, **kwargs)

Therefore, with the active adapter set to default, inference uses:

score.modules_to_save["default"].weight

rather than:

score.original_module.weight

This was verified directly:

inputs = torch.ones(2, score.original_module.in_features)

with torch.no_grad():
    baseline = score(inputs).clone()

    # Changing the original, inactive copy does not change forward output.
    score.original_module.weight.zero_()
    after_zeroing_original = score(inputs).clone()

    # Changing the active adapter copy does change forward output.
    score.modules_to_save["default"].weight.zero_()
    after_zeroing_active = score(inputs).clone()

print(torch.equal(baseline, after_zeroing_original))
# True

print(torch.equal(after_zeroing_original, after_zeroing_active))
# False

As a result, the automatically loaded model uses a newly initialized classification head instead of the head trained with the adapter.

Expected behavior

For a checkpoint whose PEFT configuration contains:

{
  "modules_to_save": ["score"]
}

CrossEncoder(checkpoint_path) should restore the checkpoint tensor:

base_model.model.score.weight

into the active PEFT module:

score.modules_to_save["default"].weight

After loading, this should hold:

torch.equal(
    score.modules_to_save["default"].weight.detach().cpu().float(),
    saved_score_weight,
)
# True

The automatic Sentence Transformers loading path should produce the same model parameters and inference behavior as loading the base model and adapter through PEFT directly.

Control experiment: the PEFT checkpoint is valid

The same checkpoint restores the trained classification head correctly when using PEFT directly:

from pathlib import Path

import torch
from peft import PeftModel
from safetensors import safe_open
from transformers import AutoModelForSequenceClassification

base_model_path = Path("path/to/base-model")
adapter_path = Path("path/to/adapter-checkpoint")

with safe_open(
    adapter_path / "adapter_model.safetensors",
    framework="pt",
    device="cpu",
) as source:
    saved_score_weight = source.get_tensor(
        "base_model.model.score.weight"
    ).float()

base_model = AutoModelForSequenceClassification.from_pretrained(
    base_model_path,
    num_labels=1,
    torch_dtype=torch.float32,
)

model = PeftModel.from_pretrained(base_model, adapter_path)

active_score_weight = (
    model.base_model.model.score.modules_to_save["default"]
    .weight.detach()
    .cpu()
    .float()
)

print(torch.equal(active_score_weight, saved_score_weight))
# True

print((active_score_weight - saved_score_weight).abs().max().item())
# 0.0

This indicates that the adapter checkpoint format is valid and that PEFT can restore it correctly. The failure appears specific to Sentence Transformers' automatic adapter-loading path.

Suggested fix

The automatic loading path should avoid treating adapter_model.safetensors as an ordinary full-model checkpoint before the PEFT adapter structure has been created.

A possible approach is:

1. Load the base sequence-classification model from the adapter configuration's base_model_name_or_path.

2. Detect that the checkpoint is a PEFT adapter checkpoint, for example by the presence of adapter_config.json.

3. Load the adapter using PEFT's native API:

   python复制

   PeftModel.from_pretrained(base_model, adapter_checkpoint_path)

4. Construct or return the CrossEncoder around that PEFT-wrapped model.

Alternatively, if Sentence Transformers must retain its existing loading flow, it should translate the saved PEFT modules_to_save key:

base_model.model.score.weight

to the key expected after wrapper creation:

score.modules_to_save.default.weight

and ensure the PEFT wrapper is installed before applying the adapter state dict.

A regression test should train or construct a small sequence-classification LoRA adapter with modules_to_save=["score"], save it, reload it through CrossEncoder(checkpoint_path), and verify:

loaded_score.modules_to_save["default"].weight == saved_adapter_score_weight

and identical inference logits between the Sentence Transformers automatic-loading path and the equivalent PEFT-native loading path.