Kernels:

---

jburtoft

/

gte-encoder-neuron-kernels

Like 0

"""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"))
# (L, batch): batch 32 is the main throughput shape; batch 3 exercises the LNC2 query split.
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()
    # count trace-time calls of each fused op, so a silent fallback to the stock modules fails the test
    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                                       # passages per query -> 64 pairs
    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)

    # padding isolation: keep sequence 0, replace every neighbour with other pairs of different lengths
    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