jinaai

/

xlm-roberta-flash-implementation

like 34

Follow

Jina AI 1.96k

Running any model that uses this implementation on CPU (where flash-attn is
unavailable, so get_use_flash_attn() returns False and the native path is
taken) produces a different last_hidden_state on every fresh process, and a
substantial fraction of fresh loads return an output that is entirely NaN.
The model weights are byte-identical across loads, so this is not random weight
init — it is an uninitialized-memory read from a torch.empty(...) buffer in
the forward path.

Environment

• transformers 5.3.0, torch 2.9.0 (CPU), numpy 2.4.6
• model: jinaai/jina-colbert-v2 (and the base jinaai/xlm-roberta-flash-implementation), revision 845308d
• flash-attn not installed (CPU) → native attention path
• Also reproduces via pylate 1.5.0 (jina-colbert-v2) with the same signature.

Minimal reproduction (bare transformers, no pylate)

import hashlib, subprocess, sys
def worker():
    import numpy as np, torch
    from transformers import AutoModel, AutoTokenizer
    name = "jinaai/jina-colbert-v2"  # or any model on this implementation
    tok = AutoTokenizer.from_pretrained(name)
    m = AutoModel.from_pretrained(name, trust_remote_code=True).eval()
    enc = tok("post-market surveillance requirements", return_tensors="pt")
    with torch.no_grad():
        hs = m(**enc).last_hidden_state
    a = hs.to(torch.float64).numpy().ravel()
    print("nan=%d/%d  hash=%s" % (int(np.isnan(a).sum()), a.size,
          hashlib.sha256(np.round(np.nan_to_num(a),6).tobytes()).hexdigest()[:12]))
if "--worker" in sys.argv: worker()
else:
    for _ in range(4):
        print(subprocess.run([sys.executable, __file__, "--worker"],
              capture_output=True, text=True).stdout.strip())

Each fresh process prints a different hash; some print nan=<all>. A correct
implementation would print the same hash with nan=0 every time.

Root cause (isolated empirically)

Fresh-process matrix on jina-colbert-v2 (CPU, encoder-only):

• Weights are identical every load — full state_dict hash is byte-stable →
  not random weight init.
• No RNG seed changes anything — seeding torch, numpy, and Python
  random all still scatter → the source is not RNG.
• torch.use_deterministic_algorithms(True) "fixes" the scatter, but only via
  its fill_uninitialized_memory side effect — setting
  torch.utils.deterministic.fill_uninitialized_memory = False while leaving
  deterministic mode on brings the scatter back. And because the fill value is
  NaN, the "determinized" output is itself all-NaN.
• Monkeypatching torch.empty to any fixed fill collapses the 3-way scatter to
  a single value → confirms the output depends on the contents of a
  torch.empty buffer that is read before being fully written.

So an uninitialized torch.empty(...) buffer in the native (non-flash) forward
is read while still holding heap garbage: its bytes vary per process (scatter),
are reused within a process (byte-stable within a run), are immune to seeding,
and are NaN-filled under deterministic mode. Candidate sites that allocate with
torch.empty and fill conditionally/partially: embedding.py (the
adapter_mask branch), the pooler in modeling_xlm_roberta.py, and the
attention/MLP output buffers on the native path.

Impact

On CPU, every fresh model load yields a different embedding for the same input,
and a large fraction of loads yield all-NaN embeddings. For retrieval
(jina-colbert-v2 / PLAID) this means results silently reshuffle across process
restarts, or the model contributes only NaN scores — irreproducible and
unreliable in production.

Suggested fix

Initialize the offending buffer(s) fully — e.g. torch.zeros(...) (or ensure the
subsequent writes cover every element) instead of torch.empty(...) on the
native path. Happy to help pinpoint the exact allocation and open a PR if useful.
Is there a reason the native path uses torch.empty here rather than a
zero-initialized buffer?