handwoven8588/CodeRankEmbed-flash-attn

CodeRankEmbed-flash-attn

A bf16 quantization of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed) with flash-attention built into a custom `modeling_hf_nomic_bert.py` shipped in this repo. It is not a finetune — the weights are the original CodeRankEmbed weights cast to bf16 (no further training) — and attention runs via flash_attn varlen instead of the original eager O(seq²) path.

Why

nomic-ai/CodeRankEmbed loads through trust_remote_code, and its attention path is eager only — activation memory grows as batch × heads × seq², which OOMs at large batches even though the model is only 137M params. flash_attn's varlen path computes the same attention in O(N) memory by packing unpadded sequences, so the large batches that OOM the eager path run comfortably — with parity embeddings (no quality change). This repo ships that flash path in the modeling file itself, so no runtime patching or post-load hooks are needed.

Behavior

• —Loads bf16 by default. flash_attn requires half precision and the model runs bf16 in any real serving setup, so the weights are stored bf16 and config.json declares torch_dtype: bfloat16. The upstream custom from_pretrained silently dropped torch_dtype and always loaded fp32; the copy in this repo honors it, so the model loads bf16 natively, like any normal HF model. Pass torch_dtype=torch.float32 to load fp32 (note: the stored weights are bf16-precision, so this only widens the dtype, not the precision).
• —CUDA + `flash_attn` → flash-varlen path (fast, low VRAM). Attention tensors are cast to bf16 internally as a safety net; with the default bf16 load this is a no-op.
• —CPU, or no `flash_attn` → the original eager attention algorithm runs unchanged. Because the model is bf16, eager runs in bf16 here — numerically equivalent to the flash path, just without its memory and throughput wins. The model loads and encodes on any host; the forward selects the path automatically (_FLASH_AVAILABLE and hidden_states.is_cuda).

Usage

Identical to the original. The query prompt must include the task-instruction prefix "Represent this query for searching relevant code: "; documents need no prefix.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("handwoven8588/CodeRankEmbed-flash-attn", trust_remote_code=True)
queries = ["Represent this query for searching relevant code: Calculate the n-th factorial"]
codes   = ["def fact(n):\n    if n < 0:\n        raise ValueError\n    return 1 if n == 0 else n * fact(n - 1)"]

q = model.encode(queries, normalize_embeddings=True)
d = model.encode(codes,   normalize_embeddings=True)

Parity & performance

The weights are the original CodeRankEmbed weights (bf16-cast), so embeddings match the fp32 original to within bf16 precision. Measured on an RTX 3090 Ti, flash_attn 2.8.3, CLS pooling, L2-normalized, batch size 64:

Same 512-snippet corpus (~236k tokens, 20–900 tokens each), batch size 64. The flash-varlen path also scales to far larger batches (bs=256 fits in ~7.6 GB) where eager O(seq²) blows up. Parity (>0.997) reproduced on both an RTX 4060 and the 3090 Ti.

What changed vs the source repo

1. 1.Weights: fp32 → bf16. flash_attn only accepts half precision and the model runs bf16 in any real serving configuration, so the weights are stored bf16 and (via the load fix below) arrive bf16 — which is simply how this model is used, and removes the need for a post-load dtype cast. Parity-neutral (cosine 0.9986 vs the fp32 original); the smaller download is incidental, not the reason.
2. 2.`from_pretrained` dtype fix: the upstream custom from_pretrained instantiated the model fp32 and load_state_dict-ed the checkpoint into fp32 params, ignoring `torch_dtype`. The copy here adds the standard transformers dtype resolution (explicit arg → config.torch_dtype → checkpoint dtype) so the model loads in its declared dtype.
3. 3.Flash-varlen forward: NomicBertAttention.forward gains an unpad → flash_attn_varlen_qkvpacked_func(causal=False) → repad branch (CUDA + flashattn); the original eager block is kept as the fallback. `NomicBertModel.forward` skips `getextendedattentionmask on the flash path. Rotary embeddings are applied to the dense [B, S, 3, H, D]` tensor before unpadding — the correctness keystone.

License & attribution

MIT — same license as nomic-ai/CodeRankEmbed (see NOTICE). The weights, tokenizer, and the bulk of the modeling file are a verbatim derivative of nomic-ai/CodeRankEmbed; the modeling file derives from Tri Dao's BERT implementation, and CodeRankEmbed was trained by the CoRNStack team (Suresh et al., 2025). Cite their work:

@misc{suresh2025cornstackhighqualitycontrastivedata,
  title  = {CoRNStack: High-Quality Contrastive Data for Text and Code Retrieval},
  author = {Suresh, K N Q and Wang, Xiang and Khan, Saqib and others},
  year   = {2025},
}