Fix 2-GPU test_model_load_utils hang; test import and fixture cleanup

tests/_test_utils/fs_utils.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Filesystem helpers for tests."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+
+def _manifest(root: Path) -> dict[str, tuple[int, int]]:
+    """``relative path -> (size, mtime_ns)`` for every file below ``root``."""
+    return {
+        str(p.relative_to(root)): (p.stat().st_size, p.stat().st_mtime_ns)
+        for p in root.rglob("*")
+        if p.is_file() and not p.is_symlink()
+    }
+
+
+@contextmanager
+def assert_unmodified_tree(path: Path | str) -> Iterator[Path]:
+    """Fail if anything under ``path`` is added, removed, or rewritten inside the ``with``.
+
+    For session/module-scoped model-directory fixtures: a test that writes into a shared
+    directory silently changes what every later test sees. Comparing a file manifest on
+    teardown catches that. ``chmod``-ing the tree read-only would report at the write rather
+    than at teardown, but it only works for an unprivileged user -- root has
+    ``CAP_DAC_OVERRIDE`` and writes straight through the permission bits, and the CI
+    containers run as root.
+    """
+    path = Path(path)
+    before = _manifest(path)
+    yield path
+    if not path.exists():
+        raise AssertionError(f"shared fixture directory {path} was deleted by a test")
+    after = _manifest(path)
+    added = sorted(after.keys() - before.keys())
+    removed = sorted(before.keys() - after.keys())
+    changed = sorted(k for k in before.keys() & after.keys() if before[k] != after[k])
+    if added or removed or changed:
+        raise AssertionError(
+            f"shared fixture directory {path} was modified by a test "
+            f"(added={added}, removed={removed}, changed={changed}); "
+            "copy it into the test's own tmp_path instead of writing into the shared tree"
+        )

tests/_test_utils/onnx/quantization/utils.py

@@ -16,11 +16,18 @@
 import onnx_graphsurgeon as gs
 
 
-def assert_nodes_are_quantized(nodes):
+def assert_nodes_are_quantized(nodes, *, ignore_identity_inputs: bool = False):
+    """Assert every variable input of ``nodes`` is produced by a DequantizeLinear.
+
+    ``ignore_identity_inputs`` skips inputs fed by an ``Identity`` node, for graphs where the
+    quantizer legitimately leaves such a passthrough in place (e.g. concat elimination).
+    """
     for node in nodes:
         for inp_idx, inp in enumerate(node.inputs):
             if isinstance(inp, gs.Variable):
                 producer = node.i(inp_idx)
+                if ignore_identity_inputs and producer and producer.op == "Identity":
+                    continue
                 # Quantized path may include a Cast right after DQ
                 if producer and producer.op == "Cast":
                     producer = producer.i(0)

tests/_test_utils/torch/quantization/attention.py

@@ -0,0 +1,37 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared attention-quantization fixtures for the unit and gpu attention tests."""
+
+import pytest
+
+pytest.importorskip("transformers")
+
+from transformers import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaAttention
+
+from modelopt.torch.quantization.plugins.huggingface import _QuantAttention
+
+
+def make_quant_attention(hidden_size=128, num_q_heads=4, num_kv_heads=2):
+    """A single ``_QuantAttention``-converted Llama attention layer, pinned to the sdpa impl."""
+    config = LlamaConfig(
+        hidden_size=hidden_size,
+        num_attention_heads=num_q_heads,
+        num_key_value_heads=num_kv_heads,
+    )
+    quant_attention = _QuantAttention.convert(LlamaAttention(config, layer_idx=0))
+    quant_attention.config._attn_implementation = "sdpa"
+    return quant_attention

tests/_test_utils/torch/quantization/offload.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared helpers for accelerate-offloaded and layerwise-calibration quantization tests."""
+
+import copy
+
+import torch
+from _test_utils.torch.transformers_models import create_tiny_llama_dir
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from transformers import AutoConfig, AutoModelForCausalLM
+
+
+def make_tiny_llama_and_inputs(tmp_path, num_hidden_layers=3):
+    """Tiny LLaMA checkpoint dir + its config + a GPU token batch sized for its vocab."""
+    tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers)
+    config = AutoConfig.from_pretrained(tiny_llama_dir)
+    inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda()
+    return tiny_llama_dir, config, inputs
+
+
+def make_cpu_offloaded_model(tmp_path, num_hidden_layers=3):
+    """Tiny LLaMA with layer 0 offloaded to CPU via accelerate.
+
+    Returns ``(model, config, tiny_llama_dir, inputs)``; ``inputs`` is a GPU token batch
+    sized for the model's vocab.
+    """
+    tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers)
+    config = AutoConfig.from_pretrained(tiny_llama_dir)
+
+    with init_empty_weights():
+        model = AutoModelForCausalLM.from_config(config)
+
+    device_map = {
+        n: 0
+        for n, m in model.named_modules()
+        if "layers" not in n or n.split("layers.")[-1].isdigit()
+    }
+    device_map["model.layers.0"] = "cpu"
+
+    model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map)
+    inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda()
+    return model, config, tiny_llama_dir, inputs
+
+
+def make_layerwise_cfg(base_cfg):
+    """Copy of ``base_cfg`` with ``layerwise=True`` set on its algorithm field."""
+    cfg = copy.deepcopy(base_cfg)
+    algo = cfg.get("algorithm", "max")
+    if isinstance(algo, str):
+        cfg["algorithm"] = {"method": algo, "layerwise": True}
+    else:
+        algo["layerwise"] = True
+    return cfg
+
+
+def make_layerwise_checkpoint_cfg(base_cfg, checkpoint_dir):
+    """``make_layerwise_cfg`` plus a ``layerwise_checkpoint_dir``."""
+    cfg = make_layerwise_cfg(base_cfg)
+    cfg["algorithm"]["layerwise_checkpoint_dir"] = checkpoint_dir
+    return cfg

tests/_test_utils/torch/quantization/quant_utils.py

@@ -15,6 +15,8 @@
 
 import torch
 
+from modelopt.torch.quantization.nn import TensorQuantizer
+
 
 def quant(x, amax, num_bits=8, fake=False, narrow_range=True):
     """Quantize x using torch."""
@@ -32,3 +34,23 @@ def quant(x, amax, num_bits=8, fake=False, narrow_range=True):
 
 def get_model_size(model):
     return sum([p.element_size() * p.nelement() for p in model.parameters()])
+
+
+def nvfp4_static_amax_dtypes(model):
+    """Map of ``module name -> amax dtype`` for every NVFP4 static quantizer with an amax."""
+    return {
+        name: module.amax.dtype
+        for name, module in model.named_modules()
+        if isinstance(module, TensorQuantizer)
+        and module.is_nvfp4_static
+        and module.amax is not None
+    }
+
+
+def assert_nvfp4_static_amaxes_fp32(amax_dtypes, model_dtype, label):
+    """NVFP4 static amaxes must stay fp32 regardless of the model's own dtype."""
+    assert amax_dtypes, f"{label}: expected NVFP4 static amaxes for model dtype {model_dtype}"
+    assert all(amax_dtype == torch.float32 for amax_dtype in amax_dtypes.values()), (
+        f"{label}: expected all NVFP4 static amaxes to be fp32 for model dtype {model_dtype}, "
+        f"got {amax_dtypes}"
+    )

tests/_test_utils/torch/speculative/dflash.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared DFlash test config, used by the unit and gpu speculative-decoding tests."""
+
+from copy import deepcopy
+
+from modelopt.torch.speculative.config import DFLASH_DEFAULT_CFG
+
+DFLASH_BLOCK_SIZE = 4
+DFLASH_NUM_DRAFT_LAYERS = 2
+
+
+def get_dflash_config(
+    block_size: int = DFLASH_BLOCK_SIZE,
+    num_layers: int = DFLASH_NUM_DRAFT_LAYERS,
+    offline: bool | None = None,
+):
+    """DFlash config sized for a tiny model: no torch.compile, token 0 as the mask token.
+
+    ``offline`` is only written when set, so callers that don't care keep the default.
+    """
+    config = deepcopy(DFLASH_DEFAULT_CFG["config"])
+    config["dflash_block_size"] = block_size
+    config["dflash_use_torch_compile"] = False
+    config["dflash_mask_token_id"] = 0  # use token 0 as mask for tiny model
+    config["dflash_architecture_config"] = {"num_hidden_layers": num_layers}
+    if offline is not None:
+        config["dflash_offline"] = offline
+    return config

tests/_test_utils/torch/transformers_models.py

@@ -35,6 +35,7 @@
     GptOssConfig,
     LlamaConfig,
     LlamaForSequenceClassification,
+    MixtralConfig,
     NemotronConfig,
     PreTrainedModel,
     Qwen3Config,
@@ -566,6 +567,25 @@ def create_tiny_gpt_oss_dir(
     )
 
 
+##### MIXTRAL #####
+def get_tiny_mixtral(**config_kwargs) -> PreTrainedModel:
+    set_seed(SEED)
+    kwargs = {
+        "dtype": torch.bfloat16,
+        "hidden_size": 32,
+        "intermediate_size": 32,
+        "num_hidden_layers": 2,
+        "num_attention_heads": 4,
+        "num_key_value_heads": 2,
+        "num_local_experts": 4,
+        "num_experts_per_tok": 2,
+        "max_position_embeddings": 32,
+        "vocab_size": 32,
+    }
+    kwargs.update(config_kwargs)
+    return AutoModelForCausalLM.from_config(MixtralConfig(**kwargs))
+
+
 ##### LLAMA #####
 def get_tiny_llama(**config_kwargs) -> PreTrainedModel:
     set_seed(SEED)
@@ -671,10 +691,6 @@ def get_tiny_bert(**config_kwargs) -> PreTrainedModel:
     return AutoModelForQuestionAnswering.from_config(BertConfig(**kwargs))
 
 
-def create_tiny_bert_dir(tmp_path: Path | str, **config_kwargs) -> Path:
-    return _create_tiny_llm_dir(Path(tmp_path) / "tiny_bert", get_tiny_bert, **config_kwargs)
-
-
 ##### ViT (vision) #####
 def get_tiny_vit(**config_kwargs) -> PreTrainedModel:
     set_seed(SEED)

tests/conftest.py

@@ -20,6 +20,7 @@
 import pytest
 import torch
 import torch.distributed as dist
+from _test_utils.fs_utils import assert_unmodified_tree
 from _test_utils.torch.distributed.utils import init_process
 
 import modelopt.torch.opt as mto
@@ -91,16 +92,7 @@ def pytest_collection_modifyitems(config, items):
         item.add_marker(pytest.mark.timeout(_DEFAULT_TIMEOUT[group]))
 
 
-@pytest.fixture
-def tiny_tokenizer():
-    """Real tiny HF tokenizer (vocab=128) shared across unit and gpu test lanes."""
-    # Lazy import: transformers_models.py runs ``pytest.importorskip("transformers")``
-    # at module load, which we don't want to trigger at conftest import time.
-    from _test_utils.torch.transformers_models import get_tiny_tokenizer
-
-    return get_tiny_tokenizer()
-
-
+# General Fixtures #################################################################################
 @pytest.fixture
 def skip_on_windows():
     if platform.system() == "Windows":
@@ -161,3 +153,29 @@ def enable_hf_checkpointing():
 def project_root_path(request: pytest.FixtureRequest) -> Path:
     """Fixture providing the project root path for tests."""
     return Path(request.config.rootpath)
+
+
+# Transformers Models Fixtures #####################################################################
+@pytest.fixture
+def tiny_tokenizer():
+    """Real tiny HF tokenizer (vocab=128) shared across unit and gpu test lanes."""
+    # Lazy import: transformers_models.py runs ``pytest.importorskip("transformers")``
+    # at module load, which we don't want to trigger at conftest import time.
+    from _test_utils.torch.transformers_models import get_tiny_tokenizer
+
+    return get_tiny_tokenizer()
+
+
+@pytest.fixture(scope="session")
+def tiny_wan22_path(tmp_path_factory):
+    """Tiny Wan 2.2 pipeline dir, built once per session (the build is the expensive part).
+
+    Shared by the gpu sparse-attention tests and the diffusers example tests.
+    """
+    # Lazy import for the same reason as ``tiny_tokenizer``: diffusers_models.py pulls in
+    # transformers at module load.
+    from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir
+
+    pipeline_dir = create_tiny_wan22_pipeline_dir(tmp_path_factory.mktemp("tiny_wan22"))
+    with assert_unmodified_tree(pipeline_dir) as path:
+        yield str(path)