feat: add C2LLM embedding model (codefuse-ai/C2LLM-7B) - #77
Open
StartupBros wants to merge 1 commit into
Open
Conversation
A Qwen2 decoder backbone with a PMA (Pooling by Multihead Attention) head
in place of mean/last-token pooling: a single learned seed vector attends
over the token states, so the pooled vector is a learned query over the
sequence rather than a positional pick.
Five details in the reference are load-bearing and none is the obvious
choice, so each is called out in the module docstring:
* the attention scale is sqrt(dim_V) -- the full compressed width (3584),
not the per-head width (112)
* the softmax is re-masked afterwards (A = A * pad_mask), so rows over
padded input deliberately do not sum to 1
* the residual adds the PROJECTED seed, not the raw seed. The model uses
PMA_v2/MAB_POST_v2; the PMA/MAB_POST classes directly above them in the
same file add the raw seed instead
* the output feed-forward is ReLU, unlike the SiLU backbone
* pma_norm is false, so the reference returns an unnormalized vector;
text_embeds is normalized for the library contract and pooler_output
keeps the raw vector
The backbone is Qwen2, not Qwen3: q/k/v carry biases and there are no
per-head q_norm/k_norm tensors.
Verified against the real checkpoint on an M3 Ultra. The upstream reference
cannot be imported for an end-to-end comparison (modeling_c2llm.py declares
deepspeed and peft imports that transformers' check_imports enforces), so
each half was compared separately, which also localizes any failure:
PMA head -> a verbatim transcription of MAB_POST_v2: cosine 1.0 on both
rows, including a left-padded one
backbone -> transformers' own Qwen2Model loaded from the same tensors:
41 unpadded positions, min cosine 1.0, and 0 missing /
0 unexpected tensors on load
Parameter count is 7,122,031,616, which is the checkpoint's 7,667,028,992
minus the 152064x3584 lm_head that sanitize drops.
Serving-verified: loaded under oMLX 0.6.3 on an M3 Ultra and returned a
3584-dimensional vector over /v1/embeddings.
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.
Adds support for
codefuse-ai/C2LLM-7B(model_type: c2llm), currently tied for #1 on the MTEB Code leaderboard.What it is
A Qwen2 decoder backbone with a PMA (Pooling by Multihead Attention) head instead of mean/last-token pooling. A single learned seed vector attends over the token states, so the pooled vector is a learned query over the sequence rather than a positional pick.
The backbone is Qwen2, not Qwen3 — q/k/v carry biases and there are no per-head
q_norm/k_normtensors.Parity details worth reviewing
Five things in the reference are load-bearing and none is the obvious choice. Each is documented in the module docstring:
sqrt(dim_V)— the full compressed width (3584), not the per-head width (112)A = A * pad_mask), so rows over padded input deliberately do not sum to 1PMA_v2/MAB_POST_v2; thePMA/MAB_POSTclasses directly above them in the same upstream file add the raw seed insteadpma_normis false, so the reference returns an unnormalized vector.text_embedsis normalized to match the library contract;pooler_outputkeeps the raw vectorVerification
Against the real checkpoint on an M3 Ultra. The upstream reference cannot be imported for an end-to-end comparison —
modeling_c2llm.pydeclaresdeepspeedandpeftimports that transformers'check_importsenforces — so each half was compared separately, which also localizes any failure:MAB_POST_v2Qwen2Model, same tensorsParameter count is 7,122,031,616 — the checkpoint's 7,667,028,992 minus the 152064x3584
lm_headthatsanitizedrops.Also serving-verified end to end: loaded under oMLX 0.6.3 and returned a 3584-dimensional vector over
/v1/embeddings.