File size: 3,657 Bytes
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 87 88 89 90 91 92 93 94 95 96 97 98 99 | """End-to-end CPU check of the torch glue (patch_model, buffer layouts, mask clamp, fallbacks) with the NKI
kernels executed by nki.simulate. torch_neuronx is stubbed; the model is truncated to 2 layers to keep the
simulator time down. Simulator only: NOT device evidence.
Needs torch + transformers + nki + ml_dtypes in one venv.
"""
import os
import sys
import types
import numpy as np
import pytest
torch = pytest.importorskip("torch")
transformers = pytest.importorskip("transformers")
nki = pytest.importorskip("nki")
ml_dtypes = pytest.importorskip("ml_dtypes")
pytestmark = [pytest.mark.sim, pytest.mark.cpu]
MODEL = os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-multilingual-reranker-base")
class _Sim:
"""kernel[lnc](*torch_args) -> nki.simulate on NumPy copies, result back as a torch tensor."""
def __init__(self, k):
self.k = k
def __getitem__(self, lnc):
def run(*args):
np_args = []
for a in args:
if isinstance(a, torch.Tensor):
a = a.detach().contiguous()
if a.dtype == torch.bfloat16:
a = a.view(torch.int16).numpy().view(ml_dtypes.bfloat16)
else:
a = a.numpy()
np_args.append(a)
y = np.asarray(nki.simulate(self.k[lnc])(*np_args))
if y.dtype == ml_dtypes.bfloat16:
return torch.from_numpy(y.view(np.int16).copy()).view(torch.bfloat16)
return torch.from_numpy(y.copy())
return run
@pytest.fixture(scope="module")
def pkg():
if "torch_neuronx" not in sys.modules or not hasattr(sys.modules["torch_neuronx"], "synchronize"):
stub = types.ModuleType("torch_neuronx")
stub.nki_op = lambda name, mutates_args=None: (lambda f: f)
sys.modules["torch_neuronx"] = stub
else:
pytest.skip("real torch_neuronx present; use test_device.py")
from conftest import load_package
p = load_package()
to = sys.modules["gte_encoder_neuron_kernels.torch_ops"]
for n in ("res_ln_kernel", "gated_mlp_ln_kernel", "attn_block_kernel"):
setattr(to, n, _Sim(getattr(to, n)))
return p
def _model(n_layers):
m = transformers.AutoModelForSequenceClassification.from_pretrained(
MODEL, trust_remote_code=True, attn_implementation="eager").eval()
m.new.encoder.layer = m.new.encoder.layer[:n_layers]
return m
@pytest.mark.parametrize("L,lens", [(128, [128, 37])])
def test_patched_model_matches_fp32(pkg, L, lens):
tok = transformers.AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
ref_m = _model(2)
torch.manual_seed(0)
B = len(lens)
ids = torch.randint(1000, 30000, (B, L))
msk = torch.zeros(B, L, dtype=torch.long)
for i, n in enumerate(lens):
msk[i, :n] = 1
ids[:, 0] = tok.cls_token_id
with torch.no_grad():
ref = ref_m(ids, msk).logits.float()
hs_ref = ref_m.new(ids, msk).last_hidden_state.float()
m = _model(2)
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)
assert pkg.patch_model(m) == 2
with torch.no_grad():
hs = m.new(ids, msk).last_hidden_state.float()
out = m(ids, msk).logits.float()
valid = msk.bool()
rel = (hs - hs_ref)[valid].norm() / hs_ref[valid].norm()
print(f"hidden rel_l2 {rel:.2e}, logits {out.flatten().tolist()} vs fp32 {ref.flatten().tolist()}")
assert rel < 2e-2
assert (out - ref).abs().max() < 0.1
|