Add BAAI/bge-m3 sparse (lexical_weights) support via --pooling m3_sparse#899
Open
SpongeRobert wants to merge 5 commits into
Open
Add BAAI/bge-m3 sparse (lexical_weights) support via --pooling m3_sparse#899SpongeRobert wants to merge 5 commits into
SpongeRobert wants to merge 5 commits into
Conversation
added 5 commits
July 16, 2026 21:06
Reference output for the bge-m3 sparse head, verified bit-exact against FlagEmbedding BGEM3FlagModel on three texts including Russian. Consumed by the insta snapshot tests for --pooling m3_sparse. Includes the generator script and the fp16 sparse_linear.pt caveat.
bge-m3 sparse is a standalone `sparse_linear` head shipped as
`sparse_linear.pt` next to the weights, applied to the same hidden states the
(already correct) dense bge-m3 path produces. Unlike SPLADE it is indexed by
the *input* token ids rather than by vocabulary argmax, so the pooled vector is
sparse by construction.
The head scatters `relu(sparse_linear(h))` into vocabulary space keyed by input
token id, deduplicating by max and dropping {cls, eos, pad, unk}, which the
router's existing `sparsify` then turns into the same {index, value} pairs
SPLADE produces — so `/embed_sparse` and its wire format are untouched.
Notes:
- `sparse_linear.pt` stores fp16 weights; the VarBuilder dtype casts them to
the hidden-state dtype, without which the matmul dtype-mismatches.
- Special token ids are resolved from the tokenizer files rather than
hardcoded, mirroring FlagEmbedding's `_process_token_weights`.
- `Pool` gained `Copy`; the variant is spelled `m3_sparse` on the CLI (matching
the model card and `Display`) with a `m3-sparse` alias for consistency with
the other variants.
- Backends and models that cannot serve it reject it at load with a message
instead of reaching an `unreachable!()`. The flash path is not ported yet.
Runs the three golden texts (including Russian) through the candle CPU path and asserts the index set matches exactly (Jaccard 1.0) with values within 1e-3. Values are not bit-exact by design: the golden was captured against exact GeLU while candle's bge-m3 uses the GeLU + tanh approximation, which lands ~6e-4 away. That is the real candle/FlagEmbedding delta, not a defect. The batch is built locally rather than via `common::batch`: XLM-RoBERTa position ids start at `pad_token_id + 1` and the shared helper starts at 0, which is only self-consistent with the snapshots it produced. Against a real FlagEmbedding reference the offset matters — without it values drift by ~0.13 while the index set still matches. Also covers that a model without a sparse head is rejected, not panicked on.
The SEN-03 `Pool::M3Sparse` head only served on the CPU `bert.rs` path; production bge-m3 runs on the GPU/flash backbone, so port the head there. - Move the head arithmetic (`token_weights`) and the scatter-max (`scatter_sequence`) onto `BertM3Head` so the padded and flash paths share one implementation and cannot drift; `bert.rs::forward` now calls both. - flash_bert: replace the `bail!` stub with a real `M3Sparse` branch. Flash packs the batch into one varlen sequence, so scatter each request's token slice (delimited by `cumulative_seq_lengths`) into vocab space and return the dense `[n, vocab_size]` tensor the router's `sparsify` expects — wire format and router untouched. `load` rejects non-XLM-R; `load_roberta` loads `sparse_linear.pt` via the threaded `model_path`. Position ids carry XLM-R's offset from the real tokenization pipeline, so the flash path needs no special handling. GPU parity smoke (A100, flash fp16) vs the FlagEmbedding golden incl. Russian: index Jaccard 1.0, max|Δ| = 6.4e-4 (threshold 1e-3). Dense `--pooling cls` regression unchanged (1024-dim, head matches golden).
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
1 participant
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.
Summary
Adds BAAI/bge-m3
sparse(lexical_weights) support to the candle backend via a new pooling mode--pooling m3_sparse, addressing the core ask of #141. Dense bge-m3 already works in TEI (--pooling cls); this applies the model's trainedsparse_linearhead over the same hidden states, so the existing/embed_sparseendpoint returns bge-m3 sparse vectors with no wire-format change.Closes the sparse half of #141. (Follows up my design comment above.)
Motivation
bge-m3 sparse is a separately trained
sparse_linearlayer (Linear(1024→1), shipped assparse_linear.pt), not an MLM head — so--pooling spladecan't produce it (load_robertahas nolm_head.decoder). Today the README/model support reflects "bge-m3, no sparse". This PR closes that gap while reusing the existing SPLADE plumbing.Design
Pool::M3Sparse(backends/core/src/lib.rs), CLI--pooling m3_sparse(kebab aliasm3-sparsefor CLI consistency).BertM3Head(backends/candle/src/models/bert.rs): loadssparse_linear.ptviaVarBuilder::from_pth(candle pickle; fp16→f32), computesrelu(sparse_linear(hidden)).squeeze(-1), then scatter-max over the input token ids into a densevocab_sizevector, dropping special tokens and non-positive weights. The router's existingsparsifyclosure turns that into{index, value}pairs — routers/types untouched. Unlike SPLADE, indices are the input token ids (not vocab argmax), soinput_idsare threaded into pooling.bert.rs) and CUDA/flash (flash_bert.rs) share onetoken_weights+scatter_sequenceimplementation, so the two paths can't diverge. Flash usescumulative_seq_lengthsto scatter per-sequence.core/src/download.rsfetchessparse_linear.ptwhenpool == M3Sparse.ort,python) rejectm3_sparsethe same way they rejectSplade.Correctness / parity
Verified bit-exact against FlagEmbedding (
BGEM3FlagModel) on a golden fixture including a Russian sentence:The residual ~1e-4 is expected: the golden reference uses transformers' exact GeLU, whereas TEI logs
hidden_act=gelu → GeLU+tanh approximation— index sets are identical. Test + fixture added underbackends/candle/tests/(runs incargo test).Points I'd like maintainer input on
These are the questions from my design comment, now concrete in code:
--pooling m3_sparserather than auto-detecting fromsparse_linear.ptpresence. Happy to switch to auto-detect if you prefer.is_splade()→ renamedis_sparse()(gates/embed_sparse,embed_all,normalize— semantics shared by both sparse poolings).download_artifacts(...)gains apool: Option<Pool>arg (to fetchsparse_linear.pt).Poolnow derivesCopy.I can shape these differently to minimize churn if you'd rather.
Notes
v1.9.3release tag; glad to rebase ontomainon request.m3_sparsepooling option.Disclosure: implementation and this PR were prepared with AI assistance; the parity results above were run and verified by me.