File size: 3,917 Bytes
25e8424 fb58c6c 25e8424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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 # noqa: E402
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": # v1.0 behaviour: fused attention block only up to 512, nkilib cte above
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) # compile + warm-up
torch_neuronx.synchronize()
for _ in range(5):
cm(input_ids=ids, attention_mask=msk)
torch_neuronx.synchronize()
bar = os.environ.get("BARRIER_DIR") # multi-process DP: start timing together
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
|