Skip to content

Fix distillation with integer teacher scores - #1613

Open
zhechen-cs wants to merge 1 commit into
FlagOpen:masterfrom
zhechen-cs:fix/integer-teacher-scores
Open

zhechen-cs wants to merge 1 commit into
FlagOpen:masterfrom
zhechen-cs:fix/integer-teacher-scores

Conversation

@zhechen-cs

Copy link
Copy Markdown

Integer pos_scores and neg_scores are accepted by the training dataset, but a batch containing only integer scores becomes a Long tensor in the embedder forward pass. Softmax then raises before the distillation loss can be computed.

Construct teacher scores as float32 in both the shared embedder forward and BGE-M3's separate forward. Add CPU regressions using a small XLM-R encoder for the base embedder and M3's dense and unified modes, covering both distillation losses and comparing losses and parameter gradients with equivalent floating-point scores.

Fixes #1471.

Validation with PyTorch 2.14.1 and Transformers 5.19.0:

  • On unmodified fd1a2bd, the new test file had 6 failures at softmax for integer scores and 6 passing controls.
  • python -m pytest tests/test_embedder_teacher_scores.py tests/test_finetune_trainer_compat.py -q: 16 passed.
  • Ruff error checks on the changed files, formatting check on the new test, and git diff --check passed.

The tests use randomly initialized weights on CPU. No GPU or cross-device training was run.

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.

m3模型训练报错源码需要修改,不然会报"host_softmax" not implemented for 'Long'

1 participant