jinaai

/

xlm-roberta-flash-implementation

like 34

Follow

Jina AI 2.03k

Hello!

Heads up, this PR is agent-generated, but human reviewed.

Pull Request overview

• Recompute RotaryEmbedding.inv_freq when the cos/sin cache is (re)built, so it survives meta-device loading on transformers>=5.

Details

inv_freq is registered as a non-persistent buffer computed in __init__, so it is not part of any checkpoint. transformers>=5 builds the model on the meta device and materialises non-persistent buffers without re-running __init__, so the buffer ends up holding whatever happened to be in memory:

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("jinaai/jina-colbert-v2", trust_remote_code=True)
rot = model.encoder.layers[0].mixer.rotary_emb

expected = 1.0 / (10000.0 ** (torch.arange(0, rot.dim, 2, dtype=torch.float32) / rot.dim))
print(rot.inv_freq.flatten()[:5].tolist())
print((rot.inv_freq.float().cpu() - expected).abs().max().item())

# transformers 4.43.2 -> [1.0, 0.7498942, 0.5623413, 0.4216965, 0.3162278],  max|diff| 0.0
# transformers 5.15.0 -> [5.45325056e+08, 1.335e-42, 0.0, 0.0, 0.0],         max|diff| 5.45e+08

The value differs on every process, so results are not reproducible between runs and occasionally come out as NaN. Within a single process it is perfectly deterministic, which makes it easy to mistake for a tokenization or pooling problem.

Two things hide it:

• The existing self-heal branch is guarded on dtype (if self.inv_freq.dtype != torch.float32: ... further down in _update_cos_sin_cache). The buffer is materialised as fp32 regardless of the load dtype, so that branch never fires under fp32, bf16, or "auto".
• The module-level _init_weights added in "Fixup post init (for v5 remote compatibility)" (#60) only covers nn.Linear and nn.Embedding, so it does not reach this buffer.

Impact

Positional information is effectively destroyed. Measured on jinaai/jina-colbert-v2 with MultiVectorNanoBEIREvaluator (full NanoBEIR, 13 datasets, default bfloat16 load):

So the model card's claim of outperforming ColBERTv2 is currently not reproducible on transformers v5, and becomes reproducible again with this change. A paired A/B on three subsets, same script and datasets, with inv_freq as the only difference:

I only measured this on jina-colbert-v2, but the code path is shared by every model that points its auto_map at this repository.

The change

_compute_inv_freq is already called here when the rotary base changes. This PR simply calls it whenever the cache is (re)built, rather than trusting the value from __init__:

             self._seq_len_cached = seqlen
-            if rotary_base_changed:
-                self.inv_freq = self._compute_inv_freq(device=device)
+            self.inv_freq = self._compute_inv_freq(device=device)

The rotary_base_changed parameter is still needed and still used, since it is one of the conditions that forces the cache rebuild in the enclosing if.

This is a no-op on older transformers, where __init__ already produced the same value: it is just recomputed instead of read back. The cost is a 32 element arange plus pow, and only when the cache is rebuilt, so it is not on the hot path.

Verification

With this patch applied and no other changes, jina-colbert-v2 output matches a transformers==4.43.2 reference (same token ids, fp32, CPU) to 2.4e-07 on the token embeddings, with identical MaxSim scores. Repeated runs are now bit-identical.

One related note

scale a few lines below is registered the same way (persistent=False, computed in __init__) and would have the same problem when XPos is enabled. No published config sets rotary_emb_scale_base, so it is None in practice and I left it out to keep this diff minimal. Happy to include it if you would like.

Please let me know if you have any questions or feedback!

• Tom Aarsen