diff --git a/examples/mamba/pulse/uv.lock b/examples/mamba/pulse/uv.lock index 138fb4e48..bf04b3a61 100644 --- a/examples/mamba/pulse/uv.lock +++ b/examples/mamba/pulse/uv.lock @@ -1008,7 +1008,7 @@ wheels = [ [[package]] name = "torch-to-nnef" -version = "0.24.2" +version = "0.24.3" source = { editable = "../../../" } dependencies = [ { name = "mako" }, diff --git a/examples/speech_enhancement/dpdfnet/uv.lock b/examples/speech_enhancement/dpdfnet/uv.lock index 5e163c0e6..350447653 100644 --- a/examples/speech_enhancement/dpdfnet/uv.lock +++ b/examples/speech_enhancement/dpdfnet/uv.lock @@ -749,7 +749,7 @@ wheels = [ [[package]] name = "torch-to-nnef" -version = "0.24.2" +version = "0.24.3" source = { editable = "../../../" } dependencies = [ { name = "mako" }, diff --git a/examples/tts/pocket_tts/uv.lock b/examples/tts/pocket_tts/uv.lock index b30dc5aa1..a515c421d 100644 --- a/examples/tts/pocket_tts/uv.lock +++ b/examples/tts/pocket_tts/uv.lock @@ -1413,7 +1413,7 @@ wheels = [ [[package]] name = "torch-to-nnef" -version = "0.24.2" +version = "0.24.3" source = { editable = "../../../" } dependencies = [ { name = "mako" }, diff --git a/examples/vad/FSMN-wasm/uv.lock b/examples/vad/FSMN-wasm/uv.lock index 429bb0e4f..4d7ee9f60 100644 --- a/examples/vad/FSMN-wasm/uv.lock +++ b/examples/vad/FSMN-wasm/uv.lock @@ -859,7 +859,7 @@ wheels = [ [[package]] name = "torch-to-nnef" -version = "0.24.2" +version = "0.24.3" source = { editable = "../../../" } dependencies = [ { name = "mako" }, diff --git a/packages/llm/tests/test_gpt_oss_handler.py b/packages/llm/tests/test_gpt_oss_handler.py new file mode 100644 index 000000000..d5f823e33 --- /dev/null +++ b/packages/llm/tests/test_gpt_oss_handler.py @@ -0,0 +1,158 @@ +from types import SimpleNamespace + +import torch +from torch import nn + +from torch_to_nnef_llm.models.handlers import GptOssArchitectureHandler +from torch_to_nnef_llm.models.handlers.registry import get_handler + + +class FakeGptOssExpert(nn.Module): + def __init__(self, experts_implementation): + super().__init__() + self.config = SimpleNamespace( + _experts_implementation=experts_implementation + ) + + +class FakeGptOssModel(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(_experts_implementation="grouped_mm") + self.grouped = FakeGptOssExpert("grouped_mm") + self.eager = FakeGptOssExpert("eager") + + +def test_gpt_oss_handler_registered(): + assert get_handler("gpt_oss") is GptOssArchitectureHandler + + +def test_gpt_oss_handler_uses_traceable_experts_implementation(): + model = FakeGptOssModel() + + GptOssArchitectureHandler().prepare_model_for_export(model) + + assert model.config._experts_implementation == "batched_mm" + assert model.grouped.config._experts_implementation == "batched_mm" + assert model.eager.config._experts_implementation == "eager" + + +NEG = torch.finfo(torch.float32).min + + +def _wrapper(*, sliding_window=4, layer_types=None, force_causal_mask=True): + if layer_types is None: + layer_types = ["sliding_attention", "full_attention"] + return SimpleNamespace( + model=SimpleNamespace( + config=SimpleNamespace( + sliding_window=sliding_window, + layer_types=layer_types, + ) + ), + force_causal_mask=force_causal_mask, + with_dyn_cache=False, + ) + + +def _inputs(*, seq_length, past_length, n_layers=2): + input_ids = torch.zeros(1, seq_length, dtype=torch.long) + kv = [torch.zeros(1, 2, past_length, 8) for _ in range(2 * n_layers)] + return (input_ids, *kv) + + +def _visible(mask): + """Boolean [S, K] view of an additive mask: True where attention is kept.""" + return mask[0, 0] == 0.0 + + +def test_full_and_sliding_masks_are_distinct(): + handler = GptOssArchitectureHandler() + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=8, past_length=0), + wrapper=_wrapper(sliding_window=4), + ) + mapping = ctx.model_inputs["attention_mask"] + + assert isinstance(mapping, dict) + assert set(mapping) == {"full_attention", "sliding_attention"} + # The whole point of the handler: these must not be the same tensor. + assert not torch.equal( + mapping["full_attention"], mapping["sliding_attention"] + ) + + +def test_full_mask_is_plain_causal(): + handler = GptOssArchitectureHandler() + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=4, past_length=0), + wrapper=_wrapper(sliding_window=2), + ) + visible = _visible(ctx.model_inputs["attention_mask"]["full_attention"]) + expected = torch.tril(torch.ones(4, 4, dtype=torch.bool)) + assert torch.equal(visible, expected) + + +def test_sliding_mask_keeps_exactly_window_keys(): + handler = GptOssArchitectureHandler() + window = 3 + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=6, past_length=0), + wrapper=_wrapper(sliding_window=window), + ) + visible = _visible(ctx.model_inputs["attention_mask"]["sliding_attention"]) + + for q in range(6): + kept = visible[q].nonzero().flatten().tolist() + # Query q sees keys in (q - window, q], clipped at 0. + assert kept == list(range(max(0, q - window + 1), q + 1)) + assert len(kept) <= window + + +def test_sliding_mask_accounts_for_past_length(): + handler = GptOssArchitectureHandler() + window, past = 4, 10 + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=2, past_length=past), + wrapper=_wrapper(sliding_window=window), + ) + visible = _visible(ctx.model_inputs["attention_mask"]["sliding_attention"]) + + assert visible.shape == (2, past + 2) + # First new token sits at absolute position 10 and sees keys 7..10. + assert visible[0].nonzero().flatten().tolist() == [7, 8, 9, 10] + assert visible[1].nonzero().flatten().tolist() == [8, 9, 10, 11] + + +def test_masks_are_additive_with_neg_inf_where_hidden(): + handler = GptOssArchitectureHandler() + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=3, past_length=0), + wrapper=_wrapper(sliding_window=2), + ) + mask = ctx.model_inputs["attention_mask"]["sliding_attention"] + assert mask.shape == (1, 1, 3, 3) + assert mask.dtype == torch.float32 + # Row 2 with window 2 hides key 0 only. + assert mask[0, 0, 2, 0] == NEG + assert mask[0, 0, 2, 1] == 0.0 + assert mask[0, 0, 2, 2] == 0.0 + + +def test_full_attention_only_model_keeps_base_single_mask(): + """No sliding layers means the base handler's single mask is correct.""" + handler = GptOssArchitectureHandler() + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=4, past_length=0), + wrapper=_wrapper(sliding_window=0, layer_types=["full_attention"] * 2), + ) + assert isinstance(ctx.model_inputs["attention_mask"], torch.Tensor) + + +def test_no_mask_mapping_when_causal_mask_not_forced(): + handler = GptOssArchitectureHandler() + ctx = handler.build_forward_inputs( + inputs=_inputs(seq_length=4, past_length=0), + wrapper=_wrapper(force_causal_mask=False), + ) + assert ctx.model_inputs["attention_mask"] is None diff --git a/packages/llm/tests/test_sdpa_attention_export.py b/packages/llm/tests/test_sdpa_attention_export.py index 5aaf6baf6..fd16f1c56 100644 --- a/packages/llm/tests/test_sdpa_attention_export.py +++ b/packages/llm/tests/test_sdpa_attention_export.py @@ -162,6 +162,16 @@ def test_reify_sdpa_operator_rejects_eager_attention(): exporter._resolve_attn_implementation("eager", True) +def test_existing_export_test_dir_is_allowed_for_existing_export_root( + tmp_path, +): + export_dir = tmp_path / "export" + test_dir = export_dir / "tests" + test_dir.mkdir(parents=True) + + assert exporter._ensure_export_test_dir(export_dir, True) == test_dir + + def test_dump_llm_routes_reified_sdpa_to_loader(monkeypatch, tmp_path): captured_load = {} captured_dump = {} @@ -214,6 +224,42 @@ def fake_load(*args, **kwargs): assert captured_load["experts_implementation"] == "batched_mm" +def test_dump_llm_routes_reified_sdpa_input_upcast_policy( + monkeypatch, tmp_path +): + captured_inference_target = {} + + class _Exporter: + # non-multimodal model_type so dump_llm's multimodal guard passes + model_infos = SimpleNamespace(conf=SimpleNamespace(model_type="fake")) + + def build_inference_target(self, **kwargs): + captured_inference_target.update(kwargs) + return object() + + def dump_with_inference_target(self, **_kwargs): + return None + + def dump(self, **kwargs): + return exporter.LLMExporter.dump(self, **kwargs) + + def fake_load(*_args, **_kwargs): + return _Exporter() + + monkeypatch.setattr(exporter.LLMExporter, "load", staticmethod(fake_load)) + exporter.dump_llm( + "fake/model", + export_dirpath=tmp_path / "export", + reify_sdpa_operator=True, + upcast_reified_sdpa_inputs_to_f32=False, + ) + + assert captured_inference_target["reify_sdpa_operator"] is True + assert ( + captured_inference_target["upcast_reified_sdpa_inputs_to_f32"] is False + ) + + def test_cli_reify_sdpa_operator_implies_sdpa_attention(monkeypatch, tmp_path): captured = {} diff --git a/packages/llm/torch_to_nnef_llm/exporter.py b/packages/llm/torch_to_nnef_llm/exporter.py index 6d8113d60..0ac879e3c 100644 --- a/packages/llm/torch_to_nnef_llm/exporter.py +++ b/packages/llm/torch_to_nnef_llm/exporter.py @@ -139,6 +139,14 @@ def _resolve_attn_implementation( return attn_implementation +def _ensure_export_test_dir( + export_dirpath: Path, ignore_already_exist_dir: bool +) -> Path: + test_dir = export_dirpath / "tests" + test_dir.mkdir(parents=True, exist_ok=ignore_already_exist_dir) + return test_dir + + #: Default number of retries for transient Hugging Face download failures. DEFAULT_HF_DOWNLOAD_N_RETRIES = 5 @@ -791,8 +799,9 @@ def export_model( inference_target.dynamic_axes = dynamic_axes # Add io.npz test in exproted dir for dbg purpose - test_dir = export_dirpath / "tests" - test_dir.mkdir(parents=True) + test_dir = _ensure_export_test_dir( + export_dirpath, ignore_already_exist_dir + ) if check_inference_modes: self._dump_modes_json( @@ -863,6 +872,7 @@ def dump(self, **kwargs): "force_f32_linear_accumulator", "force_f32_normalization", "reify_sdpa_operator", + "upcast_reified_sdpa_inputs_to_f32", "tract_check_io_tolerance", ] if key in kwargs @@ -884,6 +894,7 @@ def build_inference_target( force_f32_linear_accumulator: T.Optional[bool] = None, force_f32_normalization: T.Optional[bool] = None, reify_sdpa_operator: T.Optional[bool] = None, + upcast_reified_sdpa_inputs_to_f32: T.Optional[bool] = None, tract_check_io_tolerance: TractCheckTolerance = LM_CHECK_TOLERANCE, compression_method: T.Optional[str] = None, compression_registry: T.Optional[str] = None, @@ -930,6 +941,10 @@ def build_inference_target( if reify_sdpa_operator is not None: inference_target.reify_sdpa_operator = reify_sdpa_operator + if upcast_reified_sdpa_inputs_to_f32 is not None: + inference_target.upcast_reified_sdpa_inputs_to_f32 = ( + upcast_reified_sdpa_inputs_to_f32 + ) if ( self.is_half_precision_model diff --git a/packages/llm/torch_to_nnef_llm/loader.py b/packages/llm/torch_to_nnef_llm/loader.py index 8e699cec8..e69907159 100644 --- a/packages/llm/torch_to_nnef_llm/loader.py +++ b/packages/llm/torch_to_nnef_llm/loader.py @@ -655,6 +655,7 @@ def _from_pretrained( @require_extra_decorator(extra=T2NExtra.LLM_TRACT, module="transformers") +# pylint: disable-next=too-many-branches def load_model( hf_model_slug: T.Optional[str] = None, local_dir: T.Optional[Path] = None, diff --git a/packages/llm/torch_to_nnef_llm/models/handlers/__init__.py b/packages/llm/torch_to_nnef_llm/models/handlers/__init__.py index 4fa656568..56c0fe03b 100644 --- a/packages/llm/torch_to_nnef_llm/models/handlers/__init__.py +++ b/packages/llm/torch_to_nnef_llm/models/handlers/__init__.py @@ -16,6 +16,7 @@ Gemma4VideoEncoderHandler, Gemma4VisionEncoderHandler, ) +from .gpt_oss import GptOssArchitectureHandler from .idefics3_vl import ( Idefics3ArchitectureHandler, Idefics3VisionEncoderHandler, @@ -53,6 +54,7 @@ "Gemma4ArchitectureHandler", "Gemma4VideoEncoderHandler", "Gemma4VisionEncoderHandler", + "GptOssArchitectureHandler", "IOSpec", "Idefics3ArchitectureHandler", "Idefics3VisionEncoderHandler", diff --git a/packages/llm/torch_to_nnef_llm/models/handlers/gpt_oss.py b/packages/llm/torch_to_nnef_llm/models/handlers/gpt_oss.py new file mode 100644 index 000000000..3c5a61c7d --- /dev/null +++ b/packages/llm/torch_to_nnef_llm/models/handlers/gpt_oss.py @@ -0,0 +1,108 @@ +import typing as T + +import torch + +from .base import StateContext +from .default import DefaultArchitectureHandler +from .registry import register_handler + + +@register_handler +class GptOssArchitectureHandler(DefaultArchitectureHandler): + """Handler for GPT-OSS causal decoder models.""" + + ARCH_NAMES = ("gpt_oss",) + + def prepare_model_for_export(self, model) -> None: + # torch grouped_mm currently only has a BF16 fake/meta path. T2N exports + # GPT-OSS MoE blocks as tract_moe_ffn, but the root torch.jit trace must + # still run first; use a traceable HF expert implementation for export. + for module in model.modules(): + config = getattr(module, "config", None) + if ( + config is not None + and getattr(config, "_experts_implementation", None) + == "grouped_mm" + ): + config._experts_implementation = "batched_mm" + + @staticmethod + def _build_mask_mapping( + *, + seq_length: int, + past_length: int, + sliding_window: int, + device: torch.device, + ) -> T.Dict[str, torch.Tensor]: + """Additive causal and causal-plus-window masks over ``S+P`` keys. + + Built from token positions with arange comparisons rather than a baked + triangular constant, so the graph stays correct for any ``(S, P)`` at + inference time. + """ + # Float arithmetic throughout (comparisons cast straight to the mask + # dtype, AND via multiply) rather than boolean `&` on tensors: t2n's + # shape inference re-runs bitwise ops with a float placeholder and + # fails with "bitwise_and not implemented for Float". Same reason as + # the note in `gemma3_vl._build_mask_mapping`. + dtype = torch.float32 + total_length = seq_length + past_length + q_pos = torch.arange( + past_length, total_length, device=device + ).unsqueeze(1) + k_pos = torch.arange(total_length, device=device).unsqueeze(0) + + visible_full = (k_pos <= q_pos).to(dtype) + # Query at absolute position q sees keys in (q - window, q], matching + # `masking_utils.sliding_window_overlay`'s `kv_idx > q_idx - window`. + within_window = (k_pos > (q_pos - sliding_window)).to(dtype) + visible_sliding = visible_full * within_window + + neg = torch.finfo(dtype).min + + def to_additive(visible: torch.Tensor) -> torch.Tensor: + return ((1.0 - visible) * neg).unsqueeze(0).unsqueeze(0) + + return { + "full_attention": to_additive(visible_full), + "sliding_attention": to_additive(visible_sliding), + } + + def build_forward_inputs( + self, + *, + inputs: T.Tuple[torch.Tensor, ...], + wrapper, + ) -> StateContext: + """Pass per-layer masks so sliding layers keep their window. + + GPT-OSS alternates ``sliding_attention`` and ``full_attention`` layers. + The base handler hands the model a single 4D causal mask, and + ``masking_utils._preprocess_mask_arguments`` returns any 4D mask as-is, + so both ``create_causal_mask`` and ``create_sliding_window_causal_mask`` + early-exit with that same tensor. The model's mask mapping then holds + the unwindowed mask under both keys and every layer attends over the + whole context. + + That is invisible below the window, where the two masks agree, and + degrades output as the sequence grows past it. + """ + ctx = super().build_forward_inputs(inputs=inputs, wrapper=wrapper) + if not getattr(wrapper, "force_causal_mask", False): + return ctx + + config = wrapper.model.config + sliding_window = int(getattr(config, "sliding_window", 0) or 0) + layer_types = tuple(getattr(config, "layer_types", ()) or ()) + # Nothing to correct when the model is uniformly full-attention: the + # base handler's single causal mask is already right. + if sliding_window <= 0 or "sliding_attention" not in layer_types: + return ctx + + ctx.model_inputs["attention_mask"] = self._build_mask_mapping( + seq_length=inputs[0].shape[1], + past_length=inputs[1].shape[2], + sliding_window=sliding_window, + device=inputs[0].device, + ) + return ctx diff --git a/packages/nemo-asr/pyproject.toml b/packages/nemo-asr/pyproject.toml index 6c4739bd4..9ff0e00a1 100644 --- a/packages/nemo-asr/pyproject.toml +++ b/packages/nemo-asr/pyproject.toml @@ -126,10 +126,23 @@ commands_pre = [ "-c", "real=$(command -v cmake) && mkdir -p {envtmpdir}/bin && printf '#!/bin/bash\\nexec \"%s\" -DCMAKE_POLICY_VERSION_MINIMUM=3.5 \"$@\"\\n' \"$real\" > {envtmpdir}/bin/cmake && chmod +x {envtmpdir}/bin/cmake", ], + # Install torch/torchaudio from the CPU-only index FIRST and alone, so they + # resolve to the CPU builds (no ~GBs of nvidia CUDA deps). [ "bash", "-c", - "export PATH={envtmpdir}/bin:$PATH && uv pip install -i https://download.pytorch.org/whl/cpu torch==2.6.* torchaudio==2.6.* 'nemo_toolkit[asr]==2.7.2'", + "uv pip install -i https://download.pytorch.org/whl/cpu torch==2.6.* torchaudio==2.6.*", + ], + # Then install nemo's deps from PyPI (torch is already installed above and + # satisfies the constraint, so it stays the CPU build). Pin numpy<2.5: a + # py3.13-capable numba (>=0.61, here 0.65.x) requires numpy<2.5, but the + # test dependency-group pre-installs numpy 2.5.0; without this constraint uv + # keeps 2.5.0 and backtracks to the ancient numba 0.53.1 / llvmlite 0.36.0, + # which cannot build on Python 3.13. numpy 2.4.x satisfies matplotlib too. + [ + "bash", + "-c", + "export PATH={envtmpdir}/bin:$PATH && uv pip install 'nemo_toolkit[asr]==2.7.2' 'numpy<2.5'", ], [ "uv", diff --git a/scripts/export_moe_test_asset.py b/scripts/export_moe_test_asset.py new file mode 100644 index 000000000..64b6b6867 --- /dev/null +++ b/scripts/export_moe_test_asset.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Export a tiny Qwen3-MoE NNEF model + reference I/O for tract tests. + +Produces a directory with graph.nnef + .dat weight files + io.npz. + +Usage: + cd /path/to/t2n_main + .venv/bin/python scripts/export_moe_test_asset.py /path/to/output_dir +""" +import logging +import os +import shutil +import sys +import tempfile +from pathlib import Path + +import numpy as np +import torch +from torch import nn +from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeConfig, + Qwen3MoeSparseMoeBlock, +) + +from torch_to_nnef import TractNNEF, export_model_to_nnef +from torch_to_nnef.exceptions import T2NError +from torch_to_nnef.inference_target.tract import TractCli + + +class Qwen3TinyMoE(nn.Module): + """Minimal Qwen3 MoE block: 4 experts, top-2, hidden=16, intermediate=32.""" + + def __init__(self): + super().__init__() + cfg = Qwen3MoeConfig( + hidden_size=16, + moe_intermediate_size=32, + num_experts=4, + num_experts_per_tok=2, + hidden_act="silu", + # real Qwen3-MoE checkpoints renormalize the top-k gates; the + # adapter rejects norm_topk_prob=False (the config default). + norm_topk_prob=True, + ) + self.moe = Qwen3MoeSparseMoeBlock(cfg) + # transformers MoE experts allocate weights with torch.empty + # (uninitialized); a standalone block must be seeded or the reference + # output is NaN. + torch.manual_seed(0) + with torch.no_grad(): + for p in self.moe.parameters(): + nn.init.normal_(p, std=0.02) + + def forward(self, x): + # Qwen3MoeSparseMoeBlock expects [batch, seq, hidden] + return self.moe(x) + + +def _inference_target() -> TractNNEF: + # Prefer a tract build that has tract_moe_ffn (the op is unreleased); fall + # back to the latest official version (export still writes the asset even + # if the post-export IO check then fails). + tract_path = os.environ.get("T2N_TEST_TRACT_PATH") + if tract_path: + cli = TractCli(Path(tract_path)) + return TractNNEF( + cli.version, specific_tract_binary_path=Path(tract_path) + ) + return TractNNEF.latest() + + +def main(): + if len(sys.argv) < 2: + print(f"Usage: {sys.argv[0]} ") + sys.exit(1) + + output_dir = Path(sys.argv[1]) + output_dir.mkdir(parents=True, exist_ok=True) + + torch.manual_seed(42) + model = Qwen3TinyMoE() + model.eval() + + # Test input: [batch=1, seq=3, hidden=16] + torch.manual_seed(123) + test_input = torch.randn(1, 3, 16) + + with torch.no_grad(): + ref_output = model(test_input) + assert torch.isfinite(ref_output).all(), "reference output is not finite" + + with tempfile.TemporaryDirectory() as tmpdir: + export_path = Path(tmpdir) / "qwen3_moe_tiny.nnef" + try: + exported = export_model_to_nnef( + model=model, + args=(test_input,), + file_path_export=export_path, + inference_target=_inference_target(), + input_names=["input_0"], + output_names=["output_0"], + compression_level=None, + log_level=logging.INFO, + ) + except T2NError as e: + exported = export_path + if not exported.exists(): + raise RuntimeError(f"Export failed: {e}") from e + print(f"WARNING: post-export validation failed (expected): {e}") + + nnef_dir = Path(exported) + for item in nnef_dir.iterdir(): + dest = output_dir / item.name + if item.is_dir(): + shutil.copytree(item, dest, dirs_exist_ok=True) + else: + shutil.copy2(item, dest) + + np.savez( + output_dir / "io.npz", + input_0=test_input.numpy(), + output_0=ref_output.numpy(), + ) + + print(f"Exported to: {output_dir}") + total = 0 + for f in sorted(output_dir.rglob("*")): + if f.is_file(): + size = f.stat().st_size + total += size + print(f" {f.relative_to(output_dir)}: {size} bytes") + print(f" TOTAL: {total} bytes ({total / 1024:.1f} KB)") + + +if __name__ == "__main__": + main() diff --git a/tests/test_binary_int_float_promotion.py b/tests/test_binary_int_float_promotion.py index 92a2b076c..96b7b396f 100644 --- a/tests/test_binary_int_float_promotion.py +++ b/tests/test_binary_int_float_promotion.py @@ -106,6 +106,20 @@ def forward(self, ys: torch.Tensor) -> torch.Tensor: return torch.atan2(ys, self.xs) +class _MixedDtypeBroadcastedAffine(torch.nn.Module): + """RMSNorm-style `weight[D] * hidden[B, S, D]` after dtype promotion.""" + + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter( + torch.linspace(0.5, 1.5, steps=8).to(torch.float16) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + 1e-5) + return self.weight * x + + @pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) def test_add_sub_mixed_dtype(inference_target): """`float + int_buf` and `float - int_buf`.""" @@ -147,3 +161,10 @@ def test_atan2_mixed_dtype(inference_target): """`atan2(float, int_buf)`.""" ys = torch.randn(4) check_model_io_test(_MixedAtan2(), ys, inference_target) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_promoted_operand_broadcast_keeps_operand_shape(inference_target): + """A casted rank-1 operand must still broadcast on trailing axes.""" + x = torch.randn(1, 3, 8) + check_model_io_test(_MixedDtypeBroadcastedAffine(), x, inference_target) diff --git a/tests/test_f16.py b/tests/test_f16.py index 83391246a..1927abf3c 100644 --- a/tests/test_f16.py +++ b/tests/test_f16.py @@ -103,6 +103,9 @@ def check_contains_f32_upcast_attn(inference_target, path): assert "tract_transformers_sdpa(" in graph_content if inference_target.force_attention_inner_in_f32: assert "acc_datum_type = 'f32'" in graph_content + assert "datum_type = 'f32'" in graph_content + assert "to = 'f32'" in graph_content + assert "to = 'f16'" in graph_content elif inference_target.force_attention_inner_in_f32: assert ( "fragment scaled_dot_product_attention_3d_f16_df32(" @@ -112,6 +115,48 @@ def check_contains_f32_upcast_attn(inference_target, path): assert "fragment scaled_dot_product_attention_3d_f16(" in graph_content +def check_contains_f32_accum_sdpa_without_input_upcast(inference_target, path): + assert path.exists() + graph_content = _read_graph_nnef_from_archive(path) + assert "tract_transformers_sdpa(" in graph_content + assert "acc_datum_type = 'f32'" in graph_content + assert "datum_type = 'f16'" in graph_content + assert "sdpa_q_f32" not in graph_content + assert "sdpa_k_f32" not in graph_content + assert "sdpa_v_f32" not in graph_content + + +@pytest.mark.skipif( + condition=torch_version() < "2.2.0", + reason="torch older than 2.2.0 lack too much half operator support on CPU", +) +def test_reified_sdpa_can_keep_f16_inputs_with_f32_accum(): + targets = [ + deepcopy(inf) + for inf in TRACT_INFERENCES_TO_TESTS_APPROX + if inf.version >= "0.22.0" + ] + if not targets: + pytest.skip("No tract target with reified SDPA support") + + inference_target = targets[0] + inference_target.check_io = False + inference_target.force_attention_inner_in_f32 = True + inference_target.reify_sdpa_operator = True + inference_target.upcast_reified_sdpa_inputs_to_f32 = False + + check_model_io_test( + model=TernaryPrimitive(op=F.scaled_dot_product_attention), + test_input=( + torch.arange(12).reshape(1, 3, 4).half(), + torch.arange(12).reshape(1, 3, 4).half(), + torch.arange(12).reshape(1, 3, 4).half(), + ), + inference_target=inference_target, + callback_post_export=check_contains_f32_accum_sdpa_without_input_upcast, + ) + + @pytest.mark.parametrize( "id,test_input,model,inference_target", attn_test_suite.test_samples, diff --git a/tests/test_moe.py b/tests/test_moe.py new file mode 100644 index 000000000..2d100b9fc --- /dev/null +++ b/tests/test_moe.py @@ -0,0 +1,622 @@ +"""Tests for MoE FFN export to tract_moe_ffn operator.""" + +import os +import tarfile +import tempfile +from copy import deepcopy +from functools import partial +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from torch_to_nnef.inference_target import TractNNEF +from torch_to_nnef.nnef_io.tensor import DatBinHeader +from torch_to_nnef.op.custom_extractors import MoEFFN +from torch_to_nnef.op.custom_extractors.moe import _GraniteMoEAdapter +from torch_to_nnef.tensor.offload import OffloadedTensor +from torch_to_nnef.tensor.opaque import OpaqueTensorRef +from torch_to_nnef.tensor.quant import ( + fp_to_tract_q4_0_with_min_max_calibration, +) + +from .utils import ( + TRACT_INFERENCES_TO_TESTS_APPROX, + check_model_io_test, + skipif_limited_offload_support, + skipif_unsupported_qtensor, +) + + +class MoEFFNWrapper(nn.Module): + def __init__(self, num_experts, d_model, d_hidden, k=2, activation="silu"): + super().__init__() + self.moe = MoEFFN( + num_experts=num_experts, + d_model=d_model, + d_hidden=d_hidden, + k=k, + activation=activation, + ) + + def forward(self, x): + return self.moe(x) + + +class MoEFFNWithBiasWrapper(nn.Module): + def __init__(self, num_experts, d_model, d_hidden, k=2): + super().__init__() + self.moe = MoEFFN( + num_experts=num_experts, + d_model=d_model, + d_hidden=d_hidden, + k=k, + bias=True, + ) + + def forward(self, x): + return self.moe(x) + + +def _skip_if_unsupported(inference_target): + # tract_moe_ffn first ships in tract 0.23.4; releases 0.23.0..0.23.3 are + # already out without it, so default CI (official versions) must skip. + # An explicitly provided tract (T2N_TEST_TRACT_PATH / + # T2N_TEST_TRACT_VERSION) is trusted to have the op regardless of its + # reported version, so the locally built dev binary (e.g. 0.23.2-pre) + # still runs these. + if not isinstance(inference_target, TractNNEF): + pytest.skip("MoE export requires a tract inference target") + explicit = ( + "T2N_TEST_TRACT_PATH" in os.environ + or "T2N_TEST_TRACT_VERSION" in os.environ + ) + if not explicit and inference_target.version < "0.23.4": + pytest.skip( + "tract_moe_ffn first ships in tract 0.23.4; set " + "T2N_TEST_TRACT_PATH to a build that has the op to run these" + ) + + +def _init_moe_weights(module, seed=0): + """Initialize a freshly built MoE block. + + transformers MoE experts allocate weights with `torch.empty` (uninitialized + garbage) and rely on the model's `_init_weights`. A standalone block skips + that, so we seed small finite values to get a meaningful reference. + """ + torch.manual_seed(seed) + with torch.no_grad(): + # Small finite values for every parameter, including biases, so the + # bias paths are exercised numerically rather than left at zero. + for p in module.parameters(): + nn.init.normal_(p, std=0.02) + + +def _read_graph_from_archive(path): + with tarfile.open(path, "r:*") as tf: + for member in tf.getmembers(): + if member.name.endswith("graph.nnef"): + return tf.extractfile(member).read().decode("utf-8") + raise AssertionError("graph.nnef not found in NNEF archive") + + +def _assert_moe_expert_weights_q40( + inference_target, + path, + expected_count, + expected_shapes=None, +): + """Check split expert tensors were exported as tract Q40 values.""" + if not isinstance(inference_target, TractNNEF): + return + expected_dtype = ( + DatBinHeader.TractCustomTypes.Q40 + if inference_target.version >= "0.21.11" + else DatBinHeader.TractCustomTypes.Q40_LEGACY + ) + with tempfile.TemporaryDirectory() as td, tarfile.open(path, "r:*") as tf: + members = [ + m + for m in tf.getmembers() + if m.name.endswith(("_w1.dat", "_w2.dat", "_w3.dat")) + ] + assert len(members) == expected_count, [m.name for m in members] + for member in members: + tf.extract(member, td) + header = DatBinHeader.from_dat(Path(td) / member.name) + assert header.torch_dtype_or_custom == expected_dtype + if expected_shapes is not None: + suffix = member.name.rsplit("_", 1)[-1].removesuffix(".dat") + assert header.dims == expected_shapes[suffix] + + +def _assert_moe_expert_weight_shapes( + inference_target, + path, + expected_count, + expected_shapes, +): + """Check split expert tensor shapes independently of their storage dtype.""" + if not isinstance(inference_target, TractNNEF): + return + with tempfile.TemporaryDirectory() as td, tarfile.open(path, "r:*") as tf: + members = [ + m + for m in tf.getmembers() + if m.name.endswith(("_w1.dat", "_w2.dat", "_w3.dat")) + ] + assert len(members) == expected_count, [m.name for m in members] + for member in members: + tf.extract(member, td) + header = DatBinHeader.from_dat(Path(td) / member.name) + suffix = member.name.rsplit("_", 1)[-1].removesuffix(".dat") + assert header.dims == expected_shapes[suffix] + + +def _assert_graph_contains(inference_target, path, fragment): + if not isinstance(inference_target, TractNNEF): + return + graph = _read_graph_from_archive(path) + assert fragment in graph + + +def _assert_graph_not_contains(inference_target, path, fragment): + if not isinstance(inference_target, TractNNEF): + return + graph = _read_graph_from_archive(path) + assert fragment not in graph + + +def _opaque_ref(tensor, tmp_path, name): + offloaded = OffloadedTensor.from_original_tensor( + tensor, + name, + offload_dir=tmp_path, + ) + return OpaqueTensorRef( + torch.empty(tuple(tensor.shape), dtype=tensor.dtype, device="meta"), + offloaded, + ) + + +@skipif_limited_offload_support +def test_granite_adapter_materializes_opaque_expert_views(tmp_path): + """Granite packed expert views must not remain meta tensors.""" + input_weight = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) + output_weight = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape( + 2, 4, 3 + ) + router_weight = torch.arange(2 * 4, dtype=torch.float32).reshape(2, 4) + moe = SimpleNamespace( + input_linear=SimpleNamespace( + weight=_opaque_ref(input_weight, tmp_path, "input_linear") + ), + output_linear=SimpleNamespace( + weight=_opaque_ref(output_weight, tmp_path, "output_linear") + ), + router=SimpleNamespace( + layer=SimpleNamespace( + weight=_opaque_ref(router_weight, tmp_path, "router") + ), + top_k=2, + ), + ) + + adapter = _GraniteMoEAdapter() + + w1 = adapter.expert_w1(moe) + w2 = adapter.expert_w2(moe) + w3 = adapter.expert_w3(moe) + + assert w1.device.type != "meta" + assert w2.device.type != "meta" + assert w3.device.type != "meta" + assert torch.equal(w1, input_weight[:, :3, :].transpose(-1, -2)) + assert torch.equal(w2, output_weight.transpose(-1, -2)) + assert torch.equal(w3, input_weight[:, 3:, :].transpose(-1, -2)) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_basic(inference_target): + """Export MoEFFN with 4 experts, top-2.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWrapper(num_experts=4, d_model=16, d_hidden=32, k=2) + model.eval() + check_model_io_test( + model=model, + test_input=(torch.randn(8, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@skipif_unsupported_qtensor +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_split_experts_q40(inference_target): + """Export MoEFFN with split expert tensors quantized to tract Q40.""" + _skip_if_unsupported(inference_target) + export_target = deepcopy(inference_target) + export_target.check_io = False + model = MoEFFNWrapper(num_experts=4, d_model=32, d_hidden=64, k=2) + model.moe._t2n_quantize_moe_experts_q40 = True + model.eval() + check_model_io_test( + model=model, + test_input=(torch.randn(8, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=export_target, + callback_post_export=partial( + _assert_moe_expert_weights_q40, + expected_count=2, + ), + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_split_experts_linear_layout(inference_target): + """Export MoE experts in native linear-filter layout.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWrapper(num_experts=4, d_model=32, d_hidden=64, k=2) + model.moe._t2n_moe_expert_layout = "linear" + model.eval() + + def _assert_linear_layout(inference_target, path): + _assert_graph_contains( + inference_target, + path, + "expert_layout = 'linear'", + ) + _assert_moe_expert_weight_shapes( + inference_target, + path, + expected_count=2, + expected_shapes={ + "w1": [4, 64, 32], + "w2": [4, 32, 64], + }, + ) + + check_model_io_test( + model=model, + test_input=(torch.randn(8, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + callback_post_export=_assert_linear_layout, + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_accepts_tract_moe_ffn_layout_alias(inference_target): + """Reloaded checkpoints may label canonical expert tensors by target op.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWrapper(num_experts=4, d_model=32, d_hidden=64, k=2) + model.moe._t2n_moe_expert_layout = "tract_moe_ffn" + model.eval() + + def _assert_canonical_layout(inference_target, path): + _assert_graph_not_contains( + inference_target, + path, + "expert_layout = 'linear'", + ) + _assert_moe_expert_weight_shapes( + inference_target, + path, + expected_count=2, + expected_shapes={ + "w1": [4, 32, 64], + "w2": [4, 64, 32], + }, + ) + + check_model_io_test( + model=model, + test_input=(torch.randn(8, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + callback_post_export=_assert_canonical_layout, + ) + + +@skipif_unsupported_qtensor +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_split_experts_q40_linear_layout(inference_target): + """Export Q40 experts after independently selecting linear layout.""" + _skip_if_unsupported(inference_target) + export_target = deepcopy(inference_target) + export_target.check_io = False + model = MoEFFNWrapper(num_experts=4, d_model=32, d_hidden=64, k=2) + model.moe._t2n_quantize_moe_experts_q40 = True + model.moe._t2n_moe_expert_layout = "linear" + model.eval() + + def _assert_linear_q40(inference_target, path): + _assert_graph_contains( + inference_target, + path, + "expert_layout = 'linear'", + ) + _assert_moe_expert_weights_q40( + inference_target, + path, + expected_count=2, + expected_shapes={ + "w1": [4, 64, 32], + "w2": [4, 32, 64], + }, + ) + + check_model_io_test( + model=model, + test_input=(torch.randn(8, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=export_target, + callback_post_export=_assert_linear_q40, + ) + + +@skipif_unsupported_qtensor +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_custom_q40_quantizer(inference_target): + """Allow callers to provide a calibrated Q40 quantizer for MoE experts.""" + _skip_if_unsupported(inference_target) + export_target = deepcopy(inference_target) + export_target.check_io = False + calls = [] + + def quantizer(tensor, marker): + calls.append((tuple(tensor.shape), marker, tensor.is_contiguous())) + return fp_to_tract_q4_0_with_min_max_calibration(tensor) + + model = MoEFFNWrapper(num_experts=4, d_model=32, d_hidden=64, k=2) + model.moe._t2n_quantize_moe_experts_q40 = True + model.moe._t2n_moe_expert_layout = "linear" + model.moe._t2n_quantize_moe_experts_q40_quantizer = quantizer + model.moe._t2n_quantize_moe_experts_q40_kwargs = {"marker": "calibrated"} + model.eval() + + check_model_io_test( + model=model, + test_input=(torch.randn(8, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=export_target, + ) + + assert calls == [ + ((4, 64, 32), "calibrated", True), + ((4, 32, 64), "calibrated", True), + ] + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_with_bias(inference_target): + """Export MoEFFN with bias terms.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWithBiasWrapper(num_experts=4, d_model=16, d_hidden=32, k=2) + model.eval() + check_model_io_test( + model=model, + test_input=(torch.randn(8, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@pytest.mark.parametrize("activation", ["silu", "gelu", "relu"]) +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_activations(inference_target, activation): + """Export MoEFFN with different activations.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWrapper( + num_experts=4, d_model=16, d_hidden=32, k=1, activation=activation + ) + model.eval() + check_model_io_test( + model=model, + test_input=(torch.randn(8, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_top1(inference_target): + """Export MoEFFN with top-1 routing.""" + _skip_if_unsupported(inference_target) + model = MoEFFNWrapper(num_experts=8, d_model=32, d_hidden=64, k=1) + model.eval() + check_model_io_test( + model=model, + test_input=(torch.randn(16, 32),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +class _GptOssMoEWrapper(nn.Module): + """Wrap a single GptOssMLP block, returning only routed hidden states. + + gpt-oss exercises the op's extra features: a router bias, fused + interleaved gate/up projections with biases, and the clamped SwiGLU + activation (alpha / limit / (up + 1)). + """ + + def __init__(self, mlp): + super().__init__() + self.mlp = mlp + + def forward(self, x): + # GptOssMLP returns (hidden_states, router_scores); router_scores is + # discarded at inference, mirroring transformers' decoder layer. + return self.mlp(x)[0] + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_gpt_oss(inference_target): + """Export a tiny gpt-oss MoE block (biases + clamped SwiGLU).""" + _skip_if_unsupported(inference_target) + gpt_oss = pytest.importorskip( + "transformers.models.gpt_oss.modeling_gpt_oss", + reason="transformers too old for gpt-oss", + ) + if not hasattr(gpt_oss, "GptOssMLP"): + pytest.skip("this transformers version has no GptOssMLP") + + cfg = gpt_oss.GptOssConfig( + hidden_size=16, + intermediate_size=8, + num_local_experts=4, + num_experts_per_tok=2, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=32, + ) + mlp = gpt_oss.GptOssMLP(cfg) + _init_moe_weights(mlp) + model = _GptOssMoEWrapper(mlp.eval()).eval() + # GptOssMLP expects a 3D [batch, seq, hidden] input. + check_model_io_test( + model=model, + test_input=(torch.randn(1, 6, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +class _BlockWrapper(nn.Module): + """Wrap a transformers MoE block returning a single hidden-state tensor.""" + + def __init__(self, block): + super().__init__() + self.block = block + + def forward(self, x): + out = self.block(x) + # Older blocks return (hidden_states, router_logits); keep hidden only. + return out[0] if isinstance(out, tuple) else out + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_qwen3(inference_target): + """Export a tiny Qwen3 MoE block (concatenated gate/up, plain SwiGLU).""" + _skip_if_unsupported(inference_target) + qwen3 = pytest.importorskip( + "transformers.models.qwen3_moe.modeling_qwen3_moe", + reason="transformers too old for qwen3-moe", + ) + cfg = qwen3.Qwen3MoeConfig( + hidden_size=24, + moe_intermediate_size=8, + num_experts=4, + num_experts_per_tok=2, + hidden_act="silu", + # real Qwen3-MoE checkpoints renormalize the top-k gates + norm_topk_prob=True, + ) + block = qwen3.Qwen3MoeSparseMoeBlock(cfg) + _init_moe_weights(block) + model = _BlockWrapper(block.eval()).eval() + # The block expects a 3D [batch, seq, hidden] input. + check_model_io_test( + model=model, + test_input=(torch.randn(1, 5, 24),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_mixtral(inference_target): + """Export a tiny Mixtral MoE block (the canonical fused MoE).""" + _skip_if_unsupported(inference_target) + mixtral = pytest.importorskip( + "transformers.models.mixtral.modeling_mixtral", + reason="transformers too old for mixtral", + ) + cfg = mixtral.MixtralConfig( + hidden_size=16, + intermediate_size=32, + num_local_experts=4, + num_experts_per_tok=2, + ) + block = mixtral.MixtralSparseMoeBlock(cfg) + _init_moe_weights(block) + model = _BlockWrapper(block.eval()).eval() + check_model_io_test( + model=model, + test_input=(torch.randn(1, 5, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_olmoe(inference_target): + """Export a tiny OLMoE block (Qwen-like fused experts, no shared expert).""" + _skip_if_unsupported(inference_target) + olmoe = pytest.importorskip( + "transformers.models.olmoe.modeling_olmoe", + reason="transformers too old for olmoe", + ) + cfg = olmoe.OlmoeConfig( + hidden_size=16, + intermediate_size=32, + num_experts=4, + num_experts_per_tok=2, + # OLMoE-1B-7B ships with norm_topk_prob=False -> "softmax_all" gating. + norm_topk_prob=False, + ) + block = olmoe.OlmoeSparseMoeBlock(cfg) + _init_moe_weights(block) + model = _BlockWrapper(block.eval()).eval() + check_model_io_test( + model=model, + test_input=(torch.randn(1, 5, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) + + +@pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) +def test_moe_ffn_qwen2_shared_expert(inference_target): + """Export a tiny Qwen2 MoE block with its always-on shared expert.""" + _skip_if_unsupported(inference_target) + qwen2 = pytest.importorskip( + "transformers.models.qwen2_moe.modeling_qwen2_moe", + reason="transformers too old for qwen2-moe", + ) + cfg = qwen2.Qwen2MoeConfig( + hidden_size=16, + moe_intermediate_size=32, + shared_expert_intermediate_size=64, + num_experts=4, + num_experts_per_tok=2, + norm_topk_prob=True, + ) + block = qwen2.Qwen2MoeSparseMoeBlock(cfg) + _init_moe_weights(block) + model = _BlockWrapper(block.eval()).eval() + check_model_io_test( + model=model, + test_input=(torch.randn(1, 5, 16),), + input_names=["tokens"], + output_names=["output"], + inference_target=inference_target, + ) diff --git a/tests/test_offload.py b/tests/test_offload.py index 397608114..40fd9d3de 100644 --- a/tests/test_offload.py +++ b/tests/test_offload.py @@ -19,11 +19,14 @@ OFFLOAD_STATE_KEY, OpaqueTensor, SupportsOffloadState, + set_opaque_tensor_in_params_as_ref, ) from torch_to_nnef.tensor.quant.qtract import ( QTensorTractScaleOnly, fp_to_tract_q4_0_with_min_max_calibration, ) +from torch_to_nnef.torch_graph.ir_graph import module_tracer_into_ir_graph +from torch_to_nnef.torch_graph.ir_module_tracer import TorchModuleTracer @skipif_limited_offload_support @@ -43,6 +46,28 @@ def test_int_opaque_tensor_is_materialized_for_meta_trace(): assert torch.equal(traced, int_weight) +@skipif_limited_offload_support +def test_offloaded_linear_with_bias_traces_on_target_device(tmp_path): + model = torch.nn.Linear(3, 4).eval() + model.weight = torch.nn.Parameter( + OffloadedTensor.from_original_tensor( + model.weight.detach(), "weight", offload_dir=tmp_path + ), + requires_grad=False, + ) + model.bias = torch.nn.Parameter( + OffloadedTensor.from_original_tensor( + model.bias.detach(), "bias", offload_dir=tmp_path + ), + requires_grad=False, + ) + + set_opaque_tensor_in_params_as_ref(model) + module_tracer_into_ir_graph( + TorchModuleTracer(model, args=(torch.randn(2, 3),)) + ) + + @skipif_limited_offload_support @pytest.mark.parametrize("inference_target", TRACT_INFERENCES_TO_TESTS_APPROX) def test_offload_tensor_export_with_tract_and_conv2d(inference_target): diff --git a/tests/test_qtensor.py b/tests/test_qtensor.py index c116dfd8a..b4c6b6ffb 100644 --- a/tests/test_qtensor.py +++ b/tests/test_qtensor.py @@ -15,7 +15,6 @@ from torch_to_nnef.exceptions import ( T2NErrorTestFailed, - T2NErrorTorchJitTraceFailed, ) from torch_to_nnef.inference_target.base import InferenceTarget from torch_to_nnef.inference_target.tract import TractCheckTolerance, TractNNEF @@ -288,12 +287,13 @@ def forward(self, x): @skipif_unsupported_qtensor -def test_opaque_trace_tensor_materializes_real_data_on_non_meta_device(): +def test_opaque_trace_tensor_uses_fake_target_device_for_meta_request(): q_tensor = fp_to_tract_q4_0_with_min_max_calibration(torch.randn(8, 32)) - meta = q_tensor._to_trace_tensor("meta") - assert meta.device.type == "meta" - assert meta.shape == q_tensor.shape + traced = q_tensor._to_trace_tensor("meta") + assert traced.device == q_tensor.device + assert traced.shape == q_tensor.shape + assert getattr(traced, "fake_mode", None) is not None # a real device must expose decompressed values, not uninitialized memory, # otherwise constant-index gathers bake garbage into the exported graph. @@ -304,10 +304,11 @@ def test_opaque_trace_tensor_materializes_real_data_on_non_meta_device(): @skipif_unsupported_qtensor @skipif_no_meta_opaque_tracing -def test_opaque_meta_arithmetic_raises_actionable_error(): - # Combining a meta-traced opaque-weight view with a real tensor is - # unsupported by the meta strategy; the trace failure must carry an - # actionable hint rather than a bare device-mismatch RuntimeError. +def test_opaque_fake_trace_allows_same_device_arithmetic( + forbid_qtensor_materialization, +): + # Fake target-device placeholders keep opaque weights non-materialized + # while still allowing same-device arithmetic during trace. class ScaledViewedWeight(nn.Module): def __init__(self): super().__init__() @@ -323,10 +324,9 @@ def forward(self, x): model = ScaledViewedWeight().eval() set_opaque_tensor_in_params_as_ref(model) - with pytest.raises(T2NErrorTorchJitTraceFailed, match="meta placeholder"): - module_tracer_into_ir_graph( - TorchModuleTracer(model, args=(torch.randn(1, 32),)) - ) + module_tracer_into_ir_graph( + TorchModuleTracer(model, args=(torch.randn(1, 32),)) + ) @skipif_unsupported_qtensor diff --git a/torch_to_nnef/inference_target/tract.py b/torch_to_nnef/inference_target/tract.py index 875984168..5149b25b5 100644 --- a/torch_to_nnef/inference_target/tract.py +++ b/torch_to_nnef/inference_target/tract.py @@ -134,6 +134,7 @@ def __init__( force_linear_accumulation_in_f32: bool = False, force_norm_in_f32: bool = False, reify_sdpa_operator: T.Optional[bool] = None, + upcast_reified_sdpa_inputs_to_f32: bool = True, upsample_with_debox: bool = False, ): """Init. @@ -187,6 +188,12 @@ def __init__( as a tract operator (intead of a NNEF fragment), default false until tract v0.22.0 included then true, except if specified. Experimental feature. + upcast_reified_sdpa_inputs_to_f32: + when reifying scaled_dot_product_attention and forcing f32 + attention internals, insert f32 casts at the SDPA boundary for + f16 inputs. Disable this only when the target tract runtime can + run tract_transformers_sdpa with f16 inputs and f32 + accumulators. upsample_with_debox: use debox upsample operator instead of deconvolution. This should be faster. @@ -225,6 +232,9 @@ def __init__( if reify_sdpa_operator is None: reify_sdpa_operator = self.version > "0.22.0" self.reify_sdpa_operator = reify_sdpa_operator + self.upcast_reified_sdpa_inputs_to_f32 = ( + upcast_reified_sdpa_inputs_to_f32 + ) self.upsample_with_debox = upsample_with_debox self.dump_identity_properties = dump_identity_properties if self.feature_flags: diff --git a/torch_to_nnef/nnef_io/writer.py b/torch_to_nnef/nnef_io/writer.py index eb4f07b4b..176a47b67 100644 --- a/torch_to_nnef/nnef_io/writer.py +++ b/torch_to_nnef/nnef_io/writer.py @@ -394,12 +394,16 @@ def _write_tensors_from_operators(self, graph, folder): folder, label, self._inference_target ): LOGGER.info("written qtensor: '%s'", label) + del qtensor + gc.collect() continue while isinstance(qtensor, OffloadedTensor): qtensor = qtensor.to_base_tensor() label = op.attribs["label"] qtensor.write_in_file(folder, label, self._inference_target) LOGGER.info("written qtensor: '%s'", label) + del qtensor + gc.collect() else: filename = op.attribs["label"] + ".dat" output_data = op.output.data diff --git a/torch_to_nnef/op/aten/attn.py b/torch_to_nnef/op/aten/attn.py index 5e1e78cb8..2c04fbd87 100644 --- a/torch_to_nnef/op/aten/attn.py +++ b/torch_to_nnef/op/aten/attn.py @@ -1,6 +1,9 @@ """Attention mechanisms.""" +import numpy as np import torch +from nnef_tools.model import Operation as NOperation +from nnef_tools.model import Tensor as NTensor from torch_to_nnef.exceptions import T2NErrorNotImplemented from torch_to_nnef.inference_target.base import InferenceTarget @@ -58,7 +61,33 @@ def _broadcast_mask_query_axis( ) +def _emit_tract_cast(g, src: NTensor, name: str, to: str, dtype) -> NTensor: + """Emit a tract_core_cast intermediate with an explicit dtype.""" + out = NTensor(g, name=name, dtype=dtype, shape=tuple(src.shape)) + NOperation( + g, + type="tract_core_cast", + attribs={"to": to}, + inputs=src, + outputs=out, + ) + return out + + +def _cast_sdpa_input_to_f32(g, src: NTensor, suffix: str) -> NTensor: + if src.dtype == np.float32: + return src + return _emit_tract_cast( + g, + src, + f"{src.name}_{suffix}", + "f32", + np.float32, + ) + + @OP_REGISTRY.register() +# pylint: disable-next=too-many-branches def scaled_dot_product_attention( g, node, name_to_tensor, inference_target, **kwargs ): @@ -142,23 +171,71 @@ def scaled_dot_product_attention( ) if reify_tract_spda: + sdpa_inputs = list(inputs) + sdpa_dtype_str = dtype_str + cast_sdpa_output_back = False + if ( + inference_target.force_attention_inner_in_f32 + and inference_target.upcast_reified_sdpa_inputs_to_f32 + and query_node.dtype == torch.float16 + ): + sdpa_inputs[0] = _cast_sdpa_input_to_f32( + g, sdpa_inputs[0], "sdpa_q_f32" + ) + sdpa_inputs[1] = _cast_sdpa_input_to_f32( + g, sdpa_inputs[1], "sdpa_k_f32" + ) + sdpa_inputs[2] = _cast_sdpa_input_to_f32( + g, sdpa_inputs[2], "sdpa_v_f32" + ) + if has_masked_attn and len(sdpa_inputs) > 3: + sdpa_inputs[3] = _cast_sdpa_input_to_f32( + g, sdpa_inputs[3], "sdpa_mask_f32" + ) + sdpa_dtype_str = "f32" + cast_sdpa_output_back = True + # Define SDPA attributes attrs = { - "datum_type": dtype_str, + "datum_type": sdpa_dtype_str, "acc_datum_type": inner_dtype, "is_causal": is_causal, } if scale is not None: attrs["scale"] = scale - add_single_output_op( - g, - node, - name_to_tensor, - "tract_transformers_sdpa", - inputs=tuple(inputs), - attrs=attrs, - ) + if cast_sdpa_output_back: + output_node = node.outputs[0] + sdpa_out = NTensor( + g, + name=f"{output_node.export_name}_sdpa_f32", + dtype=np.float32, + shape=tuple(output_node.shape), + ) + NOperation( + g, + type="tract_transformers_sdpa", + attribs=attrs, + inputs=tuple(sdpa_inputs), + outputs=sdpa_out, + ) + add_single_output_op( + g, + node, + name_to_tensor, + "tract_core_cast", + inputs=sdpa_out, + attrs={"to": dtype_str}, + ) + else: + add_single_output_op( + g, + node, + name_to_tensor, + "tract_transformers_sdpa", + inputs=tuple(sdpa_inputs), + attrs=attrs, + ) return ["tract_transformers"] tmpl_fragment_name = "scaled_dot_product_attention" diff --git a/torch_to_nnef/op/aten/math.py b/torch_to_nnef/op/aten/math.py index 6a84b1f77..7f612d324 100644 --- a/torch_to_nnef/op/aten/math.py +++ b/torch_to_nnef/op/aten/math.py @@ -125,6 +125,13 @@ def div(node, op_helper, inference_target, torch_graph, **kwargs): divisor_tensor, ), output_tensor_name_suffix=suffix_div_op_output, + # When the traced output is integer we deliberately keep the + # division (and its rounding fragment) in float and cast the + # result back to int below. Letting the generic implicit-cast + # realign the f32 operands to the int output dtype would run the + # rounding fragment (e.g. `trunc`'s `select(x < 0.0, ...)`) on + # integers, which tract cannot type-resolve. + maybe_cast_align_tract=io_casting_with_dtype is None, ) if rounding_mode: @@ -489,11 +496,9 @@ def rsub(node, op_helper, torch_graph, **kwargs): for _ in [input_node, other_node] ] for idx, inp in enumerate(inputs): - inputs[idx] = op_helper.add_single_output_op_from_nnef_tensors( - node, - "tract_core_cast", - inputs=[inp], - attrs={"to": "f32"}, + inputs[idx] = op_helper.add_cast_nnef_tensor( + inp, + cast_to=np.float32, force_full_output_tensor_name=f"{inp.name}_as_f32", ) diff --git a/torch_to_nnef/op/aten/matmul.py b/torch_to_nnef/op/aten/matmul.py index 3d8ec6df4..7d7cd7786 100644 --- a/torch_to_nnef/op/aten/matmul.py +++ b/torch_to_nnef/op/aten/matmul.py @@ -1,5 +1,6 @@ import typing as T +import numpy as np import torch from nnef_tools.model import Operation as NOperation from nnef_tools.model import Tensor as NTensor @@ -696,10 +697,40 @@ def _emit_unsqueeze_intermediate( ) +def _decomposed_attention_qk_expr(node, input_node, other_node): + """Return the einsum form for a decomposed self-attention QK matmul.""" + output_name = node.outputs[0].export_name + if "selfAttn_matmul0" not in output_name: + return None + if input_node.dtype != torch.float16 or other_node.dtype != torch.float16: + return None + if input_node.rank != other_node.rank: + return None + if input_node.rank == 3: + return "bij,bjk->bik" + if input_node.rank == 4: + return "bcij,bcjk->bcik" + return None + + +def _cast_attention_qk_operand_to_f32(g, src: NTensor, name: str) -> NTensor: + """Emit an f32 cast intermediate with the operand's own shape.""" + out = NTensor(g, name=name, dtype=np.float32, shape=tuple(src.shape)) + NOperation( + g, + type="tract_core_cast", + attribs={"to": "f32"}, + inputs=src, + outputs=out, + ) + return out + + @OP_REGISTRY.register( torch_op_ids=["matmul", "bmm", "mm"] ) # since NNEF matmul does not care about rank -def matmul(g, node, name_to_tensor, op_helper, **kwargs): +# pylint: disable-next=inconsistent-return-statements +def matmul(g, node, name_to_tensor, op_helper, inference_target, **kwargs): """Map PyTorch: 'aten:matmul', 'aten:bmm', 'aten:mm' to NNEF. NNEF `matmul` requires *equal* rank on both operands; PyTorch's @@ -731,6 +762,30 @@ def matmul(g, node, name_to_tensor, op_helper, **kwargs): b_is_v = b_rank < 2 if not (a_is_v or b_is_v): + if ( + isinstance(inference_target, TractNNEF) + and inference_target.force_attention_inner_in_f32 + ): + expr = _decomposed_attention_qk_expr(node, input_node, other_node) + if expr is not None: + a_ref = _cast_attention_qk_operand_to_f32( + g, a_ref, f"{node.outputs[0].export_name}_a_f32" + ) + b_ref = _cast_attention_qk_operand_to_f32( + g, b_ref, f"{node.outputs[0].export_name}_b_f32" + ) + add_single_output_op( + g, + node, + name_to_tensor, + "tract_core_einsum", + inputs=[a_ref, b_ref], + ensure_tuple=False, + force_consistent_inputs_shapes=False, + attrs={"expr": expr, "acc": "f32", "output": "f32"}, + ) + return ["tract_core"] + # Standard case: both rank >= 2. Let the generic emitter handle # it: `force_consistent_inputs_shapes` prepends 1s to the # smaller-rank input if needed (e.g. rank-3 @ rank-2), and the diff --git a/torch_to_nnef/op/custom_extractors/__init__.py b/torch_to_nnef/op/custom_extractors/__init__.py index 87c6d54d0..375b6da92 100644 --- a/torch_to_nnef/op/custom_extractors/__init__.py +++ b/torch_to_nnef/op/custom_extractors/__init__.py @@ -18,6 +18,7 @@ ) # load default custom registries +from torch_to_nnef.op.custom_extractors.moe import MoEFFN # noqa: F401 from torch_to_nnef.op.custom_extractors.rnn import ( GRUExtractor, LSTMCellExtractor, @@ -32,4 +33,5 @@ "LSTMExtractor", "LSTMCellExtractor", "GRUExtractor", + "MoEFFN", ] diff --git a/torch_to_nnef/op/custom_extractors/moe.py b/torch_to_nnef/op/custom_extractors/moe.py new file mode 100644 index 000000000..6a0ad73c1 --- /dev/null +++ b/torch_to_nnef/op/custom_extractors/moe.py @@ -0,0 +1,1019 @@ +"""MoE FFN export to tract_moe_ffn operator (tract_transformers extension). + +Supports (transformers): +- MoEFFN (reference wrapper) +- MixtralSparseMoeBlock / MistralSparseMoeBlock +- GptOssMLP (router + expert biases, interleaved gate/up, clamped SwiGLU) +- Qwen2/Qwen3/Qwen3.5 MoE (Qwen2 & 3.5 shared expert decomposed outside) +- OlmoeSparseMoeBlock + +All variants are normalized to the same tract_moe_ffn signature: + inputs: x [T,D], wg [E,D], w1 [E,D,H], w2 [E,H,D], w3 [E,D,H], + optional biases (wg_bias, w1_bias, w3_bias, w2_bias) + attrs: k (int), activation (str), gate (softmax_topk | softmax_all | + sigmoid | raw), optional act_alpha / act_limit (clamped SwiGLU), + optional expert_layout ("canonical" | "linear") + output: y [T,D] +The default canonical expert layout matches the signature above. The optional +linear layout stores expert weights in their native linear-filter orientation: + w1/w3 [E,H,D], w2 [E,D,H] +Layout selection is independent from whether those tensors are quantized. +A shared expert (Qwen2 / Qwen3.5) is emitted as a standard NNEF subgraph +added on top of the routed output, not baked into the op. +""" + +import logging +import typing as T + +import torch +from torch import nn + +from torch_to_nnef.exceptions import ( + T2NErrorNotImplemented, + T2NErrorStrictNNEFSpec, +) +from torch_to_nnef.inference_target import TractNNEF +from torch_to_nnef.op.custom_extractors.base import ModuleInfoExtractor +from torch_to_nnef.tensor.offload import OffloadedTensor +from torch_to_nnef.tensor.opaque import OpaqueTensorRef, opaque_to_final_tensor +from torch_to_nnef.tensor.quant import ( + QTensor, + fp_to_tract_q4_0_with_min_max_calibration, +) + +LOGGER = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Weight adapters — normalize diverse layouts into unified tensors +# --------------------------------------------------------------------------- + + +class _MoEWeightAdapter: + """Base adapter: extract MoE weights from a module into canonical shapes.""" + + def __init__(self) -> None: + self._constant_cache: T.Dict[int, torch.Tensor] = {} + + def _constant(self, tensor: torch.Tensor) -> torch.Tensor: + """Materialize an export-time constant before taking tensor views.""" + key = ( + id(tensor.opaque_tensor) + if isinstance(tensor, OpaqueTensorRef) + else id(tensor) + ) + if key not in self._constant_cache: + if isinstance(tensor, OpaqueTensorRef): + tensor = tensor.opaque_tensor + self._constant_cache[key] = opaque_to_final_tensor(tensor).detach() + return self._constant_cache[key] + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + """Router weight [E, D].""" + raise T2NErrorNotImplemented() + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + """Gate projection [E, D, H] (SwiGLU gate branch).""" + raise T2NErrorNotImplemented() + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + """Down projection [E, H, D].""" + raise T2NErrorNotImplemented() + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + """Up projection [E, D, H] (SwiGLU up branch).""" + raise T2NErrorNotImplemented() + + def expert_sources(self, m: nn.Module) -> T.Sequence[torch.Tensor]: + """Raw tensors that feed expert_w1/w2/w3. + + This is used only to detect pre-quantized expert sources. Most adapters + do not need to override it because split-time quantization is normally + requested explicitly via the module marker. + """ + return () + + def top_k(self, m: nn.Module) -> int: + raise T2NErrorNotImplemented() + + def activation(self, m: nn.Module) -> str: + return "swiglu" + + def gate(self, m: nn.Module) -> str: + """How router logits become top-k gate weights (tract_moe_ffn `gate`). + + One of: "softmax_topk" (softmax over the top-k logits), "softmax_all" + (softmax over all experts, gather top-k, no renormalization), + "sigmoid" (per-expert sigmoid), "raw" (raw top-k logits). + """ + return "softmax_topk" + + # Optional biases (None unless the arch has them, e.g. gpt-oss). + def gate_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + """Router bias [E].""" + return None + + def expert_w1_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + """Gate projection bias [E, H].""" + return None + + def expert_w2_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + """Down projection bias [E, D].""" + return None + + def expert_w3_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + """Up projection bias [E, H].""" + return None + + # Optional clamped-SwiGLU params (gpt-oss). When act_limit is not None the + # op uses the clamped activation: gate.clamp(max=limit) / + # up.clamp(+-limit) / glu = gate*sigmoid(alpha*gate) / out = (up+1)*glu. + def act_alpha(self, m: nn.Module) -> T.Optional[float]: + return None + + def act_limit(self, m: nn.Module) -> T.Optional[float]: + return None + + # Optional always-on shared expert (e.g. Qwen2-MoE / Qwen3.5-MoE): + # out = routed_experts(x) + sigmoid(shared_gate(x)) * shared_mlp(x) + # Returns the raw nn.Linear weights ([out, in], used directly as NNEF + # linear filters) or None when the arch has no shared expert. + def shared_expert( + self, m: nn.Module + ) -> T.Optional[T.Dict[str, torch.Tensor]]: + return None + + +class _MoEFFNAdapter(_MoEWeightAdapter): + """Adapter for our reference MoEFFN wrapper.""" + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + return m.gate.weight.detach() + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + return m.w1.detach() + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + return m.w2.detach() + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + # MoEFFN has no w3 — duplicate w1 shape as zeros + # so the SwiGLU gate branch becomes a no-op multiply by 0+silu + # In practice MoEFFN uses simple activation, not SwiGLU. + raise T2NErrorNotImplemented( + "MoEFFN uses simple activation, not SwiGLU. " + "Use activation attr instead." + ) + + def top_k(self, m: nn.Module) -> int: + return m.k + + def activation(self, m: nn.Module) -> str: + return m.activation_name + + def gate(self, m: nn.Module) -> str: + # MoEFFN's normalize_gates softmaxes the top-k logits; otherwise it + # uses the raw top-k logits. + return "softmax_topk" if m.normalize_gates else "raw" + + +class _MixtralAdapter(_MoEWeightAdapter): + """Adapter for transformers MixtralSparseMoeBlock. + + Handles two layouts: + - Legacy (transformers <4.52): ModuleList of experts with w1/w2/w3 + w1=gate_proj, w2=down_proj, w3=up_proj + - Modern (transformers >=4.52): fused MixtralExperts with + gate_up_proj [E, 2*H, D] and down_proj [E, D, H] + """ + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + return m.gate.weight.detach() + + def _is_fused(self, m: nn.Module) -> bool: + return hasattr(m.experts, "gate_up_proj") + + def _stack_legacy(self, m: nn.Module, attr: str) -> torch.Tensor: + return torch.stack( + [self._constant(getattr(e, attr).weight) for e in m.experts] + ) + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + if self._is_fused(m): + gate_up = self._constant(m.experts.gate_up_proj) + half = gate_up.shape[1] // 2 + return gate_up[:, :half, :].transpose(-1, -2) + # legacy: w1 = gate_proj [H, D] → [E, D, H] + return self._stack_legacy(m, "w1").transpose(-1, -2) + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + if self._is_fused(m): + return self._constant(m.experts.down_proj).transpose(-1, -2) + # legacy: w2 = down_proj [D, H] → [E, H, D] + return self._stack_legacy(m, "w2").transpose(-1, -2) + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + if self._is_fused(m): + gate_up = self._constant(m.experts.gate_up_proj) + half = gate_up.shape[1] // 2 + return gate_up[:, half:, :].transpose(-1, -2) + # legacy: w3 = up_proj [H, D] → [E, D, H] + return self._stack_legacy(m, "w3").transpose(-1, -2) + + def top_k(self, m: nn.Module) -> int: + return m.top_k + + +class _GptOssAdapter(_MoEWeightAdapter): + """Adapter for transformers GPT-OSS MoE block. + + Handles the modern GptOssMLP (transformers >=4.55: a `router` plus a fused + `GptOssExperts`) and the older `GptOssSparseMoeBlock` name. gpt-oss differs + from the generic op in several ways the adapter normalizes: + - the router carries a bias and lives on `m.router` (older: `m.gate`) + - experts fuse the gate and up projections INTERLEAVED inside gate_up_proj + [E, D, 2H] (gate = [..., 0::2], up = [..., 1::2]), and carry biases + - the activation is a clamped SwiGLU with alpha / limit and a (up + 1) term + """ + + @staticmethod + def _router(m: nn.Module) -> nn.Module: + return m.router if hasattr(m, "router") else m.gate + + def _d_model(self, m: nn.Module) -> int: + return self.gate_weight(m).shape[1] + + def _gate_up(self, m: nn.Module) -> torch.Tensor: + """gate_up_proj oriented as [E, D, 2H] (contracted axis = D).""" + gu = self._constant(m.experts.gate_up_proj) + d_model = self._d_model(m) + if gu.shape[1] != d_model and gu.shape[2] == d_model: + gu = gu.transpose(-1, -2) + return gu + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + return self._router(m).weight.detach() + + def gate_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + b = getattr(self._router(m), "bias", None) + return b.detach() if b is not None else None + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + # gate (activated) branch: interleaved even columns -> [E, D, H] + return self._gate_up(m)[:, :, 0::2].contiguous() + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + # up branch: interleaved odd columns -> [E, D, H] + return self._gate_up(m)[:, :, 1::2].contiguous() + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + # down_proj is [E, H, D] already (op's w2 layout); orient defensively + dp = self._constant(m.experts.down_proj) + d_model = self._d_model(m) + if dp.shape[2] != d_model and dp.shape[1] == d_model: + dp = dp.transpose(-1, -2) + return dp.contiguous() + + def _gate_up_bias( + self, m: nn.Module + ) -> T.Tuple[T.Optional[torch.Tensor], T.Optional[torch.Tensor]]: + b = getattr(m.experts, "gate_up_proj_bias", None) + if b is None: + return None, None + b = self._constant(b) # [E, 2H] interleaved + return b[:, 0::2].contiguous(), b[:, 1::2].contiguous() + + def expert_w1_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + return self._gate_up_bias(m)[0] + + def expert_w3_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + return self._gate_up_bias(m)[1] + + def expert_w2_bias(self, m: nn.Module) -> T.Optional[torch.Tensor]: + b = getattr(m.experts, "down_proj_bias", None) + return self._constant(b) if b is not None else None + + def top_k(self, m: nn.Module) -> int: + r = self._router(m) + if hasattr(r, "top_k"): + return r.top_k + return m.top_k + + def gate(self, m: nn.Module) -> str: + # gpt-oss softmaxes the top-k router logits. + return "softmax_topk" + + def act_alpha(self, m: nn.Module) -> T.Optional[float]: + return float(getattr(m.experts, "alpha", 1.702)) + + def act_limit(self, m: nn.Module) -> T.Optional[float]: + return float(getattr(m.experts, "limit", 7.0)) + + +class _QwenMoEAdapter(_MoEWeightAdapter): + """Adapter for transformers Qwen2MoE / Qwen3.5 MoE. + + Experts are fused tensors: gate_up_proj [E, 2*H, D], down_proj [E, D, H]. + Shared expert is NOT handled here — it is decomposed outside the op. + """ + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + return m.gate.weight.detach() + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + # gate_up_proj [E, 2*H, D] → split → gate half [E, H, D] → [E, D, H] + gate_up = self._constant(m.experts.gate_up_proj) + half = gate_up.shape[1] // 2 + return gate_up[:, :half, :].transpose(-1, -2) + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + # down_proj [E, D, H] → [E, H, D] + return self._constant(m.experts.down_proj).transpose(-1, -2) + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + # gate_up_proj [E, 2*H, D] → split → up half [E, H, D] → [E, D, H] + gate_up = self._constant(m.experts.gate_up_proj) + half = gate_up.shape[1] // 2 + return gate_up[:, half:, :].transpose(-1, -2) + + def top_k(self, m: nn.Module) -> int: + # transformers <5.x: m.top_k, >=5.x: m.gate.top_k + if hasattr(m, "top_k"): + return m.top_k + return m.gate.top_k + + def shared_expert( + self, m: nn.Module + ) -> T.Optional[T.Dict[str, torch.Tensor]]: + # Qwen2-MoE / Qwen3.5-MoE add a sigmoid-gated shared expert; Qwen3-MoE + # has none (no shared_expert attribute). + if not hasattr(m, "shared_expert"): + return None + se = m.shared_expert + # The converter emits the shared expert as silu-SwiGLU; reject any + # other activation rather than silently exporting wrong maths. + act = getattr(se, "act_fn", None) + if act is not None and "silu" not in type(act).__name__.lower(): + raise T2NErrorNotImplemented( + "shared expert activation " + f"{type(act).__name__} is unsupported (only SiLU/SwiGLU)" + ) + return { + "gate_proj": se.gate_proj.weight.detach(), + "up_proj": se.up_proj.weight.detach(), + "down_proj": se.down_proj.weight.detach(), + "router": m.shared_expert_gate.weight.detach(), + } + + def gate(self, m: nn.Module) -> str: + # Qwen / OLMoE routers softmax over ALL experts then take the top-k. + # With norm_topk_prob the top-k weights are renormalized, which is + # identical to softmaxing over the top-k logits ("softmax_topk"). + # Without it, the raw softmax-over-all weights are kept ("softmax_all", + # e.g. OLMoE-1B-7B). + router = getattr(m, "gate", None) + norm = getattr(router, "norm_topk_prob", None) + if norm is None: + norm = getattr(m, "norm_topk_prob", True) + return "softmax_topk" if norm else "softmax_all" + + +class _GraniteMoEAdapter(_MoEWeightAdapter): + """Adapter for transformers GraniteMoeMoE (IBM Granite MoE). + + Same maths as Qwen (concatenated fused gate/up, softmax over top-k logits, + SiLU SwiGLU), only the attribute names differ. Older Granite exports used + `input_linear` [E, 2H, D] / `output_linear` [E, D, H] directly on the MoE + module. Current Transformers stores them under `experts.gate_up_proj` and + `experts.down_proj`, and the router weight directly under `router.weight`. + """ + + def _gate_up_proj(self, m: nn.Module) -> torch.Tensor: + if hasattr(m, "input_linear"): + return m.input_linear.weight + return m.experts.gate_up_proj + + def _down_proj(self, m: nn.Module) -> torch.Tensor: + if hasattr(m, "output_linear"): + return m.output_linear.weight + return m.experts.down_proj + + def gate_weight(self, m: nn.Module) -> torch.Tensor: + router = m.router + if hasattr(router, "layer"): + return router.layer.weight.detach() + return router.weight.detach() + + def expert_w1(self, m: nn.Module) -> torch.Tensor: + input_weight = self._constant(self._gate_up_proj(m)) + half = input_weight.shape[1] // 2 + return input_weight[:, :half, :].transpose(-1, -2) + + def expert_w3(self, m: nn.Module) -> torch.Tensor: + input_weight = self._constant(self._gate_up_proj(m)) + half = input_weight.shape[1] // 2 + return input_weight[:, half:, :].transpose(-1, -2) + + def expert_w2(self, m: nn.Module) -> torch.Tensor: + return self._constant(self._down_proj(m)).transpose(-1, -2) + + def expert_sources(self, m: nn.Module) -> T.Sequence[torch.Tensor]: + return (self._gate_up_proj(m), self._down_proj(m)) + + def top_k(self, m: nn.Module) -> int: + return m.router.top_k + + +# --------------------------------------------------------------------------- +# Adapter dispatch +# --------------------------------------------------------------------------- + +_ADAPTER_BY_CLASSNAME: T.Dict[str, T.Type[_MoEWeightAdapter]] = { + "MoEFFN": _MoEFFNAdapter, + # Mixtral / Mistral + "MixtralSparseMoeBlock": _MixtralAdapter, + "MistralSparseMoeBlock": _MixtralAdapter, + # GPT-OSS (transformers >=4.55 renamed the block to GptOssMLP) + "GptOssMLP": _GptOssAdapter, + "GptOssSparseMoeBlock": _GptOssAdapter, + # Qwen 2 / 3 / 3.5 MoE + "Qwen2MoeSparseMoeBlock": _QwenMoEAdapter, + "Qwen3MoeSparseMoeBlock": _QwenMoEAdapter, + "Qwen3_5MoeSparseMoeBlock": _QwenMoEAdapter, + # OLMoE shares the Qwen layout (fused [E, 2H, D] experts, softmax top-k + # router with norm_topk_prob, no shared expert). + "OlmoeSparseMoeBlock": _QwenMoEAdapter, + # IBM Granite MoE: same maths, different attribute names. + "GraniteMoeMoE": _GraniteMoEAdapter, + "GraniteMoeSharedMoE": _GraniteMoEAdapter, +} + + +def _get_adapter(module: nn.Module) -> _MoEWeightAdapter: + cls_name = type(module).__name__ + adapter_cls = _ADAPTER_BY_CLASSNAME.get(cls_name) + if adapter_cls is None: + raise T2NErrorNotImplemented( + f"No MoE weight adapter for '{cls_name}'. " + f"Supported: {sorted(_ADAPTER_BY_CLASSNAME.keys())}" + ) + return adapter_cls() + + +# --------------------------------------------------------------------------- +# Reference MoEFFN module (for testing / wrapping custom MoE blocks) +# --------------------------------------------------------------------------- + + +class MoEFFN(nn.Module): + """Reference MoE FFN block for export testing. + + For production models use the transformers extractors directly. + """ + + def __init__( + self, + num_experts: int, + d_model: int, + d_hidden: int, + k: int = 2, + activation: str = "silu", + normalize_gates: bool = True, + bias: bool = False, + ): + super().__init__() + self.num_experts = num_experts + self.d_model = d_model + self.d_hidden = d_hidden + self.k = k + self.activation_name = activation + self.normalize_gates = normalize_gates + + self.gate = nn.Linear(d_model, num_experts, bias=False) + self.w1 = nn.Parameter(torch.empty(num_experts, d_model, d_hidden)) + self.w2 = nn.Parameter(torch.empty(num_experts, d_hidden, d_model)) + + if bias: + self.b1 = nn.Parameter(torch.zeros(num_experts, d_hidden)) + self.b2 = nn.Parameter(torch.zeros(num_experts, d_model)) + else: + self.b1 = None + self.b2 = None + + self._init_weights() + + _activations = {"silu": nn.SiLU(), "gelu": nn.GELU(), "relu": nn.ReLU()} + if activation not in _activations: + raise ValueError(f"Unsupported activation: {activation}") + self.activation = _activations[activation] + + def _init_weights(self): + nn.init.kaiming_uniform_(self.w1) + nn.init.kaiming_uniform_(self.w2) + nn.init.xavier_uniform_(self.gate.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward: [T, D] -> [T, D].""" + t_tokens, d = x.shape + router_logits = self.gate(x) + top_k_values, top_k_indices = torch.topk(router_logits, self.k, dim=-1) + + if self.normalize_gates: + gate_weights = torch.softmax(top_k_values, dim=-1) + else: + gate_weights = top_k_values + + output = torch.zeros(t_tokens, d, device=x.device, dtype=x.dtype) + + for ki in range(self.k): + expert_indices = top_k_indices[:, ki] + weights = gate_weights[:, ki] + + for eid in range(self.num_experts): + mask = expert_indices == eid + if not mask.any(): + continue + + xi = x[mask] + h = xi @ self.w1[eid] + if self.b1 is not None: + h = h + self.b1[eid] + h = self.activation(h) + yo = h @ self.w2[eid] + if self.b2 is not None: + yo = yo + self.b2[eid] + + output[mask] += weights[mask].unsqueeze(-1) * yo + + return output + + +# --------------------------------------------------------------------------- +# Core NNEF conversion (shared by all extractors) +# --------------------------------------------------------------------------- + + +def _emit_shared_expert( + g, + name_to_tensor, + shared, + input_tensor, + routed, + out0, + x_shape, + add_weight, + mk, +): + """Graft a shared expert on top of the routed MoE output (Qwen2 / Qwen3.5). + + out = routed + sigmoid(x @ Wr^T) * down(silu(x @ Wg^T) * (x @ Wu^T)). + All weights are raw nn.Linear [out, in], used as NNEF linear filters + (linear computes x @ filter^T) so no transpose is needed. + """ + # pylint: disable-next=import-outside-toplevel + from torch_to_nnef.op import helper + + def emit(op_type, op_inputs, out_t, attribs=None, as_list=False): + # Most ops (sigmoid/mul/add) take positional inputs; tract_core_einsum + # reads its operands from a single list-valued `inputs=[...]` argument. + helper.cast_and_add_nnef_operation( + name_to_tensor=name_to_tensor, + graph=g, + type=op_type, + inputs=list(op_inputs) if as_list else tuple(op_inputs), + outputs=(out_t,), + attribs=attribs or {}, + force_consistent_inputs_shapes=False, + ) + return out_t + + def linear(x_t, weight_t, out_t): + # x @ weight^T with weight [out, in]; NNEF `linear` lowers to a matmul + # that requires equal input ranks, so use einsum to handle the 3D + # activation against the 2D weight (and accumulate in f32). + rank = len(x_shape) + if rank == 3: + expr = "bij,oj->bio" + elif rank == 2: + expr = "ij,oj->io" + else: + raise T2NErrorNotImplemented( + f"shared expert linear expects rank 2 or 3 input, got {rank}" + ) + return emit( + "tract_core_einsum", + [x_t, weight_t], + out_t, + attribs={"expr": expr, "acc": "f32", "output": ""}, + as_list=True, + ) + + d_model = x_shape[-1] + hs = shared["gate_proj"].shape[0] + hs_shape = x_shape[:-1] + [hs] + d_shape = x_shape[:-1] + [d_model] + one_shape = x_shape[:-1] + [1] + + w_gate = add_weight("se_gate_proj", shared["gate_proj"]) + w_up = add_weight("se_up_proj", shared["up_proj"]) + w_down = add_weight("se_down_proj", shared["down_proj"]) + w_router = add_weight("se_router", shared["router"]) + + gate_h = linear(input_tensor, w_gate, mk("se_gate_h", hs_shape)) + up_h = linear(input_tensor, w_up, mk("se_up_h", hs_shape)) + gate_sig = emit("sigmoid", (gate_h,), mk("se_gate_sig", hs_shape)) + silu_g = emit("mul", (gate_h, gate_sig), mk("se_silu", hs_shape)) + inter = emit("mul", (silu_g, up_h), mk("se_inter", hs_shape)) + shared_out = linear(inter, w_down, mk("se_out", d_shape)) + logit = linear(input_tensor, w_router, mk("se_logit", one_shape)) + gate = emit("sigmoid", (logit,), mk("se_gate", one_shape)) + gated = emit("mul", (gate, shared_out), mk("se_gated", d_shape)) + emit("add", (routed, gated), out0) + + +def _is_qtensor_like(tensor: torch.Tensor) -> bool: + if isinstance(tensor, OpaqueTensorRef): + tensor = tensor.opaque_tensor + if isinstance(tensor, QTensor): + return True + if not isinstance(tensor, OffloadedTensor): + return False + offloaded_type = getattr(tensor, "offloaded_tensor_type", None) + return isinstance(offloaded_type, type) and issubclass( + offloaded_type, QTensor + ) + + +def _should_quantize_expert_weights(moe, adapter) -> bool: + return bool(getattr(moe, "_t2n_quantize_moe_experts_q40", False)) or any( + _is_qtensor_like(src) for src in adapter.expert_sources(moe) + ) + + +def _maybe_quantize_expert_weight( + node, + name: str, + data: torch.Tensor, + quantize_experts_q40: bool, + expert_q40_quantizer, + expert_q40_quantizer_kwargs, +) -> torch.Tensor: + if ( + not quantize_experts_q40 + or name not in {"w1", "w2", "w3"} + or _is_qtensor_like(data) + ): + return data + q_data = expert_q40_quantizer( + data.contiguous(), + **expert_q40_quantizer_kwargs, + ) + q_data.nnef_name = f"{node.outputs[0].name}_{name}" + return q_data + + +def _expert_layout(moe) -> str: + layout = getattr(moe, "_t2n_moe_expert_layout", "canonical") + if layout == "tract_moe_ffn": + return "canonical" + if layout not in {"canonical", "linear"}: + raise T2NErrorNotImplemented( + f"unsupported MoE expert layout {layout!r}" + ) + return layout + + +def _layout_expert_weight( + name: str, + data: torch.Tensor, + expert_layout: str, +) -> torch.Tensor: + if expert_layout != "linear" or name not in {"w1", "w2", "w3"}: + return data + # Adapters normalize experts to the canonical tract_moe_ffn shapes: + # w1/w3 [E,D,H], w2 [E,H,D]. The linear layout keeps native nn.Linear + # storage instead: w1/w3 [E,H,D], w2 [E,D,H]. This is a layout contract, + # independent from whether the tensors are stored as f32/f16 or quantized. + return data.transpose(-1, -2) + + +def _expert_q40_quantizer_config( + moe, + quantize_experts_q40_percentile: float, +): + quantizer = getattr( + moe, + "_t2n_quantize_moe_experts_q40_quantizer", + fp_to_tract_q4_0_with_min_max_calibration, + ) + kwargs = dict( + getattr(moe, "_t2n_quantize_moe_experts_q40_kwargs", {}) or {} + ) + if ( + quantizer is fp_to_tract_q4_0_with_min_max_calibration + and "percentile" not in kwargs + ): + kwargs["percentile"] = quantize_experts_q40_percentile + return quantizer, kwargs + + +def _moe_attrs(adapter, moe): + attrs = { + "k": adapter.top_k(moe), + "activation": adapter.activation(moe), + "gate": adapter.gate(moe), + } + act_alpha = adapter.act_alpha(moe) + act_limit = adapter.act_limit(moe) + if act_alpha is not None: + attrs["act_alpha"] = act_alpha + if act_limit is not None: + attrs["act_limit"] = act_limit + layout = _expert_layout(moe) + if layout != "canonical": + attrs["expert_layout"] = layout + return attrs + + +def _append_optional_moe_biases(moe, adapter, is_swiglu, inputs, add_weight): + # tract maps positional inputs as: + # x, wg, w1, w2, w3, wg_bias, w1_bias, w3_bias, w2_bias. + biases = ( + (adapter.gate_bias(moe), "wg_bias"), + (adapter.expert_w1_bias(moe), "w1_bias"), + (adapter.expert_w3_bias(moe), "w3_bias"), + (adapter.expert_w2_bias(moe), "w2_bias"), + ) + if not any(bias is not None for bias, _ in biases): + return + if not is_swiglu: + raise T2NErrorNotImplemented( + "MoE biases are only supported alongside a SwiGLU (w3) gate; " + f"got biases without w3 for {type(moe).__name__}" + ) + for bias, label in biases: + if bias is None: + raise T2NErrorNotImplemented( + f"partial MoE bias set: {label} is missing while other " + "biases are present (positional mapping needs all four)" + ) + for bias, label in biases: + inputs.append(add_weight(label, bias)) + + +def _convert_moe_to_nnef(g, node, name_to_tensor, inference_target): + """Emit tract_moe_ffn for any supported MoE module.""" + if not isinstance(inference_target, TractNNEF): + raise T2NErrorStrictNNEFSpec( + "MoE FFN export requires tract inference target " + "(tract_moe_ffn is a tract extension)" + ) + + # pylint: disable-next=import-outside-toplevel + from torch_to_nnef import torch_graph as tg + + # pylint: disable-next=import-outside-toplevel + from torch_to_nnef.op import helper + + moe = node.op_ref + adapter = _get_adapter(moe) + is_swiglu = adapter.activation(moe) == "swiglu" + + quantize_experts_q40 = _should_quantize_expert_weights(moe, adapter) + quantize_experts_q40_percentile = float( + getattr(moe, "_t2n_quantize_moe_experts_q40_percentile", 1.0) + ) + expert_q40_quantizer, expert_q40_quantizer_kwargs = ( + _expert_q40_quantizer_config(moe, quantize_experts_q40_percentile) + ) + expert_layout = _expert_layout(moe) + + def _add_weight(name: str, data: torch.Tensor): + data = _layout_expert_weight(name, data, expert_layout) + data = _maybe_quantize_expert_weight( + node, + name, + data, + quantize_experts_q40, + expert_q40_quantizer, + expert_q40_quantizer_kwargs, + ) + wnode = tg.TensorVariable( + name=f"{node.outputs[0].name}_{name}", + data=data, + shape=list(data.shape), + dtype=data.dtype, + ) + return helper.get_or_add_tensor_variable_in_nnef( + g, wnode, name_to_tensor + ) + + input_tensor = helper.get_or_add_tensor_variable_in_nnef( + g, node.inputs[0], name_to_tensor + ) + + wg = _add_weight("wg", adapter.gate_weight(moe)) + w1 = _add_weight("w1", adapter.expert_w1(moe)) + w2 = _add_weight("w2", adapter.expert_w2(moe)) + + inputs = [input_tensor, wg, w1, w2] + + if is_swiglu: + w3 = _add_weight("w3", adapter.expert_w3(moe)) + inputs.append(w3) + + attrs = _moe_attrs(adapter, moe) + + _append_optional_moe_biases(moe, adapter, is_swiglu, inputs, _add_weight) + + # tract_moe_ffn is single-output (the routed hidden states). Some modules + # (e.g. gpt-oss GptOssMLP) also return router_scores as a second output, + # but transformers discards it at inference (`hidden_states, _ = mlp(...)`) + # so it has no consumers. Map only node.outputs[0] and ignore any trailing + # router output. + if len(node.outputs) > 1: + LOGGER.debug( + "%s has %d outputs; mapping output[0] and ignoring the " + "inference-unused router output(s)", + type(moe).__name__, + len(node.outputs), + ) + + # pylint: disable-next=import-outside-toplevel + from nnef_tools.model import Tensor as NTensor + + shared = adapter.shared_expert(moe) + + out0 = helper.add_tensor_variable_node_as_nnef_tensor( + g, node.outputs[0], name_to_tensor, prevent_variable=True + ) + base = node.outputs[0].name + np_dtype = input_tensor.dtype + x_shape = list(input_tensor.shape) + + def _mk(suffix, shape): + return NTensor( + g, name=f"{base}_{suffix}", dtype=np_dtype, shape=tuple(shape) + ) + + # The routed experts go to the final output directly, unless a shared + # expert must be added on top (Qwen2 / Qwen3.5), in which case they go to + # an intermediate tensor. + routed = _mk("routed", out0.shape) if shared is not None else out0 + + # tract_moe_ffn takes intentionally heterogeneous-rank inputs: x is + # [T, D] (or [B, S, D]), the gate is [E, D], and the expert weights are + # [E, D, H]. The generic rank-aligner would left-pad the lower-rank + # operands with a leading 1 to match the 3D weights, turning a 2D x into + # [1, T, D] and producing a 3D output that no longer matches the 2D + # PyTorch reference. Disable it so the op sees the ranks it expects. + helper.cast_and_add_nnef_operation( + name_to_tensor=name_to_tensor, + graph=g, + type="tract_moe_ffn", + inputs=tuple(inputs), + outputs=(routed,), + attribs=attrs, + force_consistent_inputs_shapes=False, + ) + + if shared is not None: + _emit_shared_expert( + g, + name_to_tensor, + shared, + input_tensor, + routed, + out0, + x_shape, + _add_weight, + _mk, + ) + return ["tract_transformers", "tract_core"] + + return ["tract_transformers"] + + +# --------------------------------------------------------------------------- +# Extractors (one per MODULE_CLASS, all delegate to _convert_moe_to_nnef) +# --------------------------------------------------------------------------- + + +class MoEFFNExtractor(ModuleInfoExtractor): + """Extractor for the reference MoEFFN wrapper.""" + + MODULE_CLASS = MoEFFN + + def convert_to_nnef( + self, + g, + node, + name_to_tensor, + null_ref, + torch_graph, + inference_target, + **kw, + ): + return _convert_moe_to_nnef(g, node, name_to_tensor, inference_target) + + +# --------------------------------------------------------------------------- +# Lazy registration of transformers MoE classes +# --------------------------------------------------------------------------- + + +def _try_register(import_path: str, class_name: str): + """Try to import a transformers MoE class and register an extractor.""" + try: + # pylint: disable-next=import-outside-toplevel + import importlib + + mod = importlib.import_module(import_path) + moe_cls = getattr(mod, class_name) + + def _make_convert(): + def convert_to_nnef( + self, + g, + node, + name_to_tensor, + null_ref, + torch_graph, + inference_target, + **kw, + ): + return _convert_moe_to_nnef( + g, + node, + name_to_tensor, + inference_target, + ) + + return convert_to_nnef + + # dynamically create an extractor subclass + extractor_cls = type( + f"{class_name}Extractor", + (ModuleInfoExtractor,), + { + "MODULE_CLASS": moe_cls, + "convert_to_nnef": _make_convert(), + }, + ) + # class creation triggers metaclass registration + LOGGER.debug("Registered %s extractor", class_name) + return extractor_cls + except (ImportError, AttributeError, RuntimeError) as err: + LOGGER.debug( + "Could not register optional %s.%s MoE extractor: %s", + import_path, + class_name, + err, + ) + return None + + +def _register_all_transformers_moe(): + _candidates = [ + ( + "transformers.models.mixtral.modeling_mixtral", + "MixtralSparseMoeBlock", + ), + ( + "transformers.models.gpt_oss.modeling_gpt_oss", + "GptOssMLP", + ), + ( + "transformers.models.gpt_oss.modeling_gpt_oss", + "GptOssSparseMoeBlock", + ), + ( + "transformers.models.qwen2_moe.modeling_qwen2_moe", + "Qwen2MoeSparseMoeBlock", + ), + ( + "transformers.models.qwen3_moe.modeling_qwen3_moe", + "Qwen3MoeSparseMoeBlock", + ), + ( + "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", + "Qwen3_5MoeSparseMoeBlock", + ), + ( + "transformers.models.olmoe.modeling_olmoe", + "OlmoeSparseMoeBlock", + ), + ( + "transformers.models.granitemoe.modeling_granitemoe", + "GraniteMoeMoE", + ), + ( + "transformers.models.granitemoeshared.modeling_granitemoeshared", + "GraniteMoeSharedMoE", + ), + ] + for import_path, class_name in _candidates: + _try_register(import_path, class_name) + + +_register_all_transformers_moe() diff --git a/torch_to_nnef/op/helper.py b/torch_to_nnef/op/helper.py index 5fce23fe8..4737a9e7a 100644 --- a/torch_to_nnef/op/helper.py +++ b/torch_to_nnef/op/helper.py @@ -695,6 +695,7 @@ def add_multi_output_op( attrs=None, ensure_tuple=True, output_tensor_name_suffix: str = "", + force_consistent_inputs_shapes: bool = True, ): if len(node.outputs) == 1: LOGGER.debug( @@ -721,6 +722,7 @@ def add_multi_output_op( inputs=inputs, outputs=tuple(output_tensors), attribs=attrs or {}, + force_consistent_inputs_shapes=force_consistent_inputs_shapes, ) return output_tensors @@ -822,6 +824,47 @@ def cast_element(node, accepted_none): return int_list +def add_cast_nnef_tensor( + g, + name_to_tensor, + nnef_tensor: NTensor, + cast_to: np.dtype, + force_full_output_tensor_name: T.Optional[str] = None, +) -> NTensor: + """Emit a tract dtype cast preserving the input tensor shape.""" + cast_to = np.dtype(cast_to).type + to_str = numpy_dtype_to_tract_str(cast_to) + out_name = ( + force_full_output_tensor_name or f"{nnef_tensor.name}_as_{to_str}" + ) + existing = name_to_tensor.get(out_name) + if existing is not None: + if existing.dtype == cast_to and tuple(existing.shape) == tuple( + nnef_tensor.shape + ): + return existing + raise T2NErrorConsistency( + f"cast output name collision for {out_name!r}: " + f"existing dtype/shape {existing.dtype}/{existing.shape}, " + f"requested {cast_to}/{nnef_tensor.shape}" + ) + out = NTensor( + g, + out_name, + dtype=cast_to, + shape=tuple(nnef_tensor.shape), + ) + name_to_tensor[out_name] = out + NOperation( + g, + type="tract_core_cast", + inputs=nnef_tensor, + outputs=out, + attribs={"to": to_str}, + ) + return out + + def cast_to_if_not_dtype_and_variable( g, name_to_tensor, @@ -844,16 +887,20 @@ def cast_to_if_not_dtype_and_variable( cast_to, ) cast_to = nnef_tensor.dtype - out = add_single_output_op( + # The cast output stands in for ``node``'s output (e.g. the final + # forced-cast of a div result), so it must be named after + # ``node.outputs[0]`` (as the previous ``add_single_output_op`` path + # did). Naming it after the *input* tensor would leave the graph + # output name (e.g. ``output_0``) unbound. + out_name = node.outputs[0].export_name + if suffix: + out_name += f"_{suffix}" + out = add_cast_nnef_tensor( g, - node, name_to_tensor, - "tract_core_cast", - inputs=nnef_tensor, - attrs={ - "to": numpy_dtype_to_tract_str(cast_to), - }, - output_tensor_name_suffix=suffix, + nnef_tensor, + cast_to=cast_to, + force_full_output_tensor_name=out_name, ) return out, ["tract_core"] @@ -1187,6 +1234,20 @@ def cast_to_if_not_dtype_and_variable( suffix, ) + def add_cast_nnef_tensor( + self, + nnef_tensor: NTensor, + cast_to: np.dtype, + force_full_output_tensor_name: T.Optional[str] = None, + ) -> NTensor: + return add_cast_nnef_tensor( + self.g, + self.name_to_tensor, + nnef_tensor, + cast_to, + force_full_output_tensor_name, + ) + def cast_and_add_nnef_operation(self, **kwargs): return cast_and_add_nnef_operation( graph=self.g, name_to_tensor=self.name_to_tensor, **kwargs diff --git a/torch_to_nnef/tensor/named.py b/torch_to_nnef/tensor/named.py index 2f9dd2b37..9f4b865f7 100644 --- a/torch_to_nnef/tensor/named.py +++ b/torch_to_nnef/tensor/named.py @@ -15,6 +15,29 @@ LOGGER = logging.getLogger(__name__) +def find_fake_mode(*containers): + """Return the ``FakeTensorMode`` of any fake operand, else ``None``. + + Opaque float weights are traced as ``FakeTensor`` placeholders (see + ``torch_to_nnef.tensor.opaque._fake_trace_tensor``). A fake tensor's + dispatch only works while its owning mode is active, so when such a + tensor meets a :class:`NamedTensor` (which drops out of the subclass + dispatch) the bare op raises ``TypeError: unsupported operand``. We + re-activate the mode around the op to let the two compose. + """ + try: + from torch._subclasses.fake_tensor import ( # pylint: disable=import-outside-toplevel + FakeTensor, + ) + except ImportError: + return None + for container in containers: + for value in container: + if isinstance(value, FakeTensor): + return value.fake_mode + return None + + class NamedTensor(torch.Tensor): """Tensor enriched with name attribute.""" @@ -101,12 +124,29 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): kwargs = {} with select_ctx_disable_torch_fn(): - new_args = [a.clone() if isinstance(a, cls) else a for a in args] - new_kwargs = { - k: v.clone() if isinstance(v, cls) else v - for k, v in kwargs.items() - } - ret = func(*new_args, **new_kwargs) + # An opaque float weight is traced as a ``FakeTensor`` placeholder. + # Fake dispatch (``__torch_dispatch__``) rejects a ``NamedTensor`` + # operand outright, so when a fake tensor is present we must hand + # ``func`` the plain underlying tensor rather than a re-wrapped + # ``NamedTensor`` clone, and re-activate the fake mode so the op + # resolves (plain tensors resolve fine either way under meta). + fake_mode = find_fake_mode(args, kwargs.values()) + + def _prepare(value): + if not isinstance(value, cls): + return value + cloned = value.clone() + if fake_mode is not None: + return cloned.as_subclass(torch.Tensor) + return cloned + + new_args = [_prepare(a) for a in args] + new_kwargs = {k: _prepare(v) for k, v in kwargs.items()} + if fake_mode is not None: + with fake_mode: + ret = func(*new_args, **new_kwargs) + else: + ret = func(*new_args, **new_kwargs) if func in get_default_nowrap_functions(): return ret # important modification diff --git a/torch_to_nnef/tensor/opaque.py b/torch_to_nnef/tensor/opaque.py index 9e8a948af..16db6cf6e 100644 --- a/torch_to_nnef/tensor/opaque.py +++ b/torch_to_nnef/tensor/opaque.py @@ -1,4 +1,5 @@ import abc +import contextlib import logging import typing as T import warnings @@ -29,6 +30,37 @@ # weights at export time between with Opaque and # OpaqueTensorRef. NEW_OPAQUE_TRACING_STRATEGY = torch_version() >= "2.4.0" +_TRACE_FAKE_MODE = None + + +def _fake_trace_tensor( + shape: T.Tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Create a non-materialized tensor with real device semantics.""" + global _TRACE_FAKE_MODE # pylint: disable=global-statement + try: + from torch._subclasses.fake_tensor import ( # pylint: disable=import-outside-toplevel + FakeTensorMode, + ) + except ImportError: + return torch.empty(shape, dtype=dtype, device="meta") + if _TRACE_FAKE_MODE is None: + # ``allow_non_fake_inputs`` lets a fake placeholder compose with the + # real tensors it meets during trace. Older torch (e.g. 1.13) does + # not accept it as a constructor kwarg, so fall back to setting it + # after construction, and to a plain mode if even that is absent. + try: + _TRACE_FAKE_MODE = FakeTensorMode(allow_non_fake_inputs=True) + except TypeError: + _TRACE_FAKE_MODE = FakeTensorMode() + with contextlib.suppress(AttributeError): + _TRACE_FAKE_MODE.allow_non_fake_inputs = True + if device.type == "meta": + device = torch.device("cpu") + with _TRACE_FAKE_MODE: + return torch.empty(shape, dtype=dtype, device=device) def maybe_custom_op(f): @@ -152,8 +184,10 @@ def _to_trace_tensor(self, device: str): # index_select) also request materialization to avoid baking # uninitialized memory as NNEF constants. if device == "meta" and not dtype_is_whole_number(self.dtype): - return torch.empty( - tuple(self.shape), dtype=self.dtype, device=device + return _fake_trace_tensor( + tuple(self.shape), + dtype=self.dtype, + device=self.device, ) # Materialize on the parameter's native device rather than forcing one, # so a GPU-resident trace keeps the weight co-located with its runtime @@ -331,6 +365,9 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): # pylint: disable-next=import-outside-toplevel from torch_to_nnef.tensor import NamedTensor + # pylint: disable-next=import-outside-toplevel + from torch_to_nnef.tensor.named import find_fake_mode + if not all( issubclass(cls, t) or issubclass(NamedTensor, t) for t in types ): @@ -366,7 +403,29 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): for k, v in kwargs.items() } - ret = func(*args, **kwargs) + # If expansion produced a fake placeholder, any sibling operand + # that is still a ``NamedTensor`` (e.g. an un-quantized bias) would + # reach fake ``__torch_dispatch__`` as an unknown subclass and fail + # with "Multiple dispatch failed". Hand ``func`` the plain tensors + # and re-activate the fake mode so the op resolves. + fake_mode = find_fake_mode(args, kwargs.values()) + if fake_mode is not None: + args = [ + a.as_subclass(torch.Tensor) + if isinstance(a, NamedTensor) + else a + for a in args + ] + kwargs = { + k: v.as_subclass(torch.Tensor) + if isinstance(v, NamedTensor) + else v + for k, v in kwargs.items() + } + with fake_mode: + ret = func(*args, **kwargs) + else: + ret = func(*args, **kwargs) if skip_expansion: return ret # important modification diff --git a/torch_to_nnef/tensor/quant/qtract.py b/torch_to_nnef/tensor/quant/qtract.py index 2c318b4f0..115ed0a85 100644 --- a/torch_to_nnef/tensor/quant/qtract.py +++ b/torch_to_nnef/tensor/quant/qtract.py @@ -163,13 +163,13 @@ def _iter_binary_dat_content_chunks( u8_blob: torch.Tensor, scale: torch.Tensor, post_tract_21_11: bool = False, - chunk_groups: int = 1 << 20, + chunk_groups: int = 1 << 16, ): n_bytes_per_group = 18 - tensor_flat = u8_blob.flatten() + tensor_flat = u8_blob.detach().cpu().flatten() assert tensor_flat.numel() % 32 == 0, tensor_flat.shape group_count = tensor_flat.numel() // 32 - scale_flat = scale.reshape(group_count) + scale_flat = scale.detach().cpu().reshape(group_count) for start in range(0, group_count, chunk_groups): end = min(start + chunk_groups, group_count) diff --git a/torch_to_nnef/torch_graph/jit_passes.py b/torch_to_nnef/torch_graph/jit_passes.py index bd896026a..c1c84fa7f 100644 --- a/torch_to_nnef/torch_graph/jit_passes.py +++ b/torch_to_nnef/torch_graph/jit_passes.py @@ -313,8 +313,10 @@ def _materialize_size_fold( if kind == ATEN_NUMEL: if sizes is None: return None + concrete_sizes = T.cast(T.List[int], sizes) n = 1 - for s in sizes: + # pylint: disable-next=not-an-iterable + for s in concrete_sizes: n *= int(s) return _insert_int_constant_before(graph, n, node) if kind == ATEN_LEN: