diff --git a/CMakeLists.txt b/CMakeLists.txt index aae2a6eabe..77d52b5fc9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -8,6 +8,27 @@ project(tensorrt_model_connect VERSION 0.1.0 LANGUAGES CXX) # generated build-tree RPATH entries relative to each ELF file. set(CMAKE_BUILD_RPATH_USE_ORIGIN TRUE) +# Standard CUDA installations do not always add their bin directory to PATH +# (notably bare-metal L4T images). Preserve an explicit CUDACXX selection, but +# otherwise probe the same CUDA roots used later for headers and libraries. +if(NOT CMAKE_CUDA_COMPILER AND NOT DEFINED ENV{CUDACXX}) + find_program(_trtmc_cuda_compiler + NAMES nvcc + HINTS + "$ENV{CUDA_HOME}" + "$ENV{CUDAToolkit_ROOT}" + "$ENV{CUDA_PATH}" + PATHS + /usr/local/cuda + /opt/cuda + PATH_SUFFIXES bin + ) + if(_trtmc_cuda_compiler) + set(CMAKE_CUDA_COMPILER "${_trtmc_cuda_compiler}" + CACHE FILEPATH "CUDA compiler" FORCE) + endif() +endif() + # Enable CUDA for runtime GPU kernels when a CUDA compiler is available. include(CheckLanguage) check_language(CUDA) diff --git a/tests/e2e/models/eagle_vlm/e2e_plugins/references/hf_transformers.py b/tests/e2e/models/eagle_vlm/e2e_plugins/references/hf_transformers.py index 2aedf3d64d..aec3831b62 100644 --- a/tests/e2e/models/eagle_vlm/e2e_plugins/references/hf_transformers.py +++ b/tests/e2e/models/eagle_vlm/e2e_plugins/references/hf_transformers.py @@ -754,6 +754,7 @@ def _run_reranking_ref( hf_id = case.hf_id model_ref = _resolve_cached_model_ref(hf_id) torch_dtype_expr = _torch_dtype_for_case(case) + reference_device = str(case.metadata.get("reference_device", "cuda") or "cuda") script = textwrap.dedent(f"""\ import json, torch @@ -770,7 +771,7 @@ def _run_reranking_ref( model = AutoModelForSequenceClassification.from_pretrained( model_ref, trust_remote_code=trust_remote_code, torch_dtype=torch_dtype) - device = "cuda" if torch.cuda.is_available() else "cpu" + device = {reference_device!r} model.to(device) model.eval() diff --git a/tests/e2e/models/eagle_vlm/test_eagle_vlm_hf_reference.py b/tests/e2e/models/eagle_vlm/test_eagle_vlm_hf_reference.py new file mode 100644 index 0000000000..9a57b082c3 --- /dev/null +++ b/tests/e2e/models/eagle_vlm/test_eagle_vlm_hf_reference.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Eagle-VLM Hugging Face reference tests.""" + +from __future__ import annotations + +import sys + +from tests.e2e.models.eagle_vlm.e2e_plugins.references import hf_transformers +from tests.e2e_harness.contracts import E2ECase, RunContext, StageSpec, StageOutput + + +def test_reranking_reference_uses_requested_device(monkeypatch, tmp_path) -> None: + captured: dict[str, object] = {} + + def fake_subprocess(**kwargs): + captured.update(kwargs) + return StageOutput(stage_name="full_inference", data={"scores": [1.0]}) + + monkeypatch.setattr(hf_transformers, "run_reference_subprocess", fake_subprocess) + monkeypatch.setattr( + hf_transformers, + "_resolve_cached_model_ref", + lambda model: model, + ) + case = E2ECase( + name="rerank-reference", + hf_id="org/reranker", + family="eagle_vlm", + runtime_strategy="eagle_vlm_reranking", + task_strategy="reranking", + bundle="rerank.bundle", + inputs={"prompt": "query", "documents": ["document"]}, + metadata={"reference_device": "cpu"}, + ) + context = RunContext(case=case, artifacts_dir=str(tmp_path)) + + hf_transformers.HfTransformersReference().run_stage( + case, + StageSpec("full_inference"), + context, + ) + + command = captured["command"] + assert command[:2] == [context.reference_python_path() or sys.executable, "-c"] + script = command[2] + assert "device = 'cpu'" in script + assert "torch.cuda.is_available()" not in script diff --git a/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/__init__.py b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/__init__.py new file mode 100644 index 0000000000..cc91a3223b --- /dev/null +++ b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/__init__.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Packed-attention compatibility for Lance reference environments. + +Lance imports FlashAttention's packed, dense, rotary, and padding helpers while +loading its official image-reference modules. Platforms without a qualified +FlashAttention package can explicitly select this PyTorch implementation. It is +never selected by CPU architecture inference. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import torch.nn.functional as functional +from torch.nn.attention.bias import causal_lower_right + + +def _sequence_offsets( + cumulative: torch.Tensor, + *, + total: int, + label: str, +) -> list[int]: + offsets = [int(value) for value in cumulative.detach().cpu().tolist()] + if len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != total: + raise ValueError( + f"{label} must start at zero and end at packed length {total}: " + f"{offsets}" + ) + if any(end < start for start, end in zip(offsets, offsets[1:])): + raise ValueError(f"{label} must be non-decreasing: {offsets}") + return offsets + + +def _attention_mask(query_length: int, key_length: int, causal: bool): + if not causal: + return None + return causal_lower_right(query_length, key_length) + + +def flash_attn_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + causal: bool = False, + window_size: Sequence[int] = (-1, -1), + softcap: float = 0.0, + alibi_slopes: torch.Tensor | None = None, + deterministic: bool = False, + return_attn_probs: bool = False, +) -> torch.Tensor: + """Evaluate FlashAttention's dense batch contract with PyTorch SDPA.""" + del deterministic + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("dense q, k, and v tensors must have rank 4") + if k.shape[:3] != v.shape[:3]: + raise ValueError("dense k and v tensors must have matching shapes") + if q.shape[0] != k.shape[0] or q.shape[-1] != k.shape[-1]: + raise ValueError("dense q and k batch and head dimensions must match") + if tuple(window_size) != (-1, -1): + raise NotImplementedError("L4T SDPA fallback does not support local windows") + if softcap or alibi_slopes is not None: + raise NotImplementedError( + "L4T SDPA fallback supports only Lance's dense attention contract" + ) + if return_attn_probs: + raise NotImplementedError("L4T SDPA fallback does not return attention probabilities") + + query = q.transpose(1, 2) + key = k.transpose(1, 2) + value = v.transpose(1, 2) + output = functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=_attention_mask(query.shape[-2], key.shape[-2], causal), + dropout_p=dropout_p, + scale=softmax_scale, + enable_gqa=query.shape[1] != key.shape[1], + ) + return output.transpose(1, 2) + + +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + causal: bool = False, + window_size: Sequence[int] = (-1, -1), + softcap: float = 0.0, + alibi_slopes: torch.Tensor | None = None, + deterministic: bool = False, + return_attn_probs: bool = False, + block_table: torch.Tensor | None = None, +) -> torch.Tensor: + """Evaluate FlashAttention's packed varlen contract with PyTorch SDPA.""" + del max_seqlen_q, max_seqlen_k, deterministic + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("packed q, k, and v tensors must have rank 3") + if k.shape[:2] != v.shape[:2]: + raise ValueError("packed k and v tensors must have matching token and head counts") + if q.shape[-1] != k.shape[-1]: + raise ValueError("packed q and k head dimensions must match") + if tuple(window_size) != (-1, -1): + raise NotImplementedError("L4T SDPA fallback does not support local windows") + if softcap or alibi_slopes is not None or block_table is not None: + raise NotImplementedError( + "L4T SDPA fallback supports only Lance's dense attention contract" + ) + if return_attn_probs: + raise NotImplementedError("L4T SDPA fallback does not return attention probabilities") + + query_offsets = _sequence_offsets( + cu_seqlens_q, + total=q.shape[0], + label="cu_seqlens_q", + ) + key_offsets = _sequence_offsets( + cu_seqlens_k, + total=k.shape[0], + label="cu_seqlens_k", + ) + if len(query_offsets) != len(key_offsets): + raise ValueError("packed q and k must contain the same number of sequences") + + outputs: list[torch.Tensor] = [] + for q_start, q_end, k_start, k_end in zip( + query_offsets, + query_offsets[1:], + key_offsets, + key_offsets[1:], + ): + query = q[q_start:q_end].transpose(0, 1).unsqueeze(0) + key = k[k_start:k_end].transpose(0, 1).unsqueeze(0) + value = v[k_start:k_end].transpose(0, 1).unsqueeze(0) + mask = _attention_mask(q_end - q_start, k_end - k_start, causal) + output = functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + dropout_p=dropout_p, + scale=softmax_scale, + enable_gqa=query.shape[1] != key.shape[1], + ) + outputs.append(output.squeeze(0).transpose(0, 1)) + return ( + torch.cat(outputs, dim=0) + if outputs + else q.new_empty((0, q.shape[1], v.shape[-1])) + ) + + +__all__ = ["flash_attn_func", "flash_attn_varlen_func"] diff --git a/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/bert_padding.py b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/bert_padding.py new file mode 100644 index 0000000000..0719d09d55 --- /dev/null +++ b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/bert_padding.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Padding helpers required by Transformers' FlashAttention integration.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as functional + + +def index_first_axis(values: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + """Select entries from the first axis without changing trailing dimensions.""" + return values.index_select(0, indices.to(device=values.device, dtype=torch.long)) + + +def unpad_input( + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Pack valid batch tokens and return FlashAttention-compatible indices.""" + if hidden_states.ndim < 2 or attention_mask.shape != hidden_states.shape[:2]: + raise ValueError("attention_mask must match hidden_states batch and sequence axes") + valid = attention_mask.to(dtype=torch.bool) + lengths = valid.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(valid.flatten(), as_tuple=False).flatten() + flattened = hidden_states.reshape(-1, *hidden_states.shape[2:]) + unpadded = index_first_axis(flattened, indices) + cumulative = functional.pad(torch.cumsum(lengths, dim=0, dtype=torch.int32), (1, 0)) + maximum = int(lengths.max().item()) if lengths.numel() else 0 + return unpadded, indices, cumulative, maximum + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.Tensor, + batch_size: int, + sequence_length: int, +) -> torch.Tensor: + """Restore packed token states to a zero-padded batch tensor.""" + output = hidden_states.new_zeros( + (batch_size * sequence_length, *hidden_states.shape[1:]) + ) + output.index_copy_( + 0, + indices.to(device=hidden_states.device, dtype=torch.long), + hidden_states, + ) + return output.reshape(batch_size, sequence_length, *hidden_states.shape[1:]) + + +__all__ = ["index_first_axis", "pad_input", "unpad_input"] diff --git a/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/__init__.py b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/__init__.py new file mode 100644 index 0000000000..752c31d3e4 --- /dev/null +++ b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FlashAttention layer helpers implemented for Lance's L4T reference.""" diff --git a/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/rotary.py b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/rotary.py new file mode 100644 index 0000000000..83215b847d --- /dev/null +++ b/tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat/flash_attn/layers/rotary.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Rotary embedding helper required by Lance's Qwen2.5-VL vision encoder.""" + +from __future__ import annotations + +import torch + + +def apply_rotary_emb( + values: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = False, + inplace: bool = False, + conjugate: bool = False, + seqlen_offsets: int = 0, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, +) -> torch.Tensor: + """Apply FlashAttention's rotary contract for dense Lance vision inputs.""" + del max_seqlen + if cu_seqlens is not None: + raise NotImplementedError("L4T rotary fallback does not support packed offsets") + if not isinstance(seqlen_offsets, int): + raise NotImplementedError("L4T rotary fallback requires one integer sequence offset") + if values.ndim not in (3, 4): + raise ValueError("rotary input must have rank 3 or 4") + + sequence_axis = 1 if values.ndim == 4 else 0 + sequence_length = values.shape[sequence_axis] + cos = cos[seqlen_offsets : seqlen_offsets + sequence_length] + sin = sin[seqlen_offsets : seqlen_offsets + sequence_length] + if conjugate: + sin = -sin + if values.ndim == 4: + cos = cos.unsqueeze(0).unsqueeze(2) + sin = sin.unsqueeze(0).unsqueeze(2) + else: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + if interleaved: + first = values[..., 0::2] + second = values[..., 1::2] + rotated = torch.stack( + (first * cos - second * sin, first * sin + second * cos), + dim=-1, + ).flatten(-2) + else: + first, second = values.chunk(2, dim=-1) + rotated = torch.cat( + (first * cos - second * sin, first * sin + second * cos), + dim=-1, + ) + if inplace: + values.copy_(rotated) + return values + return rotated + + +__all__ = ["apply_rotary_emb"] diff --git a/tests/e2e/models/lance/e2e_plugins/references/lance_official.py b/tests/e2e/models/lance/e2e_plugins/references/lance_official.py index fa5748a954..0c5003f48f 100644 --- a/tests/e2e/models/lance/e2e_plugins/references/lance_official.py +++ b/tests/e2e/models/lance/e2e_plugins/references/lance_official.py @@ -25,8 +25,24 @@ _IMAGE_MODEL_ALLOW_PATTERNS = [ f"{_MODEL_DIRECTORY}/**", f"{_VIT_DIRECTORY}/**", + "Wan2.2_VAE.pth", ] _IMAGE_REFERENCE_COMPAT = Path(__file__).with_name("lance_image_compat") +_IMAGE_REFERENCE_ATTENTION_COMPAT = Path(__file__).with_name( + "lance_image_attention_compat" +) +_ATTENTION_BACKEND_ENV = "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND" +_ATTENTION_BACKENDS = frozenset({"flash_attn", "torch_sdpa"}) + + +def _attention_backend(environment: dict[str, str]) -> str: + attention_backend = environment.get(_ATTENTION_BACKEND_ENV, "flash_attn").strip() + if attention_backend not in _ATTENTION_BACKENDS: + raise RuntimeError( + "unsupported Lance reference attention backend " + f"{attention_backend!r}; expected one of {sorted(_ATTENTION_BACKENDS)}" + ) + return attention_backend def _image_reference_environment( @@ -35,14 +51,82 @@ def _image_reference_environment( """Expose only the optional imports needed by upstream's image path.""" environment = dict(os.environ if base is None else base) existing = environment.get("PYTHONPATH", "").strip() + attention_backend = _attention_backend(environment) + attention_compat = ( + str(_IMAGE_REFERENCE_ATTENTION_COMPAT) + if attention_backend == "torch_sdpa" + else "" + ) environment["PYTHONPATH"] = os.pathsep.join( value - for value in (str(_IMAGE_REFERENCE_COMPAT), existing) + for value in (attention_compat, str(_IMAGE_REFERENCE_COMPAT), existing) if value ) return environment +def _image_reference_vit_path( + artifact_dir: Path, + vit_path: Path, + *, + attention_backend: str, +) -> Path: + """Create a run-owned VIT view that selects SDPA without editing HF cache.""" + if attention_backend == "flash_attn": + return vit_path + if attention_backend != "torch_sdpa": + raise RuntimeError( + "unsupported Lance reference attention backend " + f"{attention_backend!r}; expected one of {sorted(_ATTENTION_BACKENDS)}" + ) + + source_config = vit_path / "config.json" + try: + config = json.loads(source_config.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"Lance VIT checkpoint has no valid config: {source_config}" + ) from exc + if not isinstance(config, dict): + raise RuntimeError(f"Lance VIT config must contain an object: {source_config}") + + overlay = artifact_dir / "official_vit_torch_sdpa" + overlay.mkdir(parents=True, exist_ok=True) + for source in vit_path.iterdir(): + if source.name == "config.json": + continue + destination = overlay / source.name + if destination.is_symlink(): + if destination.resolve() != source.resolve(): + raise RuntimeError( + f"Lance VIT overlay points at {destination.resolve()}, " + f"expected {source.resolve()}" + ) + elif destination.exists(): + raise RuntimeError( + f"Lance VIT overlay path is not a symlink: {destination}" + ) + else: + destination.symlink_to( + source.resolve(), + target_is_directory=source.is_dir(), + ) + + config["_attn_implementation"] = "sdpa" + overlay_config = overlay / "config.json" + if overlay_config.is_symlink() or ( + overlay_config.exists() and not overlay_config.is_file() + ): + raise RuntimeError( + f"Lance VIT overlay config is not a regular file: {overlay_config}" + ) + overlay_config.write_text( + json.dumps(config, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return overlay + + def _image_reference_workspace( artifact_dir: Path, model_root: Path, @@ -69,6 +153,7 @@ def _image_reference_workspace( def _cached_model_root( model_id: str, *, + revision: str, local_files_only: bool = True, ) -> Path: path = Path(model_id) @@ -81,6 +166,7 @@ def _cached_model_root( snapshot_download( model_id, allow_patterns=_IMAGE_MODEL_ALLOW_PATTERNS, + revision=revision, local_files_only=local_files_only, ) ).resolve() @@ -158,6 +244,7 @@ def run_stage( source = _official_source() model_root = _cached_model_root( case.hf_id, + revision=case.hf_revision, local_files_only=ctx.local_files_only, ) model_path = model_root / _MODEL_DIRECTORY @@ -177,6 +264,12 @@ def run_stage( request_path = artifact_dir / "lance_x2t_request.json" result_dir = artifact_dir / "official_output" workspace = _image_reference_workspace(artifact_dir, model_root) + environment = _image_reference_environment() + reference_vit_path = _image_reference_vit_path( + artifact_dir, + vit_path, + attention_backend=_attention_backend(environment), + ) request_path.write_text( json.dumps( { @@ -195,7 +288,6 @@ def run_stage( + "\n", encoding="utf-8", ) - environment = _image_reference_environment() command = [ "env", f"--chdir={workspace}", @@ -207,7 +299,7 @@ def run_stage( "--llm_path", str(model_path), "--vit_path", - str(vit_path), + str(reference_vit_path), "--vit_type", "qwen_2_5_vl_original", "--llm_qk_norm", diff --git a/tests/e2e/models/lance/test_lance_attention_compat.py b/tests/e2e/models/lance/test_lance_attention_compat.py new file mode 100644 index 0000000000..44749dc3a3 --- /dev/null +++ b/tests/e2e/models/lance/test_lance_attention_compat.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fallback contracts for Lance's packed FlashAttention calls.""" + +from __future__ import annotations + +import importlib + +import torch +import torch.nn.functional as F + +from tests.e2e.models.lance.e2e_plugins.references.lance_image_attention_compat.flash_attn import ( + flash_attn_varlen_func, +) + + +def test_varlen_fallback_keeps_sequences_isolated() -> None: + torch.manual_seed(7) + query = torch.randn(5, 2, 4) + key = torch.randn(5, 2, 4) + value = torch.randn(5, 2, 4) + cumulative = torch.tensor([0, 2, 5], dtype=torch.int32) + + actual = flash_attn_varlen_func( + query, + key, + value, + cumulative, + cumulative, + 3, + 3, + ) + expected = torch.cat( + [ + F.scaled_dot_product_attention( + query[start:end].transpose(0, 1).unsqueeze(0), + key[start:end].transpose(0, 1).unsqueeze(0), + value[start:end].transpose(0, 1).unsqueeze(0), + ) + .squeeze(0) + .transpose(0, 1) + for start, end in ((0, 2), (2, 5)) + ] + ) + + torch.testing.assert_close(actual, expected) + + +def test_varlen_fallback_uses_bottom_right_causal_alignment() -> None: + query = torch.tensor([[[1.0, 0.0]]]) + key = torch.tensor( + [ + [[1.0, 0.0]], + [[0.0, 1.0]], + [[1.0, 1.0]], + ] + ) + value = torch.tensor([[[1.0]], [[2.0]], [[4.0]]]) + + actual = flash_attn_varlen_func( + query, + key, + value, + torch.tensor([0, 1], dtype=torch.int32), + torch.tensor([0, 3], dtype=torch.int32), + 1, + 3, + causal=True, + ) + expected = F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + attn_mask=torch.nn.attention.bias.causal_lower_right(1, 3), + ).squeeze(0).transpose(0, 1) + + torch.testing.assert_close(actual, expected) + + +def test_varlen_fallback_supports_grouped_query_attention() -> None: + torch.manual_seed(11) + query = torch.randn(3, 4, 8) + key = torch.randn(3, 2, 8) + value = torch.randn(3, 2, 8) + cumulative = torch.tensor([0, 3], dtype=torch.int32) + + actual = flash_attn_varlen_func( + query, + key, + value, + cumulative, + cumulative, + 3, + 3, + ) + expected = F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + enable_gqa=True, + ).squeeze(0).transpose(0, 1) + + torch.testing.assert_close(actual, expected) + + +def test_fallback_exposes_transformers_flash_attention_import_surface() -> None: + package = ( + "tests.e2e.models.lance.e2e_plugins.references." + "lance_image_attention_compat.flash_attn" + ) + flash_attn = importlib.import_module(package) + bert_padding = importlib.import_module(f"{package}.bert_padding") + rotary = importlib.import_module(f"{package}.layers.rotary") + + assert callable(flash_attn.flash_attn_func) + assert callable(flash_attn.flash_attn_varlen_func) + assert callable(bert_padding.index_first_axis) + assert callable(bert_padding.pad_input) + assert callable(bert_padding.unpad_input) + assert callable(rotary.apply_rotary_emb) + + +def test_dense_fallback_matches_sdpa() -> None: + package = importlib.import_module( + "tests.e2e.models.lance.e2e_plugins.references." + "lance_image_attention_compat.flash_attn" + ) + torch.manual_seed(13) + query = torch.randn(2, 5, 4, 8) + key = torch.randn(2, 5, 2, 8) + value = torch.randn(2, 5, 2, 8) + + actual = package.flash_attn_func(query, key, value) + expected = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + enable_gqa=True, + ).transpose(1, 2) + + torch.testing.assert_close(actual, expected) + + +def test_rotary_fallback_matches_non_interleaved_contract() -> None: + rotary = importlib.import_module( + "tests.e2e.models.lance.e2e_plugins.references." + "lance_image_attention_compat.flash_attn.layers.rotary" + ) + values = torch.tensor( + [[[[1.0, 2.0, 3.0, 4.0]], [[5.0, 6.0, 7.0, 8.0]]]] + ) + cos = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + sin = torch.tensor([[0.0, 1.0], [1.0, 0.0]]) + + actual = rotary.apply_rotary_emb(values, cos, sin) + expected = torch.tensor( + [[[[1.0, -4.0, 3.0, 2.0]], [[-7.0, 6.0, 5.0, 8.0]]]] + ) + + torch.testing.assert_close(actual, expected) + + +def test_padding_fallback_round_trips_valid_tokens() -> None: + padding = importlib.import_module( + "tests.e2e.models.lance.e2e_plugins.references." + "lance_image_attention_compat.flash_attn.bert_padding" + ) + values = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + mask = torch.tensor([[1, 1, 0], [0, 1, 1]], dtype=torch.bool) + + unpadded, indices, cumulative, maximum = padding.unpad_input(values, mask) + restored = padding.pad_input(unpadded, indices, 2, 3) + + torch.testing.assert_close(restored[mask], values[mask]) + torch.testing.assert_close(restored[~mask], torch.zeros_like(restored[~mask])) + assert cumulative.tolist() == [0, 2, 4] + assert maximum == 2 diff --git a/tests/e2e/models/lance/test_lance_official_reference.py b/tests/e2e/models/lance/test_lance_official_reference.py index a87686b7e6..ff29b2564b 100644 --- a/tests/e2e/models/lance/test_lance_official_reference.py +++ b/tests/e2e/models/lance/test_lance_official_reference.py @@ -5,12 +5,15 @@ from __future__ import annotations +import json import os from pathlib import Path import subprocess import sys from types import SimpleNamespace +import pytest + from tests.e2e.models.lance.e2e_plugins.references import lance_official @@ -34,6 +37,7 @@ def snapshot_download(model_id: str, **kwargs) -> str: assert ( lance_official._cached_model_root( "bytedance-research/Lance", + revision="7395315758865e6f56ab87ad06a88c7ac172f056", local_files_only=True, ) == tmp_path @@ -43,7 +47,9 @@ def snapshot_download(model_id: str, **kwargs) -> str: "allow_patterns": [ "Lance_3B/**", "Qwen2.5-VL-ViT/**", + "Wan2.2_VAE.pth", ], + "revision": "7395315758865e6f56ab87ad06a88c7ac172f056", "local_files_only": True, } @@ -72,6 +78,113 @@ def test_lance_image_reference_provides_decord_import_without_video_support() -> assert environment["PYTHONPATH"].endswith(":/existing/python/path") +def test_lance_image_reference_adds_explicit_sdpa_fallback() -> None: + environment = lance_official._image_reference_environment( + {"TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": "torch_sdpa"} + ) + + assert environment["PYTHONPATH"].split(os.pathsep)[:2] == [ + str(lance_official._IMAGE_REFERENCE_ATTENTION_COMPAT), + str(lance_official._IMAGE_REFERENCE_COMPAT), + ] + + +def test_lance_sdpa_fallback_is_importable_without_distribution_metadata() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import importlib.metadata as metadata; " + "from flash_attn import flash_attn_varlen_func; " + "assert callable(flash_attn_varlen_func); " + "\ntry:\n metadata.version('flash_attn')\n" + "except metadata.PackageNotFoundError:\n pass\n" + "else:\n raise AssertionError('fallback must not publish metadata')" + ), + ], + env=lance_official._image_reference_environment( + { + **os.environ, + "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": "torch_sdpa", + } + ), + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +def test_lance_sdpa_vit_overlay_preserves_checkpoint_and_changes_only_config( + tmp_path: Path, +) -> None: + vit_path = tmp_path / "Qwen2.5-VL-ViT" + vit_path.mkdir() + config_path = vit_path / "config.json" + config_path.write_text( + json.dumps( + { + "_attn_implementation": "flash_attention_2", + "hidden_size": 1280, + } + ), + encoding="utf-8", + ) + weights = vit_path / "vit.safetensors" + weights.write_bytes(b"weights") + + overlay = lance_official._image_reference_vit_path( + tmp_path / "artifacts", + vit_path, + attention_backend="torch_sdpa", + ) + + assert overlay != vit_path + assert json.loads((overlay / "config.json").read_text(encoding="utf-8")) == { + "_attn_implementation": "sdpa", + "hidden_size": 1280, + } + assert (overlay / "vit.safetensors").is_symlink() + assert (overlay / "vit.safetensors").resolve() == weights.resolve() + assert json.loads(config_path.read_text(encoding="utf-8"))[ + "_attn_implementation" + ] == "flash_attention_2" + + +def test_lance_native_flash_attention_uses_original_vit_checkpoint( + tmp_path: Path, +) -> None: + vit_path = tmp_path / "Qwen2.5-VL-ViT" + vit_path.mkdir() + + assert ( + lance_official._image_reference_vit_path( + tmp_path / "artifacts", + vit_path, + attention_backend="flash_attn", + ) + == vit_path + ) + + +def test_lance_image_reference_keeps_explicit_native_flash_attention() -> None: + environment = lance_official._image_reference_environment( + {"TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": "flash_attn"} + ) + + assert str(lance_official._IMAGE_REFERENCE_ATTENTION_COMPAT) not in environment[ + "PYTHONPATH" + ].split(os.pathsep) + + +def test_lance_image_reference_rejects_unknown_attention_backend() -> None: + with pytest.raises(RuntimeError, match="unsupported Lance reference attention backend"): + lance_official._image_reference_environment( + {"TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": "automatic"} + ) + + def test_lance_image_reference_keeps_upstream_visual_generation_contract( monkeypatch, tmp_path: Path, @@ -81,7 +194,13 @@ def test_lance_image_reference_keeps_upstream_visual_generation_contract( (source / "inference_lance.py").write_text("", encoding="utf-8") model_root = tmp_path / "model" (model_root / "Lance_3B").mkdir(parents=True) - (model_root / "Qwen2.5-VL-ViT").mkdir() + vit_path = model_root / "Qwen2.5-VL-ViT" + vit_path.mkdir() + (vit_path / "config.json").write_text( + '{"_attn_implementation": "flash_attention_2"}\n', + encoding="utf-8", + ) + (vit_path / "vit.safetensors").write_bytes(b"weights") image = tmp_path / "input.png" image.write_bytes(b"image") artifact_dir = tmp_path / "artifacts" @@ -89,7 +208,8 @@ def test_lance_image_reference_keeps_upstream_visual_generation_contract( captured: dict[str, list[str]] = {} captured_local_files_only: list[bool] = [] - def resolve_model(_model_id, *, local_files_only): + def resolve_model(_model_id, *, revision, local_files_only): + assert revision == "7395315758865e6f56ab87ad06a88c7ac172f056" captured_local_files_only.append(local_files_only) return model_root @@ -115,6 +235,10 @@ def run(command, **_kwargs): lambda *_args: str(artifact_dir), ) monkeypatch.setattr(lance_official.subprocess, "run", run) + monkeypatch.setenv( + "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND", + "torch_sdpa", + ) monkeypatch.setattr( lance_official, "save_full_stderr", @@ -125,6 +249,7 @@ def run(command, **_kwargs): task_strategy="vision_language_generation", inputs={"image": str(image), "prompt": "Describe the image."}, hf_id="bytedance-research/Lance", + hf_revision="7395315758865e6f56ab87ad06a88c7ac172f056", determinism={"seed": 42}, name="lance-test", ), @@ -143,6 +268,14 @@ def run(command, **_kwargs): assert command[text_template + 1] == "true" resolution = command.index("--resolution") assert command[resolution + 1] == "image_512res" + command_vit_path = Path(command[command.index("--vit_path") + 1]) + assert command_vit_path != vit_path + assert json.loads( + (command_vit_path / "config.json").read_text(encoding="utf-8") + )["_attn_implementation"] == "sdpa" + assert (command_vit_path / "vit.safetensors").resolve() == ( + vit_path / "vit.safetensors" + ).resolve() chdir_argument = next( value for value in command diff --git a/tests/e2e/models/nemotron_speech_streaming/manifests/nemotron-3.5-asr-streaming-0.6b.json b/tests/e2e/models/nemotron_speech_streaming/manifests/nemotron-3.5-asr-streaming-0.6b.json index 50a905ec1a..41f802ed92 100644 --- a/tests/e2e/models/nemotron_speech_streaming/manifests/nemotron-3.5-asr-streaming-0.6b.json +++ b/tests/e2e/models/nemotron_speech_streaming/manifests/nemotron-3.5-asr-streaming-0.6b.json @@ -25,7 +25,7 @@ "att_context_size": [56, 13] }, "metadata": { - "notes": "Reference runs NeMo's EncDecRNNTBPEModelWithPrompt via a JSONL manifest with `lang: en-US` (the field NeMo's Lhotse loader maps to supervisions[0].language); audio is converted to mono 16 kHz first. Comparator gates on WER/CER thresholds set in the sidecar." + "notes": "Reference runs NeMo's EncDecHybridRNNTCTCBPEModelWithPrompt via a JSONL manifest with `lang: en-US` (the field NeMo's Lhotse loader maps to supervisions[0].language); audio is converted to mono 16 kHz first. Comparator gates on WER/CER thresholds set in the sidecar." } } ] diff --git a/tests/e2e/models/wan2_2_ti2v/MODEL.toml b/tests/e2e/models/wan2_2_ti2v/MODEL.toml index 1f2b0db741..7745c09aa0 100644 --- a/tests/e2e/models/wan2_2_ti2v/MODEL.toml +++ b/tests/e2e/models/wan2_2_ti2v/MODEL.toml @@ -13,6 +13,7 @@ repository = "https://github.com/Wan-Video/Wan2.2.git" revision = "42bf4cfaa384bc21833865abc2f9e6c0e67233dc" relative_path = "wan2_2_ti2v/reference/Wan2.2-42bf4cfaa384" entrypoint = "wan/textimage2video.py" +environment_variable = "TRTMC_WAN_REFERENCE_REPO" suites = ["nightly"] [e2e_defaults.diffusion_media_generation] diff --git a/tests/e2e/models/wan2_2_ti2v/e2e_plugins/references/official_wan.py b/tests/e2e/models/wan2_2_ti2v/e2e_plugins/references/official_wan.py index 9471bdd4e0..b2981ab823 100644 --- a/tests/e2e/models/wan2_2_ti2v/e2e_plugins/references/official_wan.py +++ b/tests/e2e/models/wan2_2_ti2v/e2e_plugins/references/official_wan.py @@ -23,6 +23,7 @@ OFFICIAL_REVISION = "42bf4cfaa384bc21833865abc2f9e6c0e67233dc" OFFICIAL_RELATIVE_PATH = "wan2_2_ti2v/reference/Wan2.2-42bf4cfaa384" OFFICIAL_ENTRYPOINT = "wan/textimage2video.py" +OFFICIAL_SOURCE_ENV = "TRTMC_WAN_REFERENCE_REPO" UMT5_SPECIAL_TOKENS = { "pad_token": ("", 0), "eos_token": ("", 1), @@ -87,10 +88,18 @@ def _snapshot_revision(model_ref: str) -> str: def _resolve_official_source(storage_root: str | None = None) -> Path: - root_value = storage_root if storage_root is not None else os.environ.get("TRTMC_STORAGE_ROOT") - if not root_value: - raise RuntimeError("TRTMC_STORAGE_ROOT is required for the pinned official Wan reference") - source = Path(root_value) / OFFICIAL_RELATIVE_PATH + declared_source = "" if storage_root is not None else os.environ.get(OFFICIAL_SOURCE_ENV, "") + if declared_source: + source = Path(declared_source) + else: + root_value = ( + storage_root if storage_root is not None else os.environ.get("TRTMC_STORAGE_ROOT") + ) + if not root_value: + raise RuntimeError( + "TRTMC_STORAGE_ROOT is required for the pinned official Wan reference" + ) + source = Path(root_value) / OFFICIAL_RELATIVE_PATH entrypoint = source / OFFICIAL_ENTRYPOINT if source.is_symlink() or entrypoint.is_symlink() or not entrypoint.is_file(): raise RuntimeError( diff --git a/tests/e2e/models/wan2_2_ti2v/test_wan22_official_reference.py b/tests/e2e/models/wan2_2_ti2v/test_wan22_official_reference.py index 09a2f250a4..492226a060 100644 --- a/tests/e2e/models/wan2_2_ti2v/test_wan22_official_reference.py +++ b/tests/e2e/models/wan2_2_ti2v/test_wan22_official_reference.py @@ -8,6 +8,7 @@ import subprocess import struct import sys +import tomllib import zlib from pathlib import Path from types import ModuleType @@ -267,6 +268,26 @@ def test_reference_requires_the_pinned_official_source( wan22_official._resolve_official_source(str(tmp_path)) +def test_reference_uses_the_declared_checkout_when_output_cache_is_separate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + output_cache = tmp_path / "accuracy-references" + source_cache = tmp_path / "model-sources" + source = _write_official_source(source_cache) + monkeypatch.setenv("TRTMC_STORAGE_ROOT", str(output_cache)) + monkeypatch.setenv("TRTMC_WAN_REFERENCE_REPO", str(source)) + + assert wan22_official._resolve_official_source() == source + + +def test_family_declares_the_official_source_environment_variable() -> None: + manifest = tomllib.loads((_MODEL_DIR / "MODEL.toml").read_text(encoding="utf-8")) + + assert manifest["model_reference_cache"]["environment_variable"] == ( + "TRTMC_WAN_REFERENCE_REPO" + ) + + def test_reference_fails_closed_on_subprocess_error( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/model_checks/README.md b/tests/model_checks/README.md index 9611680596..f908e7f80c 100644 --- a/tests/model_checks/README.md +++ b/tests/model_checks/README.md @@ -111,8 +111,9 @@ $TRTMC_CHECK_PYTHON tools/model_checks.py run \ Use `--platform gb300`, `--platform auto-thor`, or `--platform l4t-thor`. Platform files under `tests/model_checks/platforms/` define task order and -hardware exclusions. Files under `tests/model_checks/environments/` define -paths and retention. +model exclusions. An excluded model runs no Accuracy suite or Perf entry; +Accuracy suites remain in the report as `not compared`. Files under +`tests/model_checks/environments/` define paths and retention. Checked-in environments delete Accuracy engines and Perf bundles after each binding. GB300 and Auto Thor retain their shared Hugging Face cache. L4T Thor @@ -126,7 +127,7 @@ See `benchmarks/performance/README.md` for the L4T TensorRT 11 bare-metal build. - Define an Accuracy suite in `tests/validation/workloads.yaml`. - Assign one or more suites in `tests/validation/model_workloads.yaml`. - Add the Perf entry independently in `benchmarks/performance/release.yaml`. -- Add only verified hardware exclusions under `tests/model_checks/platforms/`. +- Add platform-wide model exclusions under `tests/model_checks/platforms/`. Each Accuracy `MODEL=SUITE` binding has its own engine directory. Perf bundles remain independent from Accuracy engines. diff --git a/tests/model_checks/environments/auto-thor.yaml b/tests/model_checks/environments/auto-thor.yaml index d4eadc1ea9..7784dc592a 100644 --- a/tests/model_checks/environments/auto-thor.yaml +++ b/tests/model_checks/environments/auto-thor.yaml @@ -4,6 +4,9 @@ schema_version: trtmc.model-check-environment/v1 id: auto-thor +environment_variables: + TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND: torch_sdpa + storage: root: ${TRTMC_CHECK_STORAGE_ROOT} results_root: ${TRTMC_CHECK_STORAGE_ROOT}/results @@ -24,6 +27,7 @@ tasks: engine-retention: delete_always hf-cache-mode: shared hf-cache-retention: retain + prepare-hf-on-demand: true model-attempts: 2 model-retry-delay-seconds: 30 on-model-failure: continue diff --git a/tests/model_checks/environments/gb300.yaml b/tests/model_checks/environments/gb300.yaml index 02f6724889..8145b5b4de 100644 --- a/tests/model_checks/environments/gb300.yaml +++ b/tests/model_checks/environments/gb300.yaml @@ -4,6 +4,9 @@ schema_version: trtmc.model-check-environment/v1 id: gb300 +environment_variables: + TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND: torch_sdpa + storage: root: ${TRTMC_CHECK_STORAGE_ROOT} results_root: ${TRTMC_CHECK_STORAGE_ROOT}/results @@ -24,6 +27,7 @@ tasks: engine-retention: delete_always hf-cache-mode: shared hf-cache-retention: retain + prepare-hf-on-demand: true model-attempts: 2 model-retry-delay-seconds: 30 on-model-failure: continue diff --git a/tests/model_checks/environments/l4t-thor.yaml b/tests/model_checks/environments/l4t-thor.yaml index ac9e4bf720..5a8cdfd9c8 100644 --- a/tests/model_checks/environments/l4t-thor.yaml +++ b/tests/model_checks/environments/l4t-thor.yaml @@ -4,6 +4,13 @@ schema_version: trtmc.model-check-environment/v1 id: l4t-thor +library_dirs: + - ${TRTMC_CHECK_STORAGE_ROOT}/runtime/TensorRT-11.0.2.2/lib +executable_dirs: + - /usr/local/cuda/bin +environment_variables: + TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND: torch_sdpa + storage: root: ${TRTMC_CHECK_STORAGE_ROOT} results_root: ${TRTMC_CHECK_STORAGE_ROOT}/results @@ -24,6 +31,7 @@ tasks: engine-retention: delete_always hf-cache-mode: per_model hf-cache-retention: delete_always + prepare-hf-on-demand: true model-attempts: 2 model-retry-delay-seconds: 30 model-timeout-seconds: 21600 diff --git a/tests/model_checks/platforms/auto-thor.yaml b/tests/model_checks/platforms/auto-thor.yaml index 468e92454d..8004d7e1e5 100644 --- a/tests/model_checks/platforms/auto-thor.yaml +++ b/tests/model_checks/platforms/auto-thor.yaml @@ -10,4 +10,4 @@ execution: # Auto Thor storage is supplied by its execution environment because mount # points differ between lab installations. -unsupported: [] +excluded_models: [] diff --git a/tests/model_checks/platforms/gb300.yaml b/tests/model_checks/platforms/gb300.yaml index 6e7dcb35b9..9b1452dcf2 100644 --- a/tests/model_checks/platforms/gb300.yaml +++ b/tests/model_checks/platforms/gb300.yaml @@ -8,6 +8,4 @@ execution: task_order: [accuracy, perf] serial_tasks: false -# Keep only evidence-backed hardware exclusions here. Missing task bindings are -# configuration errors and must not be represented as hardware exclusions. -unsupported: [] +excluded_models: [] diff --git a/tests/model_checks/platforms/l4t-thor.yaml b/tests/model_checks/platforms/l4t-thor.yaml index 28a15bb31d..0f68bf0e15 100644 --- a/tests/model_checks/platforms/l4t-thor.yaml +++ b/tests/model_checks/platforms/l4t-thor.yaml @@ -11,24 +11,15 @@ execution: storage: device: /dev/nvme0n1p1 -# A disagreement is not a hardware exclusion. Keep only reproduced hardware -# limits here, scoped to the exact task binding that requires the engine. -unsupported: - - model: minimax-h3-768p - task: accuracy - binding: minimax_h3_official_profile_parity - reason: >- - The qualified single-device profile requires 180 GiB of free GPU memory, - which exceeds this platform's approximately 128 GiB unified memory. - - model: qwen-image-2512 - task: accuracy - binding: dpg_bench_diffusion_image - reason: >- - TensorRT 11.0 MMDiT engine construction exhausts the platform's 128 GiB - unified memory and can leave the build blocked in a dma-fence wait. - - model: qwen-image-edit-2511 - task: accuracy - binding: gedit_bench_image_edit - reason: >- - Equivalent Qwen Image engine profiles exhaust 128 GiB unified memory; - the edit profile uses the same engine topology with added conditioning. +excluded_models: + - deepseek-v2-lite + - flux-2-dev + - flux-2-dev-fp8 + - flux-schnell + - gpt-oss-20b + - minimax-h3-768p + - qwen-image + - qwen-image-2512 + - qwen-image-edit-2511 + - qwen3-omni-30b-a3b-instruct + - sana-wm-bidirectional diff --git a/tests/tools/test_model_checks.py b/tests/tools/test_model_checks.py index 1e80c44781..3b178e89cb 100644 --- a/tests/tools/test_model_checks.py +++ b/tests/tools/test_model_checks.py @@ -15,7 +15,7 @@ from tools import model_checks -def _platform(*, serial: bool = True, unsupported=()): +def _platform(*, serial: bool = True, excluded_models=()): return { "id": "test-platform", "source": "platform.yaml", @@ -23,7 +23,7 @@ def _platform(*, serial: bool = True, unsupported=()): "task_order": ["accuracy", "perf"], "serial_tasks": serial, }, - "unsupported": list(unsupported), + "excluded_models": list(excluded_models), } @@ -74,7 +74,7 @@ def test_plan_expands_every_accuracy_workload_and_perf_entry(): "model_count": 1, "binding_count": 4, "configured_binding_count": 4, - "unsupported_binding_count": 0, + "excluded_binding_count": 0, "blocker_count": 0, } @@ -118,20 +118,11 @@ def test_plan_can_select_distinct_accuracy_suites_per_model(): ] == [("model-a", "suite-b"), ("model-b", "suite-c")] -def test_platform_hardware_exclusion_can_target_one_accuracy_suite(): +def test_platform_model_exclusion_applies_to_every_accuracy_and_perf_binding(): plan = model_checks.resolve_plan( models=["model-a"], - tasks=["accuracy"], - platform=_platform( - unsupported=( - { - "model": "model-a", - "task": "accuracy", - "binding": "suite-b", - "reason": "suite-b exceeds unified memory", - }, - ) - ), + tasks=["accuracy", "perf"], + platform=_platform(excluded_models=("model-a",)), accuracy_catalog=_accuracy_catalog(), accuracy_workloads=(), accuracy_bindings={}, @@ -141,31 +132,54 @@ def test_platform_hardware_exclusion_can_target_one_accuracy_suite(): bindings = plan["models"][0]["tasks"]["accuracy"]["bindings"] assert [(binding["workload"], binding["status"]) for binding in bindings] == [ - ("suite-a", "configured"), - ("suite-b", "unsupported"), + ("suite-a", "excluded"), + ("suite-b", "excluded"), ] - assert plan["models"][0]["tasks"]["accuracy"]["status"] == "configured" - - -def test_platform_hardware_exclusion_must_name_a_real_binding(): - with pytest.raises(model_checks.ModelCheckError, match="unknown Accuracy binding"): - model_checks.audit_platform_unsupported( - _platform( - unsupported=( - { - "model": "model-a", - "task": "accuracy", - "binding": "missing-suite", - "reason": "hardware evidence", - }, - ) - ), + assert plan["models"][0]["tasks"]["accuracy"]["status"] == "excluded" + perf = plan["models"][0]["tasks"]["perf"]["bindings"] + assert [(binding["entry"], binding["status"]) for binding in perf] == [ + ("family-a.default", "excluded"), + ("family-a.long", "excluded"), + ] + assert plan["models"][0]["tasks"]["perf"]["status"] == "excluded" + assert plan["summary"]["excluded_binding_count"] == 4 + + +def test_platform_exclusion_must_name_a_real_model(): + with pytest.raises(model_checks.ModelCheckError, match="unknown models: missing-model"): + model_checks.audit_platform_exclusions( + _platform(excluded_models=("missing-model",)), accuracy_catalog=_accuracy_catalog(), perf_cases=_perf_cases(), ) -def test_missing_task_binding_is_a_blocker_not_hardware_unsupported(): +@pytest.mark.parametrize("legacy_field", ["unsupported", "excluded"]) +def test_platform_rejects_legacy_binding_exclusions( + tmp_path: Path, + legacy_field: str, +) -> None: + path = tmp_path / "legacy-platform.yaml" + path.write_text( + yaml.safe_dump( + { + "schema_version": model_checks.PLATFORM_SCHEMA, + "id": "legacy-platform", + "execution": { + "task_order": ["accuracy", "perf"], + "serial_tasks": True, + }, + legacy_field: [], + } + ), + encoding="utf-8", + ) + + with pytest.raises(model_checks.ModelCheckError, match="use excluded_models"): + model_checks.load_platform(str(path)) + + +def test_missing_task_binding_is_a_blocker_not_a_platform_exclusion(): plan = model_checks.resolve_plan( models=["model-b"], tasks=["accuracy", "perf"], @@ -227,7 +241,7 @@ def resolve_plan(**kwargs): "model_count": 0, "binding_count": 0, "configured_binding_count": 0, - "unsupported_binding_count": 0, + "excluded_binding_count": 0, "blocker_count": 0, }, } @@ -284,6 +298,8 @@ def test_execution_environment_preserves_command_name_and_resolves_paths( { "schema_version": model_checks.ENVIRONMENT_SCHEMA, "id": "test-platform", + "python_dirs": ["${TEST_STORAGE}/runtime/python"], + "environment_variables": {"TRTMC_MODEL_FEATURE": "enabled"}, "storage": { "root": "${TEST_STORAGE}", "results_root": "${TEST_STORAGE}/results", @@ -317,6 +333,63 @@ def test_execution_environment_preserves_command_name_and_resolves_paths( ) assert environment["tasks"]["accuracy"]["runner_python"] == "python3" assert Path(environment["tasks"]["perf"]["suite"]).is_absolute() + assert environment["python_dirs"] == [str(storage / "runtime/python")] + assert environment["environment_variables"] == { + "TRTMC_MODEL_FEATURE": "enabled" + } + + +@pytest.mark.parametrize( + "environment_variables", + [ + {"PATH": "/unmanaged"}, + {"TRTMC_VALIDATION_SOURCE_REVISION": "unmanaged"}, + {"TRTMC_MODEL_FEATURE": ""}, + {"TRTMC_MODEL_FEATURE": True}, + ], +) +def test_execution_environment_rejects_unsafe_environment_variables( + tmp_path, + monkeypatch, + environment_variables, +): + storage = tmp_path / "storage" + monkeypatch.setenv("TEST_STORAGE", str(storage)) + environment_path = tmp_path / "environment.yaml" + environment_path.write_text( + yaml.safe_dump( + { + "schema_version": model_checks.ENVIRONMENT_SCHEMA, + "id": "test-platform", + "environment_variables": environment_variables, + "storage": { + "root": "${TEST_STORAGE}", + "results_root": "${TEST_STORAGE}/results", + }, + "tasks": { + "accuracy": {"runner_python": "python3", "options": {}}, + "perf": { + "runner_python": "python3", + "suite": "benchmarks/performance/release.yaml", + "environment": ( + "benchmarks/performance/environments/gb300.yaml" + ), + }, + }, + }, + sort_keys=False, + ), + encoding="utf-8", + ) + + with pytest.raises( + model_checks.ModelCheckError, + match="environment_variables must map unmanaged TRTMC_", + ): + model_checks.load_execution_environment( + str(environment_path), + platform_id="test-platform", + ) def test_runner_executable_preserves_virtual_environment_symlink(tmp_path): @@ -345,6 +418,119 @@ def test_task_environment_uses_shared_profiles_and_allows_missing_profiles( assert os.environ["TRTMC_PYTHON_PROFILE_PREBUILT_ONLY"] == "1" +def test_task_environment_prepends_configured_runtime_libraries( + tmp_path, + monkeypatch, +): + runtime_library = tmp_path / "runtime/lib" + runtime_library.mkdir(parents=True) + monkeypatch.setenv("LD_LIBRARY_PATH", "/system/lib") + + environment = model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "library_dirs": [str(runtime_library)], + } + ) + + assert environment["LD_LIBRARY_PATH"] == ( + f"{runtime_library}{os.pathsep}/system/lib" + ) + + +def test_task_environment_prepends_configured_executable_directories( + tmp_path, + monkeypatch, +): + cuda_bin = tmp_path / "cuda/bin" + cuda_bin.mkdir(parents=True) + monkeypatch.setenv("PATH", "/usr/bin:/bin") + + environment = model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "executable_dirs": [str(cuda_bin)], + } + ) + + assert environment["PATH"] == f"{cuda_bin}{os.pathsep}/usr/bin:/bin" + + +def test_task_environment_prepends_configured_python_directories( + tmp_path, + monkeypatch, +): + runtime_python = tmp_path / "runtime/python" + runtime_python.mkdir(parents=True) + monkeypatch.setenv("PYTHONPATH", "/system/python") + + environment = model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "python_dirs": [str(runtime_python)], + } + ) + + assert environment["PYTHONPATH"] == ( + f"{runtime_python}{os.pathsep}/system/python" + ) + + +def test_task_environment_exports_checked_in_environment_variables( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND", "ambient") + + environment = model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "environment_variables": { + "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": "torch_sdpa" + }, + } + ) + + assert environment["TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND"] == "torch_sdpa" + assert os.environ["TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND"] == "ambient" + + +def test_task_environment_rejects_missing_executable_directory(tmp_path): + missing = tmp_path / "missing/bin" + + with pytest.raises(model_checks.ModelCheckError, match=str(missing)): + model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "executable_dirs": [str(missing)], + } + ) + + +def test_task_environment_rejects_missing_runtime_library_directory(tmp_path): + missing = tmp_path / "missing/lib" + + with pytest.raises(model_checks.ModelCheckError, match=str(missing)): + model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "library_dirs": [str(missing)], + } + ) + + +def test_task_environment_rejects_missing_python_directory(tmp_path): + missing = tmp_path / "missing/python" + + with pytest.raises(model_checks.ModelCheckError, match=str(missing)): + model_checks._task_environment( + { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "python_dirs": [str(missing)], + } + ) + + def test_perf_reference_contracts_come_from_selected_model_owners() -> None: plan = { "models": [ @@ -432,7 +618,7 @@ def test_selected_models_artifact_records_configured_bindings(tmp_path) -> None: "accuracy": { "bindings": [ {"model": "model-a", "status": "configured"}, - {"model": "model-b", "status": "unsupported"}, + {"model": "model-b", "status": "excluded"}, ] }, "perf": { @@ -448,7 +634,105 @@ def test_selected_models_artifact_records_configured_bindings(tmp_path) -> None: selection = model_checks._write_selected_models(plan, tmp_path) - assert selection.read_text(encoding="utf-8") == "model-a\nmodel-c\n" + assert selection.read_text(encoding="utf-8") == "model-a\nmodel-b\nmodel-c\n" + + +def test_accuracy_platform_exclusions_are_reported_without_execution(tmp_path) -> None: + plan = { + "models": [ + { + "tasks": { + "accuracy": { + "bindings": [ + { + "model": "model-a", + "workload": "suite-a", + "status": "configured", + }, + { + "model": "model-b", + "workload": "suite-b", + "status": "excluded", + "reason": "Model is excluded from platform test-platform", + }, + { + "model": "model-c", + "workload": "suite-c", + "status": "excluded", + "reason": "qualification is intentionally deferred", + }, + ] + } + } + } + ] + } + output = tmp_path / "accuracy" + + assert [ + (binding["model"], binding["workload"]) + for binding in model_checks._task_bindings(plan, "accuracy") + ] == [("model-a", "suite-a")] + written = model_checks._write_accuracy_platform_results(plan, output) + _, html_path, report = model_checks.trtmc_validate.write_report(output) + + assert [path.relative_to(output).as_posix() for path in written] == [ + "model-b/suite-b/comparison.json", + "model-c/suite-c/comparison.json", + ] + assert report["summary"]["platform_excluded"] == 2 + assert report["summary"]["not_compared"] == 2 + assert all(result["execution"]["status"] == "not_run" for result in report["results"]) + assert all("platform_exclusion" in result for result in report["results"]) + html = html_path.read_text(encoding="utf-8") + assert "2 platform excluded" in html + assert "Model is excluded from platform test-platform" in html + assert "qualification is intentionally deferred" in html + + +def test_accuracy_platform_exclusion_does_not_replace_executed_result(tmp_path) -> None: + output = tmp_path / "accuracy" + comparison = output / "model-b" / "suite-b" / "comparison.json" + comparison.parent.mkdir(parents=True) + comparison.write_text( + json.dumps( + { + "model": "model-b", + "workload": "suite-b", + "execution": {"status": "completed"}, + "validation": {"status": "passed"}, + } + ), + encoding="utf-8", + ) + plan = { + "models": [ + { + "tasks": { + "accuracy": { + "bindings": [ + { + "model": "model-b", + "workload": "suite-b", + "status": "excluded", + "reason": "Model is excluded from platform test-platform", + } + ] + } + } + } + ] + } + + with pytest.raises( + model_checks.ModelCheckError, + match="refusing to replace a previously executed Accuracy result", + ): + model_checks._write_accuracy_platform_results(plan, output) + + assert json.loads(comparison.read_text(encoding="utf-8"))["execution"] == { + "status": "completed" + } @pytest.mark.parametrize( @@ -473,6 +757,7 @@ def test_checked_in_accuracy_environment_deletes_engines_without_fixed_reserve( assert "local-files-only" not in options assert options["hf-cache-mode"] == hf_cache_mode assert options["hf-cache-retention"] == hf_cache_retention + assert options["prepare-hf-on-demand"] is True def test_l4t_accuracy_environment_bounds_each_model_attempt() -> None: @@ -484,17 +769,57 @@ def test_l4t_accuracy_environment_bounds_each_model_attempt() -> None: assert raw["tasks"]["accuracy"]["options"]["model-timeout-seconds"] == 21600 -def test_l4t_marks_minimax_profile_unsupported_by_memory_contract() -> None: - platform = model_checks.load_platform("l4t-thor") +def test_l4t_environment_selects_qualified_tensorrt_libraries() -> None: + raw = model_checks._read_yaml( + model_checks.DEFAULT_ENVIRONMENT_ROOT / "l4t-thor.yaml", + "model-check environment", + ) + + assert raw["library_dirs"] == [ + "${TRTMC_CHECK_STORAGE_ROOT}/runtime/TensorRT-11.0.2.2/lib" + ] + assert raw["executable_dirs"] == ["/usr/local/cuda/bin"] - assert any( - item["model"] == "minimax-h3-768p" - and item["task"] == "accuracy" - and item["binding"] == "minimax_h3_official_profile_parity" - and "180 GiB" in item["reason"] - for item in platform["unsupported"] + +@pytest.mark.parametrize( + ("platform", "attention_backend"), + [ + ("l4t-thor", "torch_sdpa"), + ("gb300", "torch_sdpa"), + ("auto-thor", "torch_sdpa"), + ], +) +def test_checked_in_environment_selects_lance_reference_attention_backend( + platform, + attention_backend, +) -> None: + environment = model_checks._read_yaml( + model_checks.DEFAULT_ENVIRONMENT_ROOT / f"{platform}.yaml", + "model-check environment", ) + assert environment["environment_variables"] == { + "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND": attention_backend + } + + +def test_l4t_excludes_consolidated_not_compared_models() -> None: + platform = model_checks.load_platform("l4t-thor") + + assert platform["excluded_models"] == [ + "deepseek-v2-lite", + "flux-2-dev", + "flux-2-dev-fp8", + "flux-schnell", + "gpt-oss-20b", + "minimax-h3-768p", + "qwen-image", + "qwen-image-2512", + "qwen-image-edit-2511", + "qwen3-omni-30b-a3b-instruct", + "sana-wm-bidirectional", + ] + def test_l4t_perf_environment_deletes_entry_cache_and_bundle() -> None: path = model_checks.REPOSITORY / "benchmarks/performance/environments/l4t-thor.yaml" @@ -505,7 +830,7 @@ def test_l4t_perf_environment_deletes_entry_cache_and_bundle() -> None: assert raw["execution"]["hf_cache_retention"] == "delete_always" -def test_l4t_platform_rejects_unverifiable_nvme_partition(): +def test_l4t_platform_rejects_unverifiable_data_partition(): platform = model_checks.load_platform("l4t-thor") with pytest.raises(model_checks.ModelCheckError, match="/dev/nvme0n1p1"): @@ -598,6 +923,7 @@ def test_l4t_dry_run_passes_managed_hf_cache_seed_to_accuracy(tmp_path, monkeypa command = request["commands"]["accuracy"] seed_index = command.index("--hf-cache-seed-dir") assert command[seed_index + 1] == str(seed) + assert "--local-files-only" not in command def test_run_default_output_is_concise_and_ends_with_task_summary( @@ -696,6 +1022,61 @@ def run(command, **_kwargs): assert commands[-1][-1] == "--verbose" +def test_run_forwards_exact_source_revision_to_accuracy_and_perf( + tmp_path, monkeypatch +): + storage = tmp_path / "storage" + revision = "A" * 40 + monkeypatch.setenv("TRTMC_CHECK_STORAGE_ROOT", str(storage)) + monkeypatch.setenv("TRTMC_CHECK_DATASET_ROOT", str(tmp_path / "data")) + monkeypatch.setenv("TRTMC_CHECK_BUILD_DIR", str(tmp_path / "runtime")) + monkeypatch.setenv("TRTMC_CHECK_PYTHON", sys.executable) + child_environments = [] + + def run(_command, **kwargs): + child_environments.append(kwargs["env"]) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(model_checks.subprocess, "run", run) + + assert ( + model_checks.main( + [ + "run", + "--platform", + "gb300", + "--model", + "distilgpt2", + "--revision", + revision, + "--run-id", + "revision-unit", + ] + ) + == 0 + ) + + assert len(child_environments) == 2 + for child in child_environments: + assert child["TRTMC_VALIDATION_SOURCE_REVISION"] == revision.lower() + assert child["TRTMC_PERF_SOURCE_REVISION"] == revision.lower() + + +def test_task_environment_does_not_export_symbolic_revision(tmp_path, monkeypatch): + monkeypatch.delenv("TRTMC_VALIDATION_SOURCE_REVISION", raising=False) + monkeypatch.delenv("TRTMC_PERF_SOURCE_REVISION", raising=False) + environment = { + "storage": {"python_profiles_root": str(tmp_path / "profiles")}, + "library_dirs": [], + "executable_dirs": [], + } + + child = model_checks._task_environment(environment, source_revision="HEAD") + + assert "TRTMC_VALIDATION_SOURCE_REVISION" not in child + assert "TRTMC_PERF_SOURCE_REVISION" not in child + + def test_run_resume_verifies_request_and_resumes_accuracy(tmp_path, monkeypatch): storage = tmp_path / "storage" monkeypatch.setenv("TRTMC_CHECK_STORAGE_ROOT", str(storage)) diff --git a/tests/tools/test_model_plugin_encapsulation_static.py b/tests/tools/test_model_plugin_encapsulation_static.py index 35c72f3e89..13bf07ca66 100644 --- a/tests/tools/test_model_plugin_encapsulation_static.py +++ b/tests/tools/test_model_plugin_encapsulation_static.py @@ -820,6 +820,21 @@ def _model_prefix(model: str) -> str: return "".join(part.capitalize() for part in model.split("_")) +def test_cmake_discovers_nvcc_from_cuda_roots_before_enabling_cuda() -> None: + """Keep standard CUDA installs usable when ``nvcc`` is absent from PATH.""" + text = CMAKE_ROOT.read_text(encoding="utf-8") + discovery = text.index("find_program(_trtmc_cuda_compiler") + language_probe = text.index("check_language(CUDA)") + discovery_block = text[discovery:language_probe] + + assert discovery < language_probe + assert '"$ENV{CUDA_HOME}"' in discovery_block + assert '"$ENV{CUDAToolkit_ROOT}"' in discovery_block + assert "/usr/local/cuda" in discovery_block + assert "PATH_SUFFIXES bin" in discovery_block + assert 'set(CMAKE_CUDA_COMPILER "${_trtmc_cuda_compiler}"' in discovery_block + + def test_runtime_models_do_not_include_sibling_model_folders() -> None: """Trace: ARCH-MODPLUG-001 Intent: prove each runtime model implementation includes only local model code. diff --git a/tests/tools/test_model_proof_runner.py b/tests/tools/test_model_proof_runner.py index 5a841a38c9..b8da7c0dfe 100644 --- a/tests/tools/test_model_proof_runner.py +++ b/tests/tools/test_model_proof_runner.py @@ -745,6 +745,7 @@ def test_wan22_nightly_selection_emits_only_the_pinned_source_contract( "revision": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc", "relative_path": "wan2_2_ti2v/reference/Wan2.2-42bf4cfaa384", "entrypoint": "wan/textimage2video.py", + "environment_variable": "TRTMC_WAN_REFERENCE_REPO", } diff --git a/tests/tools/test_perf_matrix.py b/tests/tools/test_perf_matrix.py index 79b8e86967..b32068389c 100644 --- a/tests/tools/test_perf_matrix.py +++ b/tests/tools/test_perf_matrix.py @@ -3212,6 +3212,25 @@ def test_lance_image_only_decord_stub_rejects_video_use() -> None: modules["decord"].VideoReader("video.mp4") +def test_lance_reference_loads_explicit_sdpa_attention_compat( + tmp_path: Path, monkeypatch +) -> None: + runner = runpy.run_path(str(REPOSITORY / "tools/lance_reference.py")) + reference_repo = tmp_path / "Lance" + reference_repo.mkdir() + (reference_repo / "inference_lance.py").write_text( + "import flash_attn\nATTENTION = flash_attn.flash_attn_varlen_func.__module__\n", + encoding="utf-8", + ) + monkeypatch.setenv("TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND", "torch_sdpa") + monkeypatch.delitem(sys.modules, "flash_attn", raising=False) + + upstream = runner["_load_upstream"](reference_repo) + + assert upstream.ATTENTION.endswith("flash_attn") + assert "lance_image_attention_compat" in sys.modules["flash_attn"].__file__ + + def test_lance_git_revision_scopes_safe_directory(tmp_path: Path, monkeypatch) -> None: runner = runpy.run_path(str(REPOSITORY / "tools/lance_reference.py")) captured: list[str] = [] diff --git a/tests/tools/test_transformers_text_reference.py b/tests/tools/test_transformers_text_reference.py index b0bde50424..452095507b 100644 --- a/tests/tools/test_transformers_text_reference.py +++ b/tests/tools/test_transformers_text_reference.py @@ -5,7 +5,10 @@ import argparse import importlib.util +import json from pathlib import Path +import sys +from types import SimpleNamespace import pytest @@ -93,3 +96,107 @@ def test_input_token_ids_preserve_native_tokenizer_framing() -> None: module = _load_reference_module() assert module._input_token_ids(torch.tensor([[2, 10, 11]])) == [2, 10, 11] + + +def test_internlm_reference_uses_pinned_independent_fast_tokenizer( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + module = _load_reference_module() + model_dir = tmp_path / "model" + model_dir.mkdir() + tokenizer_path = tmp_path / "official-tokenizer.json" + tokenizer_path.write_text("{}", encoding="utf-8") + tokenizer_config = { + "bos_token": "", + "eos_token": "", + "unk_token": "", + "pad_token": "", + "chat_template": "{{ messages }}", + "clean_up_tokenization_spaces": False, + } + (model_dir / "tokenizer_config.json").write_text( + json.dumps(tokenizer_config), + encoding="utf-8", + ) + monkeypatch.setattr( + module, + "_resolve_cached_model_ref", + lambda *_args, **_kwargs: model_dir, + ) + monkeypatch.setattr( + module, + "_resolve_internlm_reference_tokenizer_json", + lambda _model_dir, *, local_files_only: tokenizer_path, + ) + captured: dict[str, object] = {} + tokenizer = SimpleNamespace(pad_token_id=3, eos_token="") + + def fast_tokenizer(**kwargs): + captured.update(kwargs) + return tokenizer + + transformers = SimpleNamespace( + AutoTokenizer=SimpleNamespace( + from_pretrained=lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("InternLM must not use the broken slow tokenizer") + ) + ), + PreTrainedTokenizerFast=fast_tokenizer, + ) + arguments = argparse.Namespace( + model="internlm/internlm2-math-1_8b", + model_revision="", + family="internlm", + trust_remote_code=True, + local_files_only=True, + ) + + assert module._load_tokenizer(arguments, transformers) is tokenizer + assert captured == { + "tokenizer_file": str(tokenizer_path), + "model_input_names": ["input_ids", "attention_mask"], + **tokenizer_config, + } + + +def test_internlm_reference_tokenizer_honors_online_cache_preparation( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + module = _load_reference_module() + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "tokenizer.model").write_bytes(b"official-source-model") + tokenizer_path = tmp_path / "tokenizer.json" + tokenizer_path.write_text("{}", encoding="utf-8") + captured: dict[str, object] = {} + + monkeypatch.setattr( + module, + "_sha256_file", + lambda path: ( + module._INTERNLM_SOURCE_MODEL_SHA256 + if path.name == "tokenizer.model" + else module._INTERNLM_TOKENIZER_SHA256 + ), + ) + + def fake_download(**kwargs): + captured.update(kwargs) + return str(tokenizer_path) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(hf_hub_download=fake_download), + ) + + assert ( + module._resolve_internlm_reference_tokenizer_json( + model_dir, + local_files_only=False, + ) + == tokenizer_path + ) + assert captured["local_files_only"] is False diff --git a/tests/tools/test_trtmc_reference.py b/tests/tools/test_trtmc_reference.py index ffe4cb43b3..f6d9bb573b 100644 --- a/tests/tools/test_trtmc_reference.py +++ b/tests/tools/test_trtmc_reference.py @@ -615,6 +615,7 @@ def test_causal_reference_uses_native_transformers_entrypoint( "batched_mm", ) arguments.reference_family = "causal_base_continuation" + arguments.family = "internlm" manifest_path = work_dir / "manifest.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) manifest["dataset_kind"] = "mmlu_five_shot_json" @@ -657,6 +658,7 @@ def fake_run(command, **kwargs): assert command[1].endswith("tools/reference/transformers_text.py") assert command[command.index("--model-revision") + 1] == "0123456789abcdef" assert command[command.index("--experts-implementation") + 1] == "batched_mm" + assert command[command.index("--family") + 1] == "internlm" assert "validation/engine.py" not in " ".join(command) assert (work_dir / "hf_native_run.log").is_symlink() assert (work_dir / "hf_native_repro.json").is_symlink() @@ -1290,6 +1292,26 @@ def test_plugin_reference_records_nested_upstream_command(tmp_path: Path) -> Non assert "plugin_reference.py" not in " ".join(row["command"]) +def test_plugin_diffusion_failure_preserves_nested_stderr() -> None: + case = SimpleNamespace( + name="sana-wm-sample", + inputs={"seed": 42}, + determinism={"seed": 42}, + ) + output = SimpleNamespace( + timing_s=0.25, + data={ + "returncode": 1, + "num_frames": 0, + "stderr": "official SANA reference exploded", + "stdout": "", + }, + ) + + with pytest.raises(RuntimeError, match="official SANA reference exploded"): + plugin_reference._diffusion_response(case, output, "a red cube") + + def test_plugin_reference_preserves_model_owned_subprocess_command() -> None: namespace: dict[str, object] = {} @@ -1777,7 +1799,14 @@ def empty_cache() -> None: fake_torch = SimpleNamespace(cuda=FakeCuda()) class FakeModel: + def __init__(self) -> None: + self.forward_during_transcribe = None + + def forward(self, *args, **kwargs): + return args, kwargs + def transcribe(self, manifest_path, *, batch_size, verbose): + self.forward_during_transcribe = self.forward rows = [ json.loads(line) for line in Path(manifest_path) @@ -1802,12 +1831,14 @@ def transcribe(self, manifest_path, *, batch_size, verbose): assert verbose is False return [SimpleNamespace(text="CONCORD RETURNED")] + fake_model = FakeModel() + original_forward = fake_model.forward monkeypatch.setattr( speech, "_load_nemotron35_model", lambda _arguments, resolved_archive: ( fake_torch, - FakeModel(), + fake_model, ), ) monkeypatch.setattr( @@ -1841,10 +1872,157 @@ def write_fake_wav(path, _audio, _rate): ) assert len(transcribe_calls) == 1 + assert fake_model.forward_during_transcribe is not None + assert fake_model.forward_during_transcribe != original_forward + assert fake_model.forward == original_forward assert responses[0]["output_text"] == "CONCORD RETURNED" assert responses[0]["generated_token_ids"] is None +def test_nemotron35_forward_adapter_extends_prompt_by_one_frame() -> None: + prompt_tail = object() + extended_prompt = object() + + class FakePrompt: + shape = (1, 44, 8) + + def __getitem__(self, key): + assert key == (slice(None), slice(-1, None), slice(None)) + return prompt_tail + + prompt = FakePrompt() + + class FakeTorch: + @staticmethod + def cat(tensors, *, dim): + assert tensors == (prompt, prompt_tail) + assert dim == 1 + return extended_prompt + + args, kwargs = speech._extend_nemotron35_prompt_for_forward( + FakeTorch(), + (), + {"prompt": prompt, "input_signal": object()}, + ) + + assert args == () + assert kwargs["prompt"] is extended_prompt + + +def test_nemotron35_model_uses_hybrid_prompt_nemo_loader( + tmp_path: Path, + monkeypatch, +) -> None: + archive = tmp_path / "nemotron.nemo" + archive.write_bytes(b"archive") + restore_calls: list[tuple[str, object, bool, object]] = [] + device = object() + + class FakeModel: + def __init__(self) -> None: + self.eval_called = False + self.device = None + self.restore_state = None + + def load_state_dict(self, _state_dict, *, strict): + assert strict is False + return SimpleNamespace( + missing_keys=[ + "ctc_decoder.decoder_layers.0.weight", + "ctc_decoder.decoder_layers.0.bias", + ], + unexpected_keys=[], + ) + + def _set_model_restore_state(self, *, is_being_restored): + self.restore_state = is_being_restored + + def eval(self) -> None: + self.eval_called = True + + def to(self, selected_device) -> None: + self.device = selected_device + + restored = FakeModel() + + class FakeHybridPromptModel: + @staticmethod + def restore_from( + path: str, + *, + map_location, + strict, + save_restore_connector, + ): + restore_calls.append( + (path, map_location, strict, save_restore_connector) + ) + save_restore_connector.load_instance_with_state_dict( + restored, + {}, + strict, + ) + return restored + + class FakeSaveRestoreConnector: + pass + + torch_module = ModuleType("torch") + torch_module.device = lambda name: device # type: ignore[attr-defined] + models_module = ModuleType("nemo.collections.asr.models") + models_module.EncDecHybridRNNTCTCBPEModelWithPrompt = ( # type: ignore[attr-defined] + FakeHybridPromptModel + ) + asr_module = ModuleType("nemo.collections.asr") + asr_module.__path__ = [] # type: ignore[attr-defined] + asr_module.models = models_module # type: ignore[attr-defined] + collections_module = ModuleType("nemo.collections") + collections_module.__path__ = [] # type: ignore[attr-defined] + collections_module.asr = asr_module # type: ignore[attr-defined] + nemo_module = ModuleType("nemo") + nemo_module.__path__ = [] # type: ignore[attr-defined] + nemo_module.collections = collections_module # type: ignore[attr-defined] + save_restore_module = ModuleType("nemo.core.connectors.save_restore_connector") + save_restore_module.SaveRestoreConnector = ( # type: ignore[attr-defined] + FakeSaveRestoreConnector + ) + connectors_module = ModuleType("nemo.core.connectors") + connectors_module.__path__ = [] # type: ignore[attr-defined] + connectors_module.save_restore_connector = ( # type: ignore[attr-defined] + save_restore_module + ) + core_module = ModuleType("nemo.core") + core_module.__path__ = [] # type: ignore[attr-defined] + core_module.connectors = connectors_module # type: ignore[attr-defined] + nemo_module.core = core_module # type: ignore[attr-defined] + + monkeypatch.setitem(sys.modules, "torch", torch_module) + monkeypatch.setitem(sys.modules, "nemo", nemo_module) + monkeypatch.setitem(sys.modules, "nemo.collections", collections_module) + monkeypatch.setitem(sys.modules, "nemo.collections.asr", asr_module) + monkeypatch.setitem(sys.modules, "nemo.collections.asr.models", models_module) + monkeypatch.setitem(sys.modules, "nemo.core", core_module) + monkeypatch.setitem(sys.modules, "nemo.core.connectors", connectors_module) + monkeypatch.setitem( + sys.modules, + "nemo.core.connectors.save_restore_connector", + save_restore_module, + ) + + returned_torch, returned_model = speech._load_nemotron35_model( + SimpleNamespace(device="cuda"), + archive, + ) + + assert returned_torch is torch_module + assert returned_model is restored + assert len(restore_calls) == 1 + assert restore_calls[0][:3] == (str(archive), device, False) + assert restored.eval_called is True + assert restored.device is device + assert restored.restore_state is False + + def test_plugin_reference_records_actual_command_during_inference( tmp_path: Path, monkeypatch, @@ -1878,6 +2056,7 @@ def test_plugin_reference_records_actual_command_during_inference( template = SimpleNamespace( name="template", inputs={}, + metadata={}, stages=[SimpleNamespace(name="full_inference")], ) @@ -1886,6 +2065,8 @@ def run_stage(self, case, stage, context): assert case.name == "series-1" assert stage.name == "full_inference" assert context.artifacts_dir.endswith("hf_artifacts") + assert case.metadata["reference_device"] == "cpu" + assert case.metadata["reference_device_map"] == "" return SimpleNamespace( data={ "output_field": [3.0, 4.0], @@ -1924,6 +2105,8 @@ def run_stage(self, case, stage, context): str(raw_output), "--repro-metadata", str(metadata), + "--device", + "cpu", ] ) diff --git a/tests/tools/test_trtmc_validate.py b/tests/tools/test_trtmc_validate.py index 0733133001..c08d6a31fb 100644 --- a/tests/tools/test_trtmc_validate.py +++ b/tests/tools/test_trtmc_validate.py @@ -10,6 +10,7 @@ from pathlib import Path import shlex import sys +from types import SimpleNamespace import pytest @@ -672,6 +673,134 @@ def test_per_model_hf_cache_hardlinks_seed_and_deletes_only_working_copy(tmp_pat assert blob.read_text(encoding="utf-8") == "weights" +def test_per_model_hf_cache_accepts_hub_cache_as_seed(tmp_path): + seed = tmp_path / "seed-hub" + blob = seed / "models--org--model/blobs/content" + blob.parent.mkdir(parents=True) + blob.write_text("weights", encoding="utf-8") + snapshot = seed / "models--org--model/snapshots/revision" + snapshot.mkdir(parents=True) + (snapshot / "model.bin").symlink_to("../../blobs/content") + arguments = trtmc_validate.build_parser().parse_args( + [ + "--all", + "--storage-root", + str(tmp_path), + "--output", + str(tmp_path / "results"), + "--model-work-dir", + str(tmp_path / "work"), + "--engine-dir", + str(tmp_path / "engines"), + "--engine-retention", + "delete_always", + "--hf-cache-mode", + "per_model", + "--hf-cache-retention", + "delete_always", + "--hf-cache-seed-dir", + str(seed), + "--reference-cache-dir", + str(tmp_path / "references"), + ] + ) + trtmc_validate._prepare_run_directories(arguments) + + selected, _, _ = trtmc_validate._binding_resource_arguments( + arguments, + trtmc_validate.Binding("model-a", "suite-a"), + ) + + linked_blob = selected.hf_cache_dir / "hub/models--org--model/blobs/content" + linked_snapshot = ( + selected.hf_cache_dir / "hub/models--org--model/snapshots/revision/model.bin" + ) + assert linked_blob.stat().st_ino == blob.stat().st_ino + assert linked_snapshot.resolve() == linked_blob + + +def test_prepare_hf_on_demand_uses_binding_cache_and_selected_model( + tmp_path, + monkeypatch, +): + arguments = trtmc_validate.build_parser().parse_args( + [ + "--all", + "--prepare-hf-on-demand", + "--hf-python", + sys.executable, + "--output", + str(tmp_path / "results"), + ] + ) + arguments.hf_cache_dir = tmp_path / "hf-home" + arguments.hf_cache_dir.mkdir() + calls = [] + + def run(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(trtmc_validate.subprocess, "run", run) + + trtmc_validate._prepare_hf_on_demand( + trtmc_validate.Binding("model-a", "suite-a"), + arguments, + ) + + command, kwargs = calls[0] + assert command[:2] == [sys.executable, str(trtmc_validate.HF_WARM_SCRIPT)] + assert command[command.index("--models-file") + 1].endswith( + "model-a/suite-a/hf_prepare.models.txt" + ) + assert "--strict" in command + assert "--fail-fast" in command + assert kwargs["env"]["HF_HOME"] == str(arguments.hf_cache_dir) + assert ( + tmp_path / "results/model-a/suite-a/hf_prepare.models.txt" + ).read_text(encoding="utf-8") == "model-a\n" + + +def test_worker_command_propagates_on_demand_hf_preparation(tmp_path): + arguments = trtmc_validate.build_parser().parse_args( + [ + "--all", + "--prepare-hf-on-demand", + "--output", + str(tmp_path / "results"), + ] + ) + + command = trtmc_validate._worker_command( + trtmc_validate.Binding("model-a", "suite-a"), + arguments, + ) + + assert "--prepare-hf-on-demand" in command + + +def test_worker_command_propagates_hf_reference_device(tmp_path): + arguments = trtmc_validate.build_parser().parse_args( + [ + "--all", + "--hf-device", + "cpu", + "--hf-device-map", + "balanced", + "--output", + str(tmp_path / "results"), + ] + ) + + command = trtmc_validate._worker_command( + trtmc_validate.Binding("model-a", "suite-a"), + arguments, + ) + + assert command[command.index("--hf-device") + 1] == "cpu" + assert command[command.index("--hf-device-map") + 1] == "balanced" + + def test_hf_cache_seed_requires_per_model_mode(tmp_path): seed = tmp_path / "seed" seed.mkdir() @@ -1970,6 +2099,7 @@ def prepare(source, cache_root): "revision": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc", "relative_path": "wan2_2_ti2v/reference/Wan2.2-42bf4cfaa384", "entrypoint": "wan/textimage2video.py", + "environment_variable": "TRTMC_WAN_REFERENCE_REPO", }, ) lance = trtmc_validate.ensure_reference_sources( @@ -1995,7 +2125,12 @@ def prepare(source, cache_root): ) assert common.elf_reference_repo is None assert common.environment == {"TRTMC_STORAGE_ROOT": str(tmp_path)} - assert wan22.environment == {"TRTMC_STORAGE_ROOT": str(tmp_path)} + assert wan22.environment == { + "TRTMC_STORAGE_ROOT": str(tmp_path), + "TRTMC_WAN_REFERENCE_REPO": str( + tmp_path / "wan2_2_ti2v/reference/Wan2.2-42bf4cfaa384" + ), + } assert lance.environment == { "TRTMC_STORAGE_ROOT": str(tmp_path), "TRTMC_LANCE_REFERENCE_REPO": str(tmp_path / "lance/reference/Lance-4baeee086648"), @@ -2189,6 +2324,7 @@ def test_write_report_links_each_comparison(tmp_path): "validation_passed": 1, "validation_failed": 0, "validation_skipped": 0, + "platform_excluded": 0, "selected_samples": 100, } assert report["validation_status"] == "passed" @@ -3379,6 +3515,43 @@ def fake_run(args, **kwargs): ] +def test_runtime_gpu_identity_falls_back_to_cuda_runtime(monkeypatch, tmp_path): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + monkeypatch.setattr( + trtmc_validate, + "_query_nvidia_smi_gpus", + lambda: (_ for _ in ()).throw( + trtmc_validate.ValidationError("nvidia-smi was not found") + ), + ) + monkeypatch.setattr( + trtmc_validate, + "_query_cuda_runtime_gpus", + lambda: [ + { + "cuda_runtime_index": 0, + "uuid": "GPU-7d4db97a-e132-503a-816d-a252f371ec4c", + "name": "Thor", + "pci_bus_id": "0000:00:00.0", + } + ], + ) + + path = trtmc_validate.write_run_metadata(tmp_path) + metadata = json.loads(path.read_text(encoding="utf-8")) + + assert metadata["gpu_identity_source"] == "cuda-runtime" + assert metadata["gpu_devices"] == [ + { + "cuda_logical_index": 0, + "cuda_runtime_index": 0, + "uuid": "GPU-7d4db97a-e132-503a-816d-a252f371ec4c", + "name": "Thor", + "pci_bus_id": "0000:00:00.0", + } + ] + + def test_run_metadata_rejects_missing_gpu_identity(monkeypatch, tmp_path): monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") monkeypatch.setattr( @@ -3646,6 +3819,8 @@ def test_comparison_command_uses_validation_entrypoint(tmp_path): backend_dir=None, model_plugin_dir=None, cuda_visible_devices="1", + hf_device="cpu", + hf_device_map="balanced", ) command = trtmc_validate._comparison_command( @@ -3679,6 +3854,8 @@ def test_comparison_command_uses_validation_entrypoint(tmp_path): assert "--force-hf" in command assert "--require-prebuilt-bundles" in command assert "--local-files-only" in command + assert command[command.index("--hf-device") + 1] == "cpu" + assert command[command.index("--hf-device-map") + 1] == "balanced" def test_comparison_command_passes_elf_reference_checkout(tmp_path): diff --git a/tests/tools/test_validation_engine.py b/tests/tools/test_validation_engine.py index 0fc95aad6b..ecd2054444 100644 --- a/tests/tools/test_validation_engine.py +++ b/tests/tools/test_validation_engine.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +import ast import hashlib import io import json @@ -27,6 +28,22 @@ from tools.validation import engine as validation_engine +def test_native_validation_contexts_preserve_runtime_library_path() -> None: + source = Path(validation_engine.__file__).read_text(encoding="utf-8") + tree = ast.parse(source) + missing: list[int] = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + if not isinstance(node.func, ast.Name) or node.func.id != "RunContext": + continue + keywords = {item.arg for item in node.keywords if item.arg} + if "binary_path" in keywords and "ld_library_path" not in keywords: + missing.append(node.lineno) + + assert missing == [] + + def _write_bundle_config(path: Path, config: dict[str, Any]) -> None: config_data = json.dumps(config).encode("utf-8") header_data = json.dumps( @@ -2139,6 +2156,70 @@ def test_prepare_model_plugin_dataset_resolves_nested_input_assets( assert manifest["request_count"] == 1 +def test_run_model_plugin_bundle_preserves_runtime_library_path( + tmp_path: Path, + monkeypatch, +) -> None: + from tests.e2e_harness.contracts import StageOutput, StageSpec + + work_dir = tmp_path / "work" + work_dir.mkdir() + model_manifest = tmp_path / "model.json" + model_manifest.write_text("{}\n", encoding="utf-8") + (work_dir / "manifest.json").write_text( + json.dumps( + { + "dataset_kind": "model_plugin_json", + "task_eval": {"model_manifest": str(model_manifest)}, + } + ), + encoding="utf-8", + ) + (work_dir / "prompts.jsonl").write_text( + json.dumps({"sample_id": "sample-1"}) + "\n", + encoding="utf-8", + ) + case = SimpleNamespace( + bundle="model.bundle", + task_strategy="custom_strategy", + metadata={ + "model_test_dir": "tests/e2e/models/custom", + "validation_manifest_case_name": "custom-case", + }, + ) + stage = StageSpec(name="full_generation", required=True) + seen: list[str] = [] + + class Runner: + def run_stage(self, _case, _stage, context): + seen.append(context.ld_library_path) + return StageOutput(stage_name="full_generation", text="answer") + + monkeypatch.setenv("LD_LIBRARY_PATH", "/runtime/trt11:/cuda") + monkeypatch.setattr( + validation_engine, + "select_case", + lambda *_args, **_kwargs: (case, stage), + ) + monkeypatch.setattr(validation_engine, "activate_model_plugins", lambda _path: None) + monkeypatch.setattr(validation_engine, "get_runner", lambda _strategy: Runner()) + + validation_engine.run_model_plugin_bundle( + argparse.Namespace( + work_dir=str(work_dir), + bundle=str(tmp_path / "model.bundle"), + trtmc_binary="/build/trtmc", + hf_python="/venv/bin/python", + model_plugin_dir="/build/models", + predictions="bundle_predictions.json", + raw_output="bundle_raw.jsonl", + log="bundle_run.log", + ) + ) + + assert seen == ["/runtime/trt11:/cuda"] + + def test_compare_model_plugin_prediction_sets_uses_model_comparator( tmp_path: Path, monkeypatch, @@ -5156,13 +5237,14 @@ def test_run_hf_reference_subprocess_uses_hf_python(tmp_path: Path, monkeypatch) ), encoding="utf-8", ) - captured: dict[str, list[str]] = {} + captured: dict[str, object] = {} class Result: returncode = 0 - def fake_run(cmd, **_kwargs): + def fake_run(cmd, **kwargs): captured["cmd"] = cmd + captured["env"] = kwargs["env"] return Result() monkeypatch.setattr(validation_engine.subprocess, "run", fake_run) @@ -5208,6 +5290,15 @@ def fake_run(cmd, **_kwargs): assert captured["cmd"][ captured["cmd"].index("--experts-implementation") + 1 ] == "batched_mm" + assert captured["env"]["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + + monkeypatch.setenv( + validation_engine.REFERENCE_CUDA_ALLOC_CONF_ENV, + "disabled", + ) + validation_engine.run_hf_reference_subprocess(args, model, work_dir) + + assert "PYTORCH_CUDA_ALLOC_CONF" not in captured["env"] @pytest.mark.parametrize( @@ -6239,11 +6330,19 @@ def test_run_diffusion_bundle_writes_image_artifact_predictions( json.dumps({"sample_id": "partiprompts_000000", "prompt": "a red cube"}) + "\n", encoding="utf-8", ) - seen: list[tuple[str, str, str]] = [] + monkeypatch.setenv("LD_LIBRARY_PATH", "/runtime/trt11:/cuda") + seen: list[tuple[str, str, str, str]] = [] class FakeRunner: def run_stage(self, case, _stage, ctx): - seen.append((case.inputs["prompt"], ctx.binary_path, case.bundle)) + seen.append( + ( + case.inputs["prompt"], + ctx.binary_path, + case.bundle, + ctx.ld_library_path, + ) + ) frames_dir = work_dir / "fake_trt_frames" / case.name frames_dir.mkdir(parents=True) (frames_dir / "frame_0000.png").write_bytes(b"png") @@ -6294,7 +6393,14 @@ def run_stage(self, case, _stage, ctx): )) predictions = json.loads((work_dir / "bundle_predictions.json").read_text(encoding="utf-8")) - assert seen == [("a red cube", "build/trtmc", "flux-schnell-l0.bundle")] + assert seen == [ + ( + "a red cube", + "build/trtmc", + "flux-schnell-l0.bundle", + "/runtime/trt11:/cuda", + ) + ] assert predictions["responses"][0]["source"] == "bundle" assert predictions["responses"][0]["num_frames"] == 1 assert (work_dir / "bundle_run.log").read_text(encoding="utf-8") == ( diff --git a/tools/lance_reference.py b/tools/lance_reference.py index 1c6aa5ef7b..8d56b543ee 100644 --- a/tools/lance_reference.py +++ b/tools/lance_reference.py @@ -21,6 +21,14 @@ from typing import Any, Iterator, Sequence +_ATTENTION_BACKEND_ENV = "TRTMC_LANCE_REFERENCE_ATTENTION_BACKEND" +_ATTENTION_BACKENDS = frozenset({"flash_attn", "torch_sdpa"}) +_ATTENTION_COMPAT = ( + Path(__file__).resolve().parents[1] + / "tests/e2e/models/lance/e2e_plugins/references/lance_image_attention_compat" +) + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--reference-repo", required=True, type=Path) @@ -176,7 +184,21 @@ def unsupported_cpu(*_args, **_kwargs): return {"decord": decord, "decord.video_reader": video_reader} +def _configure_attention_backend() -> None: + backend = os.environ.get(_ATTENTION_BACKEND_ENV, "flash_attn").strip() + if backend not in _ATTENTION_BACKENDS: + raise RuntimeError( + "unsupported Lance reference attention backend " + f"{backend!r}; expected one of {sorted(_ATTENTION_BACKENDS)}" + ) + if backend == "torch_sdpa": + compat = str(_ATTENTION_COMPAT) + if compat not in sys.path: + sys.path.insert(0, compat) + + def _load_upstream(reference_repo: Path) -> Any: + _configure_attention_backend() try: import decord # noqa: F401 except ImportError: diff --git a/tools/model_checks.py b/tools/model_checks.py index 4b78b9b506..5fa61222b8 100644 --- a/tools/model_checks.py +++ b/tools/model_checks.py @@ -53,6 +53,16 @@ DEFAULT_PERF_SUITE = REPOSITORY / "benchmarks" / "performance" / "release.yaml" TASKS = ("accuracy", "perf") RUN_ID_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") +EXACT_GIT_REVISION_PATTERN = re.compile(r"[0-9a-fA-F]{40}") +ENVIRONMENT_VARIABLE_PATTERN = re.compile(r"TRTMC_[A-Z0-9_]+") +MANAGED_TASK_ENVIRONMENT_VARIABLES = frozenset( + { + PROFILE_ROOT_ENV, + PREBUILT_ONLY_ENV, + "TRTMC_PERF_SOURCE_REVISION", + "TRTMC_VALIDATION_SOURCE_REVISION", + } +) class ModelCheckError(ValueError): @@ -165,70 +175,35 @@ def load_platform(value: str) -> dict[str, Any]: not isinstance(required_device, str) or not Path(required_device).is_absolute() ): raise ModelCheckError(f"platform storage.device must be an absolute path: {path}") - unsupported = profile.get("unsupported", []) - if not isinstance(unsupported, list): - raise ModelCheckError(f"platform unsupported must be a list: {path}") - for index, item in enumerate(unsupported): - if not isinstance(item, Mapping): - raise ModelCheckError(f"platform unsupported[{index}] must be an object") - if not all( - isinstance(item.get(key), str) and item[key] for key in ("model", "task", "reason") - ): - raise ModelCheckError(f"platform unsupported[{index}] needs model, task, and reason") - if item["task"] not in TASKS: - raise ModelCheckError( - f"platform unsupported[{index}] has unknown task {item['task']!r}" - ) - binding = item.get("binding") - if binding is not None and (not isinstance(binding, str) or not binding.strip()): - raise ModelCheckError( - f"platform unsupported[{index}].binding must be a non-empty string" - ) + legacy_exclusions = sorted({"unsupported", "excluded"}.intersection(profile)) + if legacy_exclusions: + raise ModelCheckError( + "platform binding-level exclusions are no longer supported; use " + f"excluded_models instead of {', '.join(legacy_exclusions)}: {path}" + ) + excluded_models = profile.get("excluded_models", []) + if not isinstance(excluded_models, list) or any( + not isinstance(model, str) or not model.strip() for model in excluded_models + ): + raise ModelCheckError(f"platform excluded_models must be a list of models: {path}") + if len(excluded_models) != len(set(excluded_models)): + raise ModelCheckError(f"platform excluded_models contains duplicates: {path}") return {**profile, "source": str(path)} -def audit_platform_unsupported( +def audit_platform_exclusions( platform: Mapping[str, Any], *, accuracy_catalog: Mapping[str, Any], perf_cases: Sequence[Mapping[str, Any]], ) -> None: - accuracy_models = accuracy_catalog["models"] - perf_bindings: dict[str, set[str]] = {} - for case in perf_cases: - perf_bindings.setdefault(str(case["model"]), set()).add(str(case["id"])) - seen: set[tuple[str, str, str]] = set() - for index, item in enumerate(platform.get("unsupported", [])): - model = str(item["model"]) - task = str(item["task"]) - binding = str(item.get("binding", "") or "") - key = (model, task, binding) - if key in seen: - raise ModelCheckError( - f"platform unsupported[{index}] duplicates {model}/{task}/{binding or '*'}" - ) - seen.add(key) - if task == "accuracy": - spec = accuracy_models.get(model) - if not isinstance(spec, Mapping): - raise ModelCheckError( - f"platform unsupported[{index}] names unknown Accuracy model {model}" - ) - if binding and binding not in accuracy_catalog["sample_limits"]: - raise ModelCheckError( - f"platform unsupported[{index}] names unknown Accuracy binding " - f"{model}={binding}" - ) - else: - entries = perf_bindings.get(model, set()) - if not entries: - raise ModelCheckError( - f"platform unsupported[{index}] names unknown Perf model {model}" - ) - if binding and binding not in entries: - raise ModelCheckError( - f"platform unsupported[{index}] names unknown Perf binding {model}={binding}" - ) + known_models = set(accuracy_catalog["models"]) + known_models.update(str(case["model"]) for case in perf_cases) + unknown = sorted(set(platform.get("excluded_models", [])) - known_models) + if unknown: + raise ModelCheckError( + "platform excluded_models names unknown models: " + ", ".join(unknown) + ) def _environment_path(value: str) -> Path: @@ -283,6 +258,38 @@ def load_execution_environment(value: str, *, platform_id: str) -> dict[str, Any raise ModelCheckError( f"environment platform {environment.get('id')!r} does not match {platform_id!r}" ) + library_dirs = environment.get("library_dirs", []) + if not isinstance(library_dirs, list) or any( + not isinstance(path, str) or not path for path in library_dirs + ): + raise ModelCheckError("model-check environment library_dirs must be a list of paths") + executable_dirs = environment.get("executable_dirs", []) + if not isinstance(executable_dirs, list) or any( + not isinstance(path, str) or not path for path in executable_dirs + ): + raise ModelCheckError( + "model-check environment executable_dirs must be a list of paths" + ) + python_dirs = environment.get("python_dirs", []) + if not isinstance(python_dirs, list) or any( + not isinstance(path, str) or not path for path in python_dirs + ): + raise ModelCheckError( + "model-check environment python_dirs must be a list of paths" + ) + environment_variables = environment.get("environment_variables", {}) + if not isinstance(environment_variables, Mapping) or any( + not isinstance(name, str) + or ENVIRONMENT_VARIABLE_PATTERN.fullmatch(name) is None + or name in MANAGED_TASK_ENVIRONMENT_VARIABLES + or not isinstance(value, str) + or not value + for name, value in environment_variables.items() + ): + raise ModelCheckError( + "model-check environment environment_variables must map unmanaged " + "TRTMC_* names to non-empty strings" + ) storage = environment.get("storage") if not isinstance(storage, Mapping): raise ModelCheckError("model-check environment needs storage settings") @@ -342,6 +349,10 @@ def load_execution_environment(value: str, *, platform_id: str) -> dict[str, Any return { **environment, "source": str(path), + "library_dirs": [str(_repo_path(path)) for path in library_dirs], + "executable_dirs": [str(_repo_path(path)) for path in executable_dirs], + "python_dirs": [str(_repo_path(path)) for path in python_dirs], + "environment_variables": dict(environment_variables), "storage": { **storage, "root": str(storage_root), @@ -373,11 +384,61 @@ def load_execution_environment(value: str, *, platform_id: str) -> dict[str, Any def _task_environment( environment: Mapping[str, Any], overrides: Mapping[str, str] | None = None, + *, + source_revision: str = "", ) -> dict[str, str]: """Use a managed shared profile cache and create missing profiles on demand.""" child = os.environ.copy() child[PROFILE_ROOT_ENV] = str(environment["storage"]["python_profiles_root"]) child.pop(PREBUILT_ONLY_ENV, None) + library_dirs = [str(path) for path in environment.get("library_dirs", [])] + missing_library_dirs = [path for path in library_dirs if not Path(path).is_dir()] + if missing_library_dirs: + raise ModelCheckError( + "model-check runtime library directory does not exist: " + + ", ".join(missing_library_dirs) + ) + inherited_library_path = child.get("LD_LIBRARY_PATH", "") + if inherited_library_path: + library_dirs.append(inherited_library_path) + if library_dirs: + child["LD_LIBRARY_PATH"] = os.pathsep.join(library_dirs) + executable_dirs = [str(path) for path in environment.get("executable_dirs", [])] + missing_executable_dirs = [ + path for path in executable_dirs if not Path(path).is_dir() + ] + if missing_executable_dirs: + raise ModelCheckError( + "model-check executable directory does not exist: " + + ", ".join(missing_executable_dirs) + ) + inherited_path = child.get("PATH", "") + if inherited_path: + executable_dirs.append(inherited_path) + if executable_dirs: + child["PATH"] = os.pathsep.join(executable_dirs) + python_dirs = [str(path) for path in environment.get("python_dirs", [])] + missing_python_dirs = [path for path in python_dirs if not Path(path).is_dir()] + if missing_python_dirs: + raise ModelCheckError( + "model-check Python runtime directory does not exist: " + + ", ".join(missing_python_dirs) + ) + inherited_python_path = child.get("PYTHONPATH", "") + if inherited_python_path: + python_dirs.append(inherited_python_path) + if python_dirs: + child["PYTHONPATH"] = os.pathsep.join(python_dirs) + child.update( + { + str(name): str(value) + for name, value in environment.get("environment_variables", {}).items() + } + ) + if EXACT_GIT_REVISION_PATTERN.fullmatch(source_revision): + revision = source_revision.lower() + child["TRTMC_VALIDATION_SOURCE_REVISION"] = revision + child["TRTMC_PERF_SOURCE_REVISION"] = revision child.update(overrides or {}) return child @@ -421,20 +482,19 @@ def model_profiles_for_owners( return model_selection.normalize_models(profiles) -def _unsupported_reason( - profile: Mapping[str, Any], - *, - model: str, - task: str, - binding: str, -) -> str: - for item in profile.get("unsupported", []): - if item["model"] != model or item["task"] != task: - continue - selected_binding = str(item.get("binding", "") or "") - if not selected_binding or selected_binding == binding: - return str(item["reason"]) - return "" +def _platform_exclusion_reason(profile: Mapping[str, Any], model: str) -> str: + if model not in profile.get("excluded_models", []): + return "" + return f"Model is excluded from platform {profile['id']}" + + +def _projection_status(bindings: Sequence[Mapping[str, Any]]) -> str: + statuses = {str(binding["status"]) for binding in bindings} + if len(statuses) == 1: + status = next(iter(statuses)) + if status == "excluded": + return status + return "configured" def _accuracy_projection( @@ -466,26 +526,18 @@ def _accuracy_projection( projected = [] for binding in bindings: binding_id = binding.workload or "not-compared" - reason = _unsupported_reason( - platform, - model=model, - task="accuracy", - binding=binding_id, - ) + reason = _platform_exclusion_reason(platform, model) + status = "excluded" if reason else "configured" projected.append( { "id": f"accuracy:{model}:{binding_id}", "model": model, "workload": binding.workload, - "status": "unsupported" if reason else "configured", + "status": status, **({"reason": reason} if reason else {}), } ) - status = ( - "unsupported" - if projected and all(item["status"] == "unsupported" for item in projected) - else "configured" - ) + status = _projection_status(projected) return {"status": status, "bindings": projected} @@ -521,26 +573,18 @@ def _perf_projection( projected = [] for case in matched: entry_id = str(case["id"]) - reason = _unsupported_reason( - platform, - model=model, - task="perf", - binding=entry_id, - ) + reason = _platform_exclusion_reason(platform, model) + status = "excluded" if reason else "configured" projected.append( { "id": f"perf:{model}:{entry_id}", "model": model, "entry": entry_id, - "status": "unsupported" if reason else "configured", + "status": status, **({"reason": reason} if reason else {}), } ) - status = ( - "unsupported" - if projected and all(item["status"] == "unsupported" for item in projected) - else "configured" - ) + status = _projection_status(projected) return {"status": status, "bindings": projected} @@ -598,8 +642,8 @@ def resolve_plan( "configured_binding_count": sum( binding["status"] == "configured" for binding in bindings ), - "unsupported_binding_count": sum( - binding["status"] == "unsupported" for binding in bindings + "excluded_binding_count": sum( + binding["status"] == "excluded" for binding in bindings ), "blocker_count": blocker_count, }, @@ -631,7 +675,7 @@ def _render(plan: Mapping[str, Any]) -> str: "", f"Summary: {summary['model_count']} models, " f"{summary['binding_count']} bindings, " - f"{summary['unsupported_binding_count']} unsupported, " + f"{summary['excluded_binding_count']} excluded, " f"{summary['blocker_count']} blockers", ] ) @@ -691,7 +735,7 @@ def _resolve_request( perf_cases = perf_matrix._cases(perf_suite) perf_exclusions = perf_matrix._excluded_profiles(perf_suite) perf_matrix._validate_coverage(perf_cases, perf_exclusions) - audit_platform_unsupported( + audit_platform_exclusions( platform, accuracy_catalog=accuracy_catalog, perf_cases=perf_cases, @@ -864,6 +908,75 @@ def _task_bindings(plan: Mapping[str, Any], task: str) -> list[dict[str, Any]]: ] +def _platform_task_bindings( + plan: Mapping[str, Any], + task: str, +) -> list[dict[str, Any]]: + return [ + binding + for model in plan["models"] + for binding in model["tasks"].get(task, {}).get("bindings", []) + if binding["status"] == "excluded" + ] + + +def _accuracy_platform_result(binding: Mapping[str, Any]) -> dict[str, Any]: + reason = str(binding["reason"]) + return { + "schema_version": "trtmc.validation-result/v2", + "model": str(binding["model"]), + "workload": str(binding["workload"]), + "executor": "platform_exclusion", + "execution": {"status": "not_run", "exit_code": None}, + "comparison": { + "status": "not_run", + "mode": "platform_exclusion", + "primary_metric": None, + "metrics": {}, + "failures": [], + }, + "validation": {"status": "not_compared"}, + "not_compared_reason": reason, + "platform_exclusion": {"reason": reason}, + "reference_environment": [], + "reproduce": {}, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + + +def _write_accuracy_platform_results( + plan: Mapping[str, Any], + output: Path, +) -> list[Path]: + written: list[Path] = [] + for binding in _platform_task_bindings(plan, "accuracy"): + workload = str(binding.get("workload", "") or "") + if not workload: + raise ModelCheckError( + f"platform Accuracy binding has no workload: {binding['model']}" + ) + comparison = output / str(binding["model"]) / workload / "comparison.json" + if comparison.exists(): + try: + previous = json.loads(comparison.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ModelCheckError( + f"cannot read platform result {comparison}: {exc}" + ) from exc + if not isinstance(previous.get("platform_exclusion"), Mapping): + raise ModelCheckError( + "refusing to replace a previously executed Accuracy result with a " + f"platform exclusion: {binding['model']}={workload}" + ) + comparison.parent.mkdir(parents=True, exist_ok=True) + comparison.write_text( + json.dumps(_accuracy_platform_result(binding), indent=2) + "\n", + encoding="utf-8", + ) + written.append(comparison) + return written + + def _selected_perf_reference_contracts( plan: Mapping[str, Any], models_dir: Path, @@ -947,7 +1060,9 @@ def _write_selected_models(plan: Mapping[str, Any], run_root: Path) -> Path: { str(binding["model"]) for task in TASKS - for binding in _task_bindings(plan, task) + for model in plan["models"] + for binding in model["tasks"].get(task, {}).get("bindings", []) + if binding["status"] in {"configured", "excluded"} } ) selection = run_root / "selected-models.txt" @@ -1249,6 +1364,13 @@ def _run(arguments: argparse.Namespace) -> int: if arguments.dry_run: return 0 + accuracy_platform_results = _write_accuracy_platform_results( + plan, + run_root / "accuracy", + ) + if accuracy_platform_results: + trtmc_validate.write_report(run_root / "accuracy") + reference_environment: dict[str, str] = {} reference_contracts = _selected_perf_reference_contracts(plan, arguments.models_dir) reference_environment.update( @@ -1265,6 +1387,8 @@ def _run(arguments: argparse.Namespace) -> int: ) task_results: dict[str, int] = {} + if accuracy_platform_results and not _task_bindings(plan, "accuracy"): + task_results["accuracy"] = 0 runnable = [(task, command) for task, command in execution_commands if command is not None] for index, (task, command) in enumerate(runnable, start=1): label = _task_label(task) @@ -1275,7 +1399,11 @@ def _run(arguments: argparse.Namespace) -> int: _detailed_command(command, verbose=arguments.verbose), cwd=REPOSITORY, check=False, - env=_task_environment(environment, reference_environment), + env=_task_environment( + environment, + reference_environment, + source_revision=arguments.revision, + ), ) task_results[task] = completed.returncode status = "PASSED" if completed.returncode == 0 else "FAILED" diff --git a/tools/reference/plugin_reference.py b/tools/reference/plugin_reference.py index b432649b03..cfc6ff7457 100644 --- a/tools/reference/plugin_reference.py +++ b/tools/reference/plugin_reference.py @@ -654,9 +654,19 @@ def _diffusion_response(case: Any, output: Any, prompt: str) -> dict[str, Any]: if image_path is not None and image_path.is_file(): response["condition_image_sha256"] = _sha256_file(image_path) if response["returncode"] != 0 or response["num_frames"] < 1: + detail = "" + for stream_name in ("stderr", "stdout"): + stream = data.get(stream_name, "") + if isinstance(stream, bytes): + stream = stream.decode(errors="replace") + stream = str(stream).strip() + if stream: + detail = f"\nNested reference {stream_name} (tail):\n{stream[-8000:]}" + break raise RuntimeError( f"HF diffusion reference failed for {case.name}: " f"returncode={response['returncode']} frames={response['num_frames']}" + f"{detail}" ) return response @@ -831,6 +841,8 @@ def run(arguments: argparse.Namespace) -> None: arguments.sample_id, ) template, reference = _load_reference_plugin(manifest) + template.metadata["reference_device"] = arguments.device + template.metadata["reference_device_map"] = arguments.device_map artifacts_dir = arguments.predictions.parent / "hf_artifacts" artifacts_dir.mkdir(parents=True, exist_ok=True) command_path = arguments.predictions.parent / "hf_native_commands.jsonl" diff --git a/tools/reference/speech.py b/tools/reference/speech.py index 18ada922ed..bde7566896 100644 --- a/tools/reference/speech.py +++ b/tools/reference/speech.py @@ -28,6 +28,12 @@ "https://huggingface.co/Edresson/Speaker_Encoder_H_ASP/resolve/main/" "pytorch_model.bin" ) +NEMOTRON35_OPTIONAL_CTC_STATE_KEYS = frozenset( + { + "ctc_decoder.decoder_layers.0.bias", + "ctc_decoder.decoder_layers.0.weight", + } +) def _load_json(path: Path) -> dict[str, Any]: @@ -361,53 +367,83 @@ def _run_nemotron35_asr( target_rate = int(generation.get("sample_rate", 16000) or 16000) output_dir = arguments.predictions.parent / "hf_canary_audio" responses = [] - for prompt in prompts: - audio, _source = _audio_for_prompt(prompt, target_rate) - sample_id = str(prompt.get("sample_id", "")) - wav_path = output_dir / _safe_sample_filename(sample_id, ".wav") - _write_wav_pcm16(wav_path, audio, target_rate) - manifest_path = output_dir / _safe_sample_filename( - sample_id, - ".manifest.jsonl", - ) - language = str( - prompt.get("language") - or generation.get("language") - or "auto" + original_forward = model.forward + + def forward_with_extended_prompt(*args: Any, **kwargs: Any) -> Any: + args, kwargs = _extend_nemotron35_prompt_for_forward( + torch, + args, + kwargs, ) - manifest_path.write_text( - json.dumps( - { - "audio_filepath": str(wav_path), - "duration": len(audio) / target_rate, - "text": "", - "lang": language, - }, - ensure_ascii=False, + return original_forward(*args, **kwargs) + + model.forward = forward_with_extended_prompt + try: + for prompt in prompts: + audio, _source = _audio_for_prompt(prompt, target_rate) + sample_id = str(prompt.get("sample_id", "")) + wav_path = output_dir / _safe_sample_filename(sample_id, ".wav") + _write_wav_pcm16(wav_path, audio, target_rate) + manifest_path = output_dir / _safe_sample_filename( + sample_id, + ".manifest.jsonl", ) - + "\n", - encoding="utf-8", - ) - started = time.perf_counter() - transcriptions = model.transcribe( - str(manifest_path), - batch_size=1, - verbose=False, - ) - responses.append( - _asr_row( - prompt, - _transcription_text(transcriptions), - (time.perf_counter() - started) * 1000.0, + language = str( + prompt.get("language") + or generation.get("language") + or "auto" ) - ) - del model - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + manifest_path.write_text( + json.dumps( + { + "audio_filepath": str(wav_path), + "duration": len(audio) / target_rate, + "text": "", + "lang": language, + }, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + started = time.perf_counter() + transcriptions = model.transcribe( + str(manifest_path), + batch_size=1, + verbose=False, + ) + responses.append( + _asr_row( + prompt, + _transcription_text(transcriptions), + (time.perf_counter() - started) * 1000.0, + ) + ) + finally: + model.forward = original_forward + del model + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() return responses +def _extend_nemotron35_prompt_for_forward( + torch_module: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> tuple[tuple[Any, ...], dict[str, Any]]: + prompt = kwargs.get("prompt") + if prompt is None or prompt.shape[1] == 0: + return args, kwargs + updated_kwargs = dict(kwargs) + updated_kwargs["prompt"] = torch_module.cat( + (prompt, prompt[:, -1:, :]), + dim=1, + ) + return args, updated_kwargs + + def _resolve_nemo_archive( arguments: argparse.Namespace, *, @@ -453,12 +489,37 @@ def _load_nemotron35_model( archive: Path, ) -> tuple[Any, Any]: import torch - from nemo.collections.asr.models import EncDecRNNTBPEModelWithPrompt + from nemo.collections.asr.models import ( + EncDecHybridRNNTCTCBPEModelWithPrompt, + ) + from nemo.core.connectors.save_restore_connector import SaveRestoreConnector + + class Nemotron35SaveRestoreConnector(SaveRestoreConnector): + def load_instance_with_state_dict( + self, + instance: Any, + state_dict: Mapping[str, Any], + strict: bool, + ) -> None: + del strict + incompatible = instance.load_state_dict(state_dict, strict=False) + missing = frozenset(incompatible.missing_keys) + unexpected = frozenset(incompatible.unexpected_keys) + allowed_missing = NEMOTRON35_OPTIONAL_CTC_STATE_KEYS + if missing not in (frozenset(), allowed_missing) or unexpected: + raise RuntimeError( + "Nemotron 3.5 ASR archive state_dict mismatch: " + f"missing={sorted(missing)}, " + f"unexpected={sorted(unexpected)}" + ) + instance._set_model_restore_state(is_being_restored=False) device = torch.device(arguments.device) - model = EncDecRNNTBPEModelWithPrompt.restore_from( + model = EncDecHybridRNNTCTCBPEModelWithPrompt.restore_from( str(archive), map_location=device, + strict=False, + save_restore_connector=Nemotron35SaveRestoreConnector(), ) model.eval() if hasattr(model, "to"): diff --git a/tools/reference/transformers_text.py b/tools/reference/transformers_text.py index cdb07b7a4b..861796f7ac 100644 --- a/tools/reference/transformers_text.py +++ b/tools/reference/transformers_text.py @@ -16,6 +16,15 @@ SCHEMA_VERSION = "trtmc.native-reference-reproduction/v1" +_INTERNLM_TOKENIZER_REPO_ID = "internlm/internlm2-step-prover" +_INTERNLM_TOKENIZER_REVISION = "6c727046190546168bf3aba9a1d78d5fb325ff14" +_INTERNLM_TOKENIZER_FILENAME = "tokenizer.json" +_INTERNLM_TOKENIZER_SHA256 = ( + "1193d3a1aa3d9f74866287ca3c1f7bf64fe54dd6ecf015e751f13ebce509e411" +) +_INTERNLM_SOURCE_MODEL_SHA256 = ( + "f868398fc4e05ee1e8aeba95ddf18ddcc45b8bce55d5093bead5bbf80429b48b" +) def _load_json(path: Path) -> dict[str, Any]: @@ -38,6 +47,126 @@ def _load_jsonl(path: Path) -> list[dict[str, Any]]: return rows +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _resolve_cached_model_ref( + model_id: str, + *, + revision: str, + local_files_only: bool, +) -> Path: + model_path = Path(model_id) + if model_path.exists(): + return model_path + + try: + from huggingface_hub import snapshot_download + + kwargs: dict[str, Any] = { + "repo_id": model_id, + "local_files_only": local_files_only, + "allow_patterns": ["tokenizer.model", "tokenizer_config.json"], + } + if revision: + kwargs["revision"] = revision + return Path(snapshot_download(**kwargs)) + except Exception as exc: + raise RuntimeError( + f"could not resolve tokenizer files for {model_id!r} from the " + "Hugging Face cache" + ) from exc + + +def _resolve_internlm_reference_tokenizer_json( + model_dir: Path, + *, + local_files_only: bool, +) -> Path | None: + """Resolve InternLM's official fast tokenizer independently from the DUT.""" + source_model = model_dir / "tokenizer.model" + if ( + not source_model.is_file() + or _sha256_file(source_model) != _INTERNLM_SOURCE_MODEL_SHA256 + ): + return None + + try: + from huggingface_hub import hf_hub_download + + tokenizer_path = Path( + hf_hub_download( + repo_id=_INTERNLM_TOKENIZER_REPO_ID, + filename=_INTERNLM_TOKENIZER_FILENAME, + revision=_INTERNLM_TOKENIZER_REVISION, + local_files_only=local_files_only, + ) + ) + except Exception as exc: + raise RuntimeError( + "pinned official InternLM reference tokenizer is unavailable in " + f"the local Hugging Face cache: {_INTERNLM_TOKENIZER_REPO_ID}@" + f"{_INTERNLM_TOKENIZER_REVISION}" + ) from exc + + actual_sha256 = _sha256_file(tokenizer_path) + if actual_sha256 != _INTERNLM_TOKENIZER_SHA256: + raise RuntimeError( + "pinned official InternLM reference tokenizer SHA256 mismatch: " + f"expected {_INTERNLM_TOKENIZER_SHA256}, got {actual_sha256}: " + f"{tokenizer_path}" + ) + return tokenizer_path + + +def _load_tokenizer( + arguments: argparse.Namespace, + transformers_module: Any, +) -> Any: + tokenizer_kwargs = { + "trust_remote_code": arguments.trust_remote_code, + "local_files_only": arguments.local_files_only, + } + if arguments.model_revision: + tokenizer_kwargs["revision"] = arguments.model_revision + + if arguments.family == "internlm": + model_ref = _resolve_cached_model_ref( + arguments.model, + revision=arguments.model_revision, + local_files_only=arguments.local_files_only, + ) + tokenizer_path = _resolve_internlm_reference_tokenizer_json( + model_ref, + local_files_only=arguments.local_files_only, + ) + if tokenizer_path is not None: + tokenizer_config = _load_json(model_ref / "tokenizer_config.json") + return transformers_module.PreTrainedTokenizerFast( + tokenizer_file=str(tokenizer_path), + model_input_names=["input_ids", "attention_mask"], + bos_token=tokenizer_config.get("bos_token"), + eos_token=tokenizer_config.get("eos_token"), + unk_token=tokenizer_config.get("unk_token"), + pad_token=tokenizer_config.get("pad_token"), + chat_template=tokenizer_config.get("chat_template"), + clean_up_tokenization_spaces=tokenizer_config.get( + "clean_up_tokenization_spaces", + False, + ), + ) + + return transformers_module.AutoTokenizer.from_pretrained( + arguments.model, + **tokenizer_kwargs, + ) + + def _generation_value( arguments: argparse.Namespace, generation: Mapping[str, Any], @@ -144,6 +273,8 @@ def _reproduction_command(arguments: argparse.Namespace) -> list[str]: ] if arguments.reference_family: command.extend(["--reference-family", arguments.reference_family]) + if arguments.family: + command.extend(["--family", arguments.family]) for flag, value in ( ("--model-revision", arguments.model_revision), ("--device-map", arguments.device_map), @@ -240,16 +371,7 @@ def _load_runtime( transformers_module: Any, ) -> tuple[Any, Any, Any, bool]: transformers_module.logging.set_verbosity_error() - tokenizer_kwargs = { - "trust_remote_code": arguments.trust_remote_code, - "local_files_only": arguments.local_files_only, - } - if arguments.model_revision: - tokenizer_kwargs["revision"] = arguments.model_revision - tokenizer = transformers_module.AutoTokenizer.from_pretrained( - arguments.model, - **tokenizer_kwargs, - ) + tokenizer = _load_tokenizer(arguments, transformers_module) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token model_dtype = _model_dtype(torch_module, arguments.dtype) @@ -460,6 +582,7 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--model", required=True) parser.add_argument("--model-revision", default="") + parser.add_argument("--family", default="") parser.add_argument("--reference-family", default="") parser.add_argument("--prompts", type=Path, required=True) parser.add_argument("--answers", type=Path, required=True) diff --git a/tools/trtmc_reference.py b/tools/trtmc_reference.py index b115df7e8c..8e27a1ab39 100644 --- a/tools/trtmc_reference.py +++ b/tools/trtmc_reference.py @@ -499,7 +499,7 @@ def _native_reference_command( command.extend(["--model-revision", str(args.model_revision)]) if runner == _ELF_PREPARED_RUNNER and args.elf_reference_repo: command.extend(["--elf-reference-repo", str(args.elf_reference_repo)]) - if runner == _SPEECH_REFERENCE_RUNNER and args.family: + if runner in {_SPEECH_REFERENCE_RUNNER, _TRANSFORMERS_TEXT_RUNNER} and args.family: command.extend(["--family", str(args.family)]) if runner == _TRANSFORMERS_TEXT_RUNNER and args.experts_implementation: command.extend( diff --git a/tools/trtmc_validate.py b/tools/trtmc_validate.py index 2c929520aa..56ef8a8c2c 100644 --- a/tools/trtmc_validate.py +++ b/tools/trtmc_validate.py @@ -28,6 +28,7 @@ import tempfile import time from typing import Any, Iterable, Mapping, Sequence +import uuid as uuidlib import yaml @@ -68,6 +69,7 @@ DEFAULT_OUTPUT = REPO_ROOT / "artifacts" / "trtmc-validate" DEFAULT_ENGINE_DIR = DEFAULT_OUTPUT / "engines" DEFAULT_REFERENCE_CACHE = DEFAULT_OUTPUT / "references" +HF_WARM_SCRIPT = REPO_ROOT / "scripts" / "warm_hf_cache.py" COMMON_REFERENCE_PROFILE = "reference_common" NOT_COMPARED_DIRECTORY = "not-compared" DEFAULT_REUSED_BUNDLE_REVALIDATION_LIMIT = 1 @@ -876,6 +878,9 @@ def _comparison_command( command.extend(["--model-plugin-dir", str(arguments.model_plugin_dir)]) if arguments.cuda_visible_devices: command.extend(["--cuda-visible-devices", arguments.cuda_visible_devices]) + command.extend(["--hf-device", getattr(arguments, "hf_device", "cuda")]) + if getattr(arguments, "hf_device_map", ""): + command.extend(["--hf-device-map", arguments.hf_device_map]) if reference_sources and reference_sources.elf_reference_repo: command.extend( [ @@ -2071,6 +2076,61 @@ def _query_nvidia_smi_gpus() -> list[dict[str, Any]]: return inventory +def _query_cuda_runtime_gpus() -> list[dict[str, Any]]: + """Query identity through CUDA on platforms that do not ship nvidia-smi.""" + try: + from cuda.bindings import runtime as cudart + except ImportError as exc: + raise ValidationError(f"could not import CUDA runtime bindings: {exc}") from exc + + try: + status, count = cudart.cudaGetDeviceCount() + if int(status) != 0: + raise ValidationError(f"cudaGetDeviceCount failed with status {status}") + inventory: list[dict[str, Any]] = [] + for index in range(int(count)): + status, properties = cudart.cudaGetDeviceProperties(index) + if int(status) != 0: + raise ValidationError( + f"cudaGetDeviceProperties({index}) failed with status {status}" + ) + status, pci_bus_id = cudart.cudaDeviceGetPCIBusId(32, index) + if int(status) != 0: + raise ValidationError( + f"cudaDeviceGetPCIBusId({index}) failed with status {status}" + ) + name = properties.name + if isinstance(name, bytes): + name = name.decode("utf-8", errors="replace").rstrip("\x00") + raw_uuid = bytes(properties.uuid.bytes) + try: + gpu_uuid = f"GPU-{uuidlib.UUID(bytes=raw_uuid)}" + except (TypeError, ValueError) as exc: + raise ValidationError( + f"CUDA runtime returned an invalid UUID for device {index}" + ) from exc + if isinstance(pci_bus_id, bytes): + pci_bus_id = pci_bus_id.decode("utf-8", errors="replace") + pci_bus_id = str(pci_bus_id).rstrip("\x00 ").strip() + if not str(name).strip() or not pci_bus_id: + raise ValidationError( + f"CUDA runtime returned incomplete GPU identity for device {index}" + ) + inventory.append( + { + "cuda_runtime_index": index, + "uuid": gpu_uuid, + "name": str(name).strip(), + "pci_bus_id": pci_bus_id, + } + ) + return inventory + except ValidationError: + raise + except Exception as exc: + raise ValidationError(f"could not query CUDA runtime GPU identity: {exc}") from exc + + def _resolve_cuda_devices( cuda_visible_devices: str, inventory: Sequence[Mapping[str, Any]], @@ -2109,11 +2169,63 @@ def _resolve_cuda_devices( return resolved -def _runtime_gpu_devices(cuda_visible_devices: str) -> list[dict[str, Any]]: - return _resolve_cuda_devices( - cuda_visible_devices, - _query_nvidia_smi_gpus(), +def _resolve_cuda_runtime_devices( + cuda_visible_devices: str, + inventory: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + visible = cuda_visible_devices.strip() + if visible.lower() in {"-1", "none", "void"}: + return [] + selectors = ( + [str(device["cuda_runtime_index"]) for device in inventory] + if not visible or visible.lower() == "all" + else [selector.strip() for selector in visible.split(",")] ) + if not selectors or any(not selector for selector in selectors): + raise ValidationError("CUDA_VISIBLE_DEVICES contains an empty selector") + + resolved: list[dict[str, Any]] = [] + for logical_index, selector in enumerate(selectors): + if selector.isdigit(): + matches = [ + device + for device in inventory + if int(device["cuda_runtime_index"]) == int(selector) + ] + else: + matches = [ + device + for device in inventory + if str(device["uuid"]) == selector or str(device["uuid"]).startswith(selector) + ] + if len(matches) != 1: + raise ValidationError( + "CUDA_VISIBLE_DEVICES selector does not resolve to exactly one " + f"CUDA runtime GPU identity: {selector}" + ) + device = dict(matches[0]) + device["cuda_logical_index"] = logical_index + resolved.append(device) + return resolved + + +def _runtime_gpu_devices(cuda_visible_devices: str) -> list[dict[str, Any]]: + try: + return _resolve_cuda_devices( + cuda_visible_devices, + _query_nvidia_smi_gpus(), + ) + except ValidationError as nvidia_smi_error: + try: + return _resolve_cuda_runtime_devices( + cuda_visible_devices, + _query_cuda_runtime_gpus(), + ) + except ValidationError as cuda_runtime_error: + raise ValidationError( + "could not query GPU identity with nvidia-smi or CUDA runtime: " + f"nvidia-smi: {nvidia_smi_error}; CUDA runtime: {cuda_runtime_error}" + ) from cuda_runtime_error def write_run_metadata( @@ -2138,7 +2250,11 @@ def write_run_metadata( "hostname": platform.node(), "cuda_visible_devices": effective_cuda_visible_devices, "nvidia_visible_devices": os.environ.get("NVIDIA_VISIBLE_DEVICES", ""), - "gpu_identity_source": "nvidia-smi", + "gpu_identity_source": ( + "nvidia-smi" + if all("nvidia_smi_index" in device for device in gpu_devices) + else "cuda-runtime" + ), "gpu_devices": gpu_devices, "command": shlex.join(sys.argv), "started_at": started_at, @@ -2589,6 +2705,12 @@ def _render_comparison(result: Mapping[str, Any]) -> str: mode = str(comparison.get("mode", "") or "") not_compared_reason = str(result.get("not_compared_reason", "") or "") details = [value for value in (mode, not_compared_reason) if value] + exclusion = result.get("platform_exclusion", {}) + if isinstance(exclusion, Mapping) and exclusion: + reason = str(exclusion.get("reason", "") or "") + detail = ": ".join(value for value in ("Platform excluded", reason) if value) + if detail and detail not in details: + details.append(detail) contract = result.get("precision_contract", {}) if isinstance(contract, Mapping) and contract: base = str(contract.get("trtmc_base_precision", "") or "").upper() @@ -2738,6 +2860,10 @@ def _traffic_light_counts( return counts +def _platform_exclusion_count(results: Sequence[Mapping[str, Any]]) -> int: + return sum(isinstance(result.get("platform_exclusion"), Mapping) for result in results) + + def _traffic_light_status(result: Mapping[str, Any]) -> str: validation_status = str(result["validation"]["status"]) comparison_status = str(result["comparison"]["status"]) @@ -2823,6 +2949,10 @@ def _report_document( provenance = _report_provenance(report.get("run", {})) duration = _format_duration(report["summary"].get("duration_seconds")) duration_summary = f" · {html.escape(duration)} total duration" if duration else "" + platform_summary = "" + excluded = int(report["summary"].get("platform_excluded", 0)) + if excluded: + platform_summary = f" · {excluded} platform excluded" filters = render_report_filters( row_count=len(results), search_placeholder="Model type, operation, model, task type, or workload", @@ -2920,7 +3050,7 @@ def _report_document( {comparison_counts["disagreement"]} disagreements · {comparison_counts["not_run"]} not compared · {execution_errors} execution errors · -{report["summary"]["selected_samples"]} samples{duration_summary}
+{report["summary"]["selected_samples"]} samples{platform_summary}{duration_summary}
{html.escape(provenance)} {filters}
@@ -2938,6 +3068,7 @@ def write_report(output: Path) -> tuple[Path, Path, dict[str, Any]]: result_paths, results = _deduplicate_results(result_paths, results) validation_counts, comparison_counts, execution_errors = _report_counts(results) traffic_light_counts = _traffic_light_counts(results) + platform_excluded = _platform_exclusion_count(results) sample_counts = [ count for result in results if (count := _selected_sample_count(result)) is not None ] @@ -2964,6 +3095,7 @@ def write_report(output: Path) -> tuple[Path, Path, dict[str, Any]]: "validation_passed": validation_counts["passed"], "validation_failed": validation_counts["failed"], "validation_skipped": validation_counts["skipped"], + "platform_excluded": platform_excluded, "selected_samples": sum(sample_counts), }, "results": results, @@ -3237,6 +3369,14 @@ def build_parser() -> argparse.ArgumentParser: "the seed is never deleted" ), ) + parser.add_argument( + "--prepare-hf-on-demand", + action="store_true", + help=( + "prepare the selected model and its family-owned Hugging Face " + "dependencies immediately before each model worker" + ), + ) parser.add_argument( "--minimum-free-space-gib", type=_nonnegative_float, @@ -3250,6 +3390,16 @@ def build_parser() -> argparse.ArgumentParser: default=REPO_ROOT / "build" / "trtmc_dataset_benchmark", ) parser.add_argument("--hf-python", type=Path, default=Path(sys.executable)) + parser.add_argument( + "--hf-device", + default="cuda", + help="device used by the Hugging Face Accuracy reference", + ) + parser.add_argument( + "--hf-device-map", + default="", + help="optional Hugging Face device map for the Accuracy reference", + ) parser.add_argument("--backend-dir", type=Path) parser.add_argument("--model-plugin-dir", type=Path) parser.add_argument("--cuda-visible-devices", default="") @@ -3498,8 +3648,18 @@ def _binding_resource_arguments( if arguments.hf_cache_seed_dir is not None and not any( selected.hf_cache_dir.iterdir() ): + seed_source = arguments.hf_cache_seed_dir + seed_destination = selected.hf_cache_dir + if any( + child.name.startswith(("models--", "datasets--", "spaces--")) + for child in seed_source.iterdir() + ): + # Accept the common Hugging Face hub cache layout in addition + # to a complete HF_HOME tree. HF_HOME expects these entries + # below its hub/ directory. + seed_destination = selected.hf_cache_dir / "hub" if ( - arguments.hf_cache_seed_dir.stat().st_dev + seed_source.stat().st_dev != selected.hf_cache_dir.stat().st_dev ): raise ValidationError( @@ -3507,8 +3667,8 @@ def _binding_resource_arguments( ) try: shutil.copytree( - arguments.hf_cache_seed_dir, - selected.hf_cache_dir, + seed_source, + seed_destination, dirs_exist_ok=True, symlinks=True, copy_function=os.link, @@ -3517,7 +3677,7 @@ def _binding_resource_arguments( shutil.rmtree(selected.hf_cache_dir, ignore_errors=True) raise ValidationError( f"could not hard-link Hugging Face cache seed " - f"{arguments.hf_cache_seed_dir}: {exc}" + f"{seed_source}: {exc}" ) from exc return selected, binding_work, model_work @@ -3608,6 +3768,46 @@ def _worker_environment(arguments: argparse.Namespace) -> dict[str, str]: return environment +def _prepare_hf_on_demand( + binding: Binding, + arguments: argparse.Namespace, +) -> None: + if not arguments.prepare_hf_on_demand: + return + case_dir = _case_directory(arguments.output, binding) + case_dir.mkdir(parents=True, exist_ok=True) + selected_models = case_dir / "hf_prepare.models.txt" + selected_models.write_text(f"{binding.model}\n", encoding="utf-8") + log_path = case_dir / "hf_prepare.log" + command = [ + str(arguments.hf_python), + str(HF_WARM_SCRIPT), + "--models-file", + str(selected_models), + "--strict", + "--fail-fast", + ] + if arguments.local_files_only: + command.append("--local-only") + with log_path.open("a", encoding="utf-8") as log_file: + log_file.write(f"$ {shlex.join(command)}\n") + log_file.flush() + completed = subprocess.run( + command, + cwd=REPO_ROOT, + check=False, + env=_worker_environment(arguments), + stdout=log_file, + stderr=subprocess.STDOUT, + text=True, + ) + if completed.returncode != 0: + raise ValidationError( + f"Hugging Face preparation failed for {binding.model} " + f"with rc={completed.returncode}; see {log_path}" + ) + + def _check_free_space(arguments: argparse.Namespace, binding: Binding) -> float: root = arguments.storage_root or arguments.model_work_dir or arguments.output free_gib = shutil.disk_usage(root).free / 1024**3 @@ -3642,6 +3842,7 @@ def _worker_command( ("--trtmc-binary", arguments.trtmc_binary), ("--benchmark-binary", arguments.benchmark_binary), ("--hf-python", arguments.hf_python), + ("--hf-device", arguments.hf_device), ("--model-attempts", arguments.model_attempts), ("--model-retry-delay-seconds", arguments.model_retry_delay_seconds), ( @@ -3656,6 +3857,7 @@ def _worker_command( command.extend([option, str(value)]) for option, value in ( ("--reference-source-cache-dir", arguments.reference_source_cache_dir), + ("--hf-device-map", arguments.hf_device_map), ("--dataset-root", arguments.dataset_root), ("--backend-dir", arguments.backend_dir), ("--model-plugin-dir", arguments.model_plugin_dir), @@ -3670,6 +3872,7 @@ def _worker_command( ("--force-build", arguments.force_build), ("--no-build", arguments.no_build), ("--local-files-only", arguments.local_files_only), + ("--prepare-hf-on-demand", arguments.prepare_hf_on_demand), ): if enabled: command.append(option) @@ -4138,6 +4341,7 @@ def _run_bindings( f"\n{binding.model} / {binding.workload} / {sample_note}", flush=True, ) + _prepare_hf_on_demand(binding, binding_arguments) result = run_binding( binding, arguments=binding_arguments, diff --git a/tools/validation/engine.py b/tools/validation/engine.py index 8cd3da44f6..907bcbcb9d 100644 --- a/tools/validation/engine.py +++ b/tools/validation/engine.py @@ -76,6 +76,7 @@ DEFAULT_WAIVES = REPO_ROOT / "tests" / "e2e" / "waives.txt" REFERENCE_RUNNER = REPO_ROOT / "tools" / "trtmc_reference.py" +REFERENCE_CUDA_ALLOC_CONF_ENV = "TRTMC_REFERENCE_PYTORCH_CUDA_ALLOC_CONF" ERROR_OUTPUT_TEXT = "TensorRT Edge LLM cannot handle this request. Fails." CHOICE_LETTERS = set("ABCDEFGHIJ") GPT_OSS_MMLU_SYSTEM_PROMPT = "You are a helpful assistant. Answer with only the option letter." @@ -7219,6 +7220,7 @@ def run_time_series_bundle(args: argparse.Namespace) -> None: binary_path=str(args.trtmc_binary), hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str(getattr(args, "model_plugin_dir", "") or ""), ) @@ -7445,6 +7447,7 @@ def run_vision_bundle(args: argparse.Namespace) -> None: binary_path=str(args.trtmc_binary), hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str(getattr(args, "model_plugin_dir", "") or ""), ) @@ -7624,6 +7627,7 @@ def run_reranking_bundle(args: argparse.Namespace) -> None: binary_path=str(args.trtmc_binary), hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str(getattr(args, "model_plugin_dir", "") or ""), ) @@ -8218,6 +8222,7 @@ def run_diffusion_bundle(args: argparse.Namespace) -> None: binary_path=str(args.trtmc_binary), hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str(getattr(args, "model_plugin_dir", "") or ""), ) @@ -8355,6 +8360,7 @@ def run_diffusion_text_bundle(args: argparse.Namespace) -> None: binary_path=str(args.trtmc_binary), hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str(getattr(args, "model_plugin_dir", "") or ""), ) @@ -8732,6 +8738,7 @@ def run_model_plugin_bundle(args: argparse.Namespace) -> None: hf_python=str(getattr(args, "hf_python", "") or ""), runtime_python=str(getattr(args, "hf_python", "") or ""), reference_python=str(getattr(args, "hf_python", "") or ""), + ld_library_path=os.environ.get("LD_LIBRARY_PATH", ""), engine_dir=str(bundle_path.parent), model_plugin_dir=str( getattr(args, "model_plugin_dir", "") or "" @@ -9986,7 +9993,15 @@ def run_hf_reference_subprocess( log_path = work_dir / "hf_run.log" env = os.environ.copy() env.setdefault("TOKENIZERS_PARALLELISM", "false") - env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + allocator_conf = os.environ.get(REFERENCE_CUDA_ALLOC_CONF_ENV) + if allocator_conf is None: + env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + elif allocator_conf.strip().lower() in {"disable", "disabled", "none"}: + env.pop("PYTORCH_CUDA_ALLOC_CONF", None) + elif allocator_conf.strip(): + env["PYTORCH_CUDA_ALLOC_CONF"] = allocator_conf.strip() + else: + raise RuntimeError(f"{REFERENCE_CUDA_ALLOC_CONF_ENV} cannot be empty") with log_path.open("w", encoding="utf-8") as log_f: log_f.write(f"$ {shlex.join(cmd)}\n") log_f.flush()
Model typeOperationModelTask typeWorkloadSamplesExecution