| """Device timing: stock PyTorch Native vs patched, one model on one logical NeuronCore (single process). |
| |
| GTE_TEST_MODEL=Alibaba-NLP/gte-large-en-v1.5 NEURON_LOGICAL_NC_CONFIG=2 python tests/bench_model.py 256x32,128x32,512x8 |
| Prints and appends JSON lines (ms per batch median of N_ITERS, seq/s on one logical core) to $GTE_RESULTS. |
| Both variants use the tanh GELU so the comparison isolates the kernels. Hardware numbers only. |
| """ |
| import json |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch._dynamo as dyn |
| import torch_neuronx |
| import transformers |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from conftest import load_package |
|
|
| 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")) |
| N_ITERS = int(os.environ.get("N_ITERS", "50")) |
| SHAPES = [tuple(int(v) for v in s.split("x")) for s in (sys.argv[1] if len(sys.argv) > 1 else "256x32").split(",")] |
| VARIANTS = (sys.argv[2] if len(sys.argv) > 2 else "stock,nki").split(",") |
| dyn.config.cache_size_limit = 64 |
| pkg = load_package() |
| dev = torch.device("neuron") |
|
|
|
|
| AUTO = transformers.AutoModelForSequenceClassification if "reranker" in MODEL else transformers.AutoModel |
|
|
|
|
| def build(v): |
| m = AUTO.from_pretrained(MODEL, trust_remote_code=True, attn_implementation="eager").eval() |
| 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(torch.bfloat16) |
| if v == "nki": |
| pkg.patch_model(m) |
| elif v == "nki_cte": |
| pkg.patch_model(m, attn_lengths=(128, 256, 384, 512)) |
| return m.to(dev) |
|
|
|
|
| tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) |
| for v in VARIANTS: |
| model = build(v) |
| for L, bs in SHAPES: |
| cm = torch.compile(model, backend="neuron", dynamic=False) |
| text = " ".join(["solar wind storage grid battery policy"] * L) |
| e = tok([text] * bs, return_tensors="pt", max_length=L, padding="max_length", truncation=True) |
| ids, msk = e["input_ids"].to(dev), e["attention_mask"].to(dev) |
| with torch.no_grad(): |
| cm(input_ids=ids, attention_mask=msk) |
| torch_neuronx.synchronize() |
| for _ in range(5): |
| cm(input_ids=ids, attention_mask=msk) |
| torch_neuronx.synchronize() |
| bar = os.environ.get("BARRIER_DIR") |
| if bar: |
| n_w = int(os.environ["BARRIER_N"]) |
| open(os.path.join(bar, f"ready_{os.environ.get('NEURON_RT_VISIBLE_CORES', os.getpid())}_{L}x{bs}"), "w").close() |
| while len([f for f in os.listdir(bar) if f.endswith(f"_{L}x{bs}")]) < n_w: |
| time.sleep(0.05) |
| t_start = time.time() |
| lat = [] |
| for _ in range(N_ITERS): |
| s = time.perf_counter() |
| cm(input_ids=ids, attention_mask=msk) |
| torch_neuronx.synchronize() |
| lat.append((time.perf_counter() - s) * 1e3) |
| wall = time.time() - t_start |
| ms = float(np.median(lat)) |
| rec = dict(model=MODEL, kind="bench", variant=v, lnc=pkg.LNC, L=L, bs=bs, ms_median=round(ms, 3), |
| ms_p90=round(float(np.percentile(lat, 90)), 3), seq_per_s_one_core=round(bs * 1e3 / ms, 1), |
| iters=N_ITERS, |
| t_start=round(t_start, 2), wall_s=round(wall, 2)) |
| print(json.dumps(rec), flush=True) |
| with open(RESULTS, "a") as f: |
| f.write(json.dumps(rec) + "\n") |
| del model |
|
|
