Skip to content

feat: add C2LLM embedding model (codefuse-ai/C2LLM-7B) - #77

Open
StartupBros wants to merge 1 commit into
Blaizzy:mainfrom
StartupBros:add-c2llm-embedding
Open

feat: add C2LLM embedding model (codefuse-ai/C2LLM-7B)#77
StartupBros wants to merge 1 commit into
Blaizzy:mainfrom
StartupBros:add-c2llm-embedding

Conversation

@StartupBros

Copy link
Copy Markdown

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_norm tensors.

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:

  • 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 instantiates PMA_v2/MAB_POST_v2; the PMA/MAB_POST classes directly above them in the same upstream 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 to match the library contract; pooler_output keeps the raw vector

Verification

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:

half reference result
PMA head verbatim transcription of MAB_POST_v2 cosine 1.0 on both rows, incl. a left-padded one
backbone transformers' own Qwen2Model, same tensors min cosine 1.0 over 41 unpadded positions; 0 missing / 0 unexpected on load

Parameter count is 7,122,031,616 — the checkpoint's 7,667,028,992 minus the 152064x3584 lm_head that sanitize drops.

Also serving-verified end to end: loaded under oMLX 0.6.3 and returned a 3584-dimensional vector over /v1/embeddings.

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.
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