--- base_model: nomic-ai/CodeRankEmbed base_model_relation: quantized library_name: sentence-transformers license: mit tags: - flash-attention - code-retrieval - sentence-transformers - nomic-bert - bf16 language: - en --- # CodeRankEmbed-flash-attn A **bf16 quantization of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)** with a custom `modeling_hf_nomic_bert.py`, shipped in this repo, that **picks one of three attention implementations for the hardware it runs on.** It is not a finetune — the weights are the original CodeRankEmbed weights cast to bf16 (no further training). Two of the three replace the original's padded `O(seq²)` attention with an `O(N)` unpadded path; the third keeps the original's algorithm as the correctness reference and universal fallback. ## Why `nomic-ai/CodeRankEmbed` loads through `trust_remote_code`, and its attention path is **eager only** — activation memory grows as `batch × heads × seq²`, which runs out of memory at large batches even though the model is only 137M params. This repo adds two attention paths that compute the same attention in `O(N)` memory by packing unpadded sequences, so the large batches that the eager path cannot fit run comfortably — with **parity embeddings** (no quality change): - `torch_varlen` — `torch.nn.attention.varlen.varlen_attn`, shipped in torch itself (no extra package), available from torch **2.10.0** onward. - `flash_attn` — the `flash_attn` package's varlen kernel. Kept as a fallback for older torch builds that don't yet have `torch.nn.attention.varlen` but do have `flash_attn` installed. This package is **optional**. Both run the same FA2-family kernel and are gated by the same GPU-capability check. The modeling file ships all three paths itself, so no runtime patching or post-load hooks are needed. ## Behavior - **Three attention tiers, chosen automatically per device** (override with `NOMIC_BERT_ATTN_IMPL=torch_varlen|flash_attn|eager`): 1. **`torch_varlen`** — CUDA, compute capability sm_80+ (Ampere or newer), torch **≥ 2.10.0**. No third-party kernel needed. 2. **`flash_attn`** — CUDA, compute capability sm_80+, the `flash_attn` package importable, on a torch that doesn't yet ship `torch.nn.attention.varlen` (typically an older torch). This dependency is **optional**. 3. **`eager`** — everything else: CPU, pre-Ampere GPUs, or neither of the above available. The original's padded attention algorithm, unchanged, runs on any host. `auto` (the default) prefers `torch_varlen`, then `flash_attn`, then `eager`. Before accepting a varlen tier, `auto` runs one tiny kernel probe per device (a capability check alone can't tell whether the installed build has a kernel for the GPU's architecture, e.g. on ROCm); if the probe raises a `RuntimeError` (how a missing or unsupported kernel fails), it logs one WARNING naming the tier and the error and steps down to the next tier. An out-of-memory error, any other exception type, or a CUDA error left over from earlier work propagates instead of demoting the tier, so a transient failure never pins the device to a slower tier. A forced override (e.g. `NOMIC_BERT_ATTN_IMPL=torch_varlen`) is not probed and **raises** `RuntimeError` if that tier's precondition doesn't hold — a forced tier never falls back silently. An unrecognized override raises `ValueError`. - **See which tier engaged**: `model[0].auto_model.attention_impl` after a forward pass. The modeling file also logs the tier, and why, once at INFO. - **Padding-free on the varlen tiers.** The model unpads the batch once, right after the input ids, runs the embeddings and every encoder layer on the real tokens only, and pads back once at the end. `last_hidden_state` keeps its `[B, S, 768]` shape, with zeros at padding positions, so pooling is unchanged. - **Flattened batches (opt-in).** Loading with `model_kwargs={"attn_implementation": "flash_attention_2"}` lets sentence-transformers send each batch flattened (`[1, ΣL]` token ids with sequence boundaries, no padding at all), and the model returns one `[1, ΣL, 768]` output that sentence-transformers pools per sequence. The opt-in needs the `flash_attn` package installed, because transformers checks for it at load and raises otherwise, even though the attention itself still runs on whichever tier the model picks. On the `eager` tier the flattened batch is re-padded and runs the padded path. - **Loads bf16 by default.** `flash_attn` and `torch_varlen` both require 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 original's 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). - **`eager` runs in bf16 too** (because the stored weights are bf16), numerically equivalent to the varlen tiers, just without their memory and throughput wins. The model loads and encodes on any host regardless of which tier engages. ## 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. ```python 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) ``` With `flash_attn` installed, opt in to flattened batches: ```python model = SentenceTransformer( "handwoven8588/CodeRankEmbed-flash-attn", trust_remote_code=True, model_kwargs={"attn_implementation": "flash_attention_2"}, ) ``` On a varlen tier the encoder layers run on packed tokens only, so they compile cleanly. For about 1.2× the throughput and a third less peak memory (see Speed and memory below): ```python import torch model[0].auto_model.encoder = torch.compile(model[0].auto_model.encoder, dynamic=True) ``` `dynamic=True` matters: the number of packed tokens changes with every batch. ## Parity & performance Every number below compares this repo with the original model, [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed) (revision `3c4b608`), run as published: fp32 weights, its own modeling file, padded eager attention. Both run on the same GPU, an RTX 3090 Ti (sm_86). The original's modeling file calls `get_extended_attention_mask`, which transformers has deprecated and no longer ships by 5.17, so it was measured on transformers 5.11; this repo does not call it and loads on 5.17 too. **Corpus.** The first 16,384 `document` fields (real Python functions) of [`lightonai/cornstack`](https://huggingface.co/datasets/lightonai/cornstack)'s Python split, in shard order (`train-00000-of-00682.parquet`, then `train-00001-…`): 2,684,878 tokens, 7 to 4,746 per function, 164 mean, 92 median. Each row is one `encode()` call over all of them at the stated `batch_size=`, after one untimed warm-up of the same call. Cosine similarity is per function, on fp32-renormalized output, against the original's embedding of the same function. Peak memory is the encode's peak CUDA allocation above the loaded weights. **Load** is how the model was loaded: `default` (as in Usage) or `flattened` (the opt-in above). `torch 2.12.1+cu130`, `transformers 5.11.0`, `sentence-transformers 6.1.0`, `flash-attn 2.8.3.post1`. ### Agreement with the original model | model | attention | compiled | load | batch size | min cosine | mean cosine | | --- | --- | --- | --- | --- | --- | --- | | `nomic-ai/CodeRankEmbed` | `eager` | – | default | 4 | 1 (the reference) | 1 (the reference) | | `handwoven8588/CodeRankEmbed-flash-attn` | `eager` | – | default | 4 | 0.99902 | 0.99991 | | `handwoven8588/CodeRankEmbed-flash-attn` | `eager` | – | flattened | 4 | 0.99916 | 0.99991 | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 256 | 0.99896 | 0.99993 | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 256 | 0.99896 | 0.99993 | | `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | – | default | 256 | 0.99841 | 0.99993 | | `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | – | flattened | 256 | 0.99841 | 0.99993 | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 256 | 0.99960 | 0.99995 | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 256 | 0.99960 | 0.99995 | | `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | ✓ | default | 256 | 0.99946 | 0.99995 | | `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | ✓ | flattened | 256 | 0.99946 | 0.99995 | The original is the reference every other row is measured against, so its cosine is 1 by definition, not a measurement. The mean cosine is 0.99991 or higher in every configuration; the minimum is the one function, of 16,384, that moves most under bf16 weights and arithmetic. The batch size does not change these figures (a varlen tier gives the same five decimals at batch 32 and 256), so the largest batch that ran in every load is shown; `eager` runs at batch 4 (see below). Compiled rows sit closer to the original than uncompiled ones. ### Speed and memory | model | attention | compiled | load | batch size | peak memory | time | | --- | --- | --- | --- | --- | --- | --- | | `nomic-ai/CodeRankEmbed` | `eager` | – | default | 4 | 8.4 GiB | 68.9 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 32 | 2.5 GiB | 16.3 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 256 | 10.9 GiB | 15.1 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 1,024 | out of memory | out of memory | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 2,048 | out of memory | out of memory | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 32 | 1.4 GiB | 15.5 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 256 | 6.9 GiB | 13.6 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 1,024 | 17.0 GiB | 13.4 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 2,048 | out of memory | out of memory | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 32 | 1.6 GiB | 14.0 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 256 | 7.1 GiB | 12.8 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 1,024 | 16.3 GiB | 13.8 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 2,048 | out of memory | out of memory | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 32 | 0.9 GiB | 13.2 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 256 | 4.5 GiB | 11.4 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 1,024 | 11.1 GiB | 10.9 s | | `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 2,048 | 16.6 GiB | 10.8 s | This repo's rows are `torch_varlen`; `flash_attn` runs within 1% of it in both time and memory at every cell. The original runs at batch 4: its padded fp32 attention needs `batch × heads × seq²` memory, 8.4 GiB at batch 4 for this corpus's longest function. The same algorithm in bf16 (this repo's `eager` tier, batch 4) takes 41.2 s in 4.2 GiB; the rest of the gain comes from the unpadded path and the batches it allows. Compiled means `model[0].auto_model.encoder = torch.compile(model[0].auto_model.encoder, dynamic=True)`. - **Against the original**, the same 16,384 functions encode 5.1× faster uncompiled (13.4 s, flattened at batch 1,024) and 6.4× faster compiled (10.8 s, flattened at batch 2,048), against 68.9 s. At batch 32 this repo needs 2.5 GiB (default load) or 1.4 GiB (flattened), against the original's 8.4 GiB at batch 4. - **Memory follows the real tokens in the heaviest batch, not the number of functions.** This repo runs the encoder on real tokens only, and every out-of-memory cell failed allocating one activation over that batch's tokens: the 3,072-wide MLP hidden state, or, compiled, the 2,304-wide QKV projection. sentence-transformers length-sorts the corpus, so the default load's first batch of 1,024 holds the longest functions, about 870,000 tokens; flattened, it pairs the longest with the shortest, so its heaviest batch of 1,024 holds about 594,000. - **Flattening** cuts peak memory by 32–42% against the default load at the same batch size, and wall time by 5–21% (more at larger batches). - **Compiling** is 1.17–1.23× faster at the same batch size and cuts peak memory by 35%. It compiles one graph, with no graph breaks, once per process (the first encode takes about 5 s longer), and does not recompile as batch lengths change; the only further compiles come once a batch passes about 699,000 and 932,000 tokens, where inductor switches to 64-bit indexing for the 3,072-wide MLP activation and the 2,304-wide QKV projection. ## What changed vs the source repo 1. **Weights**: fp32 → bf16. `flash_attn` and `torch_varlen` only accept 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; the smaller download is incidental, not the reason. 2. **`from_pretrained` dtype fix**: the original's 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. **Three attention tiers**: `NomicBertAttention.forward` now selects one of three attention implementations at call time — `torch_varlen` (torch's own `torch.nn.attention.varlen.varlen_attn`, no third-party kernel, needs torch ≥ 2.10.0), `flash_attn` (the `flash_attn` package's varlen kernel, kept as a fallback for older torch builds that have it installed), and `eager` (the original's padded attention algorithm, numerically unchanged, and the default off CUDA sm_80+; it now converts a raw 2-D `[B, S]` mask to the additive form before adding it, for the case where a `device_map` split hands it a mask a varlen tier upstream never passes). On both varlen tiers `NomicBertModel.forward` unpads the batch once with torch-native `_unpad`/`_pad` helpers (a replacement for `flash_attn.bert_padding`), keeps the hidden states packed `[ΣL, 768]` through the embeddings and every layer, and pads back once at the end. Rotary embeddings rotate each packed token by its position within its own sequence, the same rotation it gets in the padded batch. On the eager tier `NomicBertModel.forward` builds the additive attention mask inline instead of calling the (now-removed-upstream) `get_extended_attention_mask` helper. Set `NOMIC_BERT_ATTN_IMPL=torch_varlen|flash_attn|eager` to force a tier (raises if it can't engage); `model[0].auto_model.attention_impl` and the one-time INFO log line report which tier engaged. 4. **`from_pretrained` forwards the Hub revision**: the original's custom `from_pretrained` fetched the weight file without the caller's `revision` (and cache/token options), so a pinned load still took the weights from the default branch. The copy here passes them through, so a pinned load fetches the weights from the pinned commit too. It also accepts transformers' `dtype=` alongside the older `torch_dtype=`, and returns the model in eval mode, as transformers' own `from_pretrained` does. 5. **Flattened input**: `NomicBertModel.forward` also accepts sentence-transformers' flattened batches (`input_ids [1, ΣL]` with per-sequence `position_ids` and `cu_seq_lens_q`), and the model declares flash-attention support so that sentence-transformers sends them when loaded with `attn_implementation="flash_attention_2"` (see Behavior). ## 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: ```bibtex @misc{suresh2025cornstackhighqualitycontrastivedata, title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking}, author={Tarun Suresh and Revanth Gangi Reddy and Yifei Xu and Zach Nussbaum and Andriy Mulyar and Brandon Duderstadt and Heng Ji}, year={2025}, eprint={2412.01007}, archivePrefix={arXiv}, primaryClass={cs.CL}, url={https://arxiv.org/abs/2412.01007}, } ```