README pass: base_model_relation=quantized, fix CPU-eager wording (bf16, not bit-identical), drop internal refs + first-person plural

README.md

@@ -1,6 +1,6 @@
 ---
-base_model:
-- nomic-ai/CodeRankEmbed
+base_model: nomic-ai/CodeRankEmbed
+base_model_relation: quantized
 library_name: sentence-transformers
 license: mit
 tags:
@@ -16,38 +16,33 @@ pipeline_tag: feature-extraction
 
 # CodeRankEmbed-flash-attn
 
-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.
+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 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**.
+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.** 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 on any host. The forward picks the path automatically.
+  `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
 
@@ -67,8 +62,9 @@ d = model.encode(codes,   normalize_embeddings=True)
 
 ## Parity & performance
 
-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:
+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:
 
 | metric | `nomic-ai/CodeRankEmbed` (fp32, eager) | this repo (bf16, flash-varlen) |
 | --- | --- | --- |
@@ -77,29 +73,26 @@ Same weights (bf16-cast) → embeddings match the fp32 original. Measured on an
 | throughput (bs=64) | 52,000 tok/s | **162,000 tok/s** (≈3.1× faster) |
 
 Same 512-snippet corpus (~236k tokens, 20–900 tokens each), batch size 64. The flash-varlen path
-additionally scales to far larger batches (bs=256 fits in ~7.6 GB) where eager O(seq²) blows up.
-The cosine is hardware-independent — the flash-varlen forward is numerically identical to the
-merged penstock #264 monkeypatch (same `unpad → flash_attn_varlen_qkvpacked_func → repad` path),
-verified at parity > 0.997 on both a 4060 and the 3090 Ti above.
+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. **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.
+1. **Weights**: fp32 → bf16 (lossy cast, parity-neutral). Halves the download.
+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. **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
+   flash_attn_varlen_qkvpacked_func(causal=False) → repad` branch (CUDA + flash_attn); the original
+   eager block is kept 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
+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: