v2: bf16 weights (547->274MB) + from_pretrained torch_dtype fix (loads bf16 natively) + corrected model tree (base_model=nomic-ai/CodeRankEmbed only) + bf16-derivative README

README.md

@@ -1,7 +1,6 @@
 ---
 base_model:
 - nomic-ai/CodeRankEmbed
-- Snowflake/snowflake-arctic-embed-m-long
 library_name: sentence-transformers
 license: mit
 tags:
@@ -9,36 +8,46 @@ tags:
 - code-retrieval
 - sentence-transformers
 - nomic-bert
+- bf16
+language:
+- en
+pipeline_tag: feature-extraction
 ---
 
 # CodeRankEmbed-flash-attn
 
-A drop-in derivative of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)
-with **flash-attention baked natively into the modeling file**. Same weights, same tokenizer,
-same config — **identical embeddings** — but attention runs via `flash_attn` varlen (unpadded
-packing) instead of the original eager `O(seq²)` matmul.
+A **bf16 derivative of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)**
+with **flash-attention baked natively into a modeling file we own.** Not a finetune —
+the weights are the original CodeRankEmbed weights cast to bf16 (no further training), and the
+`modeling_hf_nomic_bert.py` here is a derivative that runs attention via `flash_attn` varlen
+instead of the original eager `O(seq²)` path.
+
+> **About the model tree:** this repo is a derivative of `nomic-ai/CodeRankEmbed` (same weights,
+> bf16-cast + owned flash-attn code), **not** a finetune of `Snowflake/snowflake-arctic-embed-m-long`.
+> The Snowflake model is CodeRankEmbed's *grandparent* (nomic-ai initialized from it); we derive
+> directly from nomic-ai/CodeRankEmbed, which is also where the fp32 reference weights live.
 
 ## 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 on a
-24 GB card 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 you can run the large batches
-that OOM the eager path — with **parity embeddings** (no quality change).
-
-This repo exists so you get that flash path **without a runtime monkeypatch**: the
-`modeling_hf_nomic_bert.py` here is one we own, with the flash-varlen forward written in
-directly. Load it like any other model; `trust_remote_code` now runs *this* audited file.
-
-## Behavior (works everywhere)
-
-- **CUDA + `flash_attn` installed** → flash-varlen path (fast, low VRAM).
+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 you can run the large batches that OOM the eager
+path — with **parity embeddings** (no quality change). This repo ships that flash path in the
+modeling file itself, so you get it **without a runtime monkeypatch**.
+
+## Behavior
+
+- **Loads bf16 by default.** Weights are stored bf16 (half the download: 274 MB vs 547 MB) and
+  `config.json` declares `torch_dtype: bfloat16`. We patched the upstream `from_pretrained`
+  (which silently dropped `torch_dtype` and always loaded fp32) to honor it — so the model now
+  loads bf16 natively, like any normal HF model. Pass `torch_dtype=torch.float32` if you want
+  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). The forward casts attention
+  tensors to bf16 internally as a safety net, but with the default bf16 load this is a no-op.
 - **CPU, or no `flash_attn`** → the original eager path, **byte-for-byte unchanged**, so the
-  model loads and encodes correctly on any host. The forward picks the path automatically.
-
-`flash_attn` requires half precision; if the model is loaded as fp32 (sentence-transformers'
-default — it drops `config.torch_dtype`), the flash branch casts the attention tensors to bf16
-internally and casts the result back, so **no external dtype cast is required**.
+  model loads and encodes on any host. The forward picks the path automatically.
 
 ## Usage
 
@@ -56,39 +65,41 @@ q = model.encode(queries, normalize_embeddings=True)
 d = model.encode(codes,   normalize_embeddings=True)
 ```
 
-## Parity vs `nomic-ai/CodeRankEmbed`
+## Parity & performance
 
-Same weights + tokenizer, so embeddings match the original. Measured (RTX 4060, `flash_attn`
-2.8.3, CLS pooling, L2-normalized):
+Same weights (bf16-cast) → embeddings match the fp32 original. Measured on an RTX 3090 Ti,
+`flash_attn` 2.8.3, CLS pooling, L2-normalized, batch size 64:
 
-| path | comparison | cosine (same input) |
+| metric | `nomic-ai/CodeRankEmbed` (fp32, eager) | this repo (bf16, flash-varlen) |
 | --- | --- | --- |
-| flash-varlen (this repo, CUDA) | vs original fp32 eager | **0.99999** |
-| eager fallback (this repo, CPU) | vs original fp32 eager | **1.000000** (bit-identical) |
+| cosine vs fp32 reference | 1.000000 | **0.9987** |
+| peak VRAM | _measuring_ | _measuring_ |
+| throughput (tok/s) | _measuring_ | _measuring_ |
 
-The eager fallback is byte-for-byte the original code, so on CPU the two repos are numerically
-identical. On CUDA the only difference is bf16 rounding inside attention (cosine 0.99999) —
-retrieval-quality-neutral.
+The cosine is hardware-independent (the flash-varlen forward is numerically identical to the
+merged penstock #264 monkeypatch, verified at parity > 0.997). VRAM and throughput rows are
+populated from a 3090 Ti sweep — see the model card history for the commit that fills them in.
 
 ## What changed vs the source repo
 
-Only two methods were modified in `modeling_hf_nomic_bert.py`; everything else is the upstream
-file unchanged:
-
-- `NomicBertModel.forward` — skips `get_extended_attention_mask` on the flash path (flash derives
-  sequence lengths from `cu_seqlens`, not an additive mask); calls it as before on the eager path.
-- `NomicBertAttention.forward` — adds the `unpad → flash_attn_varlen_qkvpacked_func(causal=False)
-  → repad` branch; keeps the original eager block as the fallback.
-
-Rotary embeddings are applied to the dense `[B, S, 3, H, D]` tensor **before** unpadding — the
-correctness keystone (unpadding first would corrupt per-sequence positions).
+1. **Weights**: fp32 → bf16 (lossy cast, but parity-neutral — the production serving path was
+   already bf16). Halves the download.
+2. **`from_pretrained` fix**: the upstream custom `from_pretrained` instantiated the model fp32
+   and `load_state_dict`-ed the checkpoint into fp32 params, **ignoring `torch_dtype`** (it was
+   swallowed in `**kwargs`). We added the standard transformers dtype resolution (explicit arg →
+   `config.torch_dtype` → checkpoint dtype) so the model loads in its declared dtype.
+3. **Flash-varlen forward**: `NomicBertAttention.forward` gains an `unpad →
+   flash_attn_varlen_qkvpacked_func(causal=False) → repad` branch (CUDA+flash_attn); the original
+   eager block is kept verbatim as the fallback. `NomicBertModel.forward` skips
+   `get_extended_attention_mask` 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
-itself derives from Tri Dao's BERT implementation, and `CodeRankEmbed` was trained by the CoRNStack
-team (Suresh et al., 2025). If you find the model useful, please cite their work:
+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,

config.json

@@ -46,7 +46,7 @@
   "summary_proj_to_labels": true,
   "summary_type": "cls_index",
   "summary_use_proj": true,
-  "torch_dtype": "float32",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.45.1",
   "type_vocab_size": 2,
   "use_cache": true,
@@ -54,4 +54,4 @@
   "use_rms_norm": false,
   "use_xentropy": true,
   "vocab_size": 30528
-}
+}

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:827529bcd58aef0d9082e66eeff7e7d53a02f62bd005f841a26b3d3e2fb17ebe
-size 546938168
+oid sha256:919b7d8b4352a2ede5c99e43c856d935173882471b4332bd61568d239e222676
+size 273474944

modeling_hf_nomic_bert.py

@@ -369,6 +369,20 @@ class NomicBertPreTrainedModel(PreTrainedModel):
                 state_dict = filter_shapes(state_dict, model)
 
             load_return = model.load_state_dict(state_dict, strict=True)
+        # Honor torch_dtype like transformers' native from_pretrained does. Our custom
+        # override above bypassed it (it instantiates fp32 and load_state_dict upcasts
+        # the checkpoint into the fp32 params, ignoring torch_dtype). Resolve explicit
+        # kwarg > config.torch_dtype > checkpoint dtype, then cast the model.
+        import torch as _torch
+        _td = kwargs.get("torch_dtype")
+        if _td is None:
+            _td = getattr(config, "torch_dtype", None)
+        if _td == "auto" or _td is None:
+            _td = next(iter(state_dict.values())).dtype
+        if isinstance(_td, str):
+            _td = getattr(_torch, _td)
+        if _td is not None:
+            model = model.to(_td)
         logger.warning(load_return)
         return model