| """On-device accuracy for the fully patched model under PyTorch Native (torch.compile(backend="neuron")). |
| |
| * Accuracy: scores of the patched model vs the stock Hugging Face model on CPU in fp32 (exact GELU), on |
| real-text query/passage pairs with ragged lengths. Gates: Pearson >= 0.9999, and top-1 agreement per |
| query no worse than stock PyTorch Native (same tanh GELU, no kernels) on the same inputs. |
| * Padding isolation: the score of sequence 0 is bit-identical when its batch neighbours change. |
| |
| Run inside the PyTorch Native container on trn2, once per LNC: |
| NEURON_LOGICAL_NC_CONFIG=1 python -m pytest -m neuron -s tests/test_device.py |
| Results are also appended as JSON lines to $GTE_RESULTS (default tests/device_results.jsonl). |
| """ |
| import json |
| import os |
| import random |
|
|
| import numpy as np |
| import pytest |
|
|
| from conftest import load_package |
|
|
| 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-multilingual-reranker-base") |
| 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)] |
| ATTN_LENGTHS = tuple(int(v) for v in os.environ.get("GTE_ATTN_LENGTHS", "128,256,384,512,768,1024").split(",")) |
|
|
| TOPICS = { |
| "What are the benefits of renewable energy?": [ |
| "Solar and wind power generate electricity without burning fuel, so they emit no greenhouse gases while running.", |
| "Renewable projects reduce a country's dependence on imported oil and gas and keep energy spending local.", |
| "Battery storage lets grids keep using solar power after sunset, smoothing out supply.", |
| ], |
| "How do I reset my home wifi router?": [ |
| "Hold the small reset button on the back of the router for about ten seconds until the lights blink.", |
| "After a factory reset, log in to the admin page with the default password printed on the label.", |
| "Restarting the modem first and then the router often fixes a dropped internet connection.", |
| ], |
| "What are the symptoms of vitamin D deficiency?": [ |
| "Low vitamin D levels can cause tiredness, bone pain, muscle weakness and frequent infections.", |
| "People who get little sunlight are more likely to be short of vitamin D, especially in winter.", |
| "A simple blood test measures vitamin D, and supplements usually correct a deficiency within months.", |
| ], |
| "What is the best way to learn Python programming?": [ |
| "Start with small scripts that automate something you do every day, then read other people's code.", |
| "The official Python tutorial covers variables, loops, functions and classes with short examples.", |
| "Practice problems and a small personal project teach more than reading about syntax.", |
| ], |
| "How did the Roman Empire fall?": [ |
| "The Western Roman Empire weakened through civil wars, economic trouble and pressure on its borders.", |
| "In 476 the last western emperor was deposed, an event often used to mark the end of ancient Rome.", |
| "The eastern half of the empire survived for almost a thousand more years as the Byzantine Empire.", |
| ], |
| "How do vaccines protect against disease?": [ |
| "A vaccine trains the immune system to recognise a virus or bacterium without causing the illness.", |
| "After vaccination, memory cells remain ready to produce antibodies quickly if the real germ appears.", |
| "Some vaccines need booster doses because immune memory fades over time.", |
| ], |
| "How can I sleep better at night?": [ |
| "Going to bed and waking up at the same time every day helps set your internal clock.", |
| "Avoid caffeine in the afternoon and keep screens out of the bedroom in the last hour before sleep.", |
| "A cool, dark and quiet room makes it easier to fall asleep and stay asleep.", |
| ], |
| "What causes inflation?": [ |
| "Inflation rises when demand for goods and services grows faster than the economy can produce them.", |
| "Higher energy and wage costs can push companies to raise prices across the economy.", |
| "Central banks raise interest rates to slow borrowing and bring inflation back down.", |
| ], |
| } |
|
|
|
|
| def _pairs(L, n_per_query): |
| """Each query gets n_per_query passages: half on-topic, half from other topics, with varied lengths |
| (one to many sentences, so real lengths spread from short to the full L).""" |
| rng = random.Random(L) |
| qs = list(TOPICS) |
| Q, Ps = [], [] |
| for qi, q in enumerate(qs): |
| for j in range(n_per_query): |
| pool = TOPICS[q] if j % 2 == 0 else TOPICS[qs[(qi + 1 + j) % len(qs)]] |
| n_sent = rng.randint(1, max(1, L // 12)) |
| Ps.append(" ".join(rng.choice(pool) for _ in range(n_sent))) |
| Q.append(q) |
| return Q, Ps |
|
|
|
|
| def _load(dtype): |
| m = transformers.AutoModelForSequenceClassification.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 |
|
|
|
|
| def _rank_top1(s, ref, n_q, per): |
| return float(np.mean([np.argmax(s[i * per:(i + 1) * per]) == np.argmax(ref[i * per:(i + 1) * per]) |
| for i in range(n_q)])) |
|
|
|
|
| @pytest.fixture(scope="module") |
| def setup(): |
| import torch._dynamo as dyn |
| dyn.config.cache_size_limit = 64 |
| pkg = load_package() |
| |
| to = __import__("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 = _load(None) |
| dev = torch.device("neuron") |
| stock = _load(torch.bfloat16).to(dev) |
| patched = _load(torch.bfloat16) |
| n = pkg.patch_model(patched, attn_lengths=ATTN_LENGTHS) |
| patched = patched.to(dev) |
| return dict(pkg=pkg, tok=tok, cpu=cpu, dev=dev, stock=stock, patched=patched, n_layers=n, calls=calls) |
|
|
|
|
| def _run(model, ids, msk, bs, dev): |
| cm = torch.compile(model, backend="neuron", dynamic=False) |
| out = [] |
| with torch.no_grad(): |
| for i in range(0, ids.shape[0], bs): |
| out.append(cm(ids[i:i + bs].to(dev), msk[i:i + bs].to(dev)).logits[:, 0].float().cpu()) |
| return torch.cat(out).numpy(), cm |
|
|
|
|
| @pytest.mark.parametrize("L,bs", SHAPES) |
| def test_accuracy_and_isolation(setup, L, bs): |
| tok, cpu, dev = setup["tok"], setup["cpu"], setup["dev"] |
| per = 8 |
| Q, Ps = _pairs(L, per) |
| n = len(Q) - len(Q) % bs if len(Q) % bs else len(Q) |
| e = tok(Q[:n], Ps[:n], return_tensors="pt", max_length=L, padding="max_length", truncation=True) |
| ids, msk = e["input_ids"], e["attention_mask"] |
| n_q = n // per |
| with torch.no_grad(): |
| ref = torch.cat([cpu(ids[i:i + 16], msk[i:i + 16]).logits[:, 0] for i in range(0, n, 16)]).float().numpy() |
| s_stock, _ = _run(setup["stock"], ids, msk, bs, dev) |
| setup["calls"].clear() |
| s_nki, cm = _run(setup["patched"], ids, msk, bs, dev) |
| traced = dict(setup["calls"]) |
|
|
| pr_stock = float(np.corrcoef(s_stock, ref)[0, 1]) |
| pr_nki = float(np.corrcoef(s_nki, ref)[0, 1]) |
| top1_stock = _rank_top1(s_stock, ref, n_q, per) |
| top1_nki = _rank_top1(s_nki, ref, n_q, per) |
|
|
| |
| b0 = ids[:bs].clone(), msk[:bs].clone() |
| 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(b0[0].to(dev), b0[1].to(dev)).logits[0].float().cpu() |
| y1 = cm(b1[0].to(dev), b1[1].to(dev)).logits[0].float().cpu() |
| isolated = bool(torch.equal(y0, y1)) |
|
|
| lens = msk.sum(-1).numpy() |
| rec = dict(model=MODEL, lnc=setup["pkg"].LNC, L=L, bs=bs, pairs=n, real_len=[int(lens.min()), int(lens.max())], |
| layers_patched=setup["n_layers"], pearson_stock=pr_stock, pearson_nki=pr_nki, |
| top1_stock=top1_stock, top1_nki=top1_nki, |
| max_abs_err_nki=float(np.abs(s_nki - ref).max()), max_abs_err_stock=float(np.abs(s_stock - ref).max()), |
| padding_isolated=isolated, traced_ops=traced) |
| 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 and traced.get("ln", 0) >= 1, rec |
| assert pr_nki >= 0.9999, rec |
| assert top1_nki >= top1_stock, rec |
| assert isolated, rec |
|
|
