Repository navigation
Conversation
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Per-query negative groups are currently interpreted as globally concatenated passage groups in two places. With
no_in_batch_neg_flag=True, M3 distillation receives(batch_size, group_size)scores but starts each row's targets atrow * group_size. Batch size 2 and group size 3 raisesIndexError: Target 3 is out of bounds.The unified M3 ensemble has the same layout problem. Its three score matrices are already local, but they are gathered again. The passage-vector reduction introduced in #1424 makes the inferred group size one: an otherwise valid batch of 2 queries with 3 passages each gets an ensemble contrastive loss of zero, discarding the ensemble negatives. When batch size exceeds group size, it still indexes beyond the score columns. The official
m3_same_dataset.shincludes datasets with these local groups.Keep already-local score groups intact, use zero-based target columns for local M3 distillation, and retain the full passage group when combining the three M3 score branches. Globally concatenated in-batch/cross-device layouts keep their existing target offsets. This follows the conditional likelihood: each query's denominator contains its own remaining passages, and selected target columns are removed within that same group.
Validation: 47 native CPU checks pass; the original implementation fails 23 checks. Tests compare both losses and score gradients against an independent per-query conditional-likelihood calculation in float32/float64. Native tiny BERT encoders exercise dense and unified M3 training, all three heads' gradients and SGD updates, supervised/KL/M3 losses, and self-distillation. Six native in-batch objective controls pass on both versions. The complete FlagEmbedding package imports normally with Transformers 4.57.1 and PyTorch; no source extraction, module placeholders or model mocks are used.
The tests use a local synthetic vocabulary and small randomly initialized models. Pretrained checkpoints, GPU/DDP training and retrieval-quality evaluation were not run. This change does not address the distributed KL loss scaling discussed in #1604.
Related: #1410 and its merged repair #1424; the new regressions reproduce the remaining local-score layout errors on current master.