| """On-device accuracy for NewBERT embedding models (gte-*-en-v1.5, gte-multilingual-base): CLS embeddings of the |
| patched model vs the stock Hugging Face model on CPU in fp32 (exact GELU), and vs stock PyTorch Native. |
| |
| Gates: mean per-text cosine >= min(0.9999, stock bf16 mean) (patched vs CPU fp32; stock itself is below 0.9999 at L=1024); min per-text cosine within COS_MIN_SLACK of stock |
| bf16 PyTorch Native's min (an absolute 0.9999 floor on the minimum is not met by stock bf16 itself); |
| retrieval top-1 per query equal to the CPU fp32 ranking at least as often as stock Native; padding isolation |
| bit-exact; fused ops actually traced. |
| |
| GTE_TEST_MODEL=Alibaba-NLP/gte-large-en-v1.5 NEURON_LOGICAL_NC_CONFIG=2 python -m pytest -m neuron -s tests/test_device_embed.py |
| """ |
| import json |
| import os |
| import sys |
|
|
| import pytest |
|
|
| from conftest import load_package |
| from test_device import TOPICS, _pairs |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("torch_neuronx") |
| transformers = pytest.importorskip("transformers") |
| pytestmark = pytest.mark.neuron |
|
|
| MODEL = os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-large-en-v1.5") |
| RESULTS = os.environ.get("GTE_RESULTS", os.path.join(os.path.dirname(os.path.abspath(__file__)), |
| "device_results.jsonl")) |
| SHAPES = [(256, 32), (128, 32), (512, 3), (1024, 8), (768, 3)] |
| COS_MEAN_MIN = 0.9999 |
| ATTN_LENGTHS = tuple(int(v) for v in os.environ.get("GTE_ATTN_LENGTHS", "128,256,384,512,768,1024").split(",")) |
| COS_MIN_SLACK = 1e-3 |
|
|
|
|
| def _load(dtype): |
| m = transformers.AutoModel.from_pretrained(MODEL, trust_remote_code=True, attn_implementation="eager").eval() |
| if dtype is not None: |
| for mod in m.modules(): |
| for cn, c in mod.named_children(): |
| if "GELU" in type(c).__name__: |
| setattr(mod, cn, torch.nn.GELU(approximate="tanh")) |
| m = m.to(dtype) |
| return m |
|
|
|
|
| class _CLS(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m |
|
|
| def forward(self, ids, msk): |
| return self.m(input_ids=ids, attention_mask=msk).last_hidden_state[:, 0] |
|
|
|
|
| @pytest.fixture(scope="module") |
| def setup(): |
| import torch._dynamo as dyn |
| dyn.config.cache_size_limit = 64 |
| pkg = load_package() |
| to = sys.modules["gte_encoder_neuron_kernels.torch_ops"] |
| calls = {} |
| for name in ("attn_block", "gated_mlp_ln", "res_ln", "ln"): |
| def wrap(f, name=name): |
| def g(*a, **k): |
| calls[name] = calls.get(name, 0) + 1 |
| return f(*a, **k) |
| return g |
| setattr(to, name, wrap(getattr(to, name))) |
| tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) |
| cpu = _CLS(_load(None)) |
| dev = torch.device("neuron") |
| stock = _CLS(_load(torch.bfloat16)).to(dev) |
| p = _load(torch.bfloat16) |
| n = pkg.patch_model(p, attn_lengths=ATTN_LENGTHS) |
| patched = _CLS(p).to(dev) |
| return dict(pkg=pkg, tok=tok, cpu=cpu, dev=dev, stock=stock, patched=patched, n_layers=n, calls=calls) |
|
|
|
|
| def _emb(model, ids, msk, bs, dev): |
| cm = torch.compile(model, backend="neuron", dynamic=False) |
| with torch.no_grad(): |
| out = [cm(ids[i:i + bs].to(dev), msk[i:i + bs].to(dev)).float().cpu() for i in range(0, ids.shape[0], bs)] |
| return torch.cat(out), cm |
|
|
|
|
| def _cos(a, b): |
| return torch.nn.functional.cosine_similarity(a, b, dim=-1) |
|
|
|
|
| @pytest.mark.parametrize("L,bs", SHAPES) |
| def test_embed_accuracy_and_isolation(setup, L, bs): |
| tok, cpu, dev = setup["tok"], setup["cpu"], setup["dev"] |
| per = 8 |
| Q, Ps = _pairs(L, per) |
| queries = list(TOPICS) |
| texts = queries + Ps |
| n = len(texts) - len(texts) % bs |
| texts = texts[:n] |
| e = tok(texts, return_tensors="pt", max_length=L, padding="max_length", truncation=True) |
| ids, msk = e["input_ids"], e["attention_mask"] |
| with torch.no_grad(): |
| ref = torch.cat([cpu(ids[i:i + 8], msk[i:i + 8]).float() for i in range(0, n, 8)]) |
| e_stock, _ = _emb(setup["stock"], ids, msk, bs, dev) |
| setup["calls"].clear() |
| e_nki, cm = _emb(setup["patched"], ids, msk, bs, dev) |
| traced = dict(setup["calls"]) |
|
|
| cos_nki = _cos(e_nki, ref) |
| cos_stock = _cos(e_stock, ref) |
|
|
| def top1(E): |
| En = torch.nn.functional.normalize(E, dim=-1) |
| q = En[:len(queries)] |
| p = En[len(queries):] |
| return (q @ p.T).argmax(-1) |
|
|
| t_ref, t_stock, t_nki = top1(ref), top1(e_stock), top1(e_nki) |
| agree_stock = float((t_stock == t_ref).float().mean()) |
| agree_nki = float((t_nki == t_ref).float().mean()) |
|
|
| b1 = ids[:bs].clone(), msk[:bs].clone() |
| if bs > 1: |
| perm = torch.arange(n - 1, n - bs, -1) |
| b1[0][1:], b1[1][1:] = ids[perm], msk[perm] |
| with torch.no_grad(): |
| y0 = cm(ids[:bs].to(dev), msk[:bs].to(dev))[0].float().cpu() |
| y1 = cm(b1[0].to(dev), b1[1].to(dev))[0].float().cpu() |
| isolated = bool(torch.equal(y0, y1)) |
| lens = msk.sum(-1).numpy() |
| worst = torch.argsort(cos_nki)[:4].tolist() |
| rec_worst = [dict(i=int(i), len=int(lens[i]), cos_nki=float(cos_nki[i]), cos_stock=float(cos_stock[i])) |
| for i in worst] |
| rec = dict(model=MODEL, kind="embed", lnc=setup["pkg"].LNC, L=L, bs=bs, texts=n, |
| real_len=[int(lens.min()), int(lens.max())], layers_patched=setup["n_layers"], |
| cos_min_stock=float(cos_stock.min()), cos_mean_stock=float(cos_stock.mean()), |
| cos_min_nki=float(cos_nki.min()), cos_mean_nki=float(cos_nki.mean()), |
| top1_agree_stock=agree_stock, top1_agree_nki=agree_nki, padding_isolated=isolated, |
| traced_ops=traced, worst=rec_worst) |
| print(json.dumps(rec), flush=True) |
| with open(RESULTS, "a") as f: |
| f.write(json.dumps(rec) + "\n") |
| nl = setup["n_layers"] |
| if L in ATTN_LENGTHS: |
| assert traced.get("attn_block", 0) >= nl, rec |
| assert traced.get("gated_mlp_ln", 0) >= nl, rec |
| assert float(cos_nki.mean()) >= min(COS_MEAN_MIN, float(cos_stock.mean())), rec |
| assert float(cos_nki.min()) >= float(cos_stock.min()) - COS_MIN_SLACK, rec |
| assert agree_nki >= agree_stock, rec |
| assert isolated, rec |
|
|
