| """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 |
|
|
