From 20b100994edf690eabb3c6be437e3a5b577ddefa Mon Sep 17 00:00:00 2001 From: Zhe Chen Date: Sun, 11 Oct 2026 18:43:50 +0800 Subject: [PATCH] Fix distillation with integer teacher scores --- .../abc/finetune/embedder/AbsModeling.py | 2 +- .../embedder/encoder_only/m3/modeling.py | 2 +- tests/test_embedder_teacher_scores.py | 117 ++++++++++++++++++ 3 files changed, 119 insertions(+), 2 deletions(-) create mode 100644 tests/test_embedder_teacher_scores.py diff --git a/FlagEmbedding/abc/finetune/embedder/AbsModeling.py b/FlagEmbedding/abc/finetune/embedder/AbsModeling.py index 3e9e32cb3..3bff951ec 100644 --- a/FlagEmbedding/abc/finetune/embedder/AbsModeling.py +++ b/FlagEmbedding/abc/finetune/embedder/AbsModeling.py @@ -270,7 +270,7 @@ def forward( if self.training: if teacher_scores is not None: - teacher_scores = torch.tensor(teacher_scores, device=device) + teacher_scores = torch.tensor(teacher_scores, dtype=torch.float32, device=device) teacher_scores = teacher_scores.view(batch_size, -1).detach() # (batch_size, group_size) teacher_targets = F.softmax(teacher_scores, dim=-1) # (batch_size, group_size) else: diff --git a/FlagEmbedding/finetune/embedder/encoder_only/m3/modeling.py b/FlagEmbedding/finetune/embedder/encoder_only/m3/modeling.py index c3111ea5e..739690e01 100644 --- a/FlagEmbedding/finetune/embedder/encoder_only/m3/modeling.py +++ b/FlagEmbedding/finetune/embedder/encoder_only/m3/modeling.py @@ -394,7 +394,7 @@ def forward( if self.training: if teacher_scores is not None: - teacher_scores = torch.tensor(teacher_scores, device=q_dense_vecs.device) + teacher_scores = torch.tensor(teacher_scores, dtype=torch.float32, device=q_dense_vecs.device) teacher_scores = teacher_scores.view(q_dense_vecs.size(0), -1).detach() # (batch_size, group_size) teacher_targets = F.softmax(teacher_scores, dim=-1) # (batch_size, group_size) else: diff --git a/tests/test_embedder_teacher_scores.py b/tests/test_embedder_teacher_scores.py new file mode 100644 index 000000000..0abc04a6b --- /dev/null +++ b/tests/test_embedder_teacher_scores.py @@ -0,0 +1,117 @@ +"""Regression tests for integer teacher scores in embedder distillation.""" + +from copy import deepcopy + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from transformers import PreTrainedTokenizerFast, XLMRobertaConfig, XLMRobertaModel + +from FlagEmbedding.finetune.embedder.encoder_only.base.modeling import ( + BiEncoderOnlyEmbedderModel, +) +from FlagEmbedding.finetune.embedder.encoder_only.m3.modeling import ( + EncoderOnlyEmbedderM3Model, +) + + +@pytest.fixture(params=["base", "m3_dense", "m3_unified"]) +def embedder(request): + torch.manual_seed(42) + config = XLMRobertaConfig( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + max_position_embeddings=16, + hidden_dropout_prob=0.0, + attention_probs_dropout_prob=0.0, + ) + encoder = XLMRobertaModel(config) + if request.param == "base": + return BiEncoderOnlyEmbedderModel(encoder) + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer( + WordLevel( + { + "": 0, + "": 1, + "": 2, + "": 3, + }, + unk_token="", + ) + ), + cls_token="", + pad_token="", + eos_token="", + unk_token="", + ) + return EncoderOnlyEmbedderM3Model( + base_model={ + "model": encoder, + "sparse_linear": torch.nn.Linear(config.hidden_size, 1), + "colbert_linear": torch.nn.Linear(config.hidden_size, config.hidden_size), + }, + tokenizer=tokenizer, + unified_finetuning=request.param == "m3_unified", + ) + + +@pytest.mark.parametrize("kd_loss_type", ["kl_div", "m3_kd_loss"]) +@pytest.mark.parametrize("teacher_scores", [[2, 0, -1, 3], [2, 0.5, -1, 3]]) +def test_integer_teacher_scores_match_float_loss_and_gradients( + embedder, kd_loss_type, teacher_scores +): + embedder.kd_loss_type = kd_loss_type + float_embedder = deepcopy(embedder) + queries = { + "input_ids": torch.tensor([[0, 4, 5, 2, 1], [0, 6, 7, 8, 2]]), + "attention_mask": torch.tensor([[1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]), + } + passages = { + "input_ids": torch.tensor( + [ + [0, 4, 9, 2, 1], + [0, 10, 11, 2, 1], + [0, 6, 12, 13, 2], + [0, 14, 15, 2, 1], + ] + ), + "attention_mask": torch.tensor( + [ + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 0], + [1, 1, 1, 1, 1], + [1, 1, 1, 1, 0], + ] + ), + } + + expected_loss = float_embedder( + queries=queries, + passages=passages, + teacher_scores=[float(score) for score in teacher_scores], + ).loss + loss = embedder( + queries=queries, + passages=passages, + teacher_scores=teacher_scores, + ).loss + assert torch.isfinite(loss) + torch.testing.assert_close(loss, expected_loss) + + loss.backward() + expected_loss.backward() + for (name, param), (expected_name, expected_param) in zip( + embedder.named_parameters(), float_embedder.named_parameters() + ): + assert name == expected_name + if expected_param.grad is None: + assert param.grad is None + else: + assert torch.isfinite(param.grad).all() + torch.testing.assert_close(param.grad, expected_param.grad)