Skip to content

fix: preserve dense pooling in encoder checkpoints - #1622

Open
emme1t wants to merge 1 commit into
FlagOpen:masterfrom
emme1t:fix/encoder-checkpoint-pooling-1326
Open

emme1t wants to merge 1 commit into
FlagOpen:masterfrom
emme1t:fix/encoder-checkpoint-pooling-1326

Conversation

@emme1t

@emme1t emme1t commented Nov 1, 2026

Copy link
Copy Markdown

Encoder-only fine-tuning checkpoints omit the Sentence Transformers pipeline configuration. Loading a CLS-pooled BGE-M3 checkpoint through SentenceTransformer therefore creates a mean-pooling pipeline and changes the resulting embeddings.

Save the dense pipeline metadata from the base encoder and M3 model save paths, including the configured pooling and optional normalization. The helper writes configuration alongside the existing weights, maps last_token to Sentence Transformers' lasttoken, and refreshes the module list on subsequent saves. M3 sparse and ColBERT weights retain their existing save path.

Fixes #1326.

Validation

Offline tests use tiny BERT and XLM-RoBERTa models, perform an optimizer step, save through the real trainer, reload with Sentence Transformers, and compare the actual embeddings. They cover all three pooling methods, normalization on/off, repeated saves, and M3 auxiliary weights.

  • Original production code: 10 failed / 4 passed for the new regression file.
  • Transformers 5.22.1 / Sentence Transformers 6.2.0: 25 passed, including existing trainer/import/padding tests.
  • Transformers 4.44.2 / Sentence Transformers 3.0.1: 14 passed for the new regression file.
  • Ruff and formatting pass for the new helper/tests; changed existing modules introduce no lint findings relative to the baseline. Compilation, dependency checks and wheel build pass.

Commands:

python -m pytest tests/test_encoder_sentence_transformers_export.py -q
python -m pytest tests/test_encoder_sentence_transformers_export.py tests/test_finetune_trainer_compat.py tests/test_imports_v5.py tests/test_tokenizer_pad_compat.py::test_pad_with_compat_normalizes_list_of_mappings tests/test_tokenizer_pad_compat.py::test_pad_with_compat_passes_mapping_inputs_through -q
uv build --wheel

The broader legacy-version run also reproduces two existing Trainer compatibility test failures: those tests directly assert the newer processing_class attribute. Both fail identically with the original production files under Transformers 4.44.2.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Some configuration files are missing after finetuning bge-m3 model

1 participant