Skip to content

Keep per-query negative groups intact in M3 training - #1609

Open
Afloat16 wants to merge 1 commit into
FlagOpen:masterfrom
Afloat16:fix/local-m3-distillation-layout
Open

Afloat16 wants to merge 1 commit into
FlagOpen:masterfrom
Afloat16:fix/local-m3-distillation-layout

Conversation

@Afloat16

@Afloat16 Afloat16 commented Oct 2, 2026

Copy link
Copy Markdown

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 at row * group_size. Batch size 2 and group size 3 raises IndexError: 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.sh includes 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.

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.

1 participant