diff --git a/.github/workflows/golden_regen.yml b/.github/workflows/golden_regen.yml index 4953d556..b9c6294e 100644 --- a/.github/workflows/golden_regen.yml +++ b/.github/workflows/golden_regen.yml @@ -77,15 +77,17 @@ jobs: if: inputs.device != 'cuda' run: pip install onnxruntime - - name: Install onnxruntime (GPU) - if: inputs.device == 'cuda' - run: pip install onnxruntime-gpu - - name: Install dependencies run: | pip install -r requirements/ci/requirements.txt pip install -e '.[testing,transformers]' + - name: Install onnxruntime (GPU) + if: inputs.device == 'cuda' + run: | + # Install after project dependencies so the CPU ORT package cannot overwrite it. + pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu + - name: Reject trust_remote_code cases run: | # Safety: never auto-run models that require trust_remote_code diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 65be966b..8213b49c 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.github/workflows/gpu_l4_golden_parity.yml @@ -59,7 +59,9 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12. + pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu + pip install --no-deps onnxruntime-genai-cuda - name: Check for golden test data id: check_golden diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 13d17d4f..ab4f6f89 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.github/workflows/gpu_l5_generation_e2e.yml @@ -59,7 +59,9 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing,transformers]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12. + pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu + pip install --no-deps onnxruntime-genai-cuda - name: Run L5 generation E2E tests env: diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 137760f4..104fc3d4 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -283,7 +283,9 @@ jobs: pip install -r requirements/ci/requirements.txt pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs pip install -e '.[testing]' - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ "onnxruntime-gpu<1.27" onnxruntime-genai-cuda + # PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12. + pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu + pip install --no-deps onnxruntime-genai-cuda - name: Run fast integration tests env: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/.github/workflows/validation_examples_gpu.yml b/.github/workflows/validation_examples_gpu.yml index 65157577..f8b273f6 100644 --- a/.github/workflows/validation_examples_gpu.yml +++ b/.github/workflows/validation_examples_gpu.yml @@ -83,7 +83,9 @@ jobs: pip install -r requirements/ci/requirements.txt pip install -e '.[testing,transformers]' pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs - pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda + # PyPI's ORT 1.27+ wheel requires CUDA 13, while these runners use CUDA 12. + pip install --pre --no-deps --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ onnxruntime-gpu + pip install --no-deps onnxruntime-genai-cuda - name: Run ${{ matrix.name }} env: diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 3799ea33..9bc5b9f8 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -9,6 +9,7 @@ "BaseModelConfig", "CausalLMConfig", "CausalLMTask", + "ComponentInfo", "DepthAnythingConfig", "EncoderConfig", "EpCapabilities", @@ -42,6 +43,7 @@ "ep_registry", "get_build_dtype", "get_ep", + "inspect_components", "models", "optimize_model", "register_ep", @@ -81,6 +83,7 @@ from mobius._constants import OPSET_VERSION from mobius._diffusers_builder import build_diffusers_pipeline from mobius._execution_providers import EpCapabilities, ep_registry, get_ep, register_ep +from mobius._inspect import ComponentInfo, inspect_components from mobius._model_package import ModelPackage from mobius._optimizations import optimize_model from mobius._registry import ( diff --git a/src/mobius/_inspect.py b/src/mobius/_inspect.py new file mode 100644 index 00000000..47c3011f --- /dev/null +++ b/src/mobius/_inspect.py @@ -0,0 +1,173 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Inspect a model's component layout without building it. + +``inspect_components`` reports the components mobius would produce for a model +(their package keys and optimization roles) **without** constructing graphs or +loading weights. External tools such as Olive use this to plan per-component +work — e.g. optimizing a VLM's ``decoder`` differently from its +``vision_encoder`` — before calling :func:`mobius.build`. + +The component names returned here are the same keys :func:`mobius.build` +produces in its :class:`~mobius._model_package.ModelPackage` (and therefore the +subfolder names ``ModelPackage.save`` writes for multi-component models). +""" + +from __future__ import annotations + +__all__ = ["ComponentInfo", "inspect_components"] + +import dataclasses +import logging + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class ComponentInfo: + """A single component of a model. + + Attributes: + name: Component name. This is the ``ModelPackage`` key mobius produces + (and the subfolder name a multi-component export is saved under). + role: Optimization role of the component, e.g. ``"decoder"``, + ``"encoder"``, or ``"embedding"``. Mobius uses this to gate fusion + passes (only ``"decoder"`` receives GQA / QKV-packing). It is the + value declared in the task's ``model_roles``. + source_paths: Runtime ``named_modules()`` paths that make up this + component inside the full HuggingFace model. These are not + checkpoint/state-dict key prefixes. A single component may map to + multiple disjoint sub-modules (e.g. a decoder assembled from + ``model.layers``, ``model.norm`` and ``lm_head``), so this is a + tuple. Empty when the component is the whole model or the layout is + unknown. Tools such as Olive use these to optimize a submodule in + place before exporting the full model. + """ + + name: str + role: str + source_paths: tuple[str, ...] = () + + +def _resolve_task_model_type_and_config( + model_id: str, task, trust_remote_code: bool +) -> tuple[str, str | None, object | None]: + """Resolve the mobius task name for a model id without building it. + + Mirrors the model_type/task resolution in :func:`mobius.build`, limited to + what is needed to pick a task and inspect class-level component metadata + (no module construction, no weight loading). + """ + import transformers + + from mobius._config_resolver import _default_task_for_model, _try_load_config_json + from mobius._registry import _detect_fallback_registration, registry + + try: + hf_config = transformers.AutoConfig.from_pretrained( + model_id, trust_remote_code=trust_remote_code + ) + except (ValueError, KeyError, OSError): + hf_config = _try_load_config_json(model_id) + if hf_config is None: + # An explicit task is enough to report component names and roles, + # but source paths require a model type and therefore stay empty. + if task is not None: + return task, None, None + raise ValueError( + f"Could not load a HuggingFace config for {model_id!r}. inspect_components supports " + "transformers/registry models; diffusers pipelines are not supported." + ) from None + + model_type = hf_config.model_type + + # model_type adjustments that affect task selection (subset of build()): + # Qwen3.5-MoE ships the same model_type for text-only and VL checkpoints. + if model_type == "qwen3_5_moe" and getattr(hf_config, "vision_config", None) is not None: + model_type = "qwen3_5_moe_vl" + # wav2vec2/hubert/wavlm with a CTC head map to the mms (CTC) registration. + if model_type in ("wav2vec2", "hubert", "wavlm"): + architectures = getattr(hf_config, "architectures", None) or [] + if any("ForCTC" in arch for arch in architectures): + model_type = "mms" + + if task is not None: + return task, model_type, hf_config + if model_type in registry: + return _default_task_for_model(model_type), model_type, hf_config + + fallback = _detect_fallback_registration(hf_config) + if fallback is not None and fallback.task is not None: + return fallback.task, model_type, hf_config + return _default_task_for_model(model_type), model_type, hf_config + + +def _get_hf_component_sources( + module_class: type, + model_type: str, + hf_config: object, +) -> dict[str, tuple[str, ...]]: + """Read runtime HuggingFace component paths from a registered model class.""" + resolver = getattr(module_class, "get_hf_component_sources", None) + if resolver is not None: + return resolver(model_type=model_type, hf_config=hf_config) + return getattr(module_class, "HF_COMPONENT_SOURCES", {}) + + +def inspect_components( + model_id: str, + task=None, + trust_remote_code: bool = False, +) -> list[ComponentInfo]: + """Return the components mobius would produce for a model. + + Args: + model_id: HuggingFace model id or local path. + task: Optional task name (e.g. ``"vision-language"``) or + :class:`~mobius.tasks.ModelTask` instance. When ``None``, the task + is auto-detected from the model type. + trust_remote_code: Whether to trust remote code when loading the + HuggingFace config. + + Returns: + A list of :class:`ComponentInfo`. Single-component models (most LLMs) + return a single entry named ``"model"``; multi-component models (VLMs, + encoder-decoders, speech models) return one entry per component. + + Raises: + ValueError: If a config cannot be resolved for ``model_id`` (e.g. a + diffusers pipeline, which is not supported). + """ + from mobius._registry import registry + from mobius.tasks import get_task + + resolved_task, model_type, hf_config = _resolve_task_model_type_and_config( + model_id, task, trust_remote_code + ) + task_obj = get_task(resolved_task) + roles = task_obj.model_roles or {} + + # Runtime HF paths are owned by the registered model class. Most classes + # declare a fixed ``HF_COMPONENT_SOURCES`` mapping; classes shared by + # several HF layouts can resolve paths from the already-loaded config. + component_sources: dict[str, tuple[str, ...]] = {} + if model_type is not None and hf_config is not None and model_type in registry: + module_class = registry.get(model_type) + component_sources = _get_hf_component_sources(module_class, model_type, hf_config) + + components = [ + ComponentInfo( + name=name, + role=role, + source_paths=tuple(component_sources.get(name, ())), + ) + for name, role in roles.items() + ] + logger.debug( + "inspect_components(%s): task=%s components=%s", + model_id, + resolved_task, + [c.name for c in components], + ) + return components diff --git a/src/mobius/_inspect_test.py b/src/mobius/_inspect_test.py new file mode 100644 index 00000000..4ec1b1d5 --- /dev/null +++ b/src/mobius/_inspect_test.py @@ -0,0 +1,579 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import mobius +from mobius._inspect import ( + ComponentInfo, + _get_hf_component_sources, + _resolve_task_model_type_and_config, + inspect_components, +) +from mobius._registry import registry +from mobius.tasks import get_task + + +def _patch_autoconfig(monkeypatch, hf_config): + import transformers + + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + staticmethod(lambda *a, **k: hf_config), + ) + + +def _fake_hf_config(model_type): + if model_type == "blip-2": + return SimpleNamespace( + model_type=model_type, + text_config=SimpleNamespace(model_type="opt"), + use_decoder_only_language_model=True, + ) + return SimpleNamespace(model_type=model_type) + + +def test_inspect_components_exported_from_package(): + assert mobius.inspect_components is inspect_components + assert mobius.ComponentInfo is ComponentInfo + + +def test_single_component_llm_returns_one_component(monkeypatch): + _patch_autoconfig( + monkeypatch, SimpleNamespace(model_type="llama", architectures=["LlamaForCausalLM"]) + ) + components = inspect_components("fake/llama") + assert components == [ComponentInfo(name="model", role="decoder")] + + +def test_vlm_returns_decoder_vision_embedding(monkeypatch): + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="llava")) + components = inspect_components("fake/llava") + assert [(c.name, c.role) for c in components] == [ + ("decoder", "decoder"), + ("vision_encoder", "encoder"), + ("embedding", "embedding"), + ] + + +def test_qwen3_vl_returns_hf_source_paths(monkeypatch): + # Runtime paths split the text backbone from the embedding component. + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="qwen3_vl")) + components = {c.name: c for c in inspect_components("fake/qwen-vl")} + + assert components["decoder"].source_paths == ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ) + assert components["vision_encoder"].source_paths == ("model.visual",) + assert components["embedding"].source_paths == ("model.language_model.embed_tokens",) + + +def test_qwen3_tts_embedders_return_shared_hf_source_paths(monkeypatch): + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="qwen3_tts")) + components = {c.name: c for c in inspect_components("fake/qwen3-tts")} + + assert components["talker_step_embedder"].source_paths == ( + "talker.model.codec_embedding", + "talker.code_predictor.model.codec_embedding", + ) + assert components["talker_prefill_embedder"].source_paths == ( + "talker.model.text_embedding", + "talker.text_projection", + "talker.model.codec_embedding", + ) + + +def test_qwen2_5_vl_uses_runtime_paths_not_checkpoint_prefixes(monkeypatch): + # Modern HF Qwen2.5-VL uses the same runtime component roots as Qwen3-VL, + # despite exposing different checkpoint weight prefixes. + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="qwen2_5_vl")) + components = {c.name: c for c in inspect_components("fake/qwen2.5-vl")} + + assert components["decoder"].source_paths == ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ) + assert components["vision_encoder"].source_paths == ("model.visual",) + assert components["embedding"].source_paths == ("model.language_model.embed_tokens",) + + +@pytest.mark.parametrize("model_type", ["idefics2", "idefics3", "smolvlm"]) +def test_shared_llava_class_resolves_idefics_runtime_layout(monkeypatch, model_type): + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type=model_type)) + components = {c.name: c for c in inspect_components(f"fake/{model_type}")} + + assert components["decoder"].source_paths[0] == "model.text_model.layers" + assert components["vision_encoder"].source_paths == ( + "model.vision_model", + "model.connector", + ) + assert components["embedding"].source_paths == ("model.text_model.embed_tokens",) + + +def test_shared_llava_class_does_not_guess_unverified_runtime_layout(monkeypatch): + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="fuyu")) + components = inspect_components("fake/fuyu") + assert all(not component.source_paths for component in components) + + +@pytest.mark.parametrize( + ("text_model_type", "decoder_only", "embedding_path"), + [ + ("opt", True, "language_model.model.decoder.embed_tokens"), + ("t5", False, "language_model.shared"), + ], +) +def test_blip2_resolves_text_runtime_layout( + monkeypatch, + text_model_type, + decoder_only, + embedding_path, +): + _patch_autoconfig( + monkeypatch, + SimpleNamespace( + model_type="blip-2", + text_config=SimpleNamespace(model_type=text_model_type), + use_decoder_only_language_model=decoder_only, + ), + ) + components = {c.name: c for c in inspect_components("fake/blip2")} + assert components["embedding"].source_paths == (embedding_path,) + + +def test_blip2_does_not_guess_unknown_text_runtime_layout(monkeypatch): + _patch_autoconfig( + monkeypatch, + SimpleNamespace( + model_type="blip-2", + text_config=SimpleNamespace(model_type="llama"), + use_decoder_only_language_model=True, + ), + ) + + assert all(not component.source_paths for component in inspect_components("fake/blip2")) + + +def test_internvl_resolves_internlm2_runtime_layout(monkeypatch): + _patch_autoconfig( + monkeypatch, + SimpleNamespace( + model_type="internvl_chat", + llm_config=SimpleNamespace(model_type="internlm2"), + architectures=["InternVLChatModel"], + ), + ) + components = {c.name: c for c in inspect_components("fake/internvl")} + + assert components["decoder"].source_paths == ( + "language_model.model.layers", + "language_model.model.norm", + "language_model.output", + ) + assert components["embedding"].source_paths == ("language_model.model.tok_embeddings",) + + +def test_internvl_does_not_guess_unknown_llm_runtime_layout(monkeypatch): + _patch_autoconfig( + monkeypatch, + SimpleNamespace( + model_type="internvl", + llm_config=SimpleNamespace(model_type="qwen3"), + architectures=[], + ), + ) + assert all(not component.source_paths for component in inspect_components("fake/internvl")) + + +def test_phi4mm_decoder_maps_to_multiple_source_paths(monkeypatch): + # A single component can span multiple disjoint HF sub-trees. + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="phi4mm")) + components = {c.name: c for c in inspect_components("fake/phi4mm")} + + assert components["decoder"].source_paths == ("model.layers", "model.norm", "lm_head") + assert components["audio_encoder"].source_paths == ( + "model.embed_tokens_extend.audio_embed", + ) + + +def test_single_component_llm_has_empty_source_paths(monkeypatch): + _patch_autoconfig( + monkeypatch, SimpleNamespace(model_type="llama", architectures=["LlamaForCausalLM"]) + ) + (component,) = inspect_components("fake/llama") + assert component.source_paths == () + + +def test_encoder_decoder_returns_two_components(monkeypatch): + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="t5")) + names = {c.name for c in inspect_components("fake/t5")} + assert names == {"encoder", "decoder"} + + +def test_explicit_task_overrides_task_but_still_resolves_source_paths(monkeypatch): + # The explicit task chooses the component layout, while the config still + # identifies the registered model class that owns its runtime HF paths. + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="qwen3_vl")) + components = inspect_components("fake/qwen-vl", task="vision-language") + assert {c.name for c in components} == {"decoder", "vision_encoder", "embedding"} + by_name = {c.name: c for c in components} + assert by_name["vision_encoder"].source_paths == ("model.visual",) + + +def test_explicit_task_without_config_returns_roles_without_source_paths(monkeypatch): + import transformers + + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + staticmethod(lambda *a, **k: (_ for _ in ()).throw(OSError("no config"))), + ) + monkeypatch.setattr( + "mobius._config_resolver._try_load_config_json", lambda *_a, **_k: None + ) + + components = inspect_components("anything", task="vision-language") + assert {c.name for c in components} == {"decoder", "vision_encoder", "embedding"} + assert all(not component.source_paths for component in components) + + +def test_qwen3_5_moe_vl_detected_from_vision_config(monkeypatch): + _patch_autoconfig( + monkeypatch, SimpleNamespace(model_type="qwen3_5_moe", vision_config=SimpleNamespace()) + ) + names = {c.name for c in inspect_components("fake/qwen-vl")} + assert "vision_encoder" in names and "decoder" in names + + +def test_ctc_architecture_uses_mms_registration(monkeypatch): + hf_config = SimpleNamespace( + model_type="wav2vec2", + architectures=["Wav2Vec2ForCTC"], + ) + _patch_autoconfig(monkeypatch, hf_config) + + task, model_type, resolved_config = _resolve_task_model_type_and_config( + "fake/wav2vec2-ctc", + task=None, + trust_remote_code=False, + ) + + assert (task, model_type, resolved_config) == ("ctc-asr", "mms", hf_config) + + +def test_unresolvable_config_raises(monkeypatch): + import transformers + + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + staticmethod(lambda *a, **k: (_ for _ in ()).throw(OSError("no config"))), + ) + monkeypatch.setattr( + "mobius._config_resolver._try_load_config_json", lambda *_a, **_k: None + ) + with pytest.raises(ValueError, match="Could not load a HuggingFace config"): + inspect_components("fake/diffusers-pipeline") + + +# Model types whose registered module class provides HF component sources. +_MULTI_COMPONENT_MODEL_TYPES = [ + # Qwen VL family + "qwen2_5_vl", + "qwen2_vl", + "qwen3_vl", + "qwen3_5_vl", + "qwen3_5_moe_vl", + # Other VLMs (3-model split) + "llava", + "idefics3", + "smolvlm", + "blip-2", + "internvl", + "internvl_chat", + "gemma3", + "mllama", + "hunyuan_vl_mot", + # Multimodal 4-model split + "phi4mm", + "gemma4", + "gemma4_unified", + # Encoder-decoder (seq2seq / speech-to-text) + "bart", + "t5", + "trocr", + "whisper", + # Speech-language / audio + "qwen3_asr", + "fun_asr", + "qwen3_tts_tokenizer_12hz", + "qwen3_tts", + "fastconformer_rnnt", +] +_EMPTY_SOURCE_PATHS = {("trocr", "encoder")} + + +@pytest.mark.parametrize("model_type", _MULTI_COMPONENT_MODEL_TYPES) +def test_hf_component_sources_keys_match_task_roles(model_type): + # Drift guard: the source-path keys provided by the module class must be + # exactly the package keys the task produces (its model_roles keys), so + # inspect_components never mislabels or drops a component. + module_class = registry.get(model_type) + sources = _get_hf_component_sources( + module_class, + model_type, + _fake_hf_config(model_type), + ) + registration = registry.get_registration(model_type) + task_name = registration.task or module_class.default_task + roles = get_task(task_name).model_roles + + assert set(sources) == set(roles), ( + f"{model_type}: HF_COMPONENT_SOURCES keys {sorted(sources)} " + f"!= model_roles keys {sorted(roles)}" + ) + for name, paths in sources.items(): + assert isinstance(paths, tuple) + if not paths: + assert (model_type, name) in _EMPTY_SOURCE_PATHS + continue + assert all(isinstance(p, str) and p for p in paths) + + +def test_inspect_does_not_instantiate_module_class(monkeypatch): + # inspect_components must read HF_COMPONENT_SOURCES off the class only — + # never construct the (expensive) module. Make __init__ blow up to prove it. + module_class = registry.get("qwen3_vl") + + def _boom(*_a, **_k): + raise AssertionError("inspect_components must not instantiate the module class") + + monkeypatch.setattr(module_class, "__init__", _boom) + _patch_autoconfig(monkeypatch, SimpleNamespace(model_type="qwen3_vl")) + + components = {c.name: c for c in inspect_components("fake/qwen-vl")} + assert components["vision_encoder"].source_paths == ("model.visual",) + + +def test_representative_vlm_source_paths_resolve_on_runtime_models(): + import torch + import transformers + + text_config = { + "vocab_size": 128, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "max_position_embeddings": 128, + "head_dim": 4, + } + qwen_vision_config = { + "depth": 1, + "hidden_size": 16, + "intermediate_size": 32, + "num_heads": 4, + "in_channels": 3, + "patch_size": 2, + "spatial_merge_size": 1, + "temporal_patch_size": 1, + "window_size": 4, + "out_hidden_size": 16, + "fullatt_block_indexes": [0], + } + vision_config = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "image_size": 8, + "patch_size": 2, + } + cases = [ + ( + "qwen2_5_vl", + transformers.Qwen2_5_VLForConditionalGeneration, + transformers.Qwen2_5_VLConfig( + text_config=text_config, + vision_config=qwen_vision_config, + ), + ), + ( + "llava", + transformers.LlavaForConditionalGeneration, + transformers.LlavaConfig( + text_config=text_config, + vision_config=vision_config, + image_token_index=127, + ), + ), + ( + "idefics3", + transformers.Idefics3ForConditionalGeneration, + transformers.Idefics3Config( + text_config={**text_config, "pad_token_id": 0}, + vision_config=vision_config, + image_token_id=127, + pad_token_id=0, + ), + ), + ] + + for model_type, hf_model_class, hf_config in cases: + with torch.device("meta"): + hf_model = hf_model_class(hf_config) + runtime_paths = {name for name, _ in hf_model.named_modules()} + sources = _get_hf_component_sources( + registry.get(model_type), + model_type, + hf_config, + ) + missing = { + f"{component}.{path}" + for component, paths in sources.items() + for path in paths + if path not in runtime_paths + } + assert not missing, f"{model_type}: missing runtime paths {sorted(missing)}" + for left_name, left_paths in sources.items(): + for right_name, right_paths in sources.items(): + if left_name >= right_name: + continue + overlaps = { + (left, right) + for left in left_paths + for right in right_paths + if left == right + or left.startswith(f"{right}.") + or right.startswith(f"{left}.") + } + assert not overlaps, ( + f"{model_type}: {left_name}/{right_name} paths overlap: {sorted(overlaps)}" + ) + + +# Representative source_paths per model family, locking in the diverse HF layouts +# (standard llava naming, T5 without an outer ``model.``, Whisper/Bart with it, +# multi-path decoders, and shared encoder for RNNT streaming). +_SOURCE_PATH_EXPECTATIONS = { + "llava": { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + }, + "idefics3": { + "decoder": ( + "model.text_model.layers", + "model.text_model.norm", + "model.text_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_model", "model.connector"), + "embedding": ("model.text_model.embed_tokens",), + }, + "gemma3": { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + }, + "mllama": { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_model", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + }, + "internvl_chat": { + "decoder": ( + "language_model.model.layers", + "language_model.model.norm", + "language_model.model.rotary_emb", + "language_model.lm_head", + ), + "vision_encoder": ("vision_model", "mlp1"), + "embedding": ("language_model.model.embed_tokens",), + }, + "blip-2": { + "decoder": ( + "language_model.model.decoder.embed_positions", + "language_model.model.decoder.final_layer_norm", + "language_model.model.decoder.layers", + "language_model.lm_head", + ), + "vision_encoder": ("vision_model", "qformer", "language_projection"), + "embedding": ("language_model.model.decoder.embed_tokens",), + }, + "hunyuan_vl_mot": { + "decoder": ( + "model.language_model.model.layers", + "model.language_model.model.norm", + "model.language_model.lm_head", + ), + "vision_encoder": ("model.visual",), + }, + "t5": {"encoder": ("encoder",), "decoder": ("decoder", "lm_head")}, + "bart": {"encoder": ("model.encoder",), "decoder": ("model.decoder", "lm_head")}, + "trocr": {"encoder": (), "decoder": ("model.decoder", "output_projection")}, + "whisper": {"encoder": ("model.encoder",), "decoder": ("model.decoder", "proj_out")}, + "qwen3_asr": { + "audio_encoder": ("thinker.audio_tower",), + "decoder": ("thinker.model.layers", "thinker.model.norm", "thinker.lm_head"), + }, + "fastconformer_rnnt": {"encoder": ("encoder",), "encoder_streaming": ("encoder",)}, + # gemma4 (towers) vs gemma4_unified (encoder-free embedders) differ in vision/audio. + "gemma4": { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.embed_vision"), + "audio_encoder": ("model.audio_tower", "model.embed_audio"), + }, + "gemma4_unified": { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_embedder", "model.embed_vision"), + "audio_encoder": ("model.embed_audio",), + }, +} + + +@pytest.mark.parametrize("model_type,expected", _SOURCE_PATH_EXPECTATIONS.items()) +def test_hf_component_sources_values(model_type, expected): + sources = _get_hf_component_sources( + registry.get(model_type), + model_type, + _fake_hf_config(model_type), + ) + for key, paths in expected.items(): + assert sources[key] == paths, f"{model_type}.{key}: {sources[key]} != {paths}" diff --git a/src/mobius/_testing/ort_inference.py b/src/mobius/_testing/ort_inference.py index 726b72b8..41c7d4b8 100644 --- a/src/mobius/_testing/ort_inference.py +++ b/src/mobius/_testing/ort_inference.py @@ -256,6 +256,10 @@ def run(self, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: return dict(zip(self._output_names, (_ort_value_to_numpy(o) for o in raw_outputs))) def close(self) -> None: + # Release ORT resources eagerly (not only at Python GC time) so + # long-running GPU test suites do not accumulate session memory. + if hasattr(self, "_session"): + del self._session self._tmpdir.cleanup() def __del__(self) -> None: diff --git a/src/mobius/_testing/ort_inference_test.py b/src/mobius/_testing/ort_inference_test.py index 86b0b479..32755368 100644 --- a/src/mobius/_testing/ort_inference_test.py +++ b/src/mobius/_testing/ort_inference_test.py @@ -13,11 +13,13 @@ from __future__ import annotations +import tempfile + import onnx_ir as ir import pytest from mobius._flags import flags -from mobius._testing.ort_inference import _MAX_EP_OPSET, _should_lower_opset +from mobius._testing.ort_inference import _MAX_EP_OPSET, OnnxModelSession, _should_lower_opset def _make_value(name: str) -> ir.Value: @@ -108,3 +110,13 @@ def test_opset_at_or_below_max_not_lowered(lowering_enabled: None) -> None: model = _model_with(_graph_with([node], opset=_MAX_EP_OPSET)) assert _should_lower_opset(model, device="cuda") is False + + +def test_close_releases_session_reference() -> None: + session = OnnxModelSession.__new__(OnnxModelSession) + session._session = object() # type: ignore[attr-defined] + session._tmpdir = tempfile.TemporaryDirectory() # type: ignore[attr-defined] + + session.close() + + assert not hasattr(session, "_session") diff --git a/src/mobius/models/bart.py b/src/mobius/models/bart.py index c79eb338..24837511 100644 --- a/src/mobius/models/bart.py +++ b/src/mobius/models/bart.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -234,6 +234,12 @@ class BartForConditionalGeneration(nn.Module): default_task = "seq2seq" category = "encoder-decoder" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": ("model.encoder",), + "decoder": ("model.decoder", "lm_head"), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/blip2.py b/src/mobius/models/blip2.py index 8d4371fc..68e8568b 100644 --- a/src/mobius/models/blip2.py +++ b/src/mobius/models/blip2.py @@ -31,7 +31,7 @@ from __future__ import annotations import re -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -192,6 +192,55 @@ class Blip2Model(nn.Module): default_task: str = "vision-language" category: str = "Multimodal" + # Runtime HF ``named_modules()`` paths for decoder-only (OPT) BLIP-2. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "language_model.model.decoder.embed_positions", + "language_model.model.decoder.final_layer_norm", + "language_model.model.decoder.layers", + "language_model.lm_head", + ), + "vision_encoder": ("vision_model", "qformer", "language_projection"), + "embedding": ("language_model.model.decoder.embed_tokens",), + } + _T5_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "language_model.encoder.block", + "language_model.encoder.final_layer_norm", + "language_model.encoder.dropout", + "language_model.decoder.block", + "language_model.decoder.final_layer_norm", + "language_model.decoder.dropout", + "language_model.lm_head", + ), + "vision_encoder": ("vision_model", "qformer", "language_projection"), + "embedding": ("language_model.shared",), + } + + @classmethod + def get_hf_component_sources( + cls, + *, + model_type: str, + hf_config: object, + ) -> dict[str, tuple[str, ...]]: + """Return runtime paths for the decoder-only or encoder-decoder LLM.""" + del model_type + text_config = getattr(hf_config, "text_config", None) + text_model_type = getattr(text_config, "model_type", None) + decoder_only = getattr(hf_config, "use_decoder_only_language_model", None) + if text_model_type == "opt": + return cls.HF_COMPONENT_SOURCES + if text_model_type == "t5": + return cls._T5_COMPONENT_SOURCES + if text_model_type is not None: + return {} + if decoder_only is True: + return cls.HF_COMPONENT_SOURCES + if decoder_only is False: + return cls._T5_COMPONENT_SOURCES + return {} + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/fun_asr.py b/src/mobius/models/fun_asr.py index fb1c9392..33b2fcf6 100644 --- a/src/mobius/models/fun_asr.py +++ b/src/mobius/models/fun_asr.py @@ -21,6 +21,8 @@ from __future__ import annotations +from typing import ClassVar + import numpy as np import onnx_ir as ir import torch @@ -503,6 +505,13 @@ class FunASRForConditionalGeneration(nn.Module): category: str = "Speech-to-Text" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "audio_encoder": ("audio_encoder", "audio_adaptor"), + "embedding": ("llm.model.embed_tokens",), + "decoder": ("llm.model.layers", "llm.model.norm", "llm.lm_head"), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/gemma3.py b/src/mobius/models/gemma3.py index dfe977fe..83f70331 100644 --- a/src/mobius/models/gemma3.py +++ b/src/mobius/models/gemma3.py @@ -12,7 +12,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -171,6 +171,18 @@ class Gemma3MultiModalModel(nn.Module): default_task: str = "vision-language" category: str = "Multimodal" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 265f8b0f..b9cc9d10 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -25,7 +25,7 @@ import dataclasses import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import numpy as np import onnx_ir as ir @@ -2967,11 +2967,11 @@ class Gemma4Model(nn.Module): Always produced: - ``decoder``: Gemma4 text decoder taking ``inputs_embeds`` - - ``vision``: SigLIP-style encoder + projector + - ``vision_encoder``: SigLIP-style encoder + projector - ``embedding``: scaled word embedding + multimodal feature fusion Added when ``config.audio is not None``: - - ``speech``: Conformer audio encoder + projection to text hidden size + - ``audio_encoder``: Conformer audio encoder + projection to text hidden size Covers all Gemma4 variants: - Vision-language (26B-A4B, 31B): ``audio=None`` @@ -2983,6 +2983,24 @@ class Gemma4Model(nn.Module): default_task: str = "gemma4" category: str = "Multimodal" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.embed_vision"), + "audio_encoder": ("model.audio_tower", "model.embed_audio"), + "embedding": ( + "model.language_model.embed_tokens", + "model.language_model.embed_tokens_per_layer", + "model.language_model.per_layer_model_projection", + "model.language_model.per_layer_projection_norm", + ), + } + def __init__(self, config: Gemma4Config): super().__init__() self.config = config @@ -3196,6 +3214,19 @@ class Gemma4UnifiedModel(nn.Module): default_task: str = "gemma4-unified" category: str = "Multimodal" + # Runtime HF sub-trees for the encoder-free unified layout. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_embedder", "model.embed_vision"), + "audio_encoder": ("model.embed_audio",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: Gemma4Config): super().__init__() self.config = config diff --git a/src/mobius/models/hunyuan_vl_mot.py b/src/mobius/models/hunyuan_vl_mot.py index bddc9610..06fa80fa 100644 --- a/src/mobius/models/hunyuan_vl_mot.py +++ b/src/mobius/models/hunyuan_vl_mot.py @@ -47,7 +47,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -469,6 +469,17 @@ class HunYuanVLMoTModel(nn.Module): default_task: str = "hunyuan-vl-mot" category: str = "Multimodal" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.model.layers", + "model.language_model.model.norm", + "model.language_model.lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/internvl.py b/src/mobius/models/internvl.py index 7d5a6e04..365272d7 100644 --- a/src/mobius/models/internvl.py +++ b/src/mobius/models/internvl.py @@ -27,7 +27,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -547,6 +547,51 @@ class InternVL2Model(nn.Module): default_task: str = "vision-language" category: str = "Multimodal" + # Runtime ``InternVLChatModel.named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "language_model.model.layers", + "language_model.model.norm", + "language_model.model.rotary_emb", + "language_model.lm_head", + ), + "vision_encoder": ("vision_model", "mlp1"), + "embedding": ("language_model.model.embed_tokens",), + } + _INTERNLM2_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "language_model.model.layers", + "language_model.model.norm", + "language_model.output", + ), + "vision_encoder": ("vision_model", "mlp1"), + "embedding": ("language_model.model.tok_embeddings",), + } + + @classmethod + def get_hf_component_sources( + cls, + *, + model_type: str, + hf_config: object, + ) -> dict[str, tuple[str, ...]]: + """Return paths for the verified OpenGVLab remote-code runtime.""" + llm_config = getattr(hf_config, "llm_config", None) + llm_model_type = getattr(llm_config, "model_type", None) + if llm_model_type == "internlm2": + return cls._INTERNLM2_COMPONENT_SOURCES + if llm_model_type in {"llama", "qwen2"}: + return cls.HF_COMPONENT_SOURCES + if llm_model_type is not None: + return {} + architectures = getattr(hf_config, "architectures", None) or () + if ( + model_type in {"internvl", "internvl2", "internvl_chat"} + or "InternVLChatModel" in architectures + ): + return cls.HF_COMPONENT_SOURCES + return {} + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/llava.py b/src/mobius/models/llava.py index b541c340..76be0a61 100644 --- a/src/mobius/models/llava.py +++ b/src/mobius/models/llava.py @@ -21,7 +21,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -194,6 +194,47 @@ class LLaVAModel(nn.Module): default_task: str = "vision-language" category: str = "Multimodal" + # Runtime HF ``named_modules()`` sub-trees for the LLaVA/PaliGemma/Mistral3 + # layout. Decoder paths exclude token embeddings, which are a separate graph. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_tower", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + } + _IDEFICS_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.text_model.layers", + "model.text_model.norm", + "model.text_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_model", "model.connector"), + "embedding": ("model.text_model.embed_tokens",), + } + + @classmethod + def get_hf_component_sources( + cls, + *, + model_type: str, + hf_config: object, + ) -> dict[str, tuple[str, ...]]: + """Return static runtime paths for verified HF layouts served by this class.""" + del hf_config + if model_type in {"llava", "llava_next", "llava_onevision", "paligemma", "mistral3"}: + return cls.HF_COMPONENT_SOURCES + if model_type in {"idefics2", "idefics3", "smolvlm"}: + return cls._IDEFICS_COMPONENT_SOURCES + # This generic mobius class is registered for additional HF families + # whose runtime trees differ. Empty paths prevent tools from slicing + # against a guessed checkpoint-style prefix. + return {} + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/mllama.py b/src/mobius/models/mllama.py index 08454f6b..8fadb5e9 100644 --- a/src/mobius/models/mllama.py +++ b/src/mobius/models/mllama.py @@ -17,7 +17,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -359,6 +359,18 @@ class MllamaCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = MllamaConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.vision_model", "model.multi_modal_projector"), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: MllamaConfig): super().__init__() self.config = config diff --git a/src/mobius/models/nemo_rnnt.py b/src/mobius/models/nemo_rnnt.py index 8c24e845..610fc86c 100644 --- a/src/mobius/models/nemo_rnnt.py +++ b/src/mobius/models/nemo_rnnt.py @@ -35,6 +35,7 @@ from __future__ import annotations import math +from typing import ClassVar import numpy as np import onnx_ir as ir @@ -972,6 +973,15 @@ class EncDecRNNTModel(nn.Module): default_task: str = "fastconformer-rnnt" category: str = "Speech-to-Text" + # Runtime HF sub-trees. ``encoder`` and ``encoder_streaming`` intentionally + # share the same encoder. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": ("encoder",), + "encoder_streaming": ("encoder",), + "decoder": ("decoder",), + "joint": ("joint",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/phi.py b/src/mobius/models/phi.py index 1e2d059e..3a62103f 100644 --- a/src/mobius/models/phi.py +++ b/src/mobius/models/phi.py @@ -11,6 +11,8 @@ from __future__ import annotations +from typing import ClassVar + import onnx_ir as ir import torch from onnxscript import OpBuilder, nn @@ -1278,6 +1280,15 @@ class Phi4MMMultiModalModel(nn.Module): default_task: str = "phi4mm-multimodal" category: str = "Multimodal" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. The speech + # component's ModelPackage key is ``audio_encoder``. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "vision_encoder": ("model.embed_tokens_extend.image_embed",), + "audio_encoder": ("model.embed_tokens_extend.audio_embed",), + "embedding": ("model.embed_tokens",), + "decoder": ("model.layers", "model.norm", "lm_head"), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/qwen35.py b/src/mobius/models/qwen35.py index 5e6bd7e7..46106e6d 100644 --- a/src/mobius/models/qwen35.py +++ b/src/mobius/models/qwen35.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -421,6 +421,18 @@ class Qwen35VL3ModelCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config @@ -603,6 +615,18 @@ class Qwen35MoEVL3ModelCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/qwen3_asr.py b/src/mobius/models/qwen3_asr.py index 8a9bb857..46ece74f 100644 --- a/src/mobius/models/qwen3_asr.py +++ b/src/mobius/models/qwen3_asr.py @@ -18,6 +18,7 @@ from __future__ import annotations import dataclasses +from typing import ClassVar import numpy as np import onnx_ir as ir @@ -533,6 +534,13 @@ class Qwen3ASRForConditionalGeneration(nn.Module): category: str = "Speech-to-Text" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "audio_encoder": ("thinker.audio_tower",), + "embedding": ("thinker.model.embed_tokens",), + "decoder": ("thinker.model.layers", "thinker.model.norm", "thinker.lm_head"), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/qwen3_tts.py b/src/mobius/models/qwen3_tts.py index f3c3efb1..b70badd3 100644 --- a/src/mobius/models/qwen3_tts.py +++ b/src/mobius/models/qwen3_tts.py @@ -3,14 +3,16 @@ """Qwen3-TTS: Text-to-speech with Qwen3-like talker and code predictor. -Architecture (4-model split for full TTS): +Architecture (five required models plus an optional speaker encoder): 1. **Talker**: Qwen3 decoder with MRoPE 3D, QK norm. Outputs logits (first code group) and last_hidden_state (for code predictor). 2. **Code Predictor**: Small 5-layer decoder with 1D RoPE. Per-step lm_heads and embeddings selected by step_index. 3. **Embedding**: text_embedding → text_projection + codec_embedding. Additive fusion of text and codec embeddings. - 4. **Speaker Encoder**: ECAPA-TDNN extracting speaker embedding from mel. + 4. **Talker Step Embedder**: Previous frame codes + text embedding. + 5. **Talker Prefill Embedder**: Prompt text ids → talker prefill embeddings. + 6. **Speaker Encoder**: Optional ECAPA-TDNN speaker embedding extraction. Config comparison (0.6B vs 1.7B): - talker.hidden_size: 1024 vs 2048 @@ -26,7 +28,7 @@ import dataclasses import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -675,13 +677,15 @@ def forward(self, op: OpBuilder, mel_input: ir.Value): class Qwen3TTSForConditionalGeneration(nn.Module): - """Qwen3-TTS composite model for full 4-model TTS pipeline. + """Qwen3-TTS composite model for the full multi-model TTS pipeline. Contains: - ``talker``: Talker decoder (inputs_embeds → logits + hidden) - ``code_predictor``: Code predictor (hidden → 15 code groups) - ``embedding``: Text + codec embedding model - - ``speaker_encoder``: ECAPA-TDNN speaker encoder + - ``talker_step_embedder``: Per-frame talker input construction + - ``talker_prefill_embedder``: Prompt prefill/trailing-text construction + - ``speaker_encoder``: Optional ECAPA-TDNN speaker encoder HuggingFace class: ``Qwen3TTSForConditionalGeneration`` """ @@ -690,6 +694,27 @@ class Qwen3TTSForConditionalGeneration(nn.Module): category: str = "Audio" config_class: type = ArchitectureConfig + # Runtime HF sub-trees for the talker/code-predictor/embedding split. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "talker": ("talker.model.layers", "talker.model.norm", "talker.codec_head"), + "code_predictor": ("talker.code_predictor",), + "embedding": ( + "talker.model.text_embedding", + "talker.text_projection", + "talker.model.codec_embedding", + ), + "talker_step_embedder": ( + "talker.model.codec_embedding", + "talker.code_predictor.model.codec_embedding", + ), + "talker_prefill_embedder": ( + "talker.model.text_embedding", + "talker.text_projection", + "talker.model.codec_embedding", + ), + "speaker_encoder": ("speaker_encoder",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/qwen3_tts_tokenizer.py b/src/mobius/models/qwen3_tts_tokenizer.py index d4e5a65f..9a41f90d 100644 --- a/src/mobius/models/qwen3_tts_tokenizer.py +++ b/src/mobius/models/qwen3_tts_tokenizer.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar from onnxscript import OpBuilder, nn @@ -645,6 +645,12 @@ class Qwen3TTSTokenizerV2Model(nn.Module): config: Architecture config. """ + # Runtime HF RVQ codec encoder/decoder sub-trees. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": ("encoder",), + "decoder": ("decoder",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.decoder = Qwen3TTSCodecDecoderModel(config) diff --git a/src/mobius/models/qwen_vl.py b/src/mobius/models/qwen_vl.py index 3b0dbb81..c0f2a185 100644 --- a/src/mobius/models/qwen_vl.py +++ b/src/mobius/models/qwen_vl.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -86,6 +86,19 @@ class Qwen25VLCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. The decoder + # paths deliberately exclude ``embed_tokens``, which belongs to embedding. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config @@ -402,6 +415,19 @@ class Qwen2VLCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. The decoder + # paths deliberately exclude ``embed_tokens``, which belongs to embedding. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config @@ -753,6 +779,19 @@ class Qwen3VL3ModelCausalLMModel(nn.Module): category: str = "Multimodal" config_class: type = ArchitectureConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. The decoder + # paths deliberately exclude ``embed_tokens``, which belongs to embedding. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "decoder": ( + "model.language_model.layers", + "model.language_model.norm", + "model.language_model.rotary_emb", + "lm_head", + ), + "vision_encoder": ("model.visual",), + "embedding": ("model.language_model.embed_tokens",), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/t5.py b/src/mobius/models/t5.py index 4ed05db7..28d7c6b0 100644 --- a/src/mobius/models/t5.py +++ b/src/mobius/models/t5.py @@ -6,7 +6,7 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar import torch from onnxscript import OpBuilder, nn @@ -441,6 +441,12 @@ class T5ForConditionalGeneration(nn.Module): default_task = "seq2seq" category = "encoder-decoder" + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": ("encoder",), + "decoder": ("decoder", "lm_head"), + } + def __init__(self, config: ArchitectureConfig): super().__init__() self.config = config diff --git a/src/mobius/models/trocr.py b/src/mobius/models/trocr.py index 3c188726..61f871a7 100644 --- a/src/mobius/models/trocr.py +++ b/src/mobius/models/trocr.py @@ -13,6 +13,8 @@ from __future__ import annotations +from typing import ClassVar + import torch from mobius.models.bart import BartForConditionalGeneration, _rename_bart_weight @@ -25,6 +27,13 @@ class TrOCRForConditionalGeneration(BartForConditionalGeneration): The vision encoder is a separate model (e.g. ViT, DeiT, BEiT). """ + # A standalone HF TrOCRForCausalLM has no encoder. Its decoder is nested + # under ``model.decoder`` and its output head is top-level. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": (), + "decoder": ("model.decoder", "output_projection"), + } + # TODO(feature): Add a VisionEncoderDecoderTask that pairs this decoder # with a vision encoder (ViT/DeiT/BEiT) for full TrOCR inference. # Currently only the decoder half is supported as a seq2seq model. diff --git a/src/mobius/models/unet_parity_test.py b/src/mobius/models/unet_parity_test.py index 95d9ab26..292872c5 100644 --- a/src/mobius/models/unet_parity_test.py +++ b/src/mobius/models/unet_parity_test.py @@ -18,6 +18,21 @@ import pytest +def _run_onnx(model, *feeds): + import onnx_ir + import onnxruntime as ort + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = str(Path(temp_dir) / "model.onnx") + onnx_ir.save(model, model_path) + session = ort.InferenceSession(model_path) + try: + return [session.run(None, feed)[0] for feed in feeds] + finally: + # Windows keeps the model mapped until the session is released. + del session + + def _remap_transformer(state_dict: dict) -> dict: out = {} for key, value in state_dict.items(): @@ -32,7 +47,6 @@ def _remap_transformer(state_dict: dict) -> dict: def test_cross_attention_block_matches_diffusers(): pytest.importorskip("diffusers") import onnx_ir - import onnxruntime as ort import torch from diffusers.models.transformers.transformer_2d import Transformer2DModel @@ -64,19 +78,12 @@ def test_cross_attention_block_matches_diffusers(): model = _make_model(graph) apply_weights(model, _remap_transformer(hf.state_dict())) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: - model_path = Path(temp_dir) / "model.onnx" - onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) - actual = session.run(None, {"hidden": hidden.numpy(), "context": context.numpy()})[0] + (actual,) = _run_onnx(model, {"hidden": hidden.numpy(), "context": context.numpy()}) assert np.abs(actual - expected).max() < 1e-4 def test_unet_matches_diffusers(): pytest.importorskip("diffusers") - import onnx_ir - import onnxruntime as ort import torch from diffusers import UNet2DConditionModel as HFUNet @@ -118,19 +125,14 @@ def test_unet_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: - model_path = Path(temp_dir) / "model.onnx" - onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) - actual = session.run( - None, - { - "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), - "encoder_hidden_states": encoder_hidden_states.numpy(), - }, - )[0] + (actual,) = _run_onnx( + model, + { + "sample": sample.numpy(), + "timestep": timestep.numpy().astype(np.int64), + "encoder_hidden_states": encoder_hidden_states.numpy(), + }, + ) assert np.abs(actual - expected).max() < 2e-4 @@ -143,8 +145,6 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): so cross-attention is present only where diffusers places it. """ pytest.importorskip("diffusers") - import onnx_ir - import onnxruntime as ort import torch from diffusers import UNet2DConditionModel as HFUNet @@ -190,19 +190,14 @@ def test_unet_sd1x_mixed_block_types_matches_diffusers(): model = DenoisingTask().build(module, config)["model"] apply_weights(model, module.preprocess_weights(dict(hf.state_dict()))) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: - model_path = Path(temp_dir) / "model.onnx" - onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) - actual = session.run( - None, - { - "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), - "encoder_hidden_states": encoder_hidden_states.numpy(), - }, - )[0] + (actual,) = _run_onnx( + model, + { + "sample": sample.numpy(), + "timestep": timestep.numpy().astype(np.int64), + "encoder_hidden_states": encoder_hidden_states.numpy(), + }, + ) assert np.abs(actual - expected).max() < 2e-4 @@ -210,8 +205,6 @@ def test_unet_lora_gate_parity(): """Runtime LoRA parity: gate=0 == diffusers base, gate=1 == diffusers+LoRA.""" pytest.importorskip("diffusers") pytest.importorskip("peft") - import onnx_ir - import onnxruntime as ort import torch from diffusers import UNet2DConditionModel as HFUNet from diffusers.utils import convert_state_dict_to_diffusers @@ -278,18 +271,16 @@ def test_unet_lora_gate_parity(): weights.update(remap_diffusers_unet_lora(lora_state, "test")) apply_weights(model, weights) - # Windows keeps the ORT model file mapped; ignore cleanup errors so the dir can be removed. - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: - model_path = Path(temp_dir) / "model.onnx" - onnx_ir.save(model, model_path) - session = ort.InferenceSession(model_path) - feed = { - "sample": sample.numpy(), - "timestep": timestep.numpy().astype(np.int64), - "encoder_hidden_states": encoder_hidden_states.numpy(), - } - off = session.run(None, {**feed, "lora_gate.test": np.array(0.0, dtype=np.float32)})[0] - on = session.run(None, {**feed, "lora_gate.test": np.array(1.0, dtype=np.float32)})[0] + feed = { + "sample": sample.numpy(), + "timestep": timestep.numpy().astype(np.int64), + "encoder_hidden_states": encoder_hidden_states.numpy(), + } + off, on = _run_onnx( + model, + {**feed, "lora_gate.test": np.array(0.0, dtype=np.float32)}, + {**feed, "lora_gate.test": np.array(1.0, dtype=np.float32)}, + ) # gate=0 disables the adapter (base); gate=1 applies it (diffusers+LoRA). assert np.abs(off - base_out).max() < 2e-4, np.abs(off - base_out).max() diff --git a/src/mobius/models/whisper.py b/src/mobius/models/whisper.py index d199e491..a8409b45 100644 --- a/src/mobius/models/whisper.py +++ b/src/mobius/models/whisper.py @@ -13,6 +13,7 @@ from __future__ import annotations import math +from typing import ClassVar import numpy as np import onnx_ir as ir @@ -184,6 +185,12 @@ class WhisperForConditionalGeneration(nn.Module): category: str = "Speech-to-Text" config_class: type = WhisperConfig + # Runtime HF ``named_modules()`` sub-trees per ONNX component. + HF_COMPONENT_SOURCES: ClassVar[dict[str, tuple[str, ...]]] = { + "encoder": ("model.encoder",), + "decoder": ("model.decoder", "proj_out"), + } + def __init__(self, config: WhisperConfig): super().__init__() self.config = config diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index a8233a79..c7c96955 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -21,6 +21,8 @@ from __future__ import annotations +from typing import ClassVar + import onnx_ir as ir from onnxscript import GraphBuilder, nn @@ -411,6 +413,15 @@ class Gemma4Task(ModelTask): pre-allocated cache (static mode). """ + #: decoder + vision + embedding, plus audio when ``config.audio`` is set. + #: ``audio_encoder`` is declared statically (it is config-gated at build time). + model_roles: ClassVar[dict[str, str]] = { + "decoder": "decoder", + "vision_encoder": "encoder", + "audio_encoder": "encoder", + "embedding": "embedding", + } + def __init__( self, *, diff --git a/testdata/cases/speech/fun-asr-nano.yaml b/testdata/cases/speech/fun-asr-nano.yaml index 3b3a444b..294635ac 100644 --- a/testdata/cases/speech/fun-asr-nano.yaml +++ b/testdata/cases/speech/fun-asr-nano.yaml @@ -1,9 +1,12 @@ -model_id: "justinchuby/Fun-ASR-Nano-2512" +model_id: "FunAudioLLM/Fun-ASR-Nano-2512" model_type: "fun_asr" revision: "main" task_type: "fun-asr-speech-language" dtype: "float32" trust_remote_code: true +ci_skip_reason: > + HF config for Fun-ASR-Nano lacks config.json model_type, so build() + auto-detection fails in CI's generic golden pipeline. inputs: audio: @@ -19,4 +22,4 @@ notes: > Fun-ASR-Nano speech recognition. 3-model split: audio_encoder (SANM encoder) + embedding (adaptor + text/audio fusion) + decoder (Qwen3-0.6B). Audio input is LFR fbank features (batch, seq_len, 560). - Weights in safetensors format at justinchuby/Fun-ASR-Nano-2512. + Weights in safetensors format at FunAudioLLM/Fun-ASR-Nano-2512.