From 58c17ea178ae585f8440864aaa65031b81bba216 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 3 Jul 2026 09:53:37 +0200 Subject: [PATCH 01/14] add openvino exporter Co-Authored-By: Claude Fable 5 --- docs/source/en/exporters.md | 78 + docs/source/en/main_classes/exporters.md | 4 + src/transformers/exporters/__init__.py | 3 +- src/transformers/exporters/auto.py | 3 + src/transformers/exporters/configs.py | 31 + src/transformers/exporters/exporter_dynamo.py | 36 +- .../exporters/exporter_dynamo.py.orig | 690 ++++++ .../exporters/exporter_executorch.py | 440 +++- src/transformers/exporters/exporter_onnx.py | 55 + .../exporters/exporter_openvino.py | 1923 +++++++++++++++++ src/transformers/exporters/utils.py | 33 +- src/transformers/masking_utils.py | 11 +- .../models/big_bird/modeling_big_bird.py | 3 +- src/transformers/models/bros/modeling_bros.py | 4 +- .../models/chameleon/modeling_chameleon.py | 33 +- src/transformers/models/fsmt/modeling_fsmt.py | 16 +- .../models/gemma3n/modeling_gemma3n.py | 2 +- .../models/gemma3n/modular_gemma3n.py | 2 +- .../modeling_gpt_neox_japanese.py | 4 +- .../models/informer/modeling_informer.py | 2 +- .../models/longt5/modeling_longt5.py | 27 +- .../mask2former/modeling_mask2former.py | 17 +- .../maskformer/modeling_maskformer_swin.py | 10 +- .../models/minimax/modeling_minimax.py | 15 +- .../models/minimax/modular_minimax.py | 15 +- .../models/mpnet/modeling_mpnet.py | 5 +- .../models/musicgen/modeling_musicgen.py | 2 +- .../modeling_musicgen_melody.py | 2 +- .../models/perceiver/modeling_perceiver.py | 22 +- .../models/phimoe/modeling_phimoe.py | 33 +- .../models/phimoe/modular_phimoe.py | 33 +- .../modeling_pp_doclayout_v3.py | 22 +- .../modular_pp_doclayout_v3.py | 22 +- .../models/prophetnet/modeling_prophetnet.py | 8 +- .../models/speecht5/modeling_speecht5.py | 2 +- .../models/swin2sr/modeling_swin2sr.py | 12 +- .../models/swinv2/modeling_swinv2.py | 12 +- .../models/t5gemma2/modeling_t5gemma2.py | 5 +- .../models/t5gemma2/modular_t5gemma2.py | 5 +- src/transformers/models/tvp/modeling_tvp.py | 4 +- .../models/unispeech/modeling_unispeech.py | 8 +- .../models/unispeech/modular_unispeech.py | 8 +- .../models/videomae/modeling_videomae.py | 161 +- .../configuration_voxtral_realtime.py | 2 + .../modeling_voxtral_realtime.py | 1 + .../modular_voxtral_realtime.py | 1 + .../models/wavlm/modeling_wavlm.py | 6 +- .../models/wavlm/modular_wavlm.py | 6 +- .../models/xcodec2/modeling_xcodec2.py | 7 +- .../models/xcodec2/modular_xcodec2.py | 7 +- .../models/xlnet/modeling_xlnet.py | 10 +- src/transformers/pytorch_utils.py | 4 +- src/transformers/testing_utils.py | 5 + src/transformers/utils/__init__.py | 1 + src/transformers/utils/import_utils.py | 5 + tests/exporters/test_export.py | 278 ++- .../test_modeling_openai_privacy_filter.py | 1 + 57 files changed, 3745 insertions(+), 412 deletions(-) create mode 100644 src/transformers/exporters/exporter_dynamo.py.orig create mode 100644 src/transformers/exporters/exporter_openvino.py diff --git a/docs/source/en/exporters.md b/docs/source/en/exporters.md index 9b2871a56e3a..7aa4459938b6 100644 --- a/docs/source/en/exporters.md +++ b/docs/source/en/exporters.md @@ -43,6 +43,7 @@ removals as we follow upstream. | ---------------------- | -------------------------- | --------------------------------------------- | | [`DynamoExporter`] | `ExportedProgram` | Any PyTorch runtime, AOT compilation | | [`OnnxExporter`] | `ONNXProgram` | Any ONNX runtime (ORT, TensorRT, OpenVINO, …) | +| [`OpenVINOExporter`] | `openvino.Model` | OpenVINO runtime (Intel CPU/GPU/NPU) | | [`ExecutorchExporter`] | `ExecutorchProgramManager` | Mobile and edge devices (ExecuTorch) | [`AutoHfExporter`] picks the right exporter from a config and [`AutoExportConfig`] picks the right @@ -72,6 +73,13 @@ pip install transformers "torch==2.12.0" "onnx==1.21.0" "onnxscript==0.7.0" onnx pip install transformers "torch==2.12.0" "executorch==1.3.1" ``` + + + +```bash +pip install transformers "torch==2.12.0" "openvino==2025.0.0" +``` + @@ -156,6 +164,30 @@ method = program.load_method("forward") outputs = method.execute(list(inputs.values())) ``` + + + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.exporters import OpenVINOExporter, OpenVINOConfig + +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") +inputs = tokenizer("Hello, world!", return_tensors="pt") + +exporter = OpenVINOExporter() +config = OpenVINOConfig(dynamic=True) +ov_model = exporter.export(model, inputs, config=config) + +ov_model.save("model.xml") + +# compile and run on CPU (or "GPU" / "NPU" if available) +import openvino as ov +compiled = ov.Core().compile_model(ov_model, "CPU") +ov_inputs = {k: v.numpy() for k, v in inputs.items()} +outputs = compiled(ov_inputs) +``` + @@ -251,6 +283,33 @@ config = ExecutorchConfig( et_program = exporter.export(model, inputs, config=config) ``` + + + +```python +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.exporters import OpenVINOExporter, OpenVINOConfig + +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") +inputs = tokenizer(["Hello, world!", "Hi"], padding=True, return_tensors="pt") + +batch = torch.export.Dim("batch", min=1, max=32) +seq = torch.export.Dim("seq", min=1, max=2048) + +exporter = OpenVINOExporter() +config = OpenVINOConfig( + dynamic_shapes={"input_ids": {0: batch, 1: seq}, "attention_mask": {0: batch, 1: seq}}, + # Emit data-dependent shape guards as runtime asserts instead of failing the export when a + # guard wouldn't hold across the explicit symbolic range — most LLMs need this under fine-grained + # ``Dim(min=, max=)`` bounds. Not needed with ``dynamic=True`` / ``Dim.AUTO``, where torch.export + # infers shape relations instead of verifying them against user-stated bounds. + prefer_deferred_runtime_asserts_over_guards=True, +) +ov_model = exporter.export(model, inputs, config=config) +``` + @@ -325,6 +384,25 @@ components = exporter.export_for_generation(model, inputs, config=config) # components = {"image_encoder": ExecutorchProgramManager, "language_model": ..., "lm_head": ..., "decode": ...} ``` + + + +```python +from transformers import AutoModelForImageTextToText, AutoProcessor +from transformers.exporters import OpenVINOExporter, OpenVINOConfig + +model = AutoModelForImageTextToText.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") +processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") +messages = [{"role": "user", "content": [{"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"}, {"type": "text", "text": "Describe this image."}]}] +text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) +inputs = processor(text=text, images=messages[0]["content"][0]["url"], return_tensors="pt").to(model.device) + +exporter = OpenVINOExporter() +config = OpenVINOConfig(dynamic=True) +components = exporter.export_for_generation(model, inputs, config=config) +# components = {"image_encoder": openvino.Model, "language_model": openvino.Model, "lm_head": openvino.Model, "decode": openvino.Model} +``` + diff --git a/docs/source/en/main_classes/exporters.md b/docs/source/en/main_classes/exporters.md index ac1f16d8a256..bf7d24a7c283 100644 --- a/docs/source/en/main_classes/exporters.md +++ b/docs/source/en/main_classes/exporters.md @@ -47,3 +47,7 @@ Learn how to use the built-in exporters in the [Exporters](../exporters) guide. ## ExecutorchConfig [[autodoc]] exporters.configs.ExecutorchConfig + +## OpenVINOConfig + +[[autodoc]] exporters.configs.OpenVINOConfig diff --git a/src/transformers/exporters/__init__.py b/src/transformers/exporters/__init__.py index 415fdb8c79ea..500976d85289 100644 --- a/src/transformers/exporters/__init__.py +++ b/src/transformers/exporters/__init__.py @@ -14,7 +14,8 @@ # limitations under the License. from .auto import AutoExportConfig, AutoHfExporter, get_hf_exporter, register_export_config, register_exporter from .base import HfExporter -from .configs import DynamoConfig, ExecutorchConfig, ExportConfigMixin, ExportFormat, OnnxConfig +from .configs import DynamoConfig, ExecutorchConfig, ExportConfigMixin, ExportFormat, OnnxConfig, OpenVINOConfig from .exporter_dynamo import DynamoExporter from .exporter_executorch import ExecutorchExporter from .exporter_onnx import OnnxExporter +from .exporter_openvino import OpenVINOExporter diff --git a/src/transformers/exporters/auto.py b/src/transformers/exporters/auto.py index 1f324d19cfbb..adef5a2928dc 100644 --- a/src/transformers/exporters/auto.py +++ b/src/transformers/exporters/auto.py @@ -22,16 +22,19 @@ from .exporter_dynamo import DynamoConfig, DynamoExporter from .exporter_executorch import ExecutorchConfig, ExecutorchExporter from .exporter_onnx import OnnxConfig, OnnxExporter +from .exporter_openvino import OpenVINOConfig, OpenVINOExporter AUTO_EXPORTER_MAPPING = { "executorch": ExecutorchExporter, + "openvino": OpenVINOExporter, "dynamo": DynamoExporter, "onnx": OnnxExporter, } AUTO_EXPORT_CONFIG_MAPPING = { "executorch": ExecutorchConfig, + "openvino": OpenVINOConfig, "dynamo": DynamoConfig, "onnx": OnnxConfig, } diff --git a/src/transformers/exporters/configs.py b/src/transformers/exporters/configs.py index 713448debacd..982751aa2aa1 100644 --- a/src/transformers/exporters/configs.py +++ b/src/transformers/exporters/configs.py @@ -28,6 +28,7 @@ class ExportFormat(Enum): """Identifies the export backend. Stored in [`ExportConfigMixin`] for serialisation round-trips.""" EXECUTORCH = "executorch" + OPENVINO = "openvino" DYNAMO = "dynamo" ONNX = "onnx" @@ -168,3 +169,33 @@ class ExecutorchConfig(DynamoConfig): export_format: ExportFormat = ExportFormat.EXECUTORCH backend: str = "xnnpack" + + +@dataclass +class OpenVINOConfig(DynamoConfig): + """ + Configuration class for exporting models to OpenVINO IR via ``openvino.convert_model``. + + Inherits all fields from [`DynamoConfig`] (`dynamic`, `strict`, `dynamic_shapes`, + `prefer_deferred_runtime_asserts_over_guards`). + + Args: + output_path (`str` or `PathLike`, *optional*): + Output path for the `.xml` file (the matching `.bin` is written alongside). When + `None` (default) the converted model is kept in memory as an ``openvino.Model``. + compress_to_fp16 (`bool`, *optional*, defaults to `True`): + Compress floating-point weights to FP16 when saving — halves on-disk size with + negligible accuracy impact on most models. Only applied when ``output_path`` is set. + stateful (`bool`, *optional*, defaults to `True`): + Fold round-tripped state tensors (KV cache, SSM states, …) into internal OV + variables (``ReadValue``/``Assign``). The runtime then carries state across + ``infer()`` calls instead of marshalling cache tensors through inputs/outputs on + every step, and a fused ``beam_idx`` input reorders state in-graph for beam search. + No-op for models without round-tripped state (encoders, prefill-only exports). + """ + + export_format: ExportFormat = ExportFormat.OPENVINO + + output_path: str | PathLike | None = None + compress_to_fp16: bool = True + stateful: bool = True diff --git a/src/transformers/exporters/exporter_dynamo.py b/src/transformers/exporters/exporter_dynamo.py index 9dee78ad15bf..152387c39b06 100644 --- a/src/transformers/exporters/exporter_dynamo.py +++ b/src/transformers/exporters/exporter_dynamo.py @@ -43,6 +43,7 @@ import importlib import inspect import sys +import types from collections.abc import MutableMapping from contextlib import contextmanager from typing import Any @@ -470,12 +471,21 @@ def _flatten_to_context(obj: Any, tensors: list) -> Any: "p": _class_to_path(cls), "v": [_flatten_to_context(i, tensors) for i in obj], } + if isinstance(obj, types.MethodType): + # Self-bound methods are handled by the `"obj"` branch below (rebound to the + # reconstructed instance); a bare bound method has no instance to rebind to. + raise TypeError("Cannot flatten a bound method not stored on its own instance for pytree context") if hasattr(obj, "__dict__"): - return { - "_t": "obj", - "p": _class_to_path(cls), - "s": {k: _flatten_to_context(v, tensors) for k, v in vars(obj).items()}, - } + state = {} + for k, v in vars(obj).items(): + if isinstance(v, types.MethodType) and v.__self__ is obj: + # Methods bound onto the instance itself (e.g. recurrent_gemma binds + # `get_seq_length` onto its `DynamicCache`) — store the underlying + # function's import path; unflatten rebinds it to the new instance. + state[k] = {"_t": "method", "f": f"{v.__func__.__module__}:{v.__func__.__qualname__}"} + else: + state[k] = _flatten_to_context(v, tensors) + return {"_t": "obj", "p": _class_to_path(cls), "s": state} raise TypeError(f"Cannot flatten {type(obj).__name__} for pytree context") @@ -522,9 +532,12 @@ def _unflatten_from_context(ctx: Any, tensors: list) -> Any: return cls(*items) # NamedTuple (requires positional args) if t == "obj": cls = _path_to_class(ctx["p"]) - state = {k: _unflatten_from_context(v, tensors) for k, v in ctx["s"].items()} instance = cls.__new__(cls) - instance.__dict__.update(state) + for k, v in ctx["s"].items(): + if type(v) is dict and v.get("_t") == "method": + instance.__dict__[k] = _path_to_class(v["f"]).__get__(instance) + else: + instance.__dict__[k] = _unflatten_from_context(v, tensors) return instance raise TypeError(f"Unknown tag {t!r} in pytree context") @@ -565,6 +578,13 @@ def _iter_subclasses(cls: type): yield from _iter_subclasses(subclass) +def is_cache_object(value: Any) -> bool: + """Whether ``value`` is a cache — a [`Cache`] instance or a model-specific class following + the ``*Cache`` naming convention (e.g. ``xLSTMCache``, ``MimiConv1dPaddingCache``), matching + what [`register_cache_pytrees_for_model`] registers as pytree nodes.""" + return isinstance(value, Cache) or type(value).__name__.endswith("Cache") + + def register_cache_pytrees_for_model(model: PreTrainedModel): """Register all relevant cache types as pytree nodes for torch.export.""" # All transformers Cache subclasses @@ -639,6 +659,8 @@ def get_auto_dynamic_shapes(inputs: Any) -> Any: "_prefill_len", # glm_image (m-rope prefill length) "cached_rotary_positional_embedding", # wav2vec2_bert, seamless_m4t, clvp "cached_sequence_length", # wav2vec2_bert, seamless_m4t, clvp + "conv1d_state", # recurrent_gemma (conv state on RecurrentGemmaRecurrentBlock) + "recurrent_states", # recurrent_gemma (RG-LRU state on RecurrentGemmaRglru) ) diff --git a/src/transformers/exporters/exporter_dynamo.py.orig b/src/transformers/exporters/exporter_dynamo.py.orig new file mode 100644 index 000000000000..654f30d4ba23 --- /dev/null +++ b/src/transformers/exporters/exporter_dynamo.py.orig @@ -0,0 +1,690 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dynamo exporter. + +Wraps `torch.export.export(strict=False)` with helpers that make Transformers +models exportable. The export pipeline uses five sections, in execution order: + +1. **Model signature patch** (`patch_forward_signature`): replaces `model.forward` + with a flat explicit signature derived from `sample_inputs` so `torch.export` does + not expand `**kwargs` into a `combined_args` bundle that mismatches `dynamic_shapes`. + This is the entry contract `torch.export` reads before tracing. +2. **Model patches** (`_PATCHES["dynamo"]` via `apply_patches("dynamo")`): reversible + class-attribute swaps applied during tracing to replace non-exportable model patterns + (data-dependent loops, in-place ops, mask checks) with export-safe equivalents. + Modeling code itself is not updated because these patches are too model-specific. +3. **Pytree registration** (`register_cache_pytrees_for_model`): flatten/unflatten + hooks (via `torch.utils._pytree.register_pytree_node`) for Cache subclasses and + custom containers so `torch.export` can trace through them. +4. **Dynamic shapes** (`get_auto_dynamic_shapes`): automatic `Dim.AUTO` inference + for all tensor and cache inputs when `DynamoConfig.dynamic=True`. +5. **Model state cleanup** (`reset_model_state`): non-Cache stateful module attributes + (`_STATEFUL_CACHE_ATTRS`) are saved on entry, set to `None` during the trace, and + restored on exit — so a previous eager forward doesn't leak into the trace and any + FakeTensors the trace planted are discarded before the next eager forward. +""" + +from __future__ import annotations + +import copy +import importlib +import inspect +import sys +from collections.abc import MutableMapping +from contextlib import ExitStack, contextmanager +from typing import Any + +from ..utils import logging +from ..utils.import_utils import is_detectron2_available, is_torch_available, torch_compilable_check +from .base import HfExporter +from .configs import DynamoConfig +from .utils import apply_patches, prepare_for_export, register_patch + + +if is_torch_available(): + import torch + from torch.export import ExportedProgram + + from ..cache_utils import Cache + from ..modeling_utils import PreTrainedModel + + +logger = logging.get_logger(__file__) + + +class DynamoExporter(HfExporter): + """Exporter that converts a [`PreTrainedModel`] to an `ExportedProgram`. + + Example: + + ```python + >>> from transformers.exporters.exporter_dynamo import DynamoExporter, DynamoConfig + + >>> exporter = DynamoExporter() + >>> exported = exporter.export(model, inputs, config=DynamoConfig(dynamic=True)) + >>> outputs = exported.module()(**inputs) + ``` + """ + + required_packages = ["torch"] + tested_versions = {"torch": "2.12.0"} + + def export( + self, + model: PreTrainedModel, + sample_inputs: MutableMapping[str, Any], + config: DynamoConfig | dict[str, Any], + ) -> ExportedProgram: + if isinstance(config, dict): + config = DynamoConfig(**config) + elif not isinstance(config, DynamoConfig): + raise TypeError(f"Expected config to be a DynamoConfig or dict, got {type(config)}") + + model, sample_inputs, output_flags = prepare_for_export(model, sample_inputs) + + dynamic_shapes = config.dynamic_shapes + if config.dynamic and dynamic_shapes is None: + dynamic_shapes = get_auto_dynamic_shapes(sample_inputs) + + register_cache_pytrees_for_model(model) + + with ( + apply_patches("dynamo"), + reset_model_state(model), + patch_model_config(model, output_flags), + patch_forward_signature(model, sample_inputs), + ): + exported_program: ExportedProgram = torch.export.export( + model, + args=(), + kwargs=copy.deepcopy(dict(sample_inputs)), + strict=config.strict, + dynamic_shapes=dynamic_shapes, + prefer_deferred_runtime_asserts_over_guards=config.prefer_deferred_runtime_asserts_over_guards, + ) + + return exported_program + + +# ── Stage 1: Model signature patch ────────────────────────────────────────── +# Replaces `model.forward` with a flat explicit signature derived from the +# inputs dict so `torch.export` does not expand `**kwargs` into a large bundle. +# `patch_model_config` lives here too — it strips output flags from the inputs +# and applies them onto `model.config` for the duration of the trace. + + +_MISSING = object() + + +@contextmanager +def _set_config_attribute(config: Any, name: str, value: Any): + """Set `config. = value` for the block; restore the original (or delete) on exit.""" + original = getattr(config, name, _MISSING) + setattr(config, name, value) + try: + yield + finally: + if original is _MISSING: + delattr(config, name) + else: + setattr(config, name, original) + + +# Output flags stripped from inputs and applied onto `model.config` for the trace. +@contextmanager +def patch_model_config(model: PreTrainedModel, output_flags: dict[str, Any]): + """Reversibly tweak `model.config` for the trace: + + - Applies `output_flags` (popped from inputs by `prepare_for_export`) onto + `model.config.` so the model picks them up via its usual ` if is + not None else self.config.` fallback. Flags the config doesn't declare are set + anyway (and deleted on exit) — submodels of a decomposed model may not declare a flag + their forward still reads via `getattr(self.config, flag, None)` (e.g. `use_cache` on + an audio-encoder config), and a flag the model never reads is a harmless no-op. + - Disables `use_mamba_kernels` on every submodel's config that declares it (mamba/jamba + kernels are not exportable). + + Originals are restored on exit. Flags whose value is `None` are skipped. + """ + with ExitStack() as stack: + if hasattr(model, "config"): + for flag, value in output_flags.items(): + if value is None: + continue + stack.enter_context(_set_config_attribute(model.config, flag, value)) + for module in model.modules(): + if hasattr(module, "config") and hasattr(module.config, "use_mamba_kernels"): + stack.enter_context(_set_config_attribute(module.config, "use_mamba_kernels", False)) + yield + + +@contextmanager +def patch_forward_signature(model: PreTrainedModel, inputs: dict[str, Any]): + """Temporarily replace `model.forward` with a flat explicit signature derived from `inputs`. + + `torch.export` infers the exported function signature from `model.forward.__signature__`. + Most transformers models use `**kwargs: Unpack[TransformersKwargs]`, which causes + `torch.export` to expand the signature into a large `combined_args` bundle that + mismatches the `dynamic_shapes` dict. This patch replaces the forward with a + minimal signature containing only the keys present in `inputs`. + """ + original_forward = model.forward + + def _flat_forward(**kwargs): + return original_forward(**kwargs) + + _flat_forward.__signature__ = inspect.Signature( + [inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None) for k in inputs] + ) + + try: + model.forward = _flat_forward + yield + finally: + model.forward = original_forward + + +# ── Stage 2: Model patches ──────────────────────────────────────────────────── +# Reversible class-attribute swaps applied during `torch.export` tracing via +# `apply_patches("dynamo")`. Each replaces a non-exportable model pattern +# (data-dependent control flow, in-place ops on views, etc.) with an +# export-safe equivalent on the owning class — every live instance sees the +# replacement until the context exits. Modeling code itself is not updated +# because these patches are too model-specific; we do strive to keep modeling +# code compliant where reasonable. +# +# Each `@register_patch("dynamo", *dotted_paths)` decorator targets one or +# more `Class.method` paths and wraps a `factory(original) -> replacement`. +# Multiple paths share the same factory when the same method shape needs to be +# swapped across several classes (e.g. `_reshaped_vision_attention_forward` +# applied to every chunked-vision attention class — see the long list below). + + +@register_patch("dynamo", "transformers.models.nllb_moe.modeling_nllb_moe.NllbMoeTop2Router._cast_classifier") +def _patch_classifier_cast(_original): + """Disable classifier dtype cast in nllb-moe (not traceable).""" + return lambda self, *args, **kwargs: None + + +@register_patch("dynamo", "torch.nn.functional.scaled_dot_product_attention") +def _patch_sdpa(original): + """Route SDPA through the MATH backend on CPU during tracing — CPU SDPA's flash/efficient + paths guard on ``Eq(batch, 1)`` (upstream https://github.com/pytorch/pytorch/issues/180202), + which trips ``GuardOnDataDependentSymNode`` whenever the batch dim comes from a data-dependent + op like ``pixel_values[bool_mask]`` (Idefics2/3 and most VLMs). The MATH decomposition has no + batch-1 dispatch, so the guard never fires. CUDA exports are left alone — the GPU kernels + don't have this guard, and we want the flash/efficient decompositions there. + """ + from torch.nn.attention import SDPBackend, sdpa_kernel + + def patch(query, *args, **kwargs): + if query.device.type == "cpu": + with sdpa_kernel(SDPBackend.MATH): + return original(query, *args, **kwargs) + return original(query, *args, **kwargs) + + return patch + + +@register_patch( + "dynamo", + # Canonical definition + the public re-export. + "transformers.utils.import_utils.is_kernels_available", + "transformers.utils.is_kernels_available", + # Local `from ...utils import is_kernels_available` rebinds in modeling modules + # — each one needs its own override since the name is looked up there. + "transformers.modeling_utils.is_kernels_available", + "transformers.models.sam3_video.modeling_sam3_video.is_kernels_available", + "transformers.models.mra.modeling_mra.is_kernels_available", + "transformers.models.rwkv.modeling_rwkv.is_kernels_available", + "transformers.models.yoso.modeling_yoso.is_kernels_available", +) +def _patch_is_kernels_available(_original): + """Force-disable the optional ``kernels`` library during export — its kernels + call into native code that ``torch.export`` cannot trace, and the pure-PyTorch + fallbacks in each model are always traceable.""" + return lambda *args, **kwargs: False + + +# --- Chunked vision/audio attention ───────────────────────────────────────── +# Sub-encoders that pack multiple variable-length sequences into one flat tensor +# with `cu_seqlens` markers fall back to `split → per-segment SDPA → cat` in the +# unpatched forward, which is a Python loop that `torch.export` can't trace. +# `_reshaped_vision_attention_forward` replaces that loop with a reshape into a +# per-segment batch followed by a single SDPA call. It handles the layout +# differences across encoders (combined `qkv` vs separate `q/k/v` vs separate +# `q_proj/k_proj/v_proj`, asymmetric `q_dim/kv_dim` split, `(cos, sin)` vs single +# rotary tensor vs none, `.proj` vs `.out_proj`, NaViT `(1, T, D)` packing, +# tuple vs single return). The `returns_tuple` flag is bound once per class at +# install time by inspecting the original `forward`'s source. +# +# NOTE: this whole stack of patches becomes unnecessary once transformers adopts a +# proper varlen-attention op (e.g. PyTorch's `torch._nested.scaled_dot_product_attention` +# or a Flex-Attention varlen kernel) — the modeling forwards can then express the +# segmented attention directly with `cu_seqlens` and trace through `torch.export` +# without this reshape-into-batch workaround. Drop this section when that lands. + + +def _reshaped_vision_attention_forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + returns_tuple: bool = False, + **kwargs, +): + """Export-safe chunked vision/audio attention: reshape segments into a batch dim, + apply rotary if provided, run one SDPA call, project, and re-emit in the original layout.""" + + # Normalise NaViT-style `(1, T, D)` packing (minicpmv4_6) to the flat `(T, D)` layout + # the rest of this wrapper assumes. The leading dim is always 1 — multi-image batches + # are packed along the sequence dim. + needs_batch_restore = hidden_states.ndim == 3 + if needs_batch_restore: + hidden_states = hidden_states.squeeze(0) + + seq_length = hidden_states.shape[0] + torch_compilable_check( + seq_length != 0, + "Chunked vision attention received an empty input.", + ) + num_segments = cu_seqlens.shape[0] - 1 + torch_compilable_check( + seq_length % num_segments == 0, + "Chunked vision attention requires uniform segment lengths during export. " + "Ensure all images have the same resolution (use do_resize=True in the processor) " + "or pad inputs to a common size.", + ) + + if hasattr(self, "qkv"): + # Grouped-query attention (q_dim != kv_dim, e.g. Exaone4.5) splits asymmetrically; + # uniform reshape into (seq, 3, num_heads, -1) only works when Q, K, V share the head count. + if hasattr(self, "q_dim") and hasattr(self, "kv_dim") and self.q_dim != self.kv_dim: + query_states, key_states, value_states = self.qkv(hidden_states).split( + [self.q_dim, self.kv_dim, self.kv_dim], dim=-1 + ) + query_states = query_states.view(seq_length, self.num_heads, self.head_dim) + key_states = key_states.view(seq_length, self.num_key_value_heads, self.head_dim) + value_states = value_states.view(seq_length, self.num_key_value_heads, self.head_dim) + else: + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).transpose(0, 1).unbind(0) + ) + else: + q_proj = getattr(self, "q_proj", getattr(self, "q", None)) + k_proj = getattr(self, "k_proj", getattr(self, "k", None)) + v_proj = getattr(self, "v_proj", getattr(self, "v", None)) + query_states = q_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) + key_states = k_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) + value_states = v_proj(hidden_states).view(seq_length, self.num_heads, self.head_dim) + + if position_embeddings is not None: + # Each vision encoder ships its own ``apply_rotary_pos_emb_vision`` in its modeling file + # (Qwen2-VL's takes (q, k, cos, sin), Qwen2.5/3-Omni's takes (x, rotary_emb), etc.). Look + # it up on the model's own module so this patch stays signature-agnostic across the + # ~19 attention classes it's installed on. + apply_rotary_pos_emb_vision = sys.modules[type(self).__module__].apply_rotary_pos_emb_vision + if isinstance(position_embeddings, (tuple, list)): + # (cos, sin) tuple convention — most VL encoders. + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + else: + # Single `rotary_pos_emb` tensor convention — Qwen2.5/3 Omni vision applies rotary per-states. + query_states = apply_rotary_pos_emb_vision(query_states.unsqueeze(0), position_embeddings).squeeze(0) + key_states = apply_rotary_pos_emb_vision(key_states.unsqueeze(0), position_embeddings).squeeze(0) + + seg_len = seq_length // num_segments + + # (seq, heads, dim) → (n_seg, seg_len, heads, dim) → (n_seg, heads, seg_len, dim) + def _to_batched(t): + return t.unflatten(0, (num_segments, seg_len)).transpose(1, 2) + + query_states = _to_batched(query_states) + key_states = _to_batched(key_states) + value_states = _to_batched(value_states) + + torch_compilable_check(query_states.shape[0] != 0, "Reshaped chunked-vision attention got zero batch.") + torch_compilable_check(query_states.shape[2] != 0, "Reshaped chunked-vision attention got zero seq.") + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + is_causal=False, + scale=self.scaling, + dropout_p=0.0 if not self.training else self.attention_dropout, + enable_gqa=getattr(self, "num_key_value_heads", self.num_heads) != self.num_heads, + ) + + # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim) + attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1).contiguous() + out_proj = self.proj if hasattr(self, "proj") else self.out_proj + attn_output = out_proj(attn_output) + + if needs_batch_restore: + attn_output = attn_output.unsqueeze(0) + + return (attn_output, None) if returns_tuple else attn_output + + +@register_patch( + "dynamo", + # Combined `qkv` + `(cos, sin)` rotary + `.proj` + "transformers.models.qwen2_vl.modeling_qwen2_vl.VisionAttention.forward", + "transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLVisionAttention.forward", + "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionAttention.forward", + "transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe.Qwen3VLMoeVisionAttention.forward", + "transformers.models.qwen3_5.modeling_qwen3_5.Qwen3_5VisionAttention.forward", + "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe.Qwen3_5MoeVisionAttention.forward", + "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeVisionAttention.forward", + "transformers.models.glm4v.modeling_glm4v.Glm4vVisionAttention.forward", + "transformers.models.glm4v_moe.modeling_glm4v_moe.Glm4vMoeVisionAttention.forward", + "transformers.models.glm_ocr.modeling_glm_ocr.GlmOcrVisionAttention.forward", + "transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe.Ernie4_5_VLMoeVisionAttention.forward", + # Asymmetric `qkv` split + `(cos, sin)` rotary + `.proj` + "transformers.models.exaone4_5.modeling_exaone4_5.Exaone4_5_VisionAttention.forward", + # Combined `qkv` + no in-attention rotary + `.proj` + "transformers.models.glm_image.modeling_glm_image.GlmImageVisionAttention.forward", + # Separate `.q` / `.k` / `.v` + single rotary tensor + `.proj` + "transformers.models.qwen2_5_omni.modeling_qwen2_5_omni.Qwen2_5OmniVisionAttention.forward", + # Separate `_proj` + `(cos, sin)` rotary + `.out_proj` (tuple return) + "transformers.models.video_llama_3.modeling_video_llama_3.VideoLlama3VisionAttention.forward", + "transformers.models.paddleocr_vl.modeling_paddleocr_vl.PaddleOCRVisionAttention.forward", + # NaViT (1, T, D) + separate `_proj` + `.out_proj` (tuple return) + "transformers.models.minicpmv4_6.modeling_minicpmv4_6.MiniCPMV4_6VisionAttention.forward", + # Audio attention: separate `_proj` + `.out_proj`, no rotary + "transformers.models.qwen2_5_omni.modeling_qwen2_5_omni.Qwen2_5OmniAudioAttention.forward", + "transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe.Qwen3OmniMoeAudioAttention.forward", + "transformers.models.qwen3_asr.modeling_qwen3_asr.Qwen3ASRAudioAttention.forward", +) +def _patch_chunked_vision_attention(original): + """Bind `returns_tuple` once per class by inspecting the original forward's source.""" + src = inspect.getsource(original) + returns_tuple = "return attn_output, attn_weight" in src or "return attn_output, None" in src + + def forward(self, *args, **kwargs): + return _reshaped_vision_attention_forward(self, *args, returns_tuple=returns_tuple, **kwargs) + + return forward + + +# ── Stage 3: Pytree registration ───────────────────────────────────────────── +# torch.export needs pytree flatten/unflatten for Cache objects and other +# custom types. The generic flattener serialises any object to a JSON-native +# context (bools, ints, strings, dicts, lists) while collecting tensors into +# a flat list — the inverse reconstructs the original object. +# +# To register a new type: it should be handled automatically by the generic +# flattener. If not, add a branch in _flatten_to_context / _unflatten_from_context. + + +def _class_to_path(cls: type) -> str: + return f"{cls.__module__}:{cls.__qualname__}" + + +def _path_to_class(path: str) -> type: + module_name, qualname = path.split(":", 1) + obj = importlib.import_module(module_name) + for part in qualname.split("."): + obj = getattr(obj, part) + return obj + + +def _flatten_to_context(obj: Any, tensors: list) -> Any: + """Single-pass: recursively build a JSON-native context while collecting tensors into `tensors`.""" + # --- Pure Python / JSON-native (exact type check — subclasses fall through to stateful objects) --- + if obj is None or type(obj) in (bool, int, float, str): + return obj + if type(obj) is list: + return [_flatten_to_context(i, tensors) for i in obj] + if type(obj) is dict: + return {k: _flatten_to_context(v, tensors) for k, v in obj.items()} + + # --- Torch objects --- + if isinstance(obj, torch.Tensor): + idx = len(tensors) + tensors.append(obj) + return {"_t": "tensor", "i": idx} + if isinstance(obj, torch.Size): + return {"_t": "size", "v": list(obj)} + if isinstance(obj, torch.device): + return {"_t": "device", "s": str(obj)} + if isinstance(obj, torch.dtype): + return {"_t": "dtype", "n": str(obj).removeprefix("torch.")} + if isinstance(obj, torch.layout): + return {"_t": "layout", "n": str(obj).removeprefix("torch.")} + if isinstance(obj, (torch.SymInt, torch.SymFloat, torch.SymBool)): + idx = len(tensors) + tensors.append(obj) + return {"_t": "sym", "i": idx} + + # --- Python types --- + if isinstance(obj, type): + return {"_t": "type", "p": _class_to_path(obj)} + + # --- Generic Python objects (by structural category) --- + cls = type(obj) + if isinstance(obj, dict): # dict subclasses (OrderedDict, etc.) + return { + "_t": "map", + "p": _class_to_path(cls), + "v": {k: _flatten_to_context(v, tensors) for k, v in obj.items()}, + } + if isinstance(obj, (tuple, list, set, frozenset)): # sequences/sets incl. NamedTuple + return { + "_t": "seq", + "p": _class_to_path(cls), + "v": [_flatten_to_context(i, tensors) for i in obj], + } + if hasattr(obj, "__dict__"): + return { + "_t": "obj", + "p": _class_to_path(cls), + "s": {k: _flatten_to_context(v, tensors) for k, v in vars(obj).items()}, + } + + raise TypeError(f"Cannot flatten {type(obj).__name__} for pytree context") + + +def _unflatten_from_context(ctx: Any, tensors: list) -> Any: + """Reconstruct an object from its JSON-native context, substituting tensor index markers.""" + # --- Pure Python / JSON-native --- + if ctx is None or type(ctx) in (bool, int, float, str): + return ctx + if type(ctx) is list: + return [_unflatten_from_context(i, tensors) for i in ctx] + if type(ctx) is dict and "_t" not in ctx: + return {k: _unflatten_from_context(v, tensors) for k, v in ctx.items()} + + # --- Torch objects --- + t = ctx["_t"] + if t == "tensor": + return tensors[ctx["i"]] + if t == "layout": + return getattr(torch, ctx["n"]) + if t == "dtype": + return getattr(torch, ctx["n"]) + if t == "device": + return torch.device(ctx["s"]) + if t == "size": + return torch.Size(ctx["v"]) + if t == "sym": + return tensors[ctx["i"]] + + # --- Python types --- + if t == "type": + return _path_to_class(ctx["p"]) + + # --- Generic Python objects --- + if t == "map": + cls = _path_to_class(ctx["p"]) + return cls({k: _unflatten_from_context(v, tensors) for k, v in ctx["v"].items()}) + if t == "seq": + cls = _path_to_class(ctx["p"]) + items = [_unflatten_from_context(i, tensors) for i in ctx["v"]] + try: + return cls(items) # tuple, list subclass, set, frozenset, etc. + except TypeError: + return cls(*items) # NamedTuple (requires positional args) + if t == "obj": + cls = _path_to_class(ctx["p"]) + state = {k: _unflatten_from_context(v, tensors) for k, v in ctx["s"].items()} + instance = cls.__new__(cls) + instance.__dict__.update(state) + return instance + + raise TypeError(f"Unknown tag {t!r} in pytree context") + + +def _pytree_flatten(obj: Any) -> tuple[list, Any]: + tensors: list = [] + context = _flatten_to_context(obj, tensors) + return tensors, context + + +def _pytree_flatten_with_keys(obj: Any): + leaves, context = _pytree_flatten(obj) + return [(torch.utils._pytree.SequenceKey(i), leaf) for i, leaf in enumerate(leaves)], context + + +def _pytree_unflatten(values, context: Any) -> Any: + return _unflatten_from_context(context, list(values)) + + +def _register_pytree_node(object_cls: type): + try: + torch.utils._pytree.register_pytree_node( + object_cls, + _pytree_flatten, + _pytree_unflatten, + serialized_type_name=_class_to_path(object_cls), + flatten_with_keys_fn=_pytree_flatten_with_keys, + ) + except ValueError as e: + if "already registered as pytree node" not in str(e): + raise + + +def _iter_subclasses(cls: type): + for subclass in cls.__subclasses__(): + yield subclass + yield from _iter_subclasses(subclass) + + +def is_cache_object(value: Any) -> bool: + """Whether ``value`` is a cache — a [`Cache`] instance or a model-specific class following + the ``*Cache`` naming convention (e.g. ``xLSTMCache``, ``MimiConv1dPaddingCache``), matching + what [`register_cache_pytrees_for_model`] registers as pytree nodes.""" + return isinstance(value, Cache) or type(value).__name__.endswith("Cache") + + +def register_cache_pytrees_for_model(model: PreTrainedModel): + """Register all relevant cache types as pytree nodes for torch.export.""" + # All transformers Cache subclasses + for cache_type in _iter_subclasses(Cache): + _register_pytree_node(cache_type) + + # Model-specific cache classes not inheriting from Cache (e.g. custom per-model caches) + for _, obj in inspect.getmembers(inspect.getmodule(model)): + if ( + inspect.isclass(obj) + and obj.__module__ == model.__class__.__module__ + and obj.__name__.endswith("Cache") + and not issubclass(obj, Cache) + ): + _register_pytree_node(obj) + + # detectron2 ImageList (used by layoutlmv2) + if is_detectron2_available() and isinstance(model, PreTrainedModel) and model.config.model_type == "layoutlmv2": + from detectron2.structures.image_list import ImageList + + _register_pytree_node(ImageList) + + +# ── Stage 4: Dynamic shapes ───────────────────────────────────────────────── +# Automatic `Dim.AUTO` inference for all tensor and cache inputs when +# `DynamoConfig.dynamic` is True and no explicit `dynamic_shapes` are provided. + + +def _auto_dynamic_shape(tensor: torch.Tensor) -> dict[int, torch.export.Dim]: + """Generate a dynamic shape with all dimensions set to Dim.AUTO for a given tensor.""" + return dict.fromkeys(range(tensor.dim()), torch.export.Dim.AUTO) + + +def get_auto_dynamic_shapes(inputs: Any) -> Any: + """Recursively build dynamic shapes for any input value. + + - Tensors → per-dimension Dim.AUTO spec. + - Scalars / None → None (no dynamic dims). + - Objects with ``__dict__`` (ModelOutput, Cache, …) → flat list of leaf specs, + matching the ``TreeSpec(list, …)`` that torch.export produces for these types. + - Lists / tuples → same container type, recursed element-wise. + - Plain dicts → recursed dict of specs. + - Everything else → None. + """ + if isinstance(inputs, torch.Tensor): + return _auto_dynamic_shape(inputs) + if inputs is None or isinstance(inputs, (int, float, bool, str)): + return None + if hasattr(inputs, "__dict__"): + leaves, _ = _pytree_flatten(inputs) + return get_auto_dynamic_shapes(leaves) + if type(inputs) in (list, tuple, set, frozenset): + return type(inputs)(get_auto_dynamic_shapes(v) for v in inputs) + if type(inputs) is dict: + return {k: get_auto_dynamic_shapes(v) for k, v in inputs.items()} + return None + + +# ── Stage 5: Model state cleanup ──────────────────────────────────────────── +# `torch.export` traces forward with FakeTensors, which can leave non-Cache stateful +# tensor attributes as FakeTensors after tracing — a follow-up eager forward then +# hits shape/dtype mismatches when it reuses the stale state. We also want stale +# eager-mode state cleared on entry so it doesn't leak into the trace. +# `reset_model_state` brackets the `torch.export.export` call: it saves every +# attribute in `_STATEFUL_CACHE_ATTRS` on every submodule, sets them to `None` for +# the trace, and restores the originals on exit (finally semantics). +# +# To register a new stateful attribute: append its name to `_STATEFUL_CACHE_ATTRS`. + +_STATEFUL_CACHE_ATTRS = ( + "_cached_decode_position_ids", # glm_image (m-rope decode position ids) + "_prefill_len", # glm_image (m-rope prefill length) + "cached_rotary_positional_embedding", # wav2vec2_bert, seamless_m4t, clvp + "cached_sequence_length", # wav2vec2_bert, seamless_m4t, clvp +) + + +@contextmanager +def reset_model_state(model: torch.nn.Module): + """Save each `_STATEFUL_CACHE_ATTRS` value, null it for the trace, restore on exit. + + FakeTensors that `torch.export` plants into these attributes during the trace are + discarded by the restore. + """ + originals = [ + (module, attr, getattr(module, attr)) + for module in model.modules() + for attr in _STATEFUL_CACHE_ATTRS + if hasattr(module, attr) + ] + for module, attr, _ in originals: + setattr(module, attr, None) + try: + yield + finally: + for module, attr, original in originals: + setattr(module, attr, original) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index bc6513555098..9c4b77f6b39a 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -347,6 +347,15 @@ def _patch_scaled_dot_product_attention(original): - enable_gqa=True - D_q != D_v (asymmetric head dims, e.g. MLA attention) - attn_mask is float (ExecuTorch CUDA SDPA only accepts bool masks) + + The MATH-path output gets an explicit ``clone(memory_format=contiguous_format)`` so + downstream strides don't depend on which SDPA layout torch picks: the pre-dispatch trace + sees a contiguous ``(N, H, L, E)`` fake output (so ``.contiguous()`` would trace to + nothing) and records downstream ``reshape``s as bare ``view`` nodes, but decomposition + re-traces SDPA via ``scaled_dot_product_flash_attention_for_cpu``, which materializes an + ``(L, N, H, E)`` buffer — invalidating those recorded views (``Cannot view a tensor with + shape/strides``). ``clone`` records unconditionally and re-executes correctly under either + layout, normalizing the strides the rest of the graph was recorded against. """ def patch(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, **kwargs): @@ -373,7 +382,7 @@ def patch(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, sca with sdpa_kernel(SDPBackend.MATH): return original( query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs - ) + ).clone(memory_format=torch.contiguous_format) return patch @@ -584,6 +593,129 @@ def patch(self, exported_program, graph_module): return patch +@register_patch("executorch", "executorch.exir.program._program.lift_constant_tensor_pass") +def _patch_lift_constant_tensor_pass(original): + """Realign ``input_specs`` with the graph placeholder order after constant lifting. + + The upstream pass picks the graph insertion point for newly lifted constant + placeholders by matching node names against ``graph_signature.user_inputs`` — + but for user inputs exported as ``ConstantArgument`` (e.g. ``input_ids=None`` + in a prefill component that runs from ``inputs_embeds``), ``user_inputs`` + holds the argument's *value* (``None``), not its name, so the match fails and + the new placeholders land *after* that input while their signature specs land + *before* it. Later positional signature rebuilds then shift every buffer arg + name by one slot, and the emitter serializes the wrong tensor for each lifted + constant (``Tensor spec has buffer of size 4, but expected nbytes of 8``). + Reordering ``input_specs`` to match the placeholders restores the invariant + the rest of the pipeline assumes. + """ + + def patch(exported_program): + exported_program = original(exported_program) + signature = exported_program.graph_signature + placeholder_names = [node.name for node in exported_program.graph.nodes if node.op == "placeholder"] + specs_by_name = {getattr(spec.arg, "name", None): spec for spec in signature.input_specs} + if ( + None not in specs_by_name + and len(specs_by_name) == len(signature.input_specs) + and sorted(specs_by_name) == sorted(placeholder_names) + ): + signature.input_specs = [specs_by_name[name] for name in placeholder_names] + return exported_program + + return patch + + +def _view_replaceable_nodes(graph_module): + """Yield ``(node, shape)`` for non-output ``view_copy`` nodes whose view shape has the same + shape dynamism as their base — the nodes ``ReplaceViewCopyWithViewPass`` may safely replace. + + ``view`` nodes share storage with their base during memory planning, so ``_ViewSpec`` + requires both to have the same ``shape_dynamism``. Models that reshape a static parameter + with input-derived dynamic dims (the ``pos_embed.reshape(1, height, width, -1)`` position- + embedding interpolation in Pvt / DepthPro / VitDet) produce a dynamic-shaped view of a + static const base, and ``_ViewSpec.__init__`` raises ``_ViewSpec is incompatible with its + base``. Those nodes must stay ``view_copy`` (an out-variant copy op, always correct) — + only the storage-sharing optimisation is skipped for them. + """ + from executorch.exir.passes.replace_view_copy_with_view_pass import _is_view_copy + from executorch.exir.tensor import determine_tensor_dynanism + + for node in graph_module.graph.nodes: + if _is_view_copy(node) and all(user.op != "output" for user in node.users): + # The view shape is node.meta["val"].shape, not node.args[1], which can contain + # an inferred -1 (same as the original pass). + shape = node.meta["val"].shape + base = node.args[0] + if determine_tensor_dynanism(shape) == base.meta["spec"].shape_dynamism: + yield node, shape + + +@register_patch( + "executorch", "executorch.exir.passes.replace_view_copy_with_view_pass.ReplaceViewCopyWithViewPass.call" +) +def _patch_replace_view_copy_with_view_call(_original): + """Replacement for ``ReplaceViewCopyWithViewPass.call`` that only replaces ``view_copy`` + nodes whose shape dynamism matches their base's — see ``_view_replaceable_nodes``.""" + from executorch.exir.passes.replace_view_copy_with_view_pass import _VIEW_OP, _ViewSpec + from torch.fx.passes.infra.pass_base import PassResult + + def patch(self, graph_module): + n_replaced = 0 + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node, shape in _view_replaceable_nodes(module): + node.target = _VIEW_OP + node.meta["spec"] = _ViewSpec(node.args[0].meta["spec"], shape) + n_replaced += 1 + module.recompile() + return PassResult(graph_module, n_replaced > 0) + + return patch + + +@register_patch( + "executorch", "executorch.exir.passes.replace_view_copy_with_view_pass.ReplaceViewCopyWithViewPass.ensures" +) +def _patch_replace_view_copy_with_view_ensures(_original): + """Companion to ``_patch_replace_view_copy_with_view_call``: the original ``ensures`` asserts + that no non-output ``view_copy`` node remains, but the patched ``call`` deliberately keeps + the ones whose shape dynamism differs from their base's.""" + + def patch(self, graph_module): + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + remaining = [node for node, _ in _view_replaceable_nodes(module)] + assert not remaining, f"view_copy nodes were not replaced with views: {remaining}" + + return patch + + +@register_patch("executorch", "torch.export.exported_program._convert_guards_to_code") +def _patch_convert_guards_to_code(_original): + """Skip stringifying ShapeEnv guards on every ``ExportedProgram`` construction. + + ``ExportedProgram.__init__`` unconditionally pretty-prints every ShapeEnv guard + into ``_guards_code``. ExecuTorch lowering constructs hundreds of intermediate + ``ExportedProgram``s (every ``_transform`` / decomposition / partition), each + re-printing the full guard set. For dynamic-shape guards with deeply nested + ``FloorDiv``/``Add`` expressions (Swin/Hiera window partitioning, BigBird + block-sparse indexing) the printer walks the *unshared* expression tree — + minutes of sympy printing per export, and on Mask2Former/Sam2 a recursion deep + enough to overflow the C stack (segfault). The strings are only consumed when + ``ExportedProgram.module()`` builds a guards fn, which torch itself force-disables + for ExecuTorch callers (``torch.export._unlift._ok_to_generate_guards_fn``), so + they are pure waste during lowering. + """ + + def patch(graph_module): + return [] + + return patch + + @register_patch( "executorch", "executorch.exir.passes.executorch_prim_ops_registry._EXECUTORCH_SYM_OPS", @@ -637,6 +769,312 @@ def patch(self, node, xnn_graph, vals_to_ids, debug_handle): return patch +@register_patch( + "executorch", "executorch.backends.transforms.remove_clone_ops.RemoveCloneOpsTransform._is_non_identity_clone" +) +def _patch_is_non_identity_clone(original): + """Keep identity clones that feed the graph output. + + XNNPACK's delegate preprocess runs ``RemoveCloneOpsTransform``, which folds identity + clones (same dim order — including ``_clone_dim_order`` of a ``permute_copy``, identity + only *after* the view-to-copy pass) onto their input. When both the clone and its input + are outputs of the delegated submodule (a value and its ``.contiguous()`` copy both + crossing the partition boundary — Clvp's mel-attention residual, PerceptionLM's + eval-mode dropout of a returned hidden state), the fold leaves + the same node twice in the output list and ``generate_node_to_external_map`` rejects the + submodule with ``Output node ... is already in the inputs``. Report output-feeding clones + as non-identity so they are kept — the partitioner only admits dim-order-preserving + clones, which XNNPACK serializes as ``XNNCopy``. + """ + + def patch(self, node): + if any(user.op == "output" for user in node.users): + return True + return original(self, node) + + return patch + + +@register_patch( + "executorch", "executorch.backends.xnnpack.partition.config.node_configs.PreluConfig.check_constraints" +) +def _patch_prelu_check_constraints(original): + """Only delegate ``prelu`` to XNNPACK when its input is 4-D. + + ``PreluConfig.check_constraints`` only verifies the weight is a parameter, but + XNNPACK's ``ChannelsLastTaggedReshapePass`` lists ``prelu`` among the ops that + require NHWC input and asserts the input can be converted (i.e. is 4-D) — + ``Attempting to convert non-NHWC compatible node to NHWC`` otherwise. Models + that apply ``nn.PReLU`` to 3-D transformer activations (dab_detr) crash there. + Rejecting the node keeps it on the portable CPU ops instead. + """ + + def patch(self, node, ep): + input_node = node.all_input_nodes[0] + val = input_node.meta.get("val") + if not (isinstance(val, torch.Tensor) and val.dim() == 4): + return False + return original(self, node, ep) + + return patch + + +@register_patch( + "executorch", + "executorch.exir.backend.canonical_partitioners.group_partitioner.GroupBasedPartitioner.propose_partitions", +) +def _patch_group_partitioner_break_quotient_cycles(original): + """Split delegated partitions that would form a dependency cycle once fused. + + ``to_backend`` fuses each partition into a single ``call_module`` node, one at a time. A + fused node conservatively depends on *all* of the partition's inputs and feeds *all* of its + outputs — even when the partition internally holds independent sub-computations. The XNNPACK + config partitioner does create such partitions: its disjoint-set grouping + (``get_matched_nodes_from_configs``) unions nodes that merely share a constant, so a single + tag can cover two independent chains. When one such partition ``A`` collapses to one node, it + introduces an edge from every ``A``-input to every ``A``-output, which can make a *previously + convex* partition ``B`` non-convex if ``B`` has nodes on both sides of ``A`` — i.e. the + quotient graph (each partition contracted to a node) has a cycle ``B -> A -> B``. + ``create_submodule_from_nodes`` then raises ``Invalid partition, found dependency cycles`` + (seen on FLAVA, where the text/image encoder streams and the multimodal encoder interleave). + + ``GroupBasedPartitioner`` only checks *pairwise merges* (``_can_merge_partitions``) against + the original graph, so it misses this fusion-induced cycle. Enforce the real fuseability + condition here: the quotient graph over the proposed partitions must be a DAG. While it isn't, + pick a multi-node partition on a detected cycle (only a multi-node partition can create the + false all-inputs-to-all-outputs edge) and split it around the other cycle members: nodes + upstream of that barrier form one partition, the rest form another, ordering the halves as + ``upstream -> barrier -> downstream``. Both halves stay delegated to XNNPACK, so nodes the + partitioner marked "do not decompose" (e.g. ``linear``) aren't left orphaned and un-lowered. + Prefer the straddling partition (whose split lands nodes on both sides); each split strictly + increases the partition count while shrinking the straddler, so this converges. Cycle-free + partitionings (every currently-passing model) are returned untouched. + """ + from torch.fx.passes.infra.partitioner import Partition + + def find_cycle(adjacency): + # Iterative DFS returning the node ids on one cycle, or None if the graph is a DAG. + color = dict.fromkeys(adjacency, 0) # 0 = unvisited, 1 = on stack, 2 = done + for root in adjacency: + if color[root] != 0: + continue + path = [root] + stack = [(root, iter(adjacency[root]))] + color[root] = 1 + while stack: + _, neighbors = stack[-1] + advanced = False + for nxt in neighbors: + if color[nxt] == 1: # back-edge into the current DFS path + return path[path.index(nxt) :] + if color[nxt] == 0: + color[nxt] = 1 + path.append(nxt) + stack.append((nxt, iter(adjacency[nxt]))) + advanced = True + break + if not advanced: + color[stack.pop()[0]] = 2 + path.pop() + return None + + def ancestors_of(barrier): + # All nodes with a path *into* the barrier set (barrier included), via reverse BFS. + seen, worklist = set(barrier), list(barrier) + while worklist: + for inp in worklist.pop().all_input_nodes: + if inp not in seen: + seen.add(inp) + worklist.append(inp) + return seen + + def split_around(cycle, victim_id, by_id): + # Split victim's nodes into (before, after) around the other cycle members: `before` = + # victim nodes upstream of the barrier, `after` = the rest. Ordered before -> barrier -> + # after, this breaks the victim's participation in the cycle. Returns None if one side + # is empty (this victim doesn't straddle the barrier). + barrier = set() + for member in cycle: + if member == victim_id: + continue + barrier.update(by_id[member].nodes if isinstance(member, int) else {member}) + ancestors = ancestors_of(barrier) + victim_nodes = list(by_id[victim_id].nodes) + before = [n for n in victim_nodes if n in ancestors] + after = [n for n in victim_nodes if n not in ancestors] + return (before, after) if before and after else None + + def patch(self): + partitions = original(self) + by_id = {p.id: p for p in partitions} + while by_id: + # Build the quotient graph: every partition contracts to its (integer) id, every + # un-delegated node stays as itself. A cycle here is exactly a fusion cycle. The + # `tag66 -> tag95` return edge can run through an un-delegated node (e.g. FLAVA's + # multimodal layer_norm), so those nodes must be vertices too — a partition-only + # quotient misses the cycle. + node_to_pid = {node: pid for pid, p in by_id.items() for node in p.nodes} + adjacency: dict = {} + for node in self.graph_module.graph.nodes: + src = node_to_pid.get(node, node) + adjacency.setdefault(src, set()) + for user in node.users: + dst = node_to_pid.get(user, user) + adjacency.setdefault(dst, set()) + if dst != src: + adjacency[src].add(dst) + cycle = find_cycle(adjacency) + if cycle is None: + break + + # Only a multi-node partition can introduce the false fusion edge (a one-node + # partition contracts to a single graph vertex, which the DAG can't put on a cycle), + # so the victim must be one of these. Prefer the straddler — the partition whose split + # lands nodes on both sides of the barrier — smallest first to minimise churn. + candidates = sorted( + (v for v in cycle if isinstance(v, int) and v in by_id and len(by_id[v].nodes) > 1), + key=lambda pid: len(by_id[pid].nodes), + ) + if not candidates: + break + chosen_id, halves = None, None + for victim_id in candidates: + halves = split_around(cycle, victim_id, by_id) + if halves is not None: + chosen_id = victim_id + break + + next_id = max(by_id) + 1 + if chosen_id is not None: + del by_id[chosen_id] + for nodes in halves: + by_id[next_id] = Partition(id=next_id, nodes=set(nodes)) + next_id += 1 + else: + # No straddler splits cleanly — dissolve the smallest candidate into single-node + # partitions, which cannot sit on a fusion cycle. Rare; keeps the loop converging. + victim = by_id.pop(candidates[0]) + for node in victim.nodes: + by_id[next_id] = Partition(id=next_id, nodes={node}) + next_id += 1 + return list(by_id.values()) + + return patch + + +@register_patch( + "executorch", + "executorch.exir.lowered_backend_module._unsafe_adjust_original_program", + "executorch.exir.backend.backend_api._unsafe_adjust_original_program", +) +def _patch_unsafe_adjust_original_program(original): + """Delete each consumed parameter/constant target at most once when adjusting the + original program after delegation. + + After a partition is lowered, ``_unsafe_adjust_original_program`` strips the params/buffers + the delegate absorbed from the top-level graph signature and state dict. Its dedup guard + (``currently_used_targets``) only skips targets still referenced by a *remaining* input spec — + it does not dedup *within* the batch being deleted. When one delegate consumes several + duplicated copies of the same shared parameter (the constant-dedup pass emits ``..._copy_1``, + ``..._copy_2`` … all keeping the original FQN as their ``target``), the loop deletes that + target on the first copy and then raises ``KeyError`` on the next. Transformers detection + models hit this because a single head is applied at every decoder layer and tied to the + encoder head — e.g. PPDocLayoutV3's ``model.decoder.class_embed.weight`` / ``bbox_embed``. + + Track the targets already removed and delete each only once (and tolerate an already-absent + key). Removing a target once is correct: its data is baked into the delegate blob, so the + repeated deletes are no-ops. The rest of the routine (graph-node erasure, output-spec and + getitem-index fixups) is unchanged, so all duplicate placeholders are still erased. + """ + from torch.export.graph_signature import InputKind + + def patch(original_program, call_delegate_node, input_specs_to_delete, output_specs_to_delete): + original_program._graph_signature.input_specs = [ + input_spec + for input_spec in original_program.graph_signature.input_specs + if input_spec.arg.name not in input_specs_to_delete + ] + + currently_used_targets = { + input_spec.target + for input_spec in original_program._graph_signature.input_specs + if input_spec.target is not None + } + + original_program._graph_signature.output_specs = [ + output_spec + for output_spec in original_program.graph_signature.output_specs + if output_spec.arg.name not in output_specs_to_delete + ] + + for node in original_program.graph.nodes: + if node.op == "placeholder": + if node.name in input_specs_to_delete: + assert len(node.users) == 0 + original_program.graph.erase_node(node) + else: + break + + deleted_targets: set = set() + for input_spec in input_specs_to_delete.values(): + input_target = input_spec.target + assert input_target is not None + # Skip targets still referenced elsewhere, and targets already removed by an earlier + # duplicate copy in this same batch (the fix: the stock routine omits this second case). + if input_target in currently_used_targets or input_target in deleted_targets: + continue + deleted_targets.add(input_target) + + if input_spec.kind == InputKind.PARAMETER: + original_program._state_dict.pop(input_target, None) + elif input_spec.kind == InputKind.BUFFER: + if input_spec.persistent: + original_program._state_dict.pop(input_target, None) + else: + original_program._constants.pop(input_target, None) + elif input_spec.kind == InputKind.CONSTANT_TENSOR: + original_program._constants.pop(input_target, None) + else: + raise RuntimeError(f"Invalid input spec {input_spec} received") + + toplevel_output_node = original_program.graph.output_node() + assert toplevel_output_node is not None + assert len(toplevel_output_node.args) == 1, ( + f"Invalid output node: {toplevel_output_node} with args {toplevel_output_node.args}" + ) + + new_output_args = [ + arg + for arg in toplevel_output_node.args[0] + if not isinstance(arg, torch.fx.Node) or arg.name not in output_specs_to_delete + ] + toplevel_output_node.args = (tuple(new_output_args),) + + getitem_idxs: list = [] + user_nodes = list(call_delegate_node.users.keys()) + for user in user_nodes: + if user.name in output_specs_to_delete: + assert user.op == "call_function" and user.target == operator.getitem + user_idx = user.args[1] + assert isinstance(user_idx, int), f"Invalid getitem type: {type(user_idx)}" + getitem_idxs.append(user_idx) + original_program.graph.erase_node(user) + + getitem_idxs.sort(reverse=True) + + user_nodes = list(call_delegate_node.users.keys()) + for user in user_nodes: + assert user.op == "call_function" and user.target == operator.getitem + user_idx = user.args[1] + assert isinstance(user_idx, int) + for i, idx in enumerate(getitem_idxs): + if user_idx > idx: + user.args = (user.args[0], user_idx - (len(getitem_idxs) - i)) + break + + return patch + + @register_patch("executorch", "executorch.backends.xnnpack.operators.node_visitor._node_visitor_dict") def _patch_squeeze_node_visitors(original): """Swap the squeeze/unsqueeze visitor entries in ``_node_visitor_dict`` with subclasses diff --git a/src/transformers/exporters/exporter_onnx.py b/src/transformers/exporters/exporter_onnx.py index ef65aaa27044..b31b46f282c3 100644 --- a/src/transformers/exporters/exporter_onnx.py +++ b/src/transformers/exporters/exporter_onnx.py @@ -245,6 +245,41 @@ def patch(self, x): return patch +@register_patch("onnx", "torch.split", "torch.Tensor.split") +def _patch_split(original): + """Expand a symbolic split size into statically-counted `narrow`s. A SymInt split size + otherwise lowers to `SplitToSequence` with a symbolic scalar `split` input, which + onnxscript's constant folder crashes on (`'NoneType' object has no attribute 'ndim'`). + """ + + def patch(input, split_size_or_sections, dim=0): + if not isinstance(split_size_or_sections, torch.SymInt): + return original(input, split_size_or_sections, dim) + split_size = split_size_or_sections + total = input.size(dim) + # `int()` specializes the chunk count at trace time, exactly like enumerating the + # list `aten.split.Tensor` returns (its meta guards on the same ceil division). + count = int((total + split_size - 1) // split_size) + return tuple( + input.narrow(dim, i * split_size, torch.sym_min(split_size, total - i * split_size)) for i in range(count) + ) + + return patch + + +@register_patch("onnx", "torch.chunk", "torch.Tensor.chunk") +def _patch_chunk(original): + """Route a symbolic chunk size through the `torch.split` patch (see `_patch_split`).""" + + def patch(input, chunks, dim=0): + chunk_size = (input.size(dim) + chunks - 1) // chunks + if not isinstance(chunk_size, torch.SymInt): + return original(input, chunks, dim) + return torch.split(input, chunk_size, dim) + + return patch + + @register_patch("onnx", "torch.randperm") def _patch_randperm(original): """Implement randperm via argsort(rand(n)) — no ONNX decomposition for aten.randperm.""" @@ -305,6 +340,26 @@ def patch(self, *args, **kwargs): return patch +@register_patch("onnx", "onnxscript.optimizer.optimize_ir") +def _patch_optimize_ir(original): + """Skip constant-folding `Resize` nodes during onnxscript optimization. + + The optimizer's constant folder evaluates foldable nodes with onnx's pure-Python + reference implementation. For `Resize` — e.g. the bicubic position-embedding + interpolation in YOLOS/SegGPT-style vision models, whose inputs are constant + initializers — that evaluation recurses per output element and takes minutes even + on tiny graphs (~4.5 min per Resize node on the YOLOS test model, vs <1 s for the + whole rest of the optimization). Keeping the Resize node in the graph costs one + native ORT kernel launch at inference instead. + """ + + def patch(model, *args, **kwargs): + kwargs.setdefault("should_fold", lambda node: False if node.op_type == "Resize" else None) + return original(model, *args, **kwargs) + + return patch + + def _patch_cummax_or_cummin(original, *, mode: str): """Decompose cummax/cummin via triangular-mask reduction (O(N^2) memory).""" diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py new file mode 100644 index 000000000000..42d499e0fb59 --- /dev/null +++ b/src/transformers/exporters/exporter_openvino.py @@ -0,0 +1,1923 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OpenVINO exporter. + +Extends [`DynamoExporter`] with the stages that turn an ``ExportedProgram`` into an +``openvino.Model``: + +1. **Torch patches** (``apply_patches("openvino")``): reversibly swap ``torch`` ops the OV + frontend can't lower (``torch.histc``, ``torch.searchsorted``, …) with decomposed + equivalents during the trace. The trace itself runs under ``torch.no_grad()`` so + modeling-internal grad regions don't become HigherOrderOp subgraphs. +2. **Dynamo trace** (inherited from [`DynamoExporter`]): signature patch, model patches, + pytree registration, dynamic shapes, state cleanup — same as for any other backend. +3. **Graph preparation**: run OV's own decomposition pass up front + (``_run_openvino_decompositions``), then repair the resulting graph in place — FX program + fixes, output-arg deduplication, per-node FX fixes, and bare-name renames. Preparing the + decomposed graph ourselves is what makes these fixes stick: handing the raw + ``ExportedProgram`` to ``convert_model`` would re-run decompositions internally and + regenerate the graph. +4. **Conversion**: the prepared module is decoded via ``TorchFXPythonDecoder`` and handed to + ``openvino.convert_model`` together with the custom ``ConversionExtension``\\ s for ops + without a built-in OV lowering. Ports are then renamed to their dotted leaf paths and + non-tensor inputs repaired. +5. **Stateful transformation** (``OpenVINOConfig.stateful``, on by default): fold + round-tripped state tensors (KV cache, SSM states, …) into internal OV variables with a + fused ``beam_idx`` reorder. Optionally written to disk via ``openvino.save_model`` when + ``OpenVINOConfig.output_path`` is set. +""" + +from __future__ import annotations + +import operator +import re +from collections.abc import MutableMapping +from typing import TYPE_CHECKING, Any + +from ..utils import logging +from ..utils.import_utils import is_openvino_available, is_torch_available +from .configs import OpenVINOConfig +from .exporter_dynamo import DynamoExporter, is_cache_object +from .exporter_onnx import disambiguate_io_names, patch_model_outputs +from .utils import ( + apply_fx_node_fixes, + apply_fx_program_fixes, + apply_patches, + get_leaf_tensors, + register_fx_node_fix, + register_fx_program_fix, + register_patch, +) + + +if is_torch_available(): + import torch + from torch.export import ExportedProgram + from torch.export.decomp_utils import CustomDecompTable + + from .. import masking_utils + + +if is_openvino_available(): + import numpy as np + import openvino + import openvino.opset14 as ov_ops + from openvino._offline_transformations import apply_make_stateful_transformation + from openvino.frontend.pytorch import ConversionExtension + from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder + from openvino.frontend.pytorch.torchdynamo.export_decompositions import ops_to_not_decompose + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + + +logger = logging.get_logger(__name__) + + +class OpenVINOExporter(DynamoExporter): + """Exporter that converts a [`PreTrainedModel`] to an OpenVINO ``openvino.Model``. + + Example: + + ```python + >>> from transformers.exporters.exporter_openvino import OpenVINOExporter, OpenVINOConfig + + >>> exporter = OpenVINOExporter() + >>> ov_model = exporter.export(model, inputs, config=OpenVINOConfig(dynamic=True)) + >>> exporter.export(model, inputs, config=OpenVINOConfig(output_path="model.xml")) + ``` + """ + + required_packages = ["torch", "openvino"] + tested_versions = {"torch": "2.12.0", "openvino": "2025.0.0"} + + def export( + self, + model: PreTrainedModel, + sample_inputs: MutableMapping[str, Any], + config: OpenVINOConfig | dict[str, Any], + ) -> openvino.Model: + if isinstance(config, dict): + config = OpenVINOConfig(**config) + elif type(config) is not OpenVINOConfig: + raise TypeError(f"Expected config to be an OpenVINOConfig or dict, got {type(config)}") + + # ``torch.no_grad()``: with grad enabled, every modeling-internal ``torch.no_grad()`` + # region (frozen towers, VQ-VAEs) traces as a ``wrap_with_set_grad_enabled`` + # HigherOrderOp subgraph, which OV's frontend can't lower. + with torch.no_grad(), patch_model_outputs(model) as (inputs_names, outputs_names), apply_patches("openvino"): + exported_program: ExportedProgram = super().export(model, sample_inputs, config=config) + + _drop_metadata_asserts(exported_program.graph_module) + # Run OV's own decomposition pass up front and decode the RESULT — handing the + # ``ExportedProgram`` to ``convert_model`` would re-run it internally, regenerating node + # names and discarding every fix applied below. + exported_program = _run_openvino_decompositions(exported_program) + apply_fx_program_fixes("openvino", exported_program) + graph_module = exported_program.module() + _deduplicate_output_args(graph_module) + apply_fx_node_fixes("openvino", graph_module) + _rename_bare_node_names(graph_module) + decoder = TorchFXPythonDecoder(graph_module, dynamic_shapes=True) + # Name every input port after its FX placeholder — OV may drop unused inputs, so all + # downstream port↔placeholder matching is done by name, never positionally. + decoder._input_signature = [n.name for n in graph_module.graph.nodes if n.op == "placeholder"] + ov_model = openvino.convert_model(decoder, extension=_OV_CONVERSION_EXTENSIONS) + _fix_non_tensor_inputs(ov_model, graph_module) + + inputs_names = [n for n in inputs_names if n in get_leaf_tensors(sample_inputs)] + inputs_names, outputs_names = disambiguate_io_names(inputs_names, outputs_names) + _rename_model_ports(ov_model, graph_module, inputs_names, outputs_names) + + if config.stateful: + _make_stateful(ov_model, exported_program, graph_module, sample_inputs, inputs_names, outputs_names) + + if config.output_path is not None: + openvino.save_model(ov_model, config.output_path, compress_to_fp16=config.compress_to_fp16) + + return ov_model + + +# ── Conversion helpers ────────────────────────────────────────────────────── +# Small helpers for ``OpenVINOExporter.export`` — extracted for readability and so each stage +# of the conversion has a single responsibility. + + +def _placeholder_for_port(port, placeholders: dict[str, Any]): + """Return the FX placeholder whose name is among ``port``'s tensor names, or ``None``. + + Every input port carries its placeholder's name (via ``decoder._input_signature``), so + port↔placeholder matching is by name — OV drops unused inputs, which would shift any + positional pairing. + """ + return next((placeholders[name] for name in port.get_names() if name in placeholders), None) + + +def _leaf_names_by_placeholder(graph_module, inputs_names: list[str]) -> dict[str, str]: + """Map each tensor placeholder's FX name to its dotted leaf-path name. + + Tensor placeholders appear in the graph in kwargs-leaf order — the same order + ``patch_model_outputs`` captured ``inputs_names`` in — so the two zip together. + """ + tensor_placeholders = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and isinstance(node.meta.get("val"), torch.Tensor) + ] + return dict(zip((node.name for node in tensor_placeholders), inputs_names)) + + +def _rename_model_ports( + ov_model: openvino.Model, + graph_module, + inputs_names: list[str], + outputs_names: list[str], +) -> None: + """Restore the dotted leaf-path names on the converted model's input/output ports. + + OV's PyTorch frontend doesn't support an ``output=`` argument, and ``input=`` only accepts + Python-identifier names (no dots) — so the dotted ``get_leaf_tensors`` form is restored + post-conversion. Input ports are matched to their leaf names through their FX placeholder; + scalar ports (no leaf) keep their placeholder name. + """ + leaf_names = _leaf_names_by_placeholder(graph_module, inputs_names) + for port in ov_model.inputs: + name = next((leaf_names[n] for n in port.get_names() if n in leaf_names), None) + if name is not None: + port.get_tensor().set_names({name}) + # A passthrough output (e.g. T5's ``encoder_last_hidden_state`` returning the + # ``encoder_outputs.last_hidden_state`` input untouched) shares its tensor with the input + # port — renaming it would clobber the input name. Give the Result its own tensor by + # routing it through a no-op ``convert_like`` first. + changed = False + for port, name in zip(ov_model.outputs, outputs_names): + tensor = port.get_tensor() + if tensor.get_names() & set(inputs_names): + result = port.get_node() + source = result.input_value(0) + copy = ov_ops.convert_like(source, source) + result.input(0).replace_source_output(copy.output(0)) + copy.output(0).get_tensor().set_names({name}) + changed = True + else: + tensor.set_names({name}) + if changed: + ov_model.validate_nodes_and_infer_types() + + +def _fix_non_tensor_inputs(ov_model: openvino.Model, graph_module) -> None: + """Repair Parameters converted from FX non-tensor placeholders. + + Non-tensor forward kwargs survive ``torch.export`` as placeholders that OV's frontend + converts to dynamic-rank Parameters — and the CPU plugin refuses to compile any Parameter + with dynamic rank. Scalars (``logits_to_keep: int``) are pinned to a static scalar shape; + ``None`` and string kwargs (e.g. ``attention_mask=None`` in SSM decode captures, Blip's + ``reduction="mean"``) produce a Parameter nothing translatable consumes, which is removed. + """ + scalar_types = {bool: openvino.Type.boolean, int: openvino.Type.i64, float: openvino.Type.f32} + placeholders = {node.name: node for node in graph_module.graph.nodes if node.op == "placeholder"} + to_remove, changed = [], False + for port in ov_model.inputs: + node = _placeholder_for_port(port, placeholders) + if node is None or not port.get_partial_shape().rank.is_dynamic: + continue + val = node.meta.get("val") + if val is None or isinstance(val, str): + to_remove.append(port.get_node()) + changed = True + elif type(val) in scalar_types: + parameter = port.get_node() + parameter.set_partial_shape(openvino.PartialShape([])) + parameter.set_element_type(scalar_types[type(val)]) + changed = True + for parameter in to_remove: + ov_model.remove_parameter(parameter) + if changed: + ov_model.validate_nodes_and_infer_types() + + +# ── Stateful transformation ───────────────────────────────────────────────── +# Folds round-tripped state tensors (KV cache, SSM conv/ssm states, …) into internal OV +# ``ReadValue``/``Assign`` variables so the runtime carries them across ``infer()`` calls +# instead of marshalling them through inputs/outputs on every step. State pairs are derived +# STRUCTURALLY: a leaf path that appears on both sides of the model was traced from the same +# cache leaf (``disambiguate_io_names`` marks exactly these collisions with ``input.`` / +# ``output.`` prefixes) — no name conventions, no per-model-type branching, and any cache +# layout (KV, SSM, sliding-window, hybrid) is covered by construction. + +_STATE_BATCH_DIM = 0 # transformers-native caches are batch-first + + +def _find_state_pairs(ov_model: openvino.Model, sample_inputs: MutableMapping[str, Any]) -> dict[str, str]: + """Return ``{input_port_name: output_port_name}`` for every round-tripped state tensor. + + A leaf path appearing on both sides is only state when it lives inside a cache object + (per [`~exporters.exporter_dynamo.is_cache_object`]) — a plain tensor kwarg the model + happens to return under the same name (e.g. Parakeet's downsampled ``attention_mask``) + is a regular output. + """ + state_roots = {key for key, value in sample_inputs.items() if is_cache_object(value)} + input_names = {name for port in ov_model.inputs for name in port.get_names()} + pairs = {} + for port in ov_model.outputs: + for name in port.get_names(): + prefix, _, path = name.partition(".") + if prefix == "output" and f"input.{path}" in input_names and path.partition(".")[0] in state_roots: + pairs[f"input.{path}"] = name + return pairs + + +def _fuse_state_reorder(ov_model: openvino.Model, state_input_names: list[str]) -> None: + """Insert a ``beam_idx`` parameter and a batch-dim ``Gather`` in front of every state input. + + Beam search reorders the cache between steps (`_reorder_cache`); once state lives inside the + model that reorder must happen inside too. The runtime passes the beam permutation as + ``beam_idx`` and the fused ``Gather`` applies it to each state variable — for greedy decoding + ``beam_idx = arange(batch)`` makes it the identity. + """ + main_input = next(port for port in ov_model.inputs if not port.get_names() & set(state_input_names)) + batch = main_input.get_partial_shape()[_STATE_BATCH_DIM] + beam_idx = ov_ops.parameter(name="beam_idx", dtype=np.int32, shape=openvino.PartialShape([batch])) + beam_idx.output(0).get_tensor().set_names({"beam_idx"}) + ov_model.add_parameters([beam_idx]) + for input_name in state_input_names: + state_port = ov_model.input(input_name) + consumers = state_port.get_target_inputs() + gather = ov_ops.gather(state_port, beam_idx, ov_ops.constant(np.int64(_STATE_BATCH_DIM))) + for consumer in consumers: + consumer.replace_source_output(gather.output(0)) + ov_model.validate_nodes_and_infer_types() + + +def _state_init_dims( + exported_program: ExportedProgram, + graph_module, + sample_inputs: MutableMapping[str, Any], + pairs: dict[str, str], + inputs_names: list[str], + outputs_names: list[str], +) -> dict[str, list]: + """Compute a per-dim init spec for every state input: ``"batch"`` (follow the main input's + batch at runtime), ``0`` (the growing dim — empty on first inference), or a concrete length. + + The growing dim is identified from the ``ExportedProgram``'s symbolic shapes: a state input + dim whose SymInt expression differs from the paired output's (e.g. ``s2`` vs ``s2 + s3``) + grows across steps; dims with identical expressions pass through unchanged and are pinned to + their length in the sample tensors (heads, head_dim, conv width, …) — deriving them from the + data rather than from ``model.config`` keeps this model-type-agnostic. + """ + leaf_names = _leaf_names_by_placeholder(graph_module, inputs_names) + input_vals = { + leaf_names[node.name]: node.meta.get("val") + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.name in leaf_names + } + # Output vals are keyed by the trace-ordered leaf names, NOT by zipping OV output ports — + # ``convert_model`` does not always preserve output order. + node_by_name = {node.name: node for node in exported_program.graph.nodes} + output_specs = [s for s in exported_program.graph_signature.output_specs if s.kind.name == "USER_OUTPUT"] + output_vals = {} + for spec, name in zip(output_specs, outputs_names): + node = node_by_name.get(getattr(spec.arg, "name", None)) + output_vals[name] = node.meta.get("val") if node is not None else None + + sample_leaves = get_leaf_tensors(sample_inputs) + init_dims: dict[str, list] = {} + for input_name, output_name in pairs.items(): + in_val, out_val = input_vals.get(input_name), output_vals.get(output_name) + sample = sample_leaves.get(input_name.partition(".")[2]) + if in_val is None or out_val is None or sample is None: + continue + dims = [] + for axis in range(in_val.ndim): + if axis == _STATE_BATCH_DIM: + dims.append("batch") + elif str(in_val.shape[axis]) == str(out_val.shape[axis]): + dims.append(int(sample.shape[axis])) + else: + dims.append(0) + # ``apply_make_stateful_transformation`` names each variable by concatenating the input + # and output tensor names — key the specs the same way for the ReadValue lookup. + init_dims[f"{input_name}{output_name}"] = dims + return init_dims + + +def _freeze_batchless_states( + ov_model: openvino.Model, + pairs: dict[str, str], + sample_inputs: MutableMapping[str, Any], +) -> None: + """Replace batch-less round-tripped tensors with baked constants (in place, updating ``pairs``). + + A round-tripped tensor without a batch dim (e.g. Cohere2's scalar ``_sliding_window_tensor``) + is config-derived, not per-sequence state — there is nothing to reorder between beams and + nothing to reset between prompts, so it becomes a graph constant instead of an OV variable. + """ + sample_leaves = get_leaf_tensors(sample_inputs) + changed = False + for input_name in list(pairs): + port = ov_model.input(input_name) + if port.get_partial_shape().rank.get_length() > _STATE_BATCH_DIM: + continue + sample = sample_leaves.get(input_name.partition(".")[2]) + if sample is None: + continue + parameter = port.get_node() + constant = ov_ops.constant(sample.cpu().numpy()) + for consumer in port.get_target_inputs(): + consumer.replace_source_output(constant.output(0)) + ov_model.remove_parameter(parameter) + del pairs[input_name] + changed = True + if changed: + ov_model.validate_nodes_and_infer_types() + + +def _build_state_initializers(ov_model: openvino.Model, init_dims: dict[str, list]) -> None: + """Give every state variable a zero-filled init expression so the runtime can materialise + empty state on the first ``infer()`` without the caller providing shapes. + + The variable's declared shape is relaxed/pinned from the same dim spec first: batch and the + growing dim go dynamic (the trace may have 0/1-specialized batch, and a static batch can't + cover the runtime-driven init), pass-through dims get their concrete sample length (a + ``[B,0,0,D]``-style degenerate init would break the state-update concat downstream). + + Exception: when a variable's update expression (its ``Assign``'s input) is fully static — + a static trace can bake the batch into a non-growing state's update while the decoder-level + shapes stay dynamic — the batch dim is pinned to the update's instead. Updating a dynamic + variable from a fully static expression makes the CPU plugin insert a Reorder between the + two descriptors, which it can't build against the variable's dynamic one. + """ + variables = {variable.get_info().variable_id: variable for variable in ov_model.get_variables()} + update_shapes = {sink.get_variable_id(): sink.input_value(0).get_partial_shape() for sink in ov_model.get_sinks()} + main_input = ov_model.inputs[0] + batch = ov_ops.gather( + ov_ops.shape_of(main_input, output_type="i64"), + ov_ops.constant([_STATE_BATCH_DIM]), + ov_ops.constant(0), + ) + for op in ov_model.get_ops(): + if op.get_type_name() != "ReadValue": + continue + update_shape = update_shapes.get(op.get_variable_id()) + update_is_static = update_shape is not None and update_shape.is_static + dims = init_dims.get(op.get_variable_id()) + if dims is None: + continue + if update_is_static: + dims = [update_shape[axis].get_length() if d == "batch" else d for axis, d in enumerate(dims)] + info = variables[op.get_variable_id()].get_info() + info.data_shape = openvino.PartialShape([-1 if d in ("batch", 0) else d for d in dims]) + variables[op.get_variable_id()].update(info) + # The growing dim's zero length is emitted as ``batch - batch`` rather than a literal + # ``[0]``: the CPU plugin fuses single-consumer init subgraphs into + # ``ReadValueWithSubgraph`` and re-infers the state descriptor from the init's static + # shape — a folded ``0`` gets baked in and seeding the state is then rejected. + zero = ov_ops.subtract(batch, batch) + shape = ov_ops.concat( + [ + batch if d == "batch" else zero if d == 0 else ov_ops.constant(np.array([d], dtype=np.int64)) + for d in dims + ], + axis=0, + ) + zero = ov_ops.constant(0.0, dtype=op.get_output_element_type(0)) + op.set_arguments([ov_ops.broadcast(zero, shape)]) + ov_model.validate_nodes_and_infer_types() + + +def _align_state_pair_types(ov_model: openvino.Model, pairs: dict[str, str]) -> None: + """Give each state pair a single CPU-friendly storage type, converting at the boundaries. + + ``Assign`` rejects an update whose type differs from the variable's (e.g. Parakeet's + streaming lengths enter as i64 but are recomputed as i32), and the CPU plugin's oneDNN + path rejects i64 state outright (xLSTM's ``seqlen_offset``). The variable stores the + output's compute type, demoted to i32 when it would be i64; the input Parameter is retyped + to match (with a ``Convert`` restoring the original type for its consumers) and the output + converted before its ``Result``. + """ + changed = False + for input_name, output_name in pairs.items(): + input_port = ov_model.input(input_name) + original_type = input_port.get_element_type() + output_port = ov_model.output(output_name) + output_type = output_port.get_element_type() + storage_type = openvino.Type.i32 if output_type == openvino.Type.i64 else output_type + if original_type != storage_type: + parameter = input_port.get_node() + consumers = input_port.get_target_inputs() + parameter.set_element_type(storage_type) + convert = ov_ops.convert(parameter, original_type) + for consumer in consumers: + consumer.replace_source_output(convert.output(0)) + changed = True + if output_type != storage_type: + result = output_port.get_node() + convert = ov_ops.convert(result.input_value(0), storage_type) + result.input(0).replace_source_output(convert.output(0)) + convert.output(0).get_tensor().set_names({output_name}) + changed = True + if changed: + ov_model.validate_nodes_and_infer_types() + + +def _make_stateful( + ov_model: openvino.Model, + exported_program: ExportedProgram, + graph_module, + sample_inputs: MutableMapping[str, Any], + inputs_names: list[str], + outputs_names: list[str], +) -> None: + """Convert round-tripped state ports into internal OV variables (in place).""" + pairs = _find_state_pairs(ov_model, sample_inputs) + if not pairs: + # Common benign case with stateful=True as the default: encoders and prefill-only + # exports have no round-tripped state. + logger.debug("No round-tripped state tensors found — leaving the model stateless.") + return + + # Init specs must be computed before freezing — removing a Parameter breaks the + # positional placeholder↔port alignment the spec derivation relies on. + init_dims = _state_init_dims(exported_program, graph_module, sample_inputs, pairs, inputs_names, outputs_names) + _freeze_batchless_states(ov_model, pairs, sample_inputs) + if not pairs: + return + + _align_state_pair_types(ov_model, pairs) + _fuse_state_reorder(ov_model, list(pairs)) + apply_make_stateful_transformation(ov_model, pairs) + _build_state_initializers(ov_model, init_dims) + _pin_state_update_shapes(ov_model) + + +def _pin_state_update_shapes(ov_model: openvino.Model) -> None: + """Reconcile each ``Assign``'s update shape with its variable's shape. + + The CPU plugin refuses to reorder dynamic descriptors into the state memory, so the two + sides must agree. When the update is fully static and the variable is not (a static trace + bakes the batch into the update while the variable was declared batch-dynamic), the variable + is pinned to the update's shape. Conversely, when shape inference leaves the update + under-specified against a static variable (olmo_hybrid's rolled conv state comes out + ``[?,32,2..]`` against a ``[?,32,3]`` variable), a ``special_zero`` Reshape pins the + statically-known dims and copies the dynamic ones from the input. Pinning a variable + refines shapes downstream (other updates read from it), so this iterates to a fixpoint. + """ + variables = {variable.get_info().variable_id: variable for variable in ov_model.get_variables()} + read_values = {op.get_variable_id(): op for op in ov_model.get_ordered_ops() if op.get_type_name() == "ReadValue"} + for _ in range(len(variables) + 1): + changed = False + for op in ov_model.get_ordered_ops(): + if op.get_type_name() != "Assign": + continue + variable = variables[op.get_variable_id()] + variable_shape = variable.get_info().data_shape + update = op.input_value(0) + update_shape = update.get_partial_shape() + if update_shape == variable_shape: + continue + if update_shape.is_static: + info = variable.get_info() + info.data_shape = update_shape + variable.update(info) + # The variable's shape must relax its init expression's, so the init gets the + # same static shape (its dims were runtime-derived but numerically identical). + read_value = read_values.get(op.get_variable_id()) + if read_value is not None and read_value.get_input_size() > 0: + target = np.array([dim.get_length() for dim in update_shape], dtype=np.int64) + pinned_init = ov_ops.reshape( + read_value.input_value(0), ov_ops.constant(target), special_zero=False + ) + read_value.set_arguments([pinned_init]) + changed = True + continue + if variable_shape.rank.is_dynamic: + continue + target = [dim.get_length() if dim.is_static else 0 for dim in variable_shape] + if all(t == 0 for t in target): + continue + pinned = ov_ops.reshape(update, ov_ops.constant(np.array(target, dtype=np.int64)), special_zero=True) + op.input(0).replace_source_output(pinned.output(0)) + changed = True + if not changed: + break + ov_model.validate_nodes_and_infer_types() + + +# ── Graph preparation ─────────────────────────────────────────────────────── +# Turns the traced `ExportedProgram` into the exact `GraphModule` handed to OV's decoder: +# decompose with OV's own table, then repair the result in place. Doing this ourselves (rather +# than letting `convert_model` decompose internally) is what makes the repairs stick. + + +_OV_NAME_OK = re.compile(r"_\d+$") + + +def _drop_metadata_asserts(graph_module) -> None: + """Drop ``_assert_tensor_metadata`` nodes before the decomposition replay. + + These asserts re-check trace-time dtypes/devices and fail the replay once other stages + have legitimately changed them. Scalar asserts are deliberately KEPT at this point — they + carry the ``torch._check`` range facts (``u >= 0``) that unbacked-symint guards inside the + replay rely on (splinter) — and are removed post-decomposition by ``_fix_drop_assert_ops``. + """ + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in list(module.graph.nodes): + if node.op == "call_function" and node.target is torch.ops.aten._assert_tensor_metadata.default: + module.graph.erase_node(node) + module.recompile() + + +def _run_openvino_decompositions(exported_program: ExportedProgram) -> ExportedProgram: + """Run the same decomposition pass ``TorchFXPythonDecoder.from_exported_program`` would. + + Decomposing up front lets the FX node fixes and the bare-name renames below operate on the + graph OV actually decodes — ``convert_model(exported_program)`` would re-run decompositions + internally, regenerating node names and silently discarding those fixes. + """ + decomp_table = CustomDecompTable() + for op in ops_to_not_decompose(): + try: + decomp_table.pop(op) + except KeyError: + pass + return exported_program.run_decompositions(decomp_table) + + +def _deduplicate_output_args(graph_module) -> None: + """Give repeated graph outputs their own node via a real ``+ 0`` op. + + Two Results sharing one OV tensor crash the translate session's ``is_number`` check: the + results-cleanup pass erases the shared tensor's numeric id on the first visit and fails + decoding the debug alias on the second. Repeats arise when decomposition collapses the + distinction between two output nodes (and ``aten.clone`` is no protection — OV folds it to + identity); ``+ 0`` is the cheapest op that survives translation as a distinct tensor + (constant-folded away afterwards by OV's optimization passes). + """ + # Ops OV translates as pass-through — their output IS their input's tensor, so an output + # arg behind one of these still aliases the underlying node. + passthrough = (torch.ops.aten.clone.default, torch.ops.aten.alias.default, torch.ops.aten.detach.default) + output_node = next(node for node in graph_module.graph.nodes if node.op == "output") + seen = set() + + def dedup(arg): + source = arg + while source.op == "call_function" and source.target in passthrough: + source = source.args[0] + if source is arg and source not in seen: + seen.add(source) + return arg + with graph_module.graph.inserting_before(output_node): + val = source.meta.get("val") + if val is not None and val.dtype == torch.bool: + copy = graph_module.graph.call_function(torch.ops.aten.logical_and.default, args=(source, source)) + else: + copy = graph_module.graph.call_function(torch.ops.aten.add.Tensor, args=(source, 0)) + copy.meta.update(source.meta) + seen.add(source) + return copy + + output_node.args = (torch.fx.node.map_arg(output_node.args[0], dedup),) + graph_module.recompile() + + +def _rename_bare_node_names(graph_module) -> None: + """Append a numeric suffix to FX node names that lack one (in every nested graph). + + OV's PyTorch frontend strips a trailing ``_`` from each tensor name to recover the op + kind, then validates the remainder — aborting with ``GeneralFailure: is_number(name)`` for + bare names (``mul``, ``clone``, ``linear``). The first node of any kind in an FX graph has + no ``_`` suffix, so the strip is a no-op and OV rejects it. HigherOrderOp bodies + (e.g. ``wrap_with_set_grad_enabled`` around frozen vision towers) are separate + ``GraphModule``\\ s with their own name counters, and their placeholders are internal closure + args (not user inputs) — everything except the top-level placeholders gets the suffix. + """ + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + is_top_level = module is graph_module + used = {n.name for n in module.graph.nodes} + for n in module.graph.nodes: + if n.op == "output" or (n.op == "placeholder" and is_top_level): + continue + if _OV_NAME_OK.search(n.name): + continue + candidate = f"{n.name}_0" + i = 0 + while candidate in used: + i += 1 + candidate = f"{n.name}_{i}" + used.discard(n.name) + used.add(candidate) + n._rename(candidate) + + +# ── FX node fixes ─────────────────────────────────────────────────────────── +# Per-node in-place rewrites applied to the `ExportedProgram` graph after the Dynamo trace +# but before `openvino.convert_model`. Each `_fix_*(gm, node) -> bool` factory is registered +# via `@register_fx_node_fix("openvino")` and returns `True` when it consumed the node +# (no further fixes run against it). Use this for OV-frontend quirks that are easier to +# repair at the FX level than to patch around at the torch op level. +# +# To add a new fix: define a `_fix_*` callable and decorate it. + + +@register_fx_node_fix("openvino") +def _fix_sym_float(gm, node): + """``torch.sym_float`` is a no-op at the OV layer (it's a Python-level SymInt→SymFloat cast). + Replace it with its input — affects deformable_detr, focalnet, mask2former, deepseek_ocr2. + """ + if node.target is not torch.sym_float: + return False + node.replace_all_uses_with(node.args[0]) + gm.graph.erase_node(node) + return True + + +@register_fx_node_fix("openvino") +def _fix_sym_min_max(gm, node): + """Rewrite ``torch.sym_min``/``torch.sym_max`` to the built-in ``min``/``max``. + + OV's FX decoder keys translations on ``str(target)``. ``torch.sym_min`` reprs to + ```` — the address varies per process so no + ``ConversionExtension`` string can match. ``min``/``max`` repr to stable + ````/````, which we register translators + for. The numeric behaviour is identical for SymInts. + """ + if node.target is torch.sym_min: + node.target = min + return True + if node.target is torch.sym_max: + node.target = max + return True + return False + + +@register_fx_program_fix("openvino") +def _fix_to_dtype_layout_in_subgraphs(exported_program): + """Rewrite ``aten.to.dtype_layout`` and ``aten.to.device`` in *submodule* FX graphs to + ``aten._to_copy`` (with just the dtype kwarg). + + OV's PyTorch frontend registers ``ConversionExtension`` handlers only for the top-level + graph; nested ``HigherOrderOp`` subgraphs (like the one wrapped by + ``wrap_with_set_grad_enabled`` in Chameleon's rotary path) still see the raw + ``aten.to.dtype_layout`` node, for which OV has no translator — resulting in a dangling + ``torch::None`` constant. Rewriting the FX target here lets OV's built-in + ``aten._to_copy.default`` handler take over (which we also override at the top level to + swallow complex-dtype casts).""" + # Walk the top-level graph AND every submodule's graph (higher-order-op subgraphs). + graphs = [exported_program.graph_module] + graphs.extend(m for _, m in exported_program.graph_module.named_children() if hasattr(m, "graph")) + for gm_or_submod in graphs: + for node in list(gm_or_submod.graph.nodes): + if node.op != "call_function": + continue + target = node.target + if target is torch.ops.aten.to.dtype: + # ``aten.to.dtype(tensor, dtype)`` — dtype is positional arg[1]. + dtype = node.args[1] if len(node.args) > 1 else node.kwargs.get("dtype") + node.target = torch.ops.aten._to_copy.default + node.args = (node.args[0],) + node.kwargs = {"dtype": dtype} if dtype is not None else {} + elif target in (torch.ops.aten.to.dtype_layout, torch.ops.aten.to.device, torch.ops.aten.to.other): + dtype = node.kwargs.get("dtype") + node.target = torch.ops.aten._to_copy.default + node.args = (node.args[0],) + node.kwargs = {"dtype": dtype} if dtype is not None else {} + gm_or_submod.recompile() + + +@register_fx_node_fix("openvino") +def _fix_drop_assert_ops(gm, node): + """Erase ``aten._assert_tensor_metadata`` / ``aten._assert_scalar`` nodes. + + ``torch.export`` inserts these as dead-code (num_users=0) runtime assertions, but OV's + frontend translates them into ``torch::None`` constants whose downstream consumers can't + drop them — causing ``OpConversionFailure``. They have no semantic effect on the model. + """ + if node.target not in (torch.ops.aten._assert_tensor_metadata.default, torch.ops.aten._assert_scalar.default): + return False + gm.graph.erase_node(node) + return True + + +@register_fx_node_fix("openvino") +def _fix_symbolic_pad(gm, node): + """Decompose ``constant_pad_nd`` with symbolic pad amounts into ``full`` + ``cat``. + + OV's translation of the inlined pad list places a symbolic amount on the wrong axis in + some graphs (mamba2's chunked scan pads seq by ``(chunk - seq % chunk) % chunk`` and the + pad lands on the state dim at runtime). Building the filler explicitly sidesteps the + list-decoding entirely; constant pads keep OV's native translation. + """ + if node.target is not torch.ops.aten.constant_pad_nd.default: + return False + x, pads = node.args[0], node.args[1] + if all(isinstance(p, int) for p in pads): + return False + value = node.args[2] if len(node.args) > 2 else 0 + val = x.meta.get("val") + if val is None: + return False + users = list(node.users) + current = x + with gm.graph.inserting_before(node): + for pair_index in range(len(pads) // 2): + dim = val.ndim - 1 - pair_index # torch pad pairs run last-dim-first + for amount, at_front in ((pads[2 * pair_index], True), (pads[2 * pair_index + 1], False)): + if isinstance(amount, int) and amount == 0: + continue + # A negative amount crops instead of padding (mamba2's conv warmup pads by + # ``kernel - seq``); symbolic amounts can be either at runtime, so build both a + # ``max(amount, 0)``-sized filler and a ``min(amount, 0)``-deep crop. + filler_size = amount if isinstance(amount, int) else gm.graph.call_function(max, args=(amount, 0)) + crop = 0 if isinstance(amount, int) else gm.graph.call_function(min, args=(amount, 0)) + if isinstance(amount, int) and amount < 0: + filler_size, crop = 0, amount + dim_size = gm.graph.call_function(torch.ops.aten.sym_size.int, args=(current, dim)) + if crop != 0: + if at_front: + start = gm.graph.call_function(operator.sub, args=(0, crop)) + current = gm.graph.call_function( + torch.ops.aten.slice.Tensor, args=(current, dim, start, dim_size) + ) + else: + end = gm.graph.call_function(operator.add, args=(dim_size, crop)) + current = gm.graph.call_function(torch.ops.aten.slice.Tensor, args=(current, dim, 0, end)) + current.meta.update(node.meta) + if not isinstance(filler_size, int) or filler_size > 0: + sizes = [ + filler_size + if i == dim + else gm.graph.call_function(torch.ops.aten.sym_size.int, args=(current, i)) + for i in range(val.ndim) + ] + filler = gm.graph.call_function( + torch.ops.aten.full.default, + args=(sizes, value), + kwargs={"dtype": val.dtype, "device": val.device}, + ) + filler.meta.update(node.meta) + operands = [filler, current] if at_front else [current, filler] + current = gm.graph.call_function(torch.ops.aten.cat.default, args=(operands, dim)) + current.meta.update(node.meta) + for user in users: + user.replace_input_with(node, current) + gm.graph.erase_node(node) + return True + + +@register_fx_node_fix("openvino") +def _fix_gather_index_extent(gm, node): + """Align ``aten.gather``'s index with the data's extent on non-axis dims. + + torch allows the index to be SMALLER than the data on non-axis dims (positions beyond the + index's extent are simply never read); OV's ``GatherElements`` requires equal shapes except + at the axis. The index is expanded to the data's extent (gathering redundantly) and the + OUTPUT narrowed back — rewriting the data input instead is fragile: decomposition re-fuses + the slice into upstream expands and the mismatch reappears (efficientloftr's fine-matching + grid gather hits this). + """ + if node.target is not torch.ops.aten.gather.default: + return False + data, dim, index = node.args[:3] + data_val, index_val = data.meta.get("val"), index.meta.get("val") + if data_val is None or index_val is None: + return False + axis = dim if dim >= 0 else dim + data_val.ndim + mismatched = [i for i in range(data_val.ndim) if i != axis and str(data_val.shape[i]) != str(index_val.shape[i])] + if not mismatched: + return False + users = list(node.users) + with gm.graph.inserting_before(node): + sizes = [ + gm.graph.call_function(torch.ops.aten.sym_size.int, args=(data, i)) if i in mismatched else -1 + for i in range(data_val.ndim) + ] + expanded = gm.graph.call_function(torch.ops.aten.expand.default, args=(index, sizes)) + expanded.meta.update(index.meta) + node.args = (data, dim, expanded) + tuple(node.args[3:]) + narrowed = node + for i in mismatched: + with gm.graph.inserting_after(narrowed): + size = gm.graph.call_function(torch.ops.aten.sym_size.int, args=(index, i)) + with gm.graph.inserting_after(size): + narrowed = gm.graph.call_function(torch.ops.aten.slice.Tensor, args=(narrowed, i, 0, size)) + narrowed.meta.update(node.meta) + for user in users: + user.replace_input_with(node, narrowed) + return True + + +@register_fx_node_fix("openvino") +def _fix_scatter_reduce(gm, node): + """Lower ``aten.scatter_reduce.two`` at the FX level — OV's frontend has no translation, + and its ``ScatterElementsUpdate`` op can't accept the ``reduce`` string as a constant input. + + Handles two patterns the MoE/SSM models use: + * ``reduce="sum", include_self=True`` → ``aten.scatter_add`` (BLT/JetMoe/NemotronH router). + * ``reduce="amax", include_self=False`` → masked-max over a one-hot expansion of ``index`` + (BLT byte-pooling). Other combinations fall through to the generic OpConversionFailure. + """ + if node.target is not torch.ops.aten.scatter_reduce.two: + return False + if len(node.args) < 5: + return False + reduce = node.args[4] + include_self = node.kwargs.get("include_self", True) + self_arg, dim, index, src = node.args[0:4] + + if reduce == "sum" and include_self is True: + with gm.graph.inserting_before(node): + new = gm.graph.call_function(torch.ops.aten.scatter_add.default, args=(self_arg, dim, index, src)) + new.meta.update(node.meta) + node.replace_all_uses_with(new) + gm.graph.erase_node(node) + return True + + if reduce == "amax" and include_self is False: + # ``self`` is initialised to zeros; each source element competes for the max at + # ``index[j]``; empty positions stay zero. Decompose to a broadcast comparison + amax: + # build a one-hot mask ``(index.unsqueeze(dim) == arange(K))``, then take the elementwise + # max of ``src`` where the mask is set (``-inf`` elsewhere), then zero any positions with + # no scatter. + self_val = self_arg.meta.get("val") + src_val = src.meta.get("val") + if self_val is None or src_val is None or not src_val.dtype.is_floating_point: + return False + ndim = self_val.ndim + d = dim if dim >= 0 else dim + ndim + k_size = self_val.shape[d] + min_value = torch.finfo(src_val.dtype).min + k_shape = [1] * (ndim + 1) + k_shape[d] = -1 + with gm.graph.inserting_before(node): + arange = gm.graph.call_function( + torch.ops.aten.arange.default, args=(k_size,), kwargs={"device": self_val.device} + ) + k_range = gm.graph.call_function(torch.ops.aten.view.default, args=(arange, k_shape)) + index_unsq = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(index, d)) + mask = gm.graph.call_function(torch.ops.aten.eq.Tensor, args=(index_unsq, k_range)) + src_unsq = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(src, d)) + masked = gm.graph.call_function(torch.ops.aten.where.ScalarOther, args=(mask, src_unsq, min_value)) + maxes = gm.graph.call_function(torch.ops.aten.amax.default, args=(masked, [d + 1])) + any_match = gm.graph.call_function(torch.ops.aten.any.dim, args=(mask, d + 1)) + result = gm.graph.call_function(torch.ops.aten.where.ScalarOther, args=(any_match, maxes, 0.0)) + result.meta.update(node.meta) + node.replace_all_uses_with(result) + gm.graph.erase_node(node) + return True + + return False + + +@register_fx_node_fix("openvino") +def _fix_empty_cat(gm, node): + """Drop ``aten.cat([empty, x], dim)`` constructed by ``DynamicLayer`` for prefill — the empty + operand is a rank-1 ``f32[0]`` from ``aten.detach_(lift_fresh_copy(...))``, which OV's torch + frontend can't broadcast against the non-empty 4D operand for a ``dim=-2`` cat (it rejects + with ``Axis -2 out of the tensor rank range [-1, 0]``). Mathematically the cat is identity + when one operand is 0-element, so replace its uses with the non-empty operand. + """ + if node.target is not torch.ops.aten.cat.default: + return False + + operands = node.args[0] + if not isinstance(operands, (list, tuple)) or len(operands) != 2: + return False + + from torch.fx.experimental.symbolic_shapes import guard_or_false + + def _is_empty(n): + val = n.meta.get("val") if hasattr(n, "meta") else None + if val is None: + return False + # ``numel() == 0`` on a compound SymInt expression trips ``GuardOnDataDependentSymNode`` + # (MinimaxM3VL — the concat operand has ``3*u0*u1*u2 + ...`` numel). Default to + # ``False`` when we can't tell — treating the cat as non-empty keeps it in the graph, + # which is always correct (the empty-cat optimisation just doesn't fire). + return guard_or_false(val.numel() == 0) + + if _is_empty(operands[0]): + keep = operands[1] + elif _is_empty(operands[1]): + keep = operands[0] + else: + return False + + node.replace_all_uses_with(keep) + gm.graph.erase_node(node) + return True + + +@register_fx_node_fix("openvino") +def _fix_empty_expand(gm, node): + """Replace ``aten.expand`` of a statically-empty tensor with an explicitly-shaped ``full``. + + OV constant-folds the ``Tile`` an expand of a lifted constant lowers to by computing + per-axis repeats ``output_dim / input_dim`` — a zero-sized dim makes that an integer + ``0 / 0``, which SIGFPEs the whole process (chmv2's dinov3 backbone expands its + ``[1, 0, C]`` ``register_tokens`` when ``num_register_tokens=0``). The expanded tensor + has no elements, so a zero-filled ``full`` is equivalent — and it translates to a + Broadcast from a scalar, which has no zero input dims and folds safely. + """ + if node.target is not torch.ops.aten.expand.default: + return False + tensor, sizes = node.args[0], node.args[1] + val = tensor.meta.get("val") if hasattr(tensor, "meta") else None + out_val = node.meta.get("val") + if val is None or out_val is None: + return False + + from torch.fx.experimental.symbolic_shapes import guard_or_false + + if not guard_or_false(val.numel() == 0): + return False + users = list(node.users) + offset = len(sizes) - val.ndim # expand aligns sizes to the input's trailing dims + with gm.graph.inserting_before(node): + full_sizes = [] + for i, size in enumerate(sizes): + if isinstance(size, int) and size == -1: # -1 keeps the input's dim + dim = val.shape[i - offset] + size = ( + int(dim) + if isinstance(dim, int) + else gm.graph.call_function(torch.ops.aten.sym_size.int, args=(tensor, i - offset)) + ) + full_sizes.append(size) + full = gm.graph.call_function( + torch.ops.aten.full.default, + args=(full_sizes, 0), + kwargs={"dtype": out_val.dtype, "device": out_val.device}, + ) + full.meta.update(node.meta) + for user in users: + user.replace_input_with(node, full) + gm.graph.erase_node(node) + return True + + +# ── Torch patches ─────────────────────────────────────────────────────────── +# Each `_patch_*(original)` factory is registered via `@register_patch("openvino", path)` +# and reversibly swaps a `torch` op the OV frontend can't lower with a decomposed +# equivalent. Reverted on exit by `apply_patches("openvino")`. +# +# To add a new patch: define a `_patch_*` factory and decorate it. + + +@register_patch("openvino", "torch.nn.functional.layer_norm") +def _patch_layer_norm(original): + """Substitute identity ``weight=ones``/``bias=zeros`` when either is ``None``. + + OV's frontend records a ``torch::None`` constant for any unwired optional, then refuses to + convert it (``None constant cannot be converted to OpenVINO opset``). LayerNorm without + affine still computes ``(x - mean) / sqrt(var + eps)``; passing identity tensors keeps the + math unchanged and gives OV concrete operands. Affects Chameleon (no-affine RMSNorm path) + and any model that calls ``F.layer_norm(..., weight=None, bias=None)``. + """ + + def patch(input, normalized_shape, weight=None, bias=None, eps=1e-5): + if weight is None: + weight = torch.ones(normalized_shape, dtype=input.dtype, device=input.device) + if bias is None: + bias = torch.zeros(normalized_shape, dtype=input.dtype, device=input.device) + return original(input, normalized_shape, weight, bias, eps) + + return patch + + +@register_patch("openvino", "torch.nn.functional.scaled_dot_product_attention") +def _patch_sdpa(original): + """Pre-expand K/V to Q's head count before calling SDPA. + + OV's ``opset13::ScaledDotProductAttention`` op rejects GQA shapes (e.g. Q=[B,4,T,D], + K/V=[B,2,T,D]) with ``Key input shape not compatible with other inputs``. Repeating K/V via + ``repeat_interleave`` on the head axis keeps the math identical and gives OV matching shapes. + """ + + def patch(query, key, value, *args, **kwargs): + q_heads, k_heads = query.shape[-3], key.shape[-3] + if q_heads != k_heads and q_heads % k_heads == 0: + reps = q_heads // k_heads + key = key.repeat_interleave(reps, dim=-3) + value = value.repeat_interleave(reps, dim=-3) + return original(query, key, value, *args, **kwargs) + + return patch + + +@register_patch("openvino", "transformers.masking_utils._vmap_expansion_sdpa") +def _patch_broadcast_mask_expansion(_original): + """Replace vmap-based mask expansion with broadcast expansion. + + OV's PyTorch frontend can't trace through ``torch.vmap`` — the input tensors look like + they "escaped" the vmap context. Same shape of fix as the ONNX exporter's. + """ + + def patch(mask_function): + def _expanded(batch_arange, head_arange, q_arange, kv_arange): + broadcasted = masking_utils._non_vmap_expansion_sdpa(batch_arange, head_arange, q_arange, kv_arange) + return mask_function(*broadcasted).expand( + batch_arange.shape[0], head_arange.shape[0], q_arange.shape[0], kv_arange.shape[0] + ) + + return _expanded + + return patch + + +@register_patch("openvino", "torch.histc") +def _patch_histc(original): + """Replace ``torch.histc`` with a deterministic ``zeros + scatter_add_`` equivalent. + + OV's PyTorch frontend has no lowering for ``aten.histc``. The MoE token-counting path uses + integer inputs (expert ids), which ``torch.histc`` doesn't support natively anyway. The + decomposition pre-allocates a ``zeros(bins)`` (static shape) and accumulates via + ``scatter_add_``, both OV-friendly primitives. + """ + + def patch(input, bins=100, min=0, max=0, *, out=None): + flat = input.reshape(-1) + if max == min == 0: + min_val = flat.min().float() + max_val = flat.max().float() + else: + min_val = torch.tensor(float(min), device=flat.device) + max_val = torch.tensor(float(max), device=flat.device) + bin_width = (max_val - min_val) / bins + idx = ((flat.float() - min_val) / bin_width).long().clamp_(0, bins - 1) + out_dtype = input.dtype if input.is_floating_point() else torch.float + counts = torch.zeros(bins, dtype=out_dtype, device=input.device) + return counts.scatter_add_(0, idx, torch.ones_like(idx, dtype=out_dtype)) + + return patch + + +@register_patch("openvino", "torch.empty_permuted") +def _patch_empty_permuted(original): + """Replace ``torch.empty_permuted(size, physical_layout, ...)`` with plain ``torch.empty(size, ...)``. + + OV's frontend has no ``aten.empty_permuted`` lowering. The op exists only to hint a memory + layout (stride) — the values are uninitialised either way, and downstream reads see the same + logical content. ``torch.empty`` is enough. + """ + + def patch(size, physical_layout, **kwargs): + return torch.empty(size, **kwargs) + + return patch + + +@register_patch("openvino", "torch.polar") +def _patch_polar(original): + """Build ``polar(abs, angle)`` as ``complex(abs*cos(angle), abs*sin(angle))``. + + OV has no ``aten.polar`` lowering. Euler's formula gives the same result through ops the + frontend already supports. + """ + + def patch(abs, angle): + return torch.complex(abs * angle.cos(), abs * angle.sin()) + + return patch + + +def _rotate_pairs(x: torch.Tensor, freqs_pairs: torch.Tensor) -> torch.Tensor: + """Complex-multiply ``x`` (viewed as ``[..., d/2, 2]`` re/im pairs) by broadcastable + ``freqs_pairs``: ``(a+bi)(c+di) = (ac-bd) + (ad+bc)i``.""" + pairs = x.float().reshape(*x.shape[:-1], -1, 2) + real = pairs[..., 0] * freqs_pairs[..., 0] - pairs[..., 1] * freqs_pairs[..., 1] + imag = pairs[..., 0] * freqs_pairs[..., 1] + pairs[..., 1] * freqs_pairs[..., 0] + return torch.stack((real, imag), dim=-1).flatten(3).type_as(x) + + +@register_patch("openvino", "transformers.models.deepseek_v2.modeling_deepseek_v2.apply_rotary_emb") +def _patch_deepseek_rotary_emb(original): + """Rewrite complex-arithmetic RoPE with the equivalent real re/im-pair math. + + The traced ``view_as_complex(x) * freqs_cis`` mixes OV's native ``ComplexTypeMark`` + representation with the ``[..., 2]`` real-pair one our ``aten.complex`` extension emits + (via the ``torch.polar`` patch) — the mul can't reconcile the two. Keeping the whole + rotation in real arithmetic confines traced complex ops to ``complex``/``view_as_real``, + which the extensions lower consistently. + """ + + def patch(xq, xk, freqs_cis): + freqs_pairs = torch.view_as_real(freqs_cis).unsqueeze(1).to(xq.device) + return _rotate_pairs(xq, freqs_pairs), _rotate_pairs(xk, freqs_pairs) + + return patch + + +@register_patch("openvino", "transformers.models.llama4.modeling_llama4.apply_rotary_emb") +def _patch_llama4_rotary_emb(original): + """Same real-pair rewrite as ``_patch_deepseek_rotary_emb`` for llama4's text RoPE.""" + + def patch(xq, xk, freqs_cis): + freqs_pairs = torch.view_as_real(freqs_cis)[:, :, None, :, :] + return _rotate_pairs(xq, freqs_pairs), _rotate_pairs(xk, freqs_pairs) + + return patch + + +@register_patch("openvino", "transformers.models.llama4.modeling_llama4.vision_apply_rotary_emb") +def _patch_llama4_vision_rotary_emb(original): + """Same real-pair rewrite as ``_patch_deepseek_rotary_emb`` for llama4's vision RoPE.""" + + def patch(query, key, freqs_ci): + freqs_pairs = torch.view_as_real(freqs_ci) + # Mirror ``reshape_for_broadcast``: keep dims 1 (seq) and -1 (d/2), plus the re/im pair. + shape = [d if i == 1 else 1 for i, d in enumerate(query.shape[:-1])] + [freqs_pairs.shape[-2], 2] + freqs_pairs = freqs_pairs.view(*shape).to(query.device) + return _rotate_pairs(query, freqs_pairs), _rotate_pairs(key, freqs_pairs) + + return patch + + +@register_patch("openvino", "torch.nn.functional.avg_pool2d") +def _patch_avg_pool2d(original): + """Clamp oversize pooling kernels to the input's spatial size. + + torch's ``ceil_mode`` pooling permits a kernel larger than the input, producing a 1×1 + output — EfficientNet's pooler (``AvgPool2d(config.hidden_dim)``) relies on it. OV's + ``AvgPool`` rejects kernels larger than the padded input. + """ + + def patch( + input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None + ): + kh, kw = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size + ph, pw = (padding, padding) if isinstance(padding, int) else padding + kh = torch.sym_min(kh, input.shape[-2] + 2 * ph) + kw = torch.sym_min(kw, input.shape[-1] + 2 * pw) + if stride is None: + stride = (kh, kw) + return original(input, (kh, kw), stride, padding, ceil_mode, count_include_pad, divisor_override) + + return patch + + +@register_patch("openvino", "torch.Tensor.unfold") +def _patch_unfold(original): + """Decompose ``Tensor.unfold(dim, size, step)`` into ``index_select`` + reshape. + + OV's ``aten.unfold`` translator builds a permutation one rank too long for 3D inputs + (PatchTST's patchification), producing an invalid Transpose. + """ + + def patch(self, dimension, size, step): + dim = dimension if dimension >= 0 else dimension + self.dim() + starts = torch.arange(0, self.shape[dim] - size + 1, step, device=self.device) + indices = (starts.unsqueeze(1) + torch.arange(size, device=self.device)).flatten() + windows = self.index_select(dim, indices).unflatten(dim, (starts.shape[0], size)) + return windows.movedim(dim + 1, -1) + + return patch + + +@register_patch("openvino", "torch.bernoulli") +def _patch_bernoulli(original): + """Strip randomness from ``torch.bernoulli`` — return ``zeros_like(p)`` during export. + + Stochastic ops have no place in an exported graph; the training-time sampling is + deterministic-zero at inference (eval mode), so the export-time substitution is correct + for the only modes that actually export. + """ + + def patch(input, *args, **kwargs): + return torch.zeros_like(input) + + return patch + + +@register_patch("openvino", "torch.randn", "torch.randn_like") +def _patch_randn(original): + """Strip randomness from ``torch.randn`` / ``torch.randn_like`` — return zeros. + + Same rationale as ``torch.bernoulli``: stochastic noise has no place in an exported graph; + the inference-time path doesn't sample, so zero is what the model would see. + """ + + def patch(*args, **kwargs): + if args and isinstance(args[0], torch.Tensor): + return torch.zeros_like(args[0]) + return torch.zeros(*args, **kwargs) + + return patch + + +@register_patch("openvino", "torch.randperm") +def _patch_randperm(original): + """Strip randomness from ``torch.randperm`` — return the identity permutation. + + Same rationale as ``torch.bernoulli`` / ``torch.randn``. + """ + + def patch(n, *, dtype=None, device=None, **kwargs): + return torch.arange(n, dtype=dtype if dtype is not None else torch.int64, device=device) + + return patch + + +@register_patch("openvino", "torch.randint") +def _patch_randint(original): + """Strip randomness from ``torch.randint`` — return zeros. + + Same rationale as ``torch.bernoulli`` / ``torch.randn``. + """ + + def patch(*args, **kwargs): + # Signatures: ``randint(high, size, ...)`` or ``randint(low, high, size, ...)``. + size = next((a for a in args if isinstance(a, (list, tuple, torch.Size))), kwargs.get("size")) + return torch.zeros(size, dtype=kwargs.get("dtype", torch.int64), device=kwargs.get("device")) + + return patch + + +@register_patch("openvino", "torch.cummax", "torch.Tensor.cummax") +def _patch_cummax(original): + """OV has no ``aten.cummax`` lowering — reuse the ONNX triangular-mask decomposition.""" + from .exporter_onnx import _patch_cummax_or_cummin + + return _patch_cummax_or_cummin(original, mode="max") + + +@register_patch("openvino", "torch.cummin", "torch.Tensor.cummin") +def _patch_cummin(original): + """OV has no ``aten.cummin`` lowering — reuse the ONNX triangular-mask decomposition.""" + from .exporter_onnx import _patch_cummax_or_cummin + + return _patch_cummax_or_cummin(original, mode="min") + + +@register_patch("openvino", "torch.searchsorted") +def _patch_searchsorted(original): + """Decompose ``torch.searchsorted`` via broadcast comparison + sum. + + OV's frontend rejects the ``aten.searchsorted.Tensor`` node when its optional inputs + (``sorter``, ``out``) trace as ``None``. Same shape of fix as the ONNX patch — for + sorted inputs the insertion index equals the count of elements satisfying the + comparison (``<`` for left, ``<=`` for right). + """ + + def patch(sorted_sequence, values, *, out_int32=False, right=False, side=None, out=None, sorter=None): + if side is not None: + right = side == "right" + if right: + mask = sorted_sequence.unsqueeze(-1) <= values.unsqueeze(-2) + else: + mask = sorted_sequence.unsqueeze(-1) < values.unsqueeze(-2) + result = mask.sum(-2) + return result.to(torch.int32) if out_int32 else result + + return patch + + +@register_patch("openvino", "torch.bincount", "torch.Tensor.bincount") +def _patch_bincount(original): + """Replace ``torch.bincount`` with ``zeros + scatter_add_`` of size ``minlength`` (or input max+1 + when unknown). + + OV's PyTorch frontend has no ``aten.bincount`` lowering — same shape of fix as + ``_patch_histc``. The static output shape ``minlength`` keeps shape inference happy. + """ + + from torch.fx.experimental.symbolic_shapes import guard_or_true + + def patch(input, weights=None, minlength=0): + flat = input.reshape(-1) + # ``flat.numel() > 0`` and ``int(flat.max().item())`` on a data-dependent SymInt trip + # ``GuardOnDataDependentSymNode`` (splinter's question-token binning). ``guard_or_true`` + # optimistically assumes non-empty — an empty ``bincount`` collapses to a zero-length + # output anyway (harmless) — and ``torch._check_is_size`` marks the ``max`` result as + # size-like so the downstream ``bins + 1 > 0`` check in AOT autograd doesn't refire. + if guard_or_true(flat.numel() > 0): + max_val = flat.max().item() + torch._check(max_val >= 0) + bin_count = max_val + 1 + else: + bin_count = 0 + bins = torch.sym_max(minlength, bin_count) + out_dtype = weights.dtype if weights is not None else torch.long + counts = torch.zeros(bins, dtype=out_dtype, device=input.device) + src = weights.reshape(-1).to(out_dtype) if weights is not None else torch.ones_like(flat, dtype=out_dtype) + return counts.scatter_add_(0, flat.long(), src) + + return patch + + +@register_patch("openvino", "torch.nn.functional.interpolate") +def _patch_interpolate(original): + """Disable antialias for ``F.interpolate(..., antialias=True)`` during OV export. + + OV's frontend has no ``aten._upsample_bilinear2d_aa`` lowering. Antialiasing is a + pre-resample low-pass filter — turning it off costs a tiny amount of image-side quality but + keeps the graph translatable. Affects siglip2 and lfm2_vl. + """ + + def patch(input, *args, **kwargs): + kwargs.pop("antialias", None) + return original(input, *args, **kwargs) + + return patch + + +@register_patch("openvino", "torch.fft.irfft") +def _patch_irfft(original): + """Compute ``irfft`` entirely in real arithmetic — split the one-sided spectrum into + real/imag planes, mirror them to the full conjugate-symmetric spectrum, and contract + against real cos/sin DFT bases. + + OV's ``DFT`` op rejects ``is_onesided=1``/``inverse=1`` together, and a complex-valued + decomposition (mirror + ``ifft``) routes complex tensors through OV's builtin ``cat`` / + ``permute`` / ``bmm`` translators, which mix OV's native ``ComplexTypeMark`` representation + with the ``[..., 2]`` real-pair one our ``aten.complex`` extension emits (same clash as + ``_patch_apply_rotary_emb``). Keeping the whole transform real confines traced complex ops + to ``complex``/``view_as_real``, which the extensions handle. + """ + + def patch(input, n=None, dim=-1, norm=None): + if n is None: + n = 2 * (input.shape[dim] - 1) + if torch.is_complex(input): + pairs = torch.view_as_real(input) + real, imag = pairs[..., 0], pairs[..., 1] + else: + real, imag = input, torch.zeros_like(input) + real = real.movedim(dim, -1) + imag = imag.movedim(dim, -1) + # Mirror to the full n-point spectrum via conjugate symmetry: X[n - k] = conj(X[k]). + mirror = slice(1, n - real.shape[-1] + 1) + real = torch.cat([real, real[..., mirror].flip(-1)], dim=-1) + imag = torch.cat([imag, -imag[..., mirror].flip(-1)], dim=-1) + # y[j] = scale * sum_k (real[k] cos(2 pi k j / n) - imag[k] sin(2 pi k j / n)) + k = torch.arange(n, device=input.device, dtype=real.dtype) + angles = 2.0 * torch.pi * k.view(-1, 1) * k / n # symmetric [n, n], no transpose needed + scale = {"forward": 1.0, "ortho": n**-0.5}.get(norm, 1.0 / n) + out = (real @ angles.cos() - imag @ angles.sin()) * scale + return out.movedim(-1, dim) + + return patch + + +@register_patch("openvino", "torch.fft.rfft") +def _patch_rfft(original): + """Replace ``rfft`` with ``fft`` + slice to the one-sided half. OV's ``DFT(is_onesided=1)`` + has no inverse-pair (see ``_patch_irfft``); using two-sided + slice gives the same result + for the forward direction. Affects audio models (wav2vec*, seamless_m4t, pop2piano).""" + + def patch(input, n=None, dim=-1, norm=None): + full = torch.fft.fft(input, n=n, dim=dim, norm=norm) + n_full = full.shape[dim] + slc = [slice(None)] * full.ndim + slc[dim] = slice(0, n_full // 2 + 1) + return full[tuple(slc)] + + return patch + + +@register_patch("openvino", "torch.fft.fft") +def _patch_fft(original): + """``torch.fft.fft`` lowers to ``aten._fft_c2c.default`` which OV's frontend doesn't + translate. Build the DFT manually from the twiddle matrix — quadratic but adequate for + audio-encoder-sized FFTs that hit this path. + """ + + def patch(input, n=None, dim=-1, norm=None): + if n is None: + n = input.shape[dim] + # Twiddle matrix W[k, j] = exp(-2j pi k j / n) — emit via complex(cos, -sin). + k = torch.arange(n, device=input.device, dtype=torch.float32) + j = k.view(-1, 1) + angles = -2.0 * torch.pi * k * j / n + twiddle = torch.complex(angles.cos(), angles.sin()) + # Move target dim to last, matmul against twiddle, move back. + x = input.to(torch.complex64) if not torch.is_complex(input) else input + x = x.movedim(dim, -1) + out = x @ twiddle.T + return out.movedim(-1, dim) + + return patch + + +@register_patch("openvino", "torch.fft.ifft") +def _patch_ifft(original): + """Inverse of ``_patch_fft`` — uses conjugate twiddle and divides by ``n``.""" + + def patch(input, n=None, dim=-1, norm=None): + if n is None: + n = input.shape[dim] + k = torch.arange(n, device=input.device, dtype=torch.float32) + j = k.view(-1, 1) + angles = 2.0 * torch.pi * k * j / n + twiddle = torch.complex(angles.cos(), angles.sin()) + x = input.to(torch.complex64) if not torch.is_complex(input) else input + x = x.movedim(dim, -1) + out = (x @ twiddle.T) / n + return out.movedim(-1, dim) + + return patch + + +@register_patch("openvino", "torch.fft.fftn") +def _patch_fftn(original): + """Multi-dim FFT decomposed as successive 1-D ``torch.fft.fft`` calls along each ``dim``. + + OV has no ``aten._fft_c2c`` lowering for N-D inputs; the iterative 1-D form composes with + our ``_patch_fft`` so each axis is translated cleanly. Affects FNet. + """ + + def patch(input, s=None, dim=None, norm=None): + dims = list(range(input.ndim)) if dim is None else list(dim) + sizes = [None] * len(dims) if s is None else list(s) + out = input + for d, n in zip(dims, sizes): + out = torch.fft.fft(out, n=n, dim=d, norm=norm) + return out + + return patch + + +@register_patch("openvino", "torch.fft.ifftn") +def _patch_ifftn(original): + """Multi-dim inverse FFT — same decomposition as ``_patch_fftn`` via ``torch.fft.ifft``.""" + + def patch(input, s=None, dim=None, norm=None): + dims = list(range(input.ndim)) if dim is None else list(dim) + sizes = [None] * len(dims) if s is None else list(s) + out = input + for d, n in zip(dims, sizes): + out = torch.fft.ifft(out, n=n, dim=d, norm=norm) + return out + + return patch + + +@register_patch("openvino", "torch.fft.rfftn") +def _patch_rfftn(original): + """Real N-D FFT — last dim uses ``rfft`` (one-sided), remaining dims use ``fft``.""" + + def patch(input, s=None, dim=None, norm=None): + dims = list(range(input.ndim)) if dim is None else list(dim) + sizes = [None] * len(dims) if s is None else list(s) + out = input + for d, n in zip(dims[:-1], sizes[:-1]): + out = torch.fft.fft(out, n=n, dim=d, norm=norm) + return torch.fft.rfft(out, n=sizes[-1], dim=dims[-1], norm=norm) + + return patch + + +@register_patch("openvino", "torch.fft.irfftn") +def _patch_irfftn(original): + """Real N-D inverse FFT — last dim uses ``irfft``, remaining dims use ``ifft``.""" + + def patch(input, s=None, dim=None, norm=None): + dims = list(range(input.ndim)) if dim is None else list(dim) + sizes = [None] * len(dims) if s is None else list(s) + out = torch.fft.irfft(input, n=sizes[-1], dim=dims[-1], norm=norm) + for d, n in zip(dims[:-1], sizes[:-1]): + out = torch.fft.ifft(out, n=n, dim=d, norm=norm) + return out.real if torch.is_complex(out) else out + + return patch + + +@register_patch("openvino", "torch.Tensor.scatter_reduce_", "torch.Tensor.scatter_reduce") +def _patch_scatter_reduce(original): + """Decompose ``scatter_reduce_(dim, index, src, reduce)`` into ``scatter_*`` variants OV + can lower. ``sum``/``amax``/``amin`` map to ``scatter_add_``/``scatter_reduce(amax)`` / + ``scatter_reduce(amin)`` already, but the ``two`` overload OV doesn't recognise has the + same algorithmic content — replace with the plain ``scatter_add_`` for ``sum`` (the only + reduce mode actually used in the failing model, BLT). + """ + + def patch(self, dim, index, src, *, reduce="sum", include_self=True): + if reduce == "sum": + if not include_self: + self.zero_() + return self.scatter_add_(dim, index, src) + return original(self, dim, index, src, reduce=reduce, include_self=include_self) + + return patch + + +# ── OpenVINO conversion extensions ────────────────────────────────────────── +# Custom OV-side translations registered in ``_OV_CONVERSION_EXTENSIONS`` and passed to +# ``openvino.convert_model(extension=...)``. Mirrors the role of ONNX's +# ``_ONNX_TRANSLATION_TABLE``: use this when an op has no equivalent torch-level decomposition. +# Each ``_convert_*(context)`` receives a ``NodeContext`` (``context.get_input(i)`` for inputs) +# and returns a list of output ports built with ``openvino.opset14`` ops. +# +# To add a new translation: implement ``_convert_*`` and append a ``ConversionExtension`` to +# ``_OV_CONVERSION_EXTENSIONS``. + + +def _convert_grouped_mm(context): + """Convert ``aten._grouped_mm`` / ``transformers.grouped_mm_fallback`` to OV ops. + + ``grouped_mm(mat_a: (M, K), mat_b: (G, K, N), offs: (G,)) -> (M, N)`` computes + ``out[offs[g-1]:offs[g]] = mat_a[offs[g-1]:offs[g]] @ mat_b[g]`` per expert ``g``. + ``G`` (number of experts) must be static at translation time, so we unroll the loop and + emit ``G`` independent ``Slice + Gather + MatMul`` triples followed by a final ``Concat``. + """ + mat_a = context.get_input(0) + mat_b = context.get_input(1) + offs = context.get_input(2) + + G = mat_b.get_partial_shape()[0].get_length() + offs_i64 = ov_ops.convert(offs, "i64") + axes_0 = ov_ops.constant(np.array([0], dtype=np.int64)) + step_1 = ov_ops.constant(np.array([1], dtype=np.int64)) + prev_end = ov_ops.constant(np.array([0], dtype=np.int64)) + + outputs = [] + for g in range(G): + g_lo = ov_ops.constant(np.array([g], dtype=np.int64)) + g_hi = ov_ops.constant(np.array([g + 1], dtype=np.int64)) + end = ov_ops.slice(offs_i64, g_lo, g_hi, step_1, axes_0) # (1,) — offs[g] + a_g = ov_ops.slice(mat_a, prev_end, end, step_1, axes_0) # (n_g, K) + w_g_3d = ov_ops.slice(mat_b, g_lo, g_hi, step_1, axes_0) # (1, K, N) + w_g = ov_ops.squeeze(w_g_3d, axes_0) # (K, N) + outputs.append(ov_ops.matmul(a_g, w_g, transpose_a=False, transpose_b=False).output(0)) + prev_end = end + + return [ov_ops.concat(outputs, axis=0).output(0)] + + +def _convert_empty_permuted(context): + """Convert ``aten.empty_permuted`` to a zero-initialised constant of the requested shape. + + ``empty_permuted`` is uninitialised — only the shape matters for downstream ops. OV has no + direct equivalent; emit a zero ``Broadcast`` of the right shape and dtype. + """ + size = context.get_input(0) + # Default to f32; in the MoE expert path the result feeds straight into integer index ops or + # gets overwritten before any read, so dtype doesn't propagate to outputs. + zero = ov_ops.constant(np.float32(0.0)) + return [ov_ops.broadcast(zero, size).output(0)] + + +def _convert_index_add(context): + """Convert ``aten.index_add(self, dim, index, source, alpha=1)`` — OV's default translator + expects 5 inputs and fails when ``alpha`` is defaulted (torch omits it from the FX call). + Emit ``ScatterElementsUpdate`` with ``sum`` reduction: expand ``index`` from a 1-D shape + ``(N,)`` to match ``source`` along all axes so per-position add works. Used by t5gemma / + t5gemma2 / speecht5 relative-attention-bias accumulation.""" + data = context.get_input(0) + dim = int(context.get_values_from_const_input(1)) + index = context.get_input(2) + source = context.get_input(3) + # Broadcast 1-D index to source's rank/shape along ``dim`` so ScatterElementsUpdate + # can consume element-wise ``source`` values. + src_shape = ov_ops.shape_of(source, output_type="i64") + # Reshape ``index`` to a shape that's ``1`` in every dim except ``dim`` — broadcast handles + # the rest. Then broadcast to source's shape explicitly to feed ScatterElementsUpdate. + ndim = source.get_partial_shape().rank.get_length() + ones = [1] * ndim + ones[dim] = -1 + index_reshaped = ov_ops.reshape( + ov_ops.convert(index, "i64"), + ov_ops.constant(np.array(ones, dtype=np.int64)), + special_zero=False, + ) + index_bcast = ov_ops.broadcast(index_reshaped, src_shape) + return [ + ov_ops.scatter_elements_update( + data, index_bcast, source, ov_ops.constant(np.int64(dim)), reduction="sum" + ).output(0) + ] + + +def _convert_view_as_real(context): + """``view_as_real(complex)`` reinterprets a complex tensor as ``[..., 2]`` real. Our + ``_convert_complex`` already represents complex tensors that way, so this is identity.""" + return [context.get_input(0)] + + +def _convert_fft_c2c(context): + """Convert ``aten._fft_c2c(self, dim, normalization, forward)`` to OV's ``DFT``/``IDFT``. + + OV's ``dft``/``idft`` expect a trailing ``[..., 2]`` real/imag pair. Our ``_convert_complex`` + produces that layout already. For models that call ``_fft_c2c`` on a real-valued tensor + (FNet, where ``torch.fft.fftn(real)`` implicitly promotes to complex), we stack a zero + imaginary component on the last dim first. We detect the input rank via partial shape and + only inject the stack when there's no trailing ``[..., 2]`` already. + """ + data = context.get_input(0) + axes = context.get_input(1) + forward = bool(context.get_values_from_const_input(3)) + # If the input doesn't already end in a 2-element axis, treat it as real and pad imag=0. + pshape = data.get_partial_shape() + needs_pair = pshape.rank.is_static and ( + not pshape[pshape.rank.get_length() - 1].is_static or pshape[pshape.rank.get_length() - 1].get_length() != 2 + ) + if needs_pair: + zeros = ov_ops.broadcast(ov_ops.constant(np.float32(0.0)), ov_ops.shape_of(data)) + data = ov_ops.concat( + [ov_ops.unsqueeze(data, ov_ops.constant(-1)), ov_ops.unsqueeze(zeros, ov_ops.constant(-1))], + axis=-1, + ) + op = ov_ops.dft if forward else ov_ops.idft + return [op(data, ov_ops.convert(axes, "i64")).output(0)] + + +def _convert_conj(context): + """Convert ``aten._conj(complex)`` — complex conjugate. With our ``[..., 2]`` real/imag + representation, this negates the imaginary part. We split into real/imag, negate imag, + and concat back. Used by manual FFT decompositions.""" + data = context.get_input(0) + # last dim is 2 — split along axis -1 into real/imag, then concat [real, -imag] + axes_neg1 = ov_ops.constant(np.array([-1], dtype=np.int64)) + real_part = ov_ops.gather(data, ov_ops.constant(np.int64(0)), axes_neg1) + imag_part = ov_ops.gather(data, ov_ops.constant(np.int64(1)), axes_neg1) + neg_imag = ov_ops.negative(imag_part) + return [ + ov_ops.concat( + [ov_ops.unsqueeze(real_part, axes_neg1), ov_ops.unsqueeze(neg_imag, axes_neg1)], + axis=-1, + ).output(0) + ] + + +def _convert_bitwise_not(context): + """Convert ``aten.bitwise_not`` — OV's default translator internally calls ``torch.sym_float`` + on the input's dynamic dims to compute output shape metadata, and that Python-level call + remains as an unconverted node in the resulting graph. Emit ``LogicalNot`` on a boolean + view of the input; ``bitwise_not`` on bool would reject with ``is_integral()`` check. + Affects deformable_detr, mask2former.""" + data = context.get_input(0) + return [ov_ops.logical_not(ov_ops.convert(data, "boolean")).output(0)] + + +def _convert_layer_norm(context): + """Convert ``aten.layer_norm(input, normalized_shape, weight, bias, eps, cudnn_enable)`` to + ``MVN + (weight * x + bias)``. OV's default translator decomposes to ``native_layer_norm`` + which returns a 3-tuple ``(out, mean, rstd)``; the unused ``mean`` / ``rstd`` outputs are + emitted as ``torch::None`` constants that fail conversion (chameleon). Emitting MVN + directly gives a single-output op with no dangling None.""" + data = context.get_input(0) + normalized_shape = context.get_values_from_const_input(1) + weight = context.get_input(2) + bias = context.get_input(3) + eps = float(context.get_values_from_const_input(4)) if context.get_input_size() > 4 else 1e-5 + ndim = data.get_partial_shape().rank.get_length() + axes_len = len(normalized_shape) if hasattr(normalized_shape, "__len__") else 1 + axes = ov_ops.constant(np.array(list(range(ndim - axes_len, ndim)), dtype=np.int64)) + normalized = ov_ops.mvn(data, axes, normalize_variance=True, eps=eps, eps_mode="inside_sqrt") + scaled = ov_ops.multiply(normalized, weight) + shifted = ov_ops.add(scaled, bias) + return [shifted.output(0)] + + +def _convert_aten_to(context): + """Convert ``aten.to.{dtype,device,dtype_layout,other}`` — emit a real ``Convert`` when the + target dtype is present, else identity. + + OV's frontend has no ``aten.to.*`` translations at all (only ``aten._to_copy.default``); + every unhandled variant falls back to a ``torch::None`` constant that fails conversion + (chameleon's rotary sub-module hits ``aten.to.dtype_layout``). ``layout`` / ``device`` + kwargs are silently dropped — OV exports are inherently device-neutral.""" + data = context.get_input(0) + if not context.has_attribute("dtype"): + # ``aten.to.device`` — just device move, no dtype. Emit identity. + return [data] + try: + dtype = context.get_attribute("dtype") + except Exception: + # Complex dtypes throw. Skip (identity) — see ``_convert_to_copy``. + return [data] + if dtype is None: + return [data] + return [ov_ops.convert(data, dtype).output(0)] + + +def _convert_to_copy(context): + """Convert ``aten._to_copy(self, dtype=..., ...)`` to an OV ``Convert``. + + OV's default translator throws (``Attribute dtype can't be converted to defined types``) + when the target dtype is ``complex64`` — no native OV complex type. Our ``_convert_complex`` + uses a ``[..., 2]`` real representation, so the complex cast is a no-op we swallow. For all + real dtypes we emit a real ``Convert`` — dropping the cast entirely regresses downstream + ops like ``aten.bitwise_and.Tensor`` that need the mask to actually be ``bool`` (cpmant, + chameleon).""" + data = context.get_input(0) + if not context.has_attribute("dtype"): + return [data] + try: + dtype = context.get_attribute("dtype") + except Exception: + # Complex dtypes throw ``Attribute dtype can't be converted to defined types``. With + # the ``[..., 2]`` real representation, the cast is a no-op. + return [data] + if dtype is None: + return [data] + return [ov_ops.convert(data, dtype).output(0)] + + +def _convert_bmm(context): + """Translate ``aten.bmm``, shielding softmax-fed ones from OV's SDPA fusion. + + The frontend-normalization fusion matches ``bmm -> softmax -> bmm`` and mis-shapes the + result when batch and heads are flattened into one dim (SpeechT5/MVP/SeamlessM4T's + relative-position eager attention): the fused op emits ``[b, b*h, q, k]`` instead of + ``[b, h, q, k]``. For a bmm consuming a ``Softmax`` output, a ``Reshape(x, ShapeOf(x))`` + no-op is appended — runtime-dependent, so normalization can't fold it away before the + fusion pass runs, and MOC's nop-elimination cleans it up afterwards. Every other bmm + translates to a plain ``MatMul``. + """ + a, b = context.get_input(0), context.get_input(1) + product = ov_ops.matmul(a, b, transpose_a=False, transpose_b=False) + if a.get_node().get_type_name() != "Softmax": + return [product.output(0)] + identity = ov_ops.reshape(product, ov_ops.shape_of(product, output_type="i64"), special_zero=False) + return [identity.output(0)] + + +def _convert_sdpa(context): + """Convert ``aten.scaled_dot_product_attention`` — wrapping OV's op with a mask-dtype fix. + + OV's ``opset13::ScaledDotProductAttention`` rejects int-typed masks. Under CUDA export + ``aten.expand`` promotes bool masks to ``i64`` during OV translation, so we insert a + ``Convert(→ boolean)`` on the mask input before instantiating the op. Q/K/V/scale pass + through unchanged.""" + q, k, v = context.get_input(0), context.get_input(1), context.get_input(2) + # A ``None`` FX arg reaches the extension as an unconverted ``PtFrameworkNode``. + mask = None + if context.get_input_size() > 3: + candidate = context.get_input(3) + if candidate.get_node().get_type_name() != "PtFrameworkNode": + mask = ov_ops.convert(candidate, "boolean") + is_causal = False + if context.get_input_size() > 5: + # ``is_causal`` is a positional FX input (arg 5), not a node attribute. + if context.get_input(5).get_node().get_type_name() == "Constant": + is_causal = bool(context.get_values_from_const_input(5)) + if mask is None: + return [ov_ops.scaled_dot_product_attention(q, k, v, causal=is_causal).output(0)] + return [ov_ops.scaled_dot_product_attention(q, k, v, mask, causal=is_causal).output(0)] + + +def _convert_complex(context): + """Convert ``aten.complex(real, imag)`` by stacking as the last dim — OV represents complex + tensors as ``[..., 2]`` real tensors via ``ComplexTypeMark``. Affects models that build + complex tensors explicitly (RoPE polar form, manual FFT decompositions).""" + real = context.get_input(0) + imag = context.get_input(1) + stacked = ov_ops.concat( + [ov_ops.unsqueeze(real, ov_ops.constant(-1)), ov_ops.unsqueeze(imag, ov_ops.constant(-1))], + axis=-1, + ) + return [stacked.output(0)] + + +# ── SymInt builtin translations ───────────────────────────────────────────── +# torch.export records Python-level math on SymInts (``a % b``, ``a // b``, ``min(a, b)``) +# as ``call_function`` nodes whose target is the Python builtin or ``torch.sym_*`` callable. +# These survive into the EP because torch never lowers them — there's no aten op that +# produces a SymInt for ``mod``/``floordiv``/etc. OV's PyTorch frontend has no translation +# for them either, so we register one per builtin keyed on its ``str(target)`` literal. +# Each translator emits an OV opset17 elementwise op; the result is a 0-d integer tensor +# that downstream shape ops (view, reshape, expand) concat into shape lists natively. + + +def _convert_sym_binop(op): + """Factory: build a 2-arg OV-op translator for SymInt binary builtins (add, mul, mod, …). + + Mixed int/float operands (e.g. ``symint - 0.5`` in deformable-attention grid math) are + promoted to the float side — OV element-wise ops require matching types. ``mod`` must map + to ``floor_mod``: Python's ``%`` is floored (``-7 % 3 == 2``) while OV's ``Mod`` truncates, + which breaks ``-seq % block``-style padding arithmetic (LongT5). + """ + + def _convert(context): + a, b = context.get_input(0), context.get_input(1) + a_type, b_type = a.get_element_type(), b.get_element_type() + if a_type != b_type: + if a_type.is_integral() and not b_type.is_integral(): + a = ov_ops.convert_like(a, b) + elif b_type.is_integral() and not a_type.is_integral(): + b = ov_ops.convert_like(b, a) + return [op(a, b).output(0)] + + return _convert + + +def _convert_sym_unop(op, *, cast_to_i64=False): + """Factory: build a 1-arg OV-op translator for SymInt unary builtins (floor, ceil, sym_float). + + ``cast_to_i64`` casts the output back to ``i64`` — Python's ``floor(x)`` / ``ceil(x)`` on a + SymFloat return an int, but OV's ``floor`` / ``ceiling`` are dtype-preserving, so a float + input yields a float output. Downstream shape ops (SequenceMark → Concat) need i64; + without the cast, mixed-dtype Concat fails ``element::Type::merge`` (focalnet).""" + + def _convert(context): + out = op(context.get_input(0)) + if cast_to_i64: + out = ov_ops.convert(out, "i64") + return [out.output(0)] + + return _convert + + +def _convert_sym_floordiv(context): + """``a // b`` over SymInts → ``floor(a / b)``, cast to i64. Used by patch/window-size + computations (focalnet, donut_swin). The i64 cast keeps the result shape-op-friendly — + downstream ``SequenceMark → Concat`` requires a uniform int dtype.""" + a, b = context.get_input(0), context.get_input(1) + return [ov_ops.convert(ov_ops.floor(ov_ops.divide(a, b)), "i64").output(0)] + + +_OV_CONVERSION_EXTENSIONS: list[Any] = [] +if is_openvino_available(): + _OV_CONVERSION_EXTENSIONS.extend( + [ + ConversionExtension("aten._grouped_mm.default", _convert_grouped_mm), + ConversionExtension("transformers.grouped_mm_fallback.default", _convert_grouped_mm), + ConversionExtension("aten.empty_permuted.default", _convert_empty_permuted), + ConversionExtension("aten.index_add.default", _convert_index_add), + ConversionExtension("aten.bmm.default", _convert_bmm), + ConversionExtension("aten.complex.default", _convert_complex), + ConversionExtension("aten.view_as_real.default", _convert_view_as_real), + ConversionExtension("aten._fft_c2c.default", _convert_fft_c2c), + ConversionExtension("aten._conj.default", _convert_conj), + ConversionExtension("aten._to_copy.default", _convert_to_copy), + ConversionExtension("aten.to.dtype", _convert_aten_to), + ConversionExtension("aten.to.dtype_layout", _convert_aten_to), + ConversionExtension("aten.to.device", _convert_aten_to), + ConversionExtension("aten.to.other", _convert_aten_to), + ConversionExtension("aten.layer_norm.default", _convert_layer_norm), + ConversionExtension("aten.scaled_dot_product_attention.default", _convert_sdpa), + ConversionExtension("aten.bitwise_not.default", _convert_bitwise_not), + # SymInt builtins — see comment block above. + ConversionExtension("", _convert_sym_binop(ov_ops.add)), + ConversionExtension("", _convert_sym_binop(ov_ops.subtract)), + ConversionExtension("", _convert_sym_binop(ov_ops.multiply)), + ConversionExtension("", _convert_sym_binop(ov_ops.divide)), + ConversionExtension("", _convert_sym_floordiv), + ConversionExtension("", _convert_sym_binop(ov_ops.floor_mod)), + ConversionExtension("", _convert_sym_binop(ov_ops.power)), + ConversionExtension("", _convert_sym_unop(ov_ops.floor, cast_to_i64=True)), + ConversionExtension("", _convert_sym_unop(ov_ops.ceiling, cast_to_i64=True)), + ConversionExtension("", _convert_sym_binop(ov_ops.minimum)), + ConversionExtension("", _convert_sym_binop(ov_ops.maximum)), + # ``torch.sym_float`` has an address-based ``str()`` (not a stable ```` + # form), so we register by its runtime str. Emits a real→f32 Convert. + ConversionExtension(str(torch.sym_float), _convert_sym_unop(lambda x: ov_ops.convert(x, "f32"))), + ] + ) diff --git a/src/transformers/exporters/utils.py b/src/transformers/exporters/utils.py index ef876702daa5..0ffd831d23e8 100644 --- a/src/transformers/exporters/utils.py +++ b/src/transformers/exporters/utils.py @@ -613,12 +613,24 @@ def decompose_prefill_decode( Reuses the full generation machinery so every architecture (decoder-only, SSM, encoder-decoder, multi-modal, …) gets correct inputs without reimplementing the loop. + Some multi-modal models (Blip2, Kosmos-2, …) override `generate()` to run their encoders + inline and delegate the generation loop to an inner language model, so the top-level + `forward()` never runs. To cover those, the decoder returned by `model.get_decoder()` is + hooked as well, and whichever module the generation loop actually called is captured + (top-level `forward()` preferred when both ran). + Returns: `dict[str, tuple[torch.nn.Module, dict]]`: - `{"prefill": (model, prefill_inputs), "decode": (model, decode_inputs)}` + `{"prefill": (module, prefill_inputs), "decode": (module, decode_inputs)}` where + `module` is `model` itself, or its decoder when `generate()` delegates to it. """ + decoder = model.get_decoder() try: - with _capture_forward(model) as calls: + with contextlib.ExitStack() as stack: + calls = stack.enter_context(_capture_forward(model)) + decoder_calls = ( + stack.enter_context(_capture_forward(decoder)) if decoder is not None and decoder is not model else [] + ) model.generate(**copy.deepcopy(inputs), max_new_tokens=2, min_new_tokens=2) except Exception as e: raise RuntimeError( @@ -627,17 +639,17 @@ def decompose_prefill_decode( f"Make sure the inputs are compatible with model.generate()." ) from e - if len(calls) < 2: + module, module_calls = (model, calls) if len(calls) >= 2 else (decoder, decoder_calls) + if len(module_calls) < 2: raise RuntimeError( - f"decompose_prefill_decode expected at least 2 calls to {type(model).__name__}.forward() " - f"during generate(max_new_tokens=2), but captured {len(calls)}. This likely means " - "generate() bypasses the top-level forward() (e.g. delegates to an inner model), " - "so prefill/decode decomposition is not supported for this architecture." + f"decompose_prefill_decode expected at least 2 forward() calls on {type(model).__name__} " + f"or its decoder during generate(max_new_tokens=2), but captured {len(calls)} on the " + f"top-level model and {len(decoder_calls)} on the decoder." ) return { - "prefill": (copy.copy(model), calls[0]), - "decode": (copy.copy(model), calls[1]), + "prefill": (copy.copy(module), module_calls[0]), + "decode": (copy.copy(module), module_calls[1]), } @@ -741,7 +753,8 @@ def decompose_for_generation( Runs `decompose_prefill_decode` to capture prefill and decode forward kwargs from a real `model.generate(**inputs, max_new_tokens=2)`. If the prefill is multi-modal (per `is_multimodal`), further splits it into one entry per submodule (vision/audio encoder, projector, language model, - `lm_head`) via `decompose_multimodal`. + `lm_head`) via `decompose_multimodal`. Models whose `generate()` delegates the loop to an inner + language model (Blip2, Kosmos-2, …) get prefill/decode captured at that inner model instead. Args: model: Generative model. Must support `model.generate(**inputs)`. diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 130b1a1dec66..c508a65a1efe 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -533,8 +533,11 @@ def sdpa_mask( batch_arange = torch.arange(batch_size, device=device) head_arange = torch.arange(1, device=device) - q_arange = torch.arange(q_length, device=device) + q_offset - kv_arange = torch.arange(kv_length, device=device) + kv_offset + # Fold the offsets into the arange bounds rather than adding them afterwards — a python + # scalar in `aten.add` gets materialized as an int64 constant by ExecuTorch's edge passes, + # whose buffer the XNNPACK lowering then downcasts without updating the spec. + q_arange = torch.arange(q_offset, q_offset + q_length, device=device) + kv_arange = torch.arange(kv_offset, kv_offset + kv_length, device=device) # Actual mask creation # Option 1: Fast non-vmap mask creation (default) @@ -632,7 +635,9 @@ def eager_mask( ) # only bidirectional masks can be skipped, otherwise we convert bool -> float if mask is not None: - min_dtype = torch.finfo(dtype).min + # A python-float `other` would be lifted as a float64 constant by `torch.export`, + # whose serialized buffer then mismatches the spec — keep both branches in `dtype`. + min_dtype = torch.tensor(torch.finfo(dtype).min, device=mask.device, dtype=dtype) # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype) return mask diff --git a/src/transformers/models/big_bird/modeling_big_bird.py b/src/transformers/models/big_bird/modeling_big_bird.py index 1343047b9f2d..906a6569a7b6 100755 --- a/src/transformers/models/big_bird/modeling_big_bird.py +++ b/src/transformers/models/big_bird/modeling_big_bird.py @@ -2497,8 +2497,7 @@ def forward( @staticmethod def prepare_question_mask(q_lengths: torch.Tensor, maxlen: int): # q_lengths -> (bz, 1) - mask = torch.arange(0, maxlen).to(q_lengths.device) - mask.unsqueeze_(0) # -> (1, maxlen) + mask = torch.arange(0, maxlen, device=q_lengths.device).unsqueeze(0) # -> (1, maxlen) mask = torch.where(mask < q_lengths, 1, 0) return mask diff --git a/src/transformers/models/bros/modeling_bros.py b/src/transformers/models/bros/modeling_bros.py index 8f025e2f37e7..bba3f6c9426c 100755 --- a/src/transformers/models/bros/modeling_bros.py +++ b/src/transformers/models/bros/modeling_bros.py @@ -817,7 +817,7 @@ def forward( subsequent_token_logits = subsequent_token_logits.masked_fill( invalid_token_mask[:, None, :], torch.finfo(subsequent_token_logits.dtype).min ) - self_token_mask = torch.eye(max_seq_length, max_seq_length + 1).to(device=device, dtype=torch.bool) + self_token_mask = torch.eye(max_seq_length, max_seq_length + 1, device=device, dtype=torch.bool) subsequent_token_logits = subsequent_token_logits.masked_fill( self_token_mask[None, :, :], torch.finfo(subsequent_token_logits.dtype).min ) @@ -941,7 +941,7 @@ def forward( batch_size, max_seq_length = attention_mask.shape device = attention_mask.device - self_token_mask = torch.eye(max_seq_length, max_seq_length + 1).to(device=device, dtype=torch.bool) + self_token_mask = torch.eye(max_seq_length, max_seq_length + 1, device=device, dtype=torch.bool) mask = bbox_first_token_mask.view(-1) bbox_first_token_mask = torch.cat( diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index 52cf77ec9218..91132ef2f47e 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -761,17 +761,23 @@ def img2bpe(self): def bpe2img_search_tensors(self): return torch.tensor(sorted(self.bpe2img.keys())), torch.tensor(sorted(self.bpe2img.values())) - @cached_property - def img2bpe_mapping_tensor(self): - mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int) - for k, v in self.img2bpe.items(): - mapping[k] = v - return mapping + def _build_img2bpe_mapping_tensor(self, device: torch.device) -> torch.Tensor: + """Build a device-local ``(max_img_id + 1,)`` int lookup table mapping VQ image ids to + BPE token ids. Cached per-device so we don't rebuild on every call.""" + if not hasattr(self, "_img2bpe_mapping_cache"): + self._img2bpe_mapping_cache: dict[torch.device, torch.Tensor] = {} + if device not in self._img2bpe_mapping_cache: + # Build on CPU via ``scatter`` then move — avoids a Python-level element-wise + # assignment loop that trips FakeTensor propagation during export. + keys = torch.tensor(list(self.img2bpe.keys()), dtype=torch.long) + values = torch.tensor(list(self.img2bpe.values()), dtype=torch.int) + mapping = torch.zeros(int(keys.max().item()) + 1, dtype=torch.int) + mapping.scatter_(0, keys, values) + self._img2bpe_mapping_cache[device] = mapping.to(device) + return self._img2bpe_mapping_cache[device] def convert_img2bpe(self, img_batch: torch.Tensor) -> torch.Tensor: - device = img_batch.device - img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")] - return img_tokens.to(device) + return self._build_img2bpe_mapping_tensor(img_batch.device)[img_batch] @auto_docstring @@ -1077,9 +1083,14 @@ def forward( slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) - # Disallow image tokens which does not include special begin-image and end-image tokens + # Disallow image tokens which does not include special begin-image and end-image tokens. + # ``logits[:, :, image_tokens] = ...`` traces as ``aten.index_put`` with ``[None, None, tensor]`` + # indices — the ``None`` values become ``torch::None`` under torch.export, which OpenVINO's + # frontend can't lower. A vocab-sized bool mask + ``masked_fill`` is equivalent and traces cleanly. image_tokens = self.model.vocabulary_mapping.image_tokens - logits[:, :, image_tokens] = torch.finfo(logits.dtype).min + image_token_mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) + image_token_mask[image_tokens] = True + logits = logits.masked_fill(image_token_mask, torch.finfo(logits.dtype).min) loss = None if labels is not None: diff --git a/src/transformers/models/fsmt/modeling_fsmt.py b/src/transformers/models/fsmt/modeling_fsmt.py index 2d364266778d..7e5e715f0c11 100644 --- a/src/transformers/models/fsmt/modeling_fsmt.py +++ b/src/transformers/models/fsmt/modeling_fsmt.py @@ -177,17 +177,6 @@ def invert_mask(attention_mask): return attention_mask.eq(0) -def triu_onnx(x, diagonal=0): - l = x.shape[0] - arange = torch.arange(l, device=x.device) - mask = arange.expand(l, l) - arange = arange.unsqueeze(-1) - if diagonal: - arange = arange + diagonal - mask = mask >= arange - return x.masked_fill(mask == 0, 0) - - def _prepare_fsmt_decoder_inputs( config, input_ids, @@ -208,8 +197,9 @@ def _prepare_fsmt_decoder_inputs( decoder_padding_mask = make_padding_mask(decoder_input_ids, pad_token_id) else: decoder_padding_mask = invert_mask(decoder_padding_mask) - causal_mask = triu_onnx(fill_with_neg_inf(torch.zeros(tgt_len, tgt_len, dtype=causal_mask_dtype)), 1).to( - device=decoder_input_ids.device + causal_mask = torch.triu( + fill_with_neg_inf(torch.zeros(tgt_len, tgt_len, dtype=causal_mask_dtype, device=decoder_input_ids.device)), + diagonal=1, ) return decoder_input_ids, decoder_padding_mask, causal_mask diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index 0ab7f0b130c6..a24a6a8a8efb 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -2173,7 +2173,7 @@ def forward( return Gemma3nModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, - past_key_values=outputs.past_key_values if use_cache else None, + past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_features if pixel_values is not None else None, diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index d07c22acd35a..4a816f708f85 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -2290,7 +2290,7 @@ def forward( return Gemma3nModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, - past_key_values=outputs.past_key_values if use_cache else None, + past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_features if pixel_values is not None else None, diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index 34fba03f5687..c6ab35493bdd 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -258,8 +258,8 @@ def _attn(self, query, key, value, attention_mask=None): batch_size, num_attention_heads, query_length, attn_head_size = query.size() key_length = key.size(-2) - query = query.view(batch_size * num_attention_heads, query_length, attn_head_size) - key = key.view(batch_size * num_attention_heads, key_length, attn_head_size) + query = query.reshape(batch_size * num_attention_heads, query_length, attn_head_size) + key = key.reshape(batch_size * num_attention_heads, key_length, attn_head_size) # [batch_size * num_heads, q_length, kv_length] attn_scores = torch.zeros( diff --git a/src/transformers/models/informer/modeling_informer.py b/src/transformers/models/informer/modeling_informer.py index 0a11533bcb59..386640cf747e 100644 --- a/src/transformers/models/informer/modeling_informer.py +++ b/src/transformers/models/informer/modeling_informer.py @@ -583,7 +583,7 @@ def forward( if top_u_sparsity_measurement is not None: # update context: copy the attention output to the context at top_u_sparsity_measurement index - dim_for_slice = torch.arange(context.size(0)).unsqueeze(-1) + dim_for_slice = torch.arange(context.size(0), device=context.device).unsqueeze(-1) context[dim_for_slice, top_u_sparsity_measurement, :] = attn_output attn_output = context diff --git a/src/transformers/models/longt5/modeling_longt5.py b/src/transformers/models/longt5/modeling_longt5.py index b3e082789e9c..83bed253e88f 100644 --- a/src/transformers/models/longt5/modeling_longt5.py +++ b/src/transformers/models/longt5/modeling_longt5.py @@ -70,9 +70,10 @@ def _split_into_blocks(x: torch.Tensor, block_len: int, dim: int) -> torch.Tenso """Split an input tensor into blocks of a given `block_len` along the given `dim`. If the dimension length is not a multiple of `block_len`, it will be padded first with selected `pad_value`. """ - # pad tensor to multiple of block_len - if x.shape[dim] % block_len != 0: - x = _pad_to_multiple(x, block_len, dim, pad_value=0) + # Pad tensor to multiple of block_len — a no-op when the length is already a multiple. + # Padding unconditionally keeps the shape arithmetic branch-free, so ``torch.export`` + # traces a single graph that holds for every sequence length. + x = _pad_to_multiple(x, block_len, dim, pad_value=0) num_blocks = x.shape[dim] // block_len output_shape = x.shape[:dim] + (num_blocks, block_len) + x.shape[(dim + 1) :] # If 0 is in output_shape, we cannot apply reshape because of incompatibility with ONNX conversion @@ -106,9 +107,9 @@ def _concatenate_3_blocks(x: torch.Tensor, block_dim: int, sequence_dim: int, pa return torch.cat(blocks_list, dim=sequence_dim) -def _make_3block_relative_position_ids(block_len: int) -> torch.Tensor: +def _make_3block_relative_position_ids(block_len: int, device: torch.device | None = None) -> torch.Tensor: """Makes 3-blocked relative position ids for local attention.""" - position_ids = torch.arange(3 * block_len, dtype=torch.int32) + position_ids = torch.arange(3 * block_len, dtype=torch.int32, device=device) center_position_ids = position_ids[block_len:-block_len] # [block_len, 3 * block_len] relative_position_ids = position_ids.unsqueeze(0) - center_position_ids.unsqueeze(1) @@ -117,10 +118,9 @@ def _make_3block_relative_position_ids(block_len: int) -> torch.Tensor: def _mask_local_attention_mask(local_attention_mask: torch.Tensor, block_len: int) -> torch.Tensor: """Mask local attention mask to enforce that tokens are not allowed to attend tokens farther than ``local_radius.""" - relative_position_ids = _make_3block_relative_position_ids(block_len) + relative_position_ids = _make_3block_relative_position_ids(block_len, device=local_attention_mask.device) locality_mask = torch.abs(relative_position_ids) < block_len locality_mask = locality_mask[None, None, :, :] - locality_mask = locality_mask.to(local_attention_mask.device) return torch.logical_and(local_attention_mask, locality_mask) @@ -156,8 +156,7 @@ def _make_global_fixed_block_ids( batch_size, seq_len = attention_mask.shape[:2] def handle_orphan_tokens(block_ids: torch.Tensor) -> torch.Tensor: - block_ends = (torch.arange(seq_len) % global_block_size) == global_block_size - 1 - block_ends = block_ends.to(block_ids.device) + block_ends = (torch.arange(seq_len, device=block_ids.device) % global_block_size) == global_block_size - 1 true_block_ends = torch.logical_and(block_ends, block_ids >= 0) full_blocks = true_block_ends.sum(-1).unsqueeze(-1).type(block_ids.dtype) - 1 block_ids = torch.where(block_ids < full_blocks, block_ids, full_blocks) @@ -183,8 +182,12 @@ def handle_orphan_tokens(block_ids: torch.Tensor) -> torch.Tensor: _sequence_block_ids_max = torch.zeros( batch_size, 0, dtype=global_block_ids.dtype, device=global_block_ids.device ) - global_segment_ids = torch.cumsum(torch.ones(batch_size, num_globals), dim=-1) - 1 - global_segment_ids = global_segment_ids.to(attention_mask.device) + # Equivalent to ``cumsum(ones(...), dim=-1) - 1``: [0, 1, …, num_globals-1] per batch. + global_segment_ids = ( + torch.arange(num_globals, dtype=attention_mask.dtype, device=attention_mask.device) + .unsqueeze(0) + .expand(batch_size, -1) + ) global_segment_ids = torch.where(global_segment_ids <= _sequence_block_ids_max, 1, 0) return global_block_ids.type(torch.int), global_segment_ids.type(torch.int) @@ -808,7 +811,7 @@ def unshape(states): # global_seq_len := seq_len // self.global_block_size # shapes: (batch_size, seq_len) & (batch_size, global_seq_len) block_ids, global_segment_ids = _make_global_fixed_block_ids( - mask if mask is not None else torch.ones(hidden_states.shape[:-1]), + mask if mask is not None else torch.ones(hidden_states.shape[:-1], device=hidden_states.device), self.global_block_size, ) # Create global inputs diff --git a/src/transformers/models/mask2former/modeling_mask2former.py b/src/transformers/models/mask2former/modeling_mask2former.py index e9623bb9acb5..ab92b3c955c3 100644 --- a/src/transformers/models/mask2former/modeling_mask2former.py +++ b/src/transformers/models/mask2former/modeling_mask2former.py @@ -978,10 +978,19 @@ def forward( ) # batch_size, num_queries, n_heads, n_levels, n_points, 2 if reference_points.shape[-1] == 2: - offset_normalizer = torch.tensor( - [[shape[1], shape[0]] for shape in spatial_shapes_list], - dtype=torch.long, - device=reference_points.device, + # ``torch.tensor([[w, h], ...], device=cuda)`` where dims are SymInts materialises + # on CPU and moves — trips FakeTensor device propagation during ``torch.export``. + # Build via ``torch.stack`` on device-side scalars instead. + offset_normalizer = torch.stack( + [ + torch.stack( + [ + torch.as_tensor(shape[1], dtype=torch.long, device=reference_points.device), + torch.as_tensor(shape[0], dtype=torch.long, device=reference_points.device), + ] + ) + for shape in spatial_shapes_list + ] ) sampling_locations = ( reference_points[:, :, None, :, None, :] diff --git a/src/transformers/models/maskformer/modeling_maskformer_swin.py b/src/transformers/models/maskformer/modeling_maskformer_swin.py index 4f41456ece07..f6c93180ecbb 100644 --- a/src/transformers/models/maskformer/modeling_maskformer_swin.py +++ b/src/transformers/models/maskformer/modeling_maskformer_swin.py @@ -457,7 +457,7 @@ def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, self.intermediate = MaskFormerSwinIntermediate(config, dim) self.output = MaskFormerSwinOutput(config, dim) - def get_attn_mask(self, input_resolution): + def get_attn_mask(self, input_resolution, device=None): """Build the cyclic-shift attention mask for shifted-window MSA; returns None when shift_size is 0. Each (h, w) position belongs to one of 9 cyclic-shift regions (3 along each axis), encoded @@ -472,8 +472,8 @@ def get_attn_mask(self, input_resolution): if self.shift_size <= 0: return None height, width = input_resolution - h_idx = torch.arange(height) - w_idx = torch.arange(width) + h_idx = torch.arange(height, device=device) + w_idx = torch.arange(width, device=device) h_region = (h_idx >= height - self.window_size).long() + (h_idx >= height - self.shift_size).long() w_region = (w_idx >= width - self.window_size).long() + (w_idx >= width - self.shift_size).long() img_mask = (h_region[None, :, None, None] * 3 + w_region[None, None, :, None]).float() @@ -511,9 +511,7 @@ def forward(self, hidden_states, input_dimensions, output_attentions=False): # partition windows hidden_states_windows = window_partition(shifted_hidden_states, self.window_size) hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels) - attn_mask = self.get_attn_mask((height_pad, width_pad)) - if attn_mask is not None: - attn_mask = attn_mask.to(hidden_states_windows.device) + attn_mask = self.get_attn_mask((height_pad, width_pad), device=hidden_states_windows.device) self_attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions) diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 620d57876e10..c23beb12c1a2 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -137,9 +137,9 @@ def __init__(self, config: MiniMaxConfig, layer_idx: int): self.layer_type = config.layer_types[layer_idx] - def get_slope_rate(self): + def get_slope_rate(self, device=None): base = 1 / (2 ** (8 / self.num_attention_heads)) - exponent = torch.arange(self.num_attention_heads) + 1 + exponent = torch.arange(self.num_attention_heads, device=device) + 1 factor = 1 - self.layer_idx / (self.num_hidden_layers - 1 + 1e-5) + 1e-5 rate = base**exponent @@ -149,7 +149,7 @@ def get_slope_rate(self): return rate def decay_factors(self, slope_rate): - block_size_range = torch.arange(self.block_size) + 1 + block_size_range = torch.arange(self.block_size, device=slope_rate.device) + 1 query_decay = torch.exp(-slope_rate * block_size_range[:, None]) key_decay = torch.exp(-slope_rate * (self.block_size - block_size_range[:, None])) @@ -188,8 +188,13 @@ def forward( attn_weights_inter = past_key_values.get_linear_cache(self.layer_idx) if attn_weights_inter is None: - attn_weights_inter = torch.zeros(batch_size, self.num_attention_heads, self.head_dim, self.head_dim).to( - value_states + attn_weights_inter = torch.zeros( + batch_size, + self.num_attention_heads, + self.head_dim, + self.head_dim, + device=value_states.device, + dtype=value_states.dtype, ) # apply attention_mask diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 3a7ed5dd1247..e8530af641d9 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -220,9 +220,9 @@ def __init__(self, config: MiniMaxConfig, layer_idx: int): self.layer_type = config.layer_types[layer_idx] - def get_slope_rate(self): + def get_slope_rate(self, device=None): base = 1 / (2 ** (8 / self.num_attention_heads)) - exponent = torch.arange(self.num_attention_heads) + 1 + exponent = torch.arange(self.num_attention_heads, device=device) + 1 factor = 1 - self.layer_idx / (self.num_hidden_layers - 1 + 1e-5) + 1e-5 rate = base**exponent @@ -232,7 +232,7 @@ def get_slope_rate(self): return rate def decay_factors(self, slope_rate): - block_size_range = torch.arange(self.block_size) + 1 + block_size_range = torch.arange(self.block_size, device=slope_rate.device) + 1 query_decay = torch.exp(-slope_rate * block_size_range[:, None]) key_decay = torch.exp(-slope_rate * (self.block_size - block_size_range[:, None])) @@ -271,8 +271,13 @@ def forward( attn_weights_inter = past_key_values.get_linear_cache(self.layer_idx) if attn_weights_inter is None: - attn_weights_inter = torch.zeros(batch_size, self.num_attention_heads, self.head_dim, self.head_dim).to( - value_states + attn_weights_inter = torch.zeros( + batch_size, + self.num_attention_heads, + self.head_dim, + self.head_dim, + device=value_states.device, + dtype=value_states.dtype, ) # apply attention_mask diff --git a/src/transformers/models/mpnet/modeling_mpnet.py b/src/transformers/models/mpnet/modeling_mpnet.py index 6cf8b83f18ae..4b79112b6e71 100644 --- a/src/transformers/models/mpnet/modeling_mpnet.py +++ b/src/transformers/models/mpnet/modeling_mpnet.py @@ -317,13 +317,12 @@ def compute_position_bias(self, x, position_ids=None, num_buckets=32): context_position = position_ids[:, :, None] memory_position = position_ids[:, None, :] else: - context_position = torch.arange(qlen, dtype=torch.long)[:, None] - memory_position = torch.arange(klen, dtype=torch.long)[None, :] + context_position = torch.arange(qlen, dtype=torch.long, device=x.device)[:, None] + memory_position = torch.arange(klen, dtype=torch.long, device=x.device)[None, :] relative_position = memory_position - context_position rp_bucket = self.relative_position_bucket(relative_position, num_buckets=num_buckets) - rp_bucket = rp_bucket.to(x.device) values = self.relative_attention_bias(rp_bucket) values = values.permute([2, 0, 1]).unsqueeze(0) values = values.expand((bsz, -1, qlen, klen)).contiguous() diff --git a/src/transformers/models/musicgen/modeling_musicgen.py b/src/transformers/models/musicgen/modeling_musicgen.py index 33a02805a38a..c6b4ff9be2de 100644 --- a/src/transformers/models/musicgen/modeling_musicgen.py +++ b/src/transformers/models/musicgen/modeling_musicgen.py @@ -140,7 +140,7 @@ def get_embedding(num_embeddings: int, embedding_dim: int): def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0): bsz, codebooks, seq_len = input_ids.size() # Create the position ids from the input token ids. - position_ids = (torch.arange(seq_len) + past_key_values_length).to(input_ids.device) + position_ids = torch.arange(seq_len, device=input_ids.device) + past_key_values_length # expand embeddings if needed if seq_len > self.weights.size(0): self.make_weights(seq_len, self.embedding_dim) diff --git a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py index a3e3f5c3251c..963cb3f66e94 100644 --- a/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py +++ b/src/transformers/models/musicgen_melody/modeling_musicgen_melody.py @@ -147,7 +147,7 @@ def get_embedding(num_embeddings: int, embedding_dim: int): def forward(self, inputs_embeds: torch.Tensor, past_key_values_length: int = 0): bsz, seq_len, _ = inputs_embeds.size() # Create the position ids from the input token ids. - position_ids = (torch.arange(seq_len) + past_key_values_length).to(inputs_embeds.device) + position_ids = torch.arange(seq_len, device=inputs_embeds.device) + past_key_values_length # expand embeddings if needed if seq_len > self.weights.size(0): self.make_weights(seq_len, self.embedding_dim) diff --git a/src/transformers/models/perceiver/modeling_perceiver.py b/src/transformers/models/perceiver/modeling_perceiver.py index 43fa45828f30..c4fe062af27b 100755 --- a/src/transformers/models/perceiver/modeling_perceiver.py +++ b/src/transformers/models/perceiver/modeling_perceiver.py @@ -2499,7 +2499,7 @@ def generate_fourier_features(pos, num_bands, max_resolution=(224, 224), concat_ return per_pos_features -def build_linear_positions(index_dims, output_range=(-1.0, 1.0)): +def build_linear_positions(index_dims, output_range=(-1.0, 1.0), device=None): """ Generate an array of position indices for an N-D input array. @@ -2508,13 +2508,17 @@ def build_linear_positions(index_dims, output_range=(-1.0, 1.0)): The shape of the index dimensions of the input array. output_range (`tuple[float]`, *optional*, defaults to `(-1.0, 1.0)`): The min and max values taken by each input index dimension. + device (`torch.device`, *optional*): + Device on which to allocate the linspace/stack. Defaults to CPU (torch default). Returns: `torch.FloatTensor` of shape `(index_dims[0], index_dims[1], .., index_dims[-1], N)`. """ def _linspace(n_xels_per_dim): - return torch.linspace(start=output_range[0], end=output_range[1], steps=n_xels_per_dim, dtype=torch.float32) + return torch.linspace( + start=output_range[0], end=output_range[1], steps=n_xels_per_dim, dtype=torch.float32, device=device + ) dim_ranges = [_linspace(n_xels_per_dim) for n_xels_per_dim in index_dims] array_index_grid = torch.meshgrid(*dim_ranges, indexing="ij") @@ -2593,7 +2597,7 @@ def forward( return position_embeddings -def _check_or_build_spatial_positions(pos, index_dims, batch_size): +def _check_or_build_spatial_positions(pos, index_dims, batch_size, device=None): """ Checks or builds spatial position features (x, y, ...). @@ -2604,12 +2608,14 @@ def _check_or_build_spatial_positions(pos, index_dims, batch_size): An iterable giving the spatial/index size of the data to be featurized. batch_size (`int`): The batch size of the data to be featurized. + device (`torch.device`, *optional*): + Device to build the spatial positions on when ``pos`` is ``None``. Returns: `torch.FloatTensor` of shape `(batch_size, prod(index_dims))` an array of position features. """ if pos is None: - pos = build_linear_positions(index_dims) + pos = build_linear_positions(index_dims, device=device) # equivalent to `torch.broadcast_to(pos[None], (batch_size,) + pos.shape)` # but `torch.broadcast_to` cannot be converted to ONNX pos = pos[None].expand((batch_size,) + pos.shape) @@ -2656,14 +2662,14 @@ def forward( dtype: torch.dtype, pos: torch.FloatTensor | None = None, ) -> torch.FloatTensor: - pos = _check_or_build_spatial_positions(pos, index_dims, batch_size) + pos = _check_or_build_spatial_positions(pos, index_dims, batch_size, device=device) fourier_pos_enc = generate_fourier_features( pos, num_bands=self.num_bands, max_resolution=self.max_resolution, concat_pos=self.concat_pos, sine_only=self.sine_only, - ).to(device=device, dtype=dtype) + ).to(dtype=dtype) return fourier_pos_enc @@ -3277,8 +3283,8 @@ def forward( if modality in self.mask_probs: mask_token = self.mask[modality].expand(batch_size, -1, -1) mask_prob = self.mask_probs[modality] - mask = torch.bernoulli(torch.full([batch_size, num_samples], mask_prob)) - mask = torch.unsqueeze(mask, dim=2).to(mask_token.device) + mask = torch.bernoulli(torch.full([batch_size, num_samples], mask_prob, device=mask_token.device)) + mask = torch.unsqueeze(mask, dim=2) output_padded = (1 - mask) * output_padded + mask * mask_token padded[modality] = output_padded diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index eda3f270d5e4..25f180050c38 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -56,10 +56,17 @@ def __init__(self, config: PhimoeConfig, device=None): self.rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + # Precompute both LongRoPE variants at init so forward stays traceable: at forward we + # blend by the runtime ``max(position_ids) + 1 > original_max_position_embeddings`` mask + # via ``torch.where``, avoiding a data-dependent Python-``bool`` on the seq_len tensor. inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) - self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + if self.rope_type != "default": + inv_freq_long, _ = self.rope_init_fn( + self.config, device, seq_len=self.config.rope_parameters["original_max_position_embeddings"] + 1 + ) + self.register_buffer("inv_freq_long", inv_freq_long, persistent=False) @staticmethod def compute_default_rope_parameters( @@ -99,16 +106,22 @@ def forward(self, x, position_ids=None, layer_type=None): f"{self.__class__.__name__} does not support layer types, but got `layer_type={layer_type}`" ) - mscale = None - seq_len = torch.max(position_ids) + 1 - if self.config.rope_parameters["rope_type"] != "default" and seq_len: - mscale = ( - self.config.rope_parameters["long_mscale"] - if seq_len > self.config.rope_parameters["original_max_position_embeddings"] - else self.config.rope_parameters["short_mscale"] + if self.rope_type != "default": + threshold = self.config.rope_parameters["original_max_position_embeddings"] + # ``max(position_ids) + 1`` (not ``shape[-1]``) so decode / sliding-window paths that + # ship a short slice of position_ids still see the true reached max — matters for + # LongRoPE scale selection. Blending via ``torch.where`` avoids the ``Python bool`` + # on a tensor that trips ``GuardOnDataDependentSymNode`` under ``torch.export``. + is_long_context = (torch.max(position_ids) + 1) > threshold + inv_freq = torch.where(is_long_context, self.inv_freq_long, self.inv_freq) + mscale = torch.where( + is_long_context, + torch.tensor(self.config.rope_parameters["long_mscale"], device=x.device, dtype=torch.float32), + torch.tensor(self.config.rope_parameters["short_mscale"], device=x.device, dtype=torch.float32), ) - inv_freq, attention_scaling = self.rope_init_fn(self.config, x.device, seq_len) - mscale = attention_scaling if mscale is None else mscale + else: + inv_freq = self.inv_freq + mscale = self.attention_scaling inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() diff --git a/src/transformers/models/phimoe/modular_phimoe.py b/src/transformers/models/phimoe/modular_phimoe.py index a4c0a081d89c..0913791c8396 100644 --- a/src/transformers/models/phimoe/modular_phimoe.py +++ b/src/transformers/models/phimoe/modular_phimoe.py @@ -49,10 +49,17 @@ def __init__(self, config: PhimoeConfig, device=None): self.rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + # Precompute both LongRoPE variants at init so forward stays traceable: at forward we + # blend by the runtime ``max(position_ids) + 1 > original_max_position_embeddings`` mask + # via ``torch.where``, avoiding a data-dependent Python-``bool`` on the seq_len tensor. inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) - self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + if self.rope_type != "default": + inv_freq_long, _ = self.rope_init_fn( + self.config, device, seq_len=self.config.rope_parameters["original_max_position_embeddings"] + 1 + ) + self.register_buffer("inv_freq_long", inv_freq_long, persistent=False) def forward(self, x, position_ids=None, layer_type=None): if layer_type is not None: @@ -60,16 +67,22 @@ def forward(self, x, position_ids=None, layer_type=None): f"{self.__class__.__name__} does not support layer types, but got `layer_type={layer_type}`" ) - mscale = None - seq_len = torch.max(position_ids) + 1 - if self.config.rope_parameters["rope_type"] != "default" and seq_len: - mscale = ( - self.config.rope_parameters["long_mscale"] - if seq_len > self.config.rope_parameters["original_max_position_embeddings"] - else self.config.rope_parameters["short_mscale"] + if self.rope_type != "default": + threshold = self.config.rope_parameters["original_max_position_embeddings"] + # ``max(position_ids) + 1`` (not ``shape[-1]``) so decode / sliding-window paths that + # ship a short slice of position_ids still see the true reached max — matters for + # LongRoPE scale selection. Blending via ``torch.where`` avoids the ``Python bool`` + # on a tensor that trips ``GuardOnDataDependentSymNode`` under ``torch.export``. + is_long_context = (torch.max(position_ids) + 1) > threshold + inv_freq = torch.where(is_long_context, self.inv_freq_long, self.inv_freq) + mscale = torch.where( + is_long_context, + torch.tensor(self.config.rope_parameters["long_mscale"], device=x.device, dtype=torch.float32), + torch.tensor(self.config.rope_parameters["short_mscale"], device=x.device, dtype=torch.float32), ) - inv_freq, attention_scaling = self.rope_init_fn(self.config, x.device, seq_len) - mscale = attention_scaling if mscale is None else mscale + else: + inv_freq = self.inv_freq + mscale = self.attention_scaling inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() diff --git a/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py b/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py index a74a7994edf6..944bf34415c0 100644 --- a/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py +++ b/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py @@ -1541,30 +1541,26 @@ def mask_to_box_coordinate(mask, dtype): x_coords = x_coords.to(dtype) y_coords = y_coords.to(dtype) + finfo_max = torch.tensor(torch.finfo(dtype).max, device=mask.device) x_coords_masked = x_coords * mask x_max = x_coords_masked.flatten(start_dim=-2).max(dim=-1).values + 1 - x_min = ( - torch.where(mask, x_coords_masked, torch.tensor(torch.finfo(dtype).max)) - .flatten(start_dim=-2) - .min(dim=-1) - .values - ) + x_min = torch.where(mask, x_coords_masked, finfo_max).flatten(start_dim=-2).min(dim=-1).values y_coords_masked = y_coords * mask y_max = y_coords_masked.flatten(start_dim=-2).max(dim=-1).values + 1 - y_min = ( - torch.where(mask, y_coords_masked, torch.tensor(torch.finfo(dtype).max)) - .flatten(start_dim=-2) - .min(dim=-1) - .values - ) + y_min = torch.where(mask, y_coords_masked, finfo_max).flatten(start_dim=-2).min(dim=-1).values unnormalized_bbox = torch.stack([x_min, y_min, x_max, y_max], dim=-1) is_mask_non_empty = torch.any(mask, dim=(-2, -1)).unsqueeze(-1) unnormalized_bbox = unnormalized_bbox * is_mask_non_empty - norm_tensor = torch.tensor([width, height, width, height], device=mask.device, dtype=dtype) + # ``torch.tensor([w, h, w, h])`` where ``w`` / ``h`` are SymInts materialises on CPU and + # then transfers, tripping FakeTensor device propagation during export. Build via + # ``torch.stack`` on already-device-side scalars instead. + width_t = torch.as_tensor(width, device=mask.device, dtype=dtype) + height_t = torch.as_tensor(height, device=mask.device, dtype=dtype) + norm_tensor = torch.stack([width_t, height_t, width_t, height_t]) normalized_bbox_xyxy = unnormalized_bbox / norm_tensor x_min_norm, y_min_norm, x_max_norm, y_max_norm = normalized_bbox_xyxy.unbind(dim=-1) diff --git a/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py b/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py index 0967ebe83f42..222dfe4d4c64 100644 --- a/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py +++ b/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py @@ -585,30 +585,26 @@ def mask_to_box_coordinate(mask, dtype): x_coords = x_coords.to(dtype) y_coords = y_coords.to(dtype) + finfo_max = torch.tensor(torch.finfo(dtype).max, device=mask.device) x_coords_masked = x_coords * mask x_max = x_coords_masked.flatten(start_dim=-2).max(dim=-1).values + 1 - x_min = ( - torch.where(mask, x_coords_masked, torch.tensor(torch.finfo(dtype).max)) - .flatten(start_dim=-2) - .min(dim=-1) - .values - ) + x_min = torch.where(mask, x_coords_masked, finfo_max).flatten(start_dim=-2).min(dim=-1).values y_coords_masked = y_coords * mask y_max = y_coords_masked.flatten(start_dim=-2).max(dim=-1).values + 1 - y_min = ( - torch.where(mask, y_coords_masked, torch.tensor(torch.finfo(dtype).max)) - .flatten(start_dim=-2) - .min(dim=-1) - .values - ) + y_min = torch.where(mask, y_coords_masked, finfo_max).flatten(start_dim=-2).min(dim=-1).values unnormalized_bbox = torch.stack([x_min, y_min, x_max, y_max], dim=-1) is_mask_non_empty = torch.any(mask, dim=(-2, -1)).unsqueeze(-1) unnormalized_bbox = unnormalized_bbox * is_mask_non_empty - norm_tensor = torch.tensor([width, height, width, height], device=mask.device, dtype=dtype) + # ``torch.tensor([w, h, w, h])`` where ``w`` / ``h`` are SymInts materialises on CPU and + # then transfers, tripping FakeTensor device propagation during export. Build via + # ``torch.stack`` on already-device-side scalars instead. + width_t = torch.as_tensor(width, device=mask.device, dtype=dtype) + height_t = torch.as_tensor(height, device=mask.device, dtype=dtype) + norm_tensor = torch.stack([width_t, height_t, width_t, height_t]) normalized_bbox_xyxy = unnormalized_bbox / norm_tensor x_min_norm, y_min_norm, x_max_norm, y_max_norm = normalized_bbox_xyxy.unbind(dim=-1) diff --git a/src/transformers/models/prophetnet/modeling_prophetnet.py b/src/transformers/models/prophetnet/modeling_prophetnet.py index d1a7265bdbba..b0b4ecaa5a6b 100644 --- a/src/transformers/models/prophetnet/modeling_prophetnet.py +++ b/src/transformers/models/prophetnet/modeling_prophetnet.py @@ -734,11 +734,10 @@ def get_main_relative_pos_embeddings( if main_relative_position_buckets is None: batch_size, sequence_length = hidden_states.shape[:2] relative_positions = ( - torch.arange(1, attn_weights.shape[-1] + 1) + torch.arange(1, attn_weights.shape[-1] + 1, device=position_ids.device) .unsqueeze(0) .unsqueeze(0) .repeat(batch_size, sequence_length, 1) - .to(position_ids.device) ) # [batch_size, sequence_length, sequence_length+1] relative_positions = relative_positions - position_ids.unsqueeze(0).repeat(batch_size, sequence_length, 1) @@ -784,11 +783,10 @@ def get_predict_relative_pos_embeddings( "`position_ids` are incorrect. They should be of the format 1 2 3 4 5 ... (key_sequence_length - 1)", ) relative_positions = ( - torch.arange(0, key_sequence_length) + torch.arange(0, key_sequence_length, device=position_ids.device) .unsqueeze(0) .unsqueeze(0) .repeat(batch_size, sequence_length, 1) - .to(position_ids.device) ) relative_positions = relative_positions - position_ids.unsqueeze(0).repeat(batch_size, sequence_length, 1) @@ -1273,7 +1271,7 @@ def forward( def compute_buffered_relative_buckets(self, position_ids): batch_size, sequence_length = position_ids.shape - position_ids = torch.arange(1, self.max_target_positions).to(position_ids.device).repeat(1, 1) + position_ids = torch.arange(1, self.max_target_positions, device=position_ids.device).repeat(1, 1) main_relative_buckets, predict_relative_buckets = compute_all_stream_relative_buckets( self.num_buckets, self.relative_max_distance, position_ids ) diff --git a/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py index f75fa9dcdcd9..436a6923eff2 100644 --- a/src/transformers/models/speecht5/modeling_speecht5.py +++ b/src/transformers/models/speecht5/modeling_speecht5.py @@ -431,7 +431,7 @@ def __init__(self, dim, max_length=1000): def forward(self, hidden_states): seq_len = hidden_states.shape[1] - pos_seq = torch.arange(0, seq_len).to(device=hidden_states.device, dtype=torch.long) + pos_seq = torch.arange(0, seq_len, device=hidden_states.device, dtype=torch.long) pos_seq = pos_seq[:, None] - pos_seq[None, :] pos_seq = torch.where(pos_seq < -self.max_length, -self.max_length, pos_seq) diff --git a/src/transformers/models/swin2sr/modeling_swin2sr.py b/src/transformers/models/swin2sr/modeling_swin2sr.py index c72616940d5f..756a4724fe7a 100644 --- a/src/transformers/models/swin2sr/modeling_swin2sr.py +++ b/src/transformers/models/swin2sr/modeling_swin2sr.py @@ -458,7 +458,7 @@ def _compute_window_shift(self, target_window_size, target_shift_size) -> tuple[ shift_size = [0 if r <= w else s for r, w, s in zip(self.input_resolution, window_size, target_shift_size)] return window_size, shift_size - def get_attn_mask(self, height, width, dtype): + def get_attn_mask(self, height, width, dtype, device=None): """Build the cyclic-shift attention mask for shifted-window MSA; returns None when shift_size is 0. Each (h, w) position belongs to one of 9 cyclic-shift regions (3 along each axis), encoded @@ -472,8 +472,8 @@ def get_attn_mask(self, height, width, dtype): """ if self.shift_size <= 0: return None - h_idx = torch.arange(height) - w_idx = torch.arange(width) + h_idx = torch.arange(height, device=device) + w_idx = torch.arange(width, device=device) h_region = (h_idx >= height - self.window_size).long() + (h_idx >= height - self.shift_size).long() w_region = (w_idx >= width - self.window_size).long() + (w_idx >= width - self.shift_size).long() img_mask = (h_region[None, :, None, None] * 3 + w_region[None, None, :, None]).to(dtype) @@ -513,9 +513,9 @@ def forward( # partition windows hidden_states_windows = window_partition(shifted_hidden_states, self.window_size) hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels) - attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype) - if attn_mask is not None: - attn_mask = attn_mask.to(hidden_states_windows.device) + attn_mask = self.get_attn_mask( + height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device + ) attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions) diff --git a/src/transformers/models/swinv2/modeling_swinv2.py b/src/transformers/models/swinv2/modeling_swinv2.py index 3b9d5ab1b01d..8a250c0be164 100644 --- a/src/transformers/models/swinv2/modeling_swinv2.py +++ b/src/transformers/models/swinv2/modeling_swinv2.py @@ -617,7 +617,7 @@ def _compute_window_shift(self, target_window_size, target_shift_size) -> tuple[ shift_size = [0 if r <= w else s for r, w, s in zip(self.input_resolution, window_size, target_shift_size)] return window_size, shift_size - def get_attn_mask(self, height, width, dtype): + def get_attn_mask(self, height, width, dtype, device=None): """Build the cyclic-shift attention mask for shifted-window MSA; returns None when shift_size is 0. Each (h, w) position belongs to one of 9 cyclic-shift regions (3 along each axis), encoded @@ -631,8 +631,8 @@ def get_attn_mask(self, height, width, dtype): """ if self.shift_size <= 0: return None - h_idx = torch.arange(height) - w_idx = torch.arange(width) + h_idx = torch.arange(height, device=device) + w_idx = torch.arange(width, device=device) h_region = (h_idx >= height - self.window_size).long() + (h_idx >= height - self.shift_size).long() w_region = (w_idx >= width - self.window_size).long() + (w_idx >= width - self.shift_size).long() img_mask = (h_region[None, :, None, None] * 3 + w_region[None, None, :, None]).to(dtype) @@ -672,9 +672,9 @@ def forward( # partition windows hidden_states_windows = window_partition(shifted_hidden_states, self.window_size) hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels) - attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype) - if attn_mask is not None: - attn_mask = attn_mask.to(hidden_states_windows.device) + attn_mask = self.get_attn_mask( + height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device + ) attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions) diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index 0085b9a31e19..bdd0edbc0ad6 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -641,7 +641,10 @@ def __init__( def forward(self, input_ids: torch.Tensor): input_embeddings = super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) - input_embeddings[input_ids == self.eoi_token_index] = self.eoi_embedding.to(input_embeddings.dtype) + # ``embeddings[bool_mask] = vector`` traces as ``aten.index_put`` with a mask list, + # which several export backends can't lower — ``torch.where`` is equivalent. + eoi_mask = (input_ids == self.eoi_token_index).unsqueeze(-1) + input_embeddings = torch.where(eoi_mask, self.eoi_embedding.to(input_embeddings.dtype), input_embeddings) return input_embeddings diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index c59546c72520..108d99ad8336 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -457,7 +457,10 @@ def __init__( def forward(self, input_ids: torch.Tensor): input_embeddings = super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) - input_embeddings[input_ids == self.eoi_token_index] = self.eoi_embedding.to(input_embeddings.dtype) + # ``embeddings[bool_mask] = vector`` traces as ``aten.index_put`` with a mask list, + # which several export backends can't lower — ``torch.where`` is equivalent. + eoi_mask = (input_ids == self.eoi_token_index).unsqueeze(-1) + input_embeddings = torch.where(eoi_mask, self.eoi_embedding.to(input_embeddings.dtype), input_embeddings) return input_embeddings diff --git a/src/transformers/models/tvp/modeling_tvp.py b/src/transformers/models/tvp/modeling_tvp.py index eeeb9ce2738c..d9f53bd457fb 100644 --- a/src/transformers/models/tvp/modeling_tvp.py +++ b/src/transformers/models/tvp/modeling_tvp.py @@ -751,9 +751,7 @@ def forward( if attention_mask is not None: # (batch_size, visual_sequence_length) visual_attention_mask = attention_mask.new_ones(visual_embedding_output.shape[:2]) - pt_mask = torch.ones(attention_mask.shape[0], 10).to( - device=attention_mask.device, dtype=attention_mask.dtype - ) + pt_mask = torch.ones(attention_mask.shape[0], 10, device=attention_mask.device, dtype=attention_mask.dtype) attention_mask = torch.cat([pt_mask, attention_mask, visual_attention_mask], dim=-1) text_prompt = self.text_prompt.expand(text_embedding_output.shape[0], -1, -1) diff --git a/src/transformers/models/unispeech/modeling_unispeech.py b/src/transformers/models/unispeech/modeling_unispeech.py index 03103760140c..ede32b34e74f 100755 --- a/src/transformers/models/unispeech/modeling_unispeech.py +++ b/src/transformers/models/unispeech/modeling_unispeech.py @@ -1138,11 +1138,13 @@ def forward( quantized_features = self.project_q(quantized_features.to(self.project_q.weight.dtype)) quantized_features = self.project_hid(quantized_features) - prob_replace_matrix = torch.empty(transformer_features.size(0), transformer_features.size(1)).fill_( - self.config.replace_prob + prob_replace_matrix = torch.full( + (transformer_features.size(0), transformer_features.size(1)), + self.config.replace_prob, + device=transformer_features.device, ) prob_replace_matrix = prob_replace_matrix.transpose(0, 1) - sampled_replace_matrix = torch.bernoulli(prob_replace_matrix).bool().to(transformer_features.device) + sampled_replace_matrix = torch.bernoulli(prob_replace_matrix).bool() sampled_replace_matrix = sampled_replace_matrix.transpose(0, 1) sampled_replace_matrix = sampled_replace_matrix.unsqueeze(-1) logits = transformer_features.masked_fill(sampled_replace_matrix, 0.0) + ( diff --git a/src/transformers/models/unispeech/modular_unispeech.py b/src/transformers/models/unispeech/modular_unispeech.py index fce518d13a65..d14e1c306c56 100644 --- a/src/transformers/models/unispeech/modular_unispeech.py +++ b/src/transformers/models/unispeech/modular_unispeech.py @@ -383,11 +383,13 @@ def forward( quantized_features = self.project_q(quantized_features.to(self.project_q.weight.dtype)) quantized_features = self.project_hid(quantized_features) - prob_replace_matrix = torch.empty(transformer_features.size(0), transformer_features.size(1)).fill_( - self.config.replace_prob + prob_replace_matrix = torch.full( + (transformer_features.size(0), transformer_features.size(1)), + self.config.replace_prob, + device=transformer_features.device, ) prob_replace_matrix = prob_replace_matrix.transpose(0, 1) - sampled_replace_matrix = torch.bernoulli(prob_replace_matrix).bool().to(transformer_features.device) + sampled_replace_matrix = torch.bernoulli(prob_replace_matrix).bool() sampled_replace_matrix = sampled_replace_matrix.transpose(0, 1) sampled_replace_matrix = sampled_replace_matrix.unsqueeze(-1) logits = transformer_features.masked_fill(sampled_replace_matrix, 0.0) + ( diff --git a/src/transformers/models/videomae/modeling_videomae.py b/src/transformers/models/videomae/modeling_videomae.py index 35bcce23c973..c85f95ff1e95 100755 --- a/src/transformers/models/videomae/modeling_videomae.py +++ b/src/transformers/models/videomae/modeling_videomae.py @@ -63,7 +63,7 @@ class VideoMAEDecoderOutput(ModelOutput): @dataclass class VideoMAEForPreTrainingOutput(ModelOutput): r""" - loss (`torch.FloatTensor` of shape `(1,)`): + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Pixel reconstruction loss. logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. @@ -532,6 +532,7 @@ def forward( self, pixel_values: torch.FloatTensor, bool_masked_pos: torch.BoolTensor, + return_loss: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> VideoMAEForPreTrainingOutput: r""" @@ -539,6 +540,9 @@ def forward( Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Each video in the batch must have the same number of masked patches. Sequence length is `(num_frames // tubelet_size) * (image_size // patch_size) ** 2`. + return_loss (`bool`, *optional*): + Whether to compute and return the pixel reconstruction loss. Defaults to `self.config.return_loss` if + that attribute is set, and `True` otherwise. Examples: ```python @@ -586,82 +590,89 @@ def forward( logits = decoder_outputs.logits loss = None - with torch.no_grad(): - # calculate the labels to be predicted - if self.config.num_channels != 3: - # Can't unnormalize with default means/stds - frames = pixel_values - else: - # first, unnormalize the frames - device = pixel_values.device - dtype = pixel_values.dtype - mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None] - std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None] - frames = pixel_values * std + mean # in [0, 1] - - batch_size, time, num_channels, height, width = frames.shape - tubelet_size, patch_size = self.config.tubelet_size, self.config.patch_size - if self.config.norm_pix_loss: - # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) - frames = frames.view( - batch_size, - time // tubelet_size, - tubelet_size, - num_channels, - height // patch_size, - patch_size, - width // patch_size, - patch_size, - ) - # step 2: move dimensions to concatenate: - frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() - # step 3: concatenate: - frames = frames.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size, - num_channels, - ) - # step 4: normalize. The authors find that the mean is about 0.48 and standard deviation is about 0.08. - frames_norm = (frames - frames.mean(dim=-2, keepdim=True)) / ( - frames.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6 - ) - # step 5: reshape to (batch_size, T//ts * H//ps * W//ps, ts * ps * ps * C) - videos_patch = frames_norm.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size * num_channels, - ) - else: + if return_loss is None: + return_loss = getattr(self.config, "return_loss", True) + if return_loss: + with torch.no_grad(): + # calculate the labels to be predicted if self.config.num_channels != 3: - raise ValueError( - "Can't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False." + # Can't unnormalize with default means/stds + frames = pixel_values + else: + # first, unnormalize the frames + device = pixel_values.device + dtype = pixel_values.dtype + mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[ + None, None, :, None, None + ] + std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[ + None, None, :, None, None + ] + frames = pixel_values * std + mean # in [0, 1] + + batch_size, time, num_channels, height, width = frames.shape + tubelet_size, patch_size = self.config.tubelet_size, self.config.patch_size + if self.config.norm_pix_loss: + # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) + frames = frames.view( + batch_size, + time // tubelet_size, + tubelet_size, + num_channels, + height // patch_size, + patch_size, + width // patch_size, + patch_size, ) - # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) - frames = frames.view( - batch_size, - time // tubelet_size, - tubelet_size, - num_channels, - height // patch_size, - patch_size, - width // patch_size, - patch_size, - ) - # step 2: move dimensions to concatenate: (batch_size, T//ts, H//ps, W//ps, ts, ps, ps, C) - frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() - # step 3: concatenate - videos_patch = frames.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size * num_channels, - ) - - batch_size, _, num_channels = videos_patch.shape - labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels) - - loss_fct = MSELoss() - loss = loss_fct(logits, labels) + # step 2: move dimensions to concatenate: + frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() + # step 3: concatenate: + frames = frames.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size, + num_channels, + ) + # step 4: normalize. The authors find that the mean is about 0.48 and standard deviation is about 0.08. + frames_norm = (frames - frames.mean(dim=-2, keepdim=True)) / ( + frames.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6 + ) + # step 5: reshape to (batch_size, T//ts * H//ps * W//ps, ts * ps * ps * C) + videos_patch = frames_norm.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size * num_channels, + ) + else: + if self.config.num_channels != 3: + raise ValueError( + "Can't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False." + ) + # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) + frames = frames.view( + batch_size, + time // tubelet_size, + tubelet_size, + num_channels, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + # step 2: move dimensions to concatenate: (batch_size, T//ts, H//ps, W//ps, ts, ps, ps, C) + frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() + # step 3: concatenate + videos_patch = frames.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size * num_channels, + ) + + batch_size, _, num_channels = videos_patch.shape + labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels) + + loss_fct = MSELoss() + loss = loss_fct(logits, labels) return VideoMAEForPreTrainingOutput( loss=loss, diff --git a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py index 3660946a1a4d..97c51b02c575 100644 --- a/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/configuration_voxtral_realtime.py @@ -113,6 +113,7 @@ class VoxtralRealtimeEncoderConfig(PreTrainedConfig): rope_parameters: RopeParameters | dict | None = None sliding_window: int = 750 head_dim: int = 64 + use_cache: bool = True def __post_init__(self, **kwargs): self.head_dim = self.head_dim if self.head_dim is not None else self.hidden_size // self.num_attention_heads @@ -170,6 +171,7 @@ class VoxtralRealtimeConfig(PreTrainedConfig): default_num_delay_tokens: int = 6 downsample_factor: int = 4 tie_word_embeddings: bool = True + use_cache: bool = True def __post_init__(self, **kwargs): if isinstance(self.audio_config, dict): diff --git a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py index bcea513fcc1e..88885e475d55 100644 --- a/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modeling_voxtral_realtime.py @@ -936,6 +936,7 @@ def get_audio_features( return audio_outputs + @merge_with_config_defaults @can_return_tuple @auto_docstring def forward( diff --git a/src/transformers/models/voxtral_realtime/modular_voxtral_realtime.py b/src/transformers/models/voxtral_realtime/modular_voxtral_realtime.py index b0db1a3195b5..3b869e765379 100644 --- a/src/transformers/models/voxtral_realtime/modular_voxtral_realtime.py +++ b/src/transformers/models/voxtral_realtime/modular_voxtral_realtime.py @@ -543,6 +543,7 @@ def get_audio_features( return audio_outputs + @merge_with_config_defaults @can_return_tuple @auto_docstring def forward( diff --git a/src/transformers/models/wavlm/modeling_wavlm.py b/src/transformers/models/wavlm/modeling_wavlm.py index 18440ebf7d25..5192d1169e07 100755 --- a/src/transformers/models/wavlm/modeling_wavlm.py +++ b/src/transformers/models/wavlm/modeling_wavlm.py @@ -241,11 +241,11 @@ def torch_multi_head_self_attention( return attn_output, attn_weights def compute_bias(self, query_length: int, key_length: int) -> torch.FloatTensor: - context_position = torch.arange(query_length, dtype=torch.long)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long)[None, :] + device = self.rel_attn_embed.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] relative_position = memory_position - context_position relative_position_bucket = self._relative_positions_bucket(relative_position) - relative_position_bucket = relative_position_bucket.to(self.rel_attn_embed.weight.device) values = self.rel_attn_embed(relative_position_bucket) values = values.permute([2, 0, 1]) return values diff --git a/src/transformers/models/wavlm/modular_wavlm.py b/src/transformers/models/wavlm/modular_wavlm.py index b3329e64913d..78136d2bf393 100644 --- a/src/transformers/models/wavlm/modular_wavlm.py +++ b/src/transformers/models/wavlm/modular_wavlm.py @@ -172,11 +172,11 @@ def torch_multi_head_self_attention( return attn_output, attn_weights def compute_bias(self, query_length: int, key_length: int) -> torch.FloatTensor: - context_position = torch.arange(query_length, dtype=torch.long)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long)[None, :] + device = self.rel_attn_embed.weight.device + context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] relative_position = memory_position - context_position relative_position_bucket = self._relative_positions_bucket(relative_position) - relative_position_bucket = relative_position_bucket.to(self.rel_attn_embed.weight.device) values = self.rel_attn_embed(relative_position_bucket) values = values.permute([2, 0, 1]) return values diff --git a/src/transformers/models/xcodec2/modeling_xcodec2.py b/src/transformers/models/xcodec2/modeling_xcodec2.py index 2e6fed319091..cec5cbafb187 100644 --- a/src/transformers/models/xcodec2/modeling_xcodec2.py +++ b/src/transformers/models/xcodec2/modeling_xcodec2.py @@ -771,12 +771,13 @@ def __init__(self, config: Xcodec2Config): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: stft_pred = self.linear(hidden_states).transpose(1, 2) magnitude, phase = stft_pred.chunk(2, dim=1) - # Cast to float32: complex exponential and irfft are not supported for fp16 (ComplexHalf) + # Cast to float32: complex tensors and irfft are not supported for fp16 (ComplexHalf) magnitude = magnitude.float() phase = phase.float() # Clamp like original: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/vq/codec_decoder_vocos.py#L138 magnitude = torch.exp(magnitude).clamp(max=1e2) - spectrogram_complex = magnitude * torch.exp(1j * phase) + # ``polar(magnitude, phase)`` is ``magnitude * exp(1j * phase)`` + spectrogram_complex = torch.polar(magnitude, phase) # Back to audio (ISTFT with manual "same" padding: torch.istft lacks a native same-padding mode, # so we use irfft + fold with explicit pre-computed padding to replicate it) @@ -797,7 +798,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: output_size=(1, output_size), kernel_size=(1, self.n_fft), stride=(1, self.hop_length), - ).squeeze()[self.padding : -self.padding] + )[0, 0, 0, self.padding : -self.padding] # Clamp as expected by original: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/vq/codec_decoder_vocos.py#L82 window_envelope = window_envelope.clamp(min=1e-11) audio = audio / window_envelope diff --git a/src/transformers/models/xcodec2/modular_xcodec2.py b/src/transformers/models/xcodec2/modular_xcodec2.py index 6716c1b2df44..687315c6f856 100644 --- a/src/transformers/models/xcodec2/modular_xcodec2.py +++ b/src/transformers/models/xcodec2/modular_xcodec2.py @@ -446,12 +446,13 @@ def __init__(self, config: Xcodec2Config): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: stft_pred = self.linear(hidden_states).transpose(1, 2) magnitude, phase = stft_pred.chunk(2, dim=1) - # Cast to float32: complex exponential and irfft are not supported for fp16 (ComplexHalf) + # Cast to float32: complex tensors and irfft are not supported for fp16 (ComplexHalf) magnitude = magnitude.float() phase = phase.float() # Clamp like original: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/vq/codec_decoder_vocos.py#L138 magnitude = torch.exp(magnitude).clamp(max=1e2) - spectrogram_complex = magnitude * torch.exp(1j * phase) + # ``polar(magnitude, phase)`` is ``magnitude * exp(1j * phase)`` + spectrogram_complex = torch.polar(magnitude, phase) # Back to audio (ISTFT with manual "same" padding: torch.istft lacks a native same-padding mode, # so we use irfft + fold with explicit pre-computed padding to replicate it) @@ -472,7 +473,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: output_size=(1, output_size), kernel_size=(1, self.n_fft), stride=(1, self.hop_length), - ).squeeze()[self.padding : -self.padding] + )[0, 0, 0, self.padding : -self.padding] # Clamp as expected by original: https://huggingface.co/HKUSTAudio/xcodec2/blob/main/vq/codec_decoder_vocos.py#L82 window_envelope = window_envelope.clamp(min=1e-11) audio = audio / window_envelope diff --git a/src/transformers/models/xlnet/modeling_xlnet.py b/src/transformers/models/xlnet/modeling_xlnet.py index ce29b5dea44b..c261ef0c2dd4 100755 --- a/src/transformers/models/xlnet/modeling_xlnet.py +++ b/src/transformers/models/xlnet/modeling_xlnet.py @@ -1090,7 +1090,9 @@ def forward( if data_mask is not None: # all mems can be attended to if mlen > 0: - mems_mask = torch.zeros([data_mask.shape[0], mlen, bsz]).to(data_mask) + mems_mask = torch.zeros( + [data_mask.shape[0], mlen, bsz], device=data_mask.device, dtype=data_mask.dtype + ) data_mask = torch.cat([mems_mask, data_mask], dim=1) if attn_mask is None: attn_mask = data_mask[:, :, :, None] @@ -1101,9 +1103,11 @@ def forward( attn_mask = (attn_mask > 0).to(dtype_float) if attn_mask is not None: - non_tgt_mask = -torch.eye(qlen).to(attn_mask) + non_tgt_mask = -torch.eye(qlen, device=attn_mask.device, dtype=attn_mask.dtype) if mlen > 0: - non_tgt_mask = torch.cat([torch.zeros([qlen, mlen]).to(attn_mask), non_tgt_mask], dim=-1) + non_tgt_mask = torch.cat( + [torch.zeros([qlen, mlen], device=attn_mask.device, dtype=attn_mask.dtype), non_tgt_mask], dim=-1 + ) non_tgt_mask = ((attn_mask + non_tgt_mask[:, :, None, None]) > 0).to(attn_mask) else: non_tgt_mask = None diff --git a/src/transformers/pytorch_utils.py b/src/transformers/pytorch_utils.py index 444c8bc457bb..ccaff43c8520 100644 --- a/src/transformers/pytorch_utils.py +++ b/src/transformers/pytorch_utils.py @@ -118,7 +118,9 @@ def __repr__(self) -> str: def forward(self, x): size_out = x.size()[:-1] + (self.nf,) - x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight) + # reshape, not view: export passes (e.g. ExecuTorch's dim-order pass) can hand this a + # non-contiguous input, which view rejects; reshape is free when contiguous. + x = torch.addmm(self.bias, x.reshape(-1, x.size(-1)), self.weight) x = x.view(size_out) return x diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py index d7b69a8cdc4b..25d8cc91cffc 100644 --- a/src/transformers/testing_utils.py +++ b/src/transformers/testing_utils.py @@ -127,6 +127,7 @@ is_onnxruntime_available, is_onnxscript_available, is_openai_available, + is_openvino_available, is_optimum_available, is_optimum_quanto_available, is_pandas_available, @@ -624,6 +625,10 @@ def require_executorch(test_case): return unittest.skipUnless(is_executorch_available(), "test requires ExecuTorch")(test_case) +def require_openvino(test_case): + return unittest.skipUnless(is_openvino_available(), "test requires OpenVINO")(test_case) + + def require_timm(test_case): """ Decorator marking a test that requires Timm. diff --git a/src/transformers/utils/__init__.py b/src/transformers/utils/__init__.py index d95af51a7df6..25b49a533260 100644 --- a/src/transformers/utils/__init__.py +++ b/src/transformers/utils/__init__.py @@ -181,6 +181,7 @@ is_onnxruntime_available, is_onnxscript_available, is_openai_available, + is_openvino_available, is_optimum_available, is_optimum_quanto_available, is_pandas_available, diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index 2be1ac3fdb61..13ce33e37dda 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -915,6 +915,11 @@ def is_onnxscript_available() -> bool: return _is_package_available("onnxscript")[0] +@lru_cache +def is_openvino_available() -> bool: + return _is_package_available("openvino")[0] + + @lru_cache def is_onnxruntime_available() -> bool: return _is_package_available("onnxruntime")[0] or _is_package_available("onnxruntime-gpu")[0] diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 18affeac92e9..aeda6ba9e892 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -24,6 +24,7 @@ from transformers.exporters.exporter_dynamo import DynamoConfig, DynamoExporter from transformers.exporters.exporter_executorch import ExecutorchConfig, ExecutorchExporter from transformers.exporters.exporter_onnx import OnnxConfig, OnnxExporter +from transformers.exporters.exporter_openvino import OpenVINOConfig, OpenVINOExporter from transformers.exporters.utils import ( decompose_for_generation, decompose_multimodal, @@ -34,6 +35,7 @@ require_executorch, require_onnxruntime, require_onnxscript, + require_openvino, set_config_for_less_flaky_test, set_model_for_less_flaky_test, slow, @@ -57,53 +59,16 @@ EXPORT_SKIPS: dict[str, dict[str, str]] = { # Every backend, every variant. "all": { - "VideoMAEForPreTraining": ( - "Computes loss even when `return_loss=False`, hitting a data-dependent guard in " - "`mse_loss`. TODO: skip loss when labels aren't provided." - ), - "OpenAIPrivacyFilterModel": ( - "`get_correct_experts_implementation` defaults to `eager` because the model is " - "sensitive to accumulation order. Eager experts forward iterates over " - "`expert_hit.nonzero()` (data-dependent shape). Users can opt into " - "`set_experts_implementation('batched_mm')` to export." - ), - "OpenAIPrivacyFilterForTokenClassification": ( - "Same root cause as `OpenAIPrivacyFilterModel` — eager experts implementation." - ), }, # Every backend, generate path only. "generate": { - "Blip2ForConditionalGeneration": ( - "`generate()` delegates to the inner language model without calling top-level " - "`forward()`, so `decompose_prefill_decode` can't capture inputs. " - "TODO: route generate through top-level `forward()`." - ), - "InstructBlipForConditionalGeneration": "Same `generate()`-delegation as Blip2.", - "InstructBlipVideoForConditionalGeneration": "Same `generate()`-delegation as Blip2.", - "Kosmos2ForConditionalGeneration": "Same `generate()`-delegation as Blip2.", "RecurrentGemmaForCausalLM": ( - "Stores recurrent/conv state as module attributes (not a `Cache` object); " - "`torch.export` can't carry that state between calls. " - "TODO: refactor to a cache-based SSM pattern (like Mamba/Mamba2)." - ), - "MoshiForConditionalGeneration": ( - "`generate()` creates `blank_user_audio_codes` outside the traced forward and " - "passes it as a kwarg; the resulting ONNX input has mismatched rank (scalar vs 3D). " - "TODO: make `blank_user_audio_codes` part of the model state." - ), - "UdopForConditionalGeneration": ( - "Exported decoder output is missing `attention_mask` vs eager — encoder-decoder " - "cross-attention mask doesn't flow through the generate decomposition correctly." - ), - "VoxtralRealtimeForConditionalGeneration": ( - "Exported prefill drops `past_key_values.*.{keys,values,_sliding_window_tensor}` " - "tensors that eager returns. Plain forward exports work. " - "TODO: align generate-decomposition path with the realtime KV-cache shape." - ), - "Gemma3nForConditionalGeneration": ( - "KV-shared layers (`num_kv_shared_layers`) reuse cache entries from earlier layers; " - "exported prefill returns only `logits` while eager surfaces the populated KV cache. " - "Same shape as Voxtral. TODO: align the generate-decomposition path." + "Keeps RG-LRU/conv state as plain module attributes (not a `Cache` passed via " + "`past_key_values`), so the state cannot be a graph input/output — the exported " + "decode step computes from zero-initialized state and its logits diverge from eager " + "(prefill exports and matches). " + "TODO: migrate the state to `LinearAttentionLayer` entries in `past_key_values` " + "(the qwen3_next / Mamba pattern)." ), }, # ONNX, every variant. @@ -116,12 +81,6 @@ }, # ONNX, generate path only. "onnx.generate": { - "ReformerModelWithLMHead": ( - "Chunked local attention exports a Constant idx that exceeds the cached-keys axis " - "length under static decode (prefill+1 token, seq=17 vs chunked axis of 16). The same " - "computation stays symbolic under dynamic so ORT can't pre-validate it. The other " - "three Reformer-local-attn ONNX variants pass." - ), }, # ONNX, dynamic-shape only. "onnx.dynamic": { @@ -143,69 +102,12 @@ # ExecuTorch — lowering failures grouped by root cause; see the first entry of each # `Same ... as` chain for the full description. "executorch": { - "BarkFineModel": ( - "ExecuTorch memory planning miscomputes the tensor spec (`buffer of size N, expected nbytes of M`) — a dtype-size mismatch in the lowered program." - ), - "ClvpModelForConditionalGeneration": ( - "A pass-through output aliases an input (`Output node is already in the inputs`)." - ), - "ColQwen2ForRetrieval": ( - "ExecuTorch dim-order lowering requires a copying view (`Cannot view a tensor ... with shape/strides`)." - ), - "DabDetrModel": ("XNNPACK partitioner: `Attempting to convert non-NHWC compatible node to NHWC`."), - "DabDetrForObjectDetection": "Same `nhwc` failure as `DabDetrModel`.", - "Ernie4_5_VLMoeModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Ernie4_5_VLMoeForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "FlavaForPreTraining": ("XNNPACK partitioner: `Invalid partition, found dependency cycles`."), - "GPT2Model": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GPT2LMHeadModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GPT2DoubleHeadsModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GPT2ForQuestionAnswering": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GPT2ForSequenceClassification": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GPT2ForTokenClassification": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Gemma3nModel": "Same `spec` failure as `BarkFineModel`.", - "Gemma3nForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "Glm46VModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Glm46VForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Glm4vModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Glm4vForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Glm4vMoeModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Glm4vMoeForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GlmImageModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GlmImageForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GlmOcrModel": "Same `view` failure as `ColQwen2ForRetrieval`.", - "GlmOcrForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", - "InstructBlipModel": "Same `spec` failure as `BarkFineModel`.", - "InstructBlipForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "InstructBlipVideoForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "InstructBlipVideoModel": "Same `spec` failure as `BarkFineModel`.", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", - "MiniMaxM3VLModel": ("Serialization rejects an i64 constant (`bad number for type int32`)."), - "MiniMaxM3SparseForConditionalGeneration": "Same `int32` failure as `MiniMaxM3VLModel`.", - "PPDocLayoutV3ForObjectDetection": ("Delegation drops a referenced weight (`KeyError` on a state-dict key)."), - "PaddleOCRVLForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "PerceptionLMModel": "Same `passthrough` failure as `ClvpModelForConditionalGeneration`.", - "PerceptionLMForConditionalGeneration": "Same `passthrough` failure as `ClvpModelForConditionalGeneration`.", - "Qwen2VLModel": "Same `spec` failure as `BarkFineModel`.", - "Qwen2VLForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "Qwen2_5OmniThinkerForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Qwen2_5_VLModel": "Same `spec` failure as `BarkFineModel`.", - "Qwen2_5_VLForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "Qwen3OmniMoeThinkerForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.", - "Qwen3_5Model": "Same `spec` failure as `BarkFineModel`.", - "Qwen3_5ForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", - "Qwen3_5ForSequenceClassification": "Same `spec` failure as `BarkFineModel`.", - "Qwen3_5ForTokenClassification": "Same `spec` failure as `BarkFineModel`.", - "Qwen3_5MoeModel": "Same `spec` failure as `BarkFineModel`.", - "Qwen3_5MoeForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", }, "executorch.generate": { - "PPFormulaNetForConditionalGeneration": ( - "ExecuTorch memory planning miscomputes the tensor spec (`buffer of size N, expected nbytes of M`) — a dtype-size mismatch in the lowered program." - ), }, "executorch.dynamic": { "BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."), @@ -216,25 +118,21 @@ "BigBirdForQuestionAnswering": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForSequenceClassification": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForTokenClassification": "Same `timeout` failure as `BigBirdModel`.", - "DepthProModel": ( - "`_ViewSpec is incompatible with its base` — mixed shape dynamism between a view and its base." - ), - "DepthProForDepthEstimation": "Same `viewspec` failure as `DepthProModel`.", "DonutSwinModel": ( "ExecuTorch memory planning overflows under unbounded dynamic shapes (`mem_offset does not fit in 64 bits`)." ), "DonutSwinForImageClassification": "Same `overflow` failure as `DonutSwinModel`.", - "Mask2FormerModel": "Same `timeout` failure as `BigBirdModel`.", - "Mask2FormerForUniversalSegmentation": "Same `timeout` failure as `BigBirdModel`.", "MaskFormerModel": "Same `timeout` failure as `BigBirdModel`.", "MaskFormerForInstanceSegmentation": "Same `timeout` failure as `BigBirdModel`.", "MaskFormerSwinModel": "Same `overflow` failure as `DonutSwinModel`.", "MaskFormerSwinBackbone": "Same `overflow` failure as `DonutSwinModel`.", "MllamaModel": "Same `overflow` failure as `DonutSwinModel`.", "MllamaForConditionalGeneration": "Same `overflow` failure as `DonutSwinModel`.", - "PvtModel": "Same `viewspec` failure as `DepthProModel`.", - "PvtForImageClassification": "Same `viewspec` failure as `DepthProModel`.", - "Sam2Model": ("Delegation drops a referenced weight (`KeyError` on a state-dict key)."), + "Sam2Model": ( + "The dropped-weight `KeyError` is fixed, but `torch.export` of the Hiera vision " + "backbone under dynamic shapes then exceeds the 1000s test timeout (same " + "Hiera-backbone dynamic-shape overrun as `Sam2VisionModel`)." + ), "Sam2VisionModel": "Same `timeout` failure as `BigBirdModel`.", "Swin2SRModel": "Same `overflow` failure as `DonutSwinModel`.", "Swin2SRForImageSuperResolution": "Same `overflow` failure as `DonutSwinModel`.", @@ -246,14 +144,16 @@ "Swinv2ForImageClassification": "Same `overflow` failure as `DonutSwinModel`.", "Swinv2ForMaskedImageModeling": "Same `overflow` failure as `DonutSwinModel`.", "Swinv2Backbone": "Same `overflow` failure as `DonutSwinModel`.", - "VitDetModel": "Same `viewspec` failure as `DepthProModel`.", - "VitDetBackbone": "Same `viewspec` failure as `DepthProModel`.", "Wav2Vec2BertForCTC": ("`flatc` schema compilation fails when serializing the program."), "Wav2Vec2BertModel": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", "Wav2Vec2BertForSequenceClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", "Wav2Vec2BertForAudioFrameClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", "Wav2Vec2BertForXVector": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", }, + # OpenVINO, every variant (currently empty — populated as we sweep). + "openvino": {}, + # OpenVINO, dynamic-shape only. + "openvino.dynamic": {}, } @@ -263,36 +163,8 @@ ONNX_DISABLE_OPTIMIZE: dict[str, dict[str, str]] = { - # Disable for every variant. - "all": { - "LayoutLMv2Model": ( - "Detectron2 FPN backbone — onnxscript optimizer drops initializers still referenced " - "by nodes, producing an invalid graph for ORT." - ), - "LayoutLMv2ForSequenceClassification": "Same as `LayoutLMv2Model`.", - "LayoutLMv2ForTokenClassification": "Same as `LayoutLMv2Model`.", - "LayoutLMv2ForQuestionAnswering": "Same as `LayoutLMv2Model`.", - "YolosModel": ( - "Optimizer takes >6 min on the YOLOS detection graph (many small Concat/Slice nodes). " - "`optimize=False` exports in 2s. TODO: revisit when onnxscript's optimizer improves." - ), - "YolosForObjectDetection": "Same as `YolosModel`.", - "PixioModel": "Same dense-small-node optimizer slowdown as YOLOS (~100–290s).", - "SegGptModel": "Same dense-small-node optimizer slowdown as YOLOS.", - "SegGptForImageSegmentation": "Same dense-small-node optimizer slowdown as YOLOS.", - }, # Disable for dynamic-shape only — static benefits from optimisation. - "dynamic": { - "ProphetNetModel": ( - "Onnxscript's `SplitToSequence` constant-folding trips `'NoneType' object has no " - "attribute 'ndim'` under dynamic shapes. Static works after the vectorized " - "`ngram_attention_bias` rewrite." - ), - "ProphetNetForConditionalGeneration": "Same `SplitToSequence` issue as `ProphetNetModel`.", - "ProphetNetDecoder": "Same `SplitToSequence` issue as `ProphetNetModel`.", - "ProphetNetForCausalLM": "Same `SplitToSequence` issue as `ProphetNetModel`.", - "ZoeDepthForDepthEstimation": "Same `SplitToSequence` issue as `ProphetNetModel`.", - }, + "dynamic": {}, } @@ -338,6 +210,66 @@ def _run_onnx_program(onnx_program, inputs) -> dict: return dict(zip(onnx_names, onnx_outputs)) +def _run_openvino_model(ov_model, inputs) -> dict: + """Compile an OpenVINO model and run it, returning outputs as a `{name: array}` dict. + + Feeds the tensor leaves that survived as input ports (stateful folding removes cache + inputs), seeds folded state variables from the sample cache leaves so outputs correspond + to the same inputs eager saw, supplies the identity `beam_idx`, and passes scalar kwargs + through under their FX placeholder names. + """ + import numpy as np + import openvino + + set_seed(1234) + compiled = openvino.compile_model(ov_model, "AUTO") + request = compiled.create_infer_request() + leaves = {path: tensor.cpu() for path, tensor in get_leaf_tensors(inputs).items()} + batch = next(iter(leaves.values())).shape[0] if leaves else 1 + + feed = {} + for port in compiled.inputs: + # Passthrough tensors carry both an input and an output name — check every alias. + for name in port.get_names(): + path = re.sub(r"^input\.", "", name) + if path in leaves: + feed[name] = leaves[path] + elif name == "beam_idx": + feed[name] = np.arange(batch, dtype=np.int32) + elif name in inputs: + feed[name] = np.array(inputs[name]) + else: + continue + break + + # Folded state variables read zeros on the first infer — seed them from the sample leaves + # (cast to the variable's dtype: the exporter may retype state, e.g. i64 lengths to i32). + # The variable id is ``input.output.``. + def _state_path(state): + return state.name[len("input.") : (len(state.name) - len("input.output.")) // 2 + len("input.")] + + for state in request.query_state(): + path = _state_path(state) + if path in leaves: + state.state = openvino.Tensor(leaves[path].numpy().astype(state.state.data.dtype, copy=False)) + + results = request.infer(feed) + outputs = {} + for port in compiled.outputs: + # Compilation may merge a named output tensor with an intermediate that kept its + # numeric id — prefer the human-readable alias over ``get_any_name``'s sorted-first. + names = sorted(port.get_names()) + name = next((n for n in names if not n.isdigit()), names[0]) + outputs[re.sub(r"^output\.", "", name)] = results[port] + + # Folded state tensors are outputs too — read them back so the returned dict covers the + # same leaves eager returns. + for state in request.query_state(): + outputs[_state_path(state)] = state.state.data.copy() + + return outputs + + def _onnx_optimize_enabled(model_class, dynamic: bool) -> bool: """Return whether onnxscript optimisation should run for this model under this shape mode. @@ -512,6 +444,33 @@ def test_onnx_export(self, dynamic): self.assertTrue(onnx_outputs, f"ONNX outputs are empty for {name}.") self.assertEqual(set(onnx_outputs.keys()), set(eager_outputs[name].keys())) + # ──────────────────── OpenVINO tests ───────────────────────── + + @slow + @DYNAMIC_EXPORT_PARAMS + @require_openvino + @pytest.mark.openvino_export_test + @pytest.mark.timeout(EXPORT_TEST_TIMEOUT) + def test_openvino_export(self, dynamic): + """Export each model class to OpenVINO IR and verify output names match eager.""" + self._skip_if_not_exportable() + exporter = OpenVINOExporter() + config = OpenVINOConfig(dynamic=dynamic) + + for model_class in self.all_model_classes: + if self._should_skip(model_class, dynamic=dynamic, backend="openvino"): + continue + + components = self._prepare_export_model_and_inputs(model_class) + eager_outputs = self._collect_eager_outputs(components) + + for name, (model, inputs) in components.items(): + with self.subTest(f"{model_class.__name__}/{name}"): + ov_model = exporter.export(model, inputs, config=config) + ov_outputs = _run_openvino_model(ov_model, inputs) + self.assertTrue(ov_outputs, f"OpenVINO outputs are empty for {name}.") + self.assertEqual(set(ov_outputs.keys()), set(eager_outputs[name].keys())) + # ──────────────────── ExecuTorch tests ─────────────────────── @slow @@ -633,6 +592,33 @@ def test_onnx_export_generate(self, dynamic): self.assertTrue(onnx_outputs, "ONNX outputs are empty.") self.assertEqual(set(onnx_outputs.keys()), set(eager_outputs[name].keys())) + # ──────────────────── OpenVINO tests ───────────────────────── + + @slow + @DYNAMIC_EXPORT_PARAMS + @require_openvino + @pytest.mark.openvino_export_test + @pytest.mark.timeout(EXPORT_TEST_TIMEOUT) + def test_openvino_export_generate(self, dynamic): + """Export prefill and decode stages to OpenVINO IR and verify output names match eager.""" + self._skip_if_not_exportable() + exporter = OpenVINOExporter() + config = OpenVINOConfig(dynamic=dynamic) + + for model_class in self.all_generative_model_classes: + if self._should_skip(model_class, generate=True, dynamic=dynamic, backend="openvino"): + continue + + components = self._prepare_export_generate_model_and_inputs(model_class) + eager_outputs = self._collect_eager_outputs(components) + + for name, (model, inputs) in components.items(): + with self.subTest(f"{model_class.__name__}/{name}"): + ov_model = exporter.export(model, inputs, config=config) + ov_outputs = _run_openvino_model(ov_model, inputs) + self.assertTrue(ov_outputs, "OpenVINO outputs are empty.") + self.assertEqual(set(ov_outputs.keys()), set(eager_outputs[name].keys())) + # ──────────────────── ExecuTorch tests ─────────────────────── @slow diff --git a/tests/models/openai_privacy_filter/test_modeling_openai_privacy_filter.py b/tests/models/openai_privacy_filter/test_modeling_openai_privacy_filter.py index 8dba6125b5fc..13b07f6cd1ca 100644 --- a/tests/models/openai_privacy_filter/test_modeling_openai_privacy_filter.py +++ b/tests/models/openai_privacy_filter/test_modeling_openai_privacy_filter.py @@ -169,6 +169,7 @@ class OpenAIPrivacyFilterModelTest(ModelTesterMixin, PipelineTesterMixin, unitte if is_torch_available() else () ) + test_torch_exportable = False # eager experts (the model's default) iterate over `nonzero()`, untraceable pipeline_model_mapping = ( { "feature-extraction": OpenAIPrivacyFilterModel, From efff4e12ab55abbb0cbe1d40c75695eafdaa795d Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 3 Jul 2026 13:45:16 +0200 Subject: [PATCH 02/14] style --- .../exporters/exporter_executorch.py | 31 ++++++++++----- .../models/informer/modular_informer.py | 2 +- tests/exporters/test_export.py | 39 +++++++------------ 3 files changed, 37 insertions(+), 35 deletions(-) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 9c4b77f6b39a..ee2e014c0f72 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -442,22 +442,33 @@ def patch(self, *sizes): "executorch.exir.passes.sym_shape_eval_pass.eval_upper_bound", ) def _patch_eval_upper_bound(original): - """Constraint-based bound, then trace hint, then ``_MAX_DIM_FLOOR``. - - Constraint propagation returns ``int_oo`` for compound expressions whose - constraints don't compose (e.g. ``((s43*s53)//s70)``) or for sums of - unbacked symbols (e.g. MoE per-expert cats ``u320+u321+...``); the - fallbacks guarantee an ``int`` so ``ConstraintBasedSymShapeEvalPass`` - doesn't raise. + """Constraint-based bound, clamped to a trace-hint-proportional cap. + + Constraint propagation misbehaves on compound expressions in two ways, and + ``ConstraintBasedSymShapeEvalPass`` needs an ``int`` in both cases: + + - It returns ``int_oo`` when constraints don't compose (e.g. ``((s43*s53)//s70)``) + or for sums of unbacked symbols (e.g. MoE per-expert cats ``u320+u321+...``). + - It returns absurdly large *finite* bounds for floordiv ratios: interval + arithmetic evaluates ``x // (x // 2)`` (window-count ratios in the Swin family, + true value 2) as ``upper(x) // lower(x // 2)``, e.g. ``513 // 1``. These + ratios appear squared in window-partition reshapes and compound across + stages, so worst-case tensor sizes reach ~2^63 bytes and ExecuTorch's memory + planner overflows (``mem_offset does not fit in 64 bits``). + + Clamp every symbolic bound to ``max(hint * _MAX_DIM_MULTIPLIER, _MAX_DIM_FLOOR)`` + — the same trace-proportional heuristic ``_fix_range_constraints`` applies to the + per-symbol ranges — so planned buffers stay proportional to the sampled inputs. """ from executorch.exir.sym_util import eval_expr def patch(maybe_symint): + if isinstance(maybe_symint, int): + return maybe_symint result = original(maybe_symint) - if isinstance(result, int): - return result hint = eval_expr(maybe_symint) - return hint if isinstance(hint, int) else _MAX_DIM_FLOOR + cap = max(hint * _MAX_DIM_MULTIPLIER, _MAX_DIM_FLOOR) if isinstance(hint, int) else _MAX_DIM_FLOOR + return min(result, cap) if isinstance(result, int) else cap return patch diff --git a/src/transformers/models/informer/modular_informer.py b/src/transformers/models/informer/modular_informer.py index 99ed7c5ebd6e..6ab6b82448a1 100644 --- a/src/transformers/models/informer/modular_informer.py +++ b/src/transformers/models/informer/modular_informer.py @@ -279,7 +279,7 @@ def forward( if top_u_sparsity_measurement is not None: # update context: copy the attention output to the context at top_u_sparsity_measurement index - dim_for_slice = torch.arange(context.size(0)).unsqueeze(-1) + dim_for_slice = torch.arange(context.size(0), device=context.device).unsqueeze(-1) context[dim_for_slice, top_u_sparsity_measurement, :] = attn_output attn_output = context diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index aeda6ba9e892..8a66deb1b8ab 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -58,8 +58,7 @@ EXPORT_SKIPS: dict[str, dict[str, str]] = { # Every backend, every variant. - "all": { - }, + "all": {}, # Every backend, generate path only. "generate": { "RecurrentGemmaForCausalLM": ( @@ -80,8 +79,7 @@ ), }, # ONNX, generate path only. - "onnx.generate": { - }, + "onnx.generate": {}, # ONNX, dynamic-shape only. "onnx.dynamic": { "GroundingDinoModel": ( @@ -107,8 +105,7 @@ "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", }, - "executorch.generate": { - }, + "executorch.generate": {}, "executorch.dynamic": { "BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."), "BigBirdForPreTraining": "Same `timeout` failure as `BigBirdModel`.", @@ -118,32 +115,26 @@ "BigBirdForQuestionAnswering": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForSequenceClassification": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForTokenClassification": "Same `timeout` failure as `BigBirdModel`.", - "DonutSwinModel": ( - "ExecuTorch memory planning overflows under unbounded dynamic shapes (`mem_offset does not fit in 64 bits`)." - ), - "DonutSwinForImageClassification": "Same `overflow` failure as `DonutSwinModel`.", "MaskFormerModel": "Same `timeout` failure as `BigBirdModel`.", "MaskFormerForInstanceSegmentation": "Same `timeout` failure as `BigBirdModel`.", - "MaskFormerSwinModel": "Same `overflow` failure as `DonutSwinModel`.", - "MaskFormerSwinBackbone": "Same `overflow` failure as `DonutSwinModel`.", - "MllamaModel": "Same `overflow` failure as `DonutSwinModel`.", - "MllamaForConditionalGeneration": "Same `overflow` failure as `DonutSwinModel`.", + "MaskFormerSwinModel": ( + "ExecuTorch memory planning overflows under unbounded dynamic shapes (`mem_offset does not fit in 64 bits`)." + ), + "MaskFormerSwinBackbone": "Same `overflow` failure as `MaskFormerSwinModel`.", + "MllamaModel": "Same `overflow` failure as `MaskFormerSwinModel`.", + "MllamaForConditionalGeneration": "Same `overflow` failure as `MaskFormerSwinModel`.", "Sam2Model": ( "The dropped-weight `KeyError` is fixed, but `torch.export` of the Hiera vision " "backbone under dynamic shapes then exceeds the 1000s test timeout (same " "Hiera-backbone dynamic-shape overrun as `Sam2VisionModel`)." ), "Sam2VisionModel": "Same `timeout` failure as `BigBirdModel`.", - "Swin2SRModel": "Same `overflow` failure as `DonutSwinModel`.", - "Swin2SRForImageSuperResolution": "Same `overflow` failure as `DonutSwinModel`.", - "SwinModel": "Same `overflow` failure as `DonutSwinModel`.", - "SwinBackbone": "Same `overflow` failure as `DonutSwinModel`.", - "SwinForImageClassification": "Same `overflow` failure as `DonutSwinModel`.", - "SwinForMaskedImageModeling": "Same `overflow` failure as `DonutSwinModel`.", - "Swinv2Model": "Same `overflow` failure as `DonutSwinModel`.", - "Swinv2ForImageClassification": "Same `overflow` failure as `DonutSwinModel`.", - "Swinv2ForMaskedImageModeling": "Same `overflow` failure as `DonutSwinModel`.", - "Swinv2Backbone": "Same `overflow` failure as `DonutSwinModel`.", + "Swin2SRModel": "Same `overflow` failure as `MaskFormerSwinModel`.", + "Swin2SRForImageSuperResolution": "Same `overflow` failure as `MaskFormerSwinModel`.", + "Swinv2Model": "Same `overflow` failure as `MaskFormerSwinModel`.", + "Swinv2ForImageClassification": "Same `overflow` failure as `MaskFormerSwinModel`.", + "Swinv2ForMaskedImageModeling": "Same `overflow` failure as `MaskFormerSwinModel`.", + "Swinv2Backbone": "Same `overflow` failure as `MaskFormerSwinModel`.", "Wav2Vec2BertForCTC": ("`flatc` schema compilation fails when serializing the program."), "Wav2Vec2BertModel": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", "Wav2Vec2BertForSequenceClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", From 17349189383120ad976f39a0a5d4c50bcf0a2576 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 09:09:12 +0200 Subject: [PATCH 03/14] more OV end ET fixes --- .../exporters/exporter_executorch.py | 24 +++++++++++ .../exporters/exporter_openvino.py | 41 ++++++++++++++----- 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index ee2e014c0f72..721e021c6ba1 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -1183,6 +1183,30 @@ def _fix_range_constraints(exported_program: ExportedProgram) -> None: ) +@register_fx_program_fix("executorch") +def _drop_runtime_asserts(exported_program: ExportedProgram) -> None: + """Drop ``_assert_scalar`` / ``_assert_tensor_metadata`` runtime asserts before lowering. + + ``_assert_scalar`` lowers a ``torch._check`` on an unbacked symint (e.g. the image-token + count in ``get_placeholder_mask``) into a ``cast_symbool_to_symint`` + ``eq`` chain whose + ``Piecewise`` result the ``_ModuleStackTracer`` used by ``to_edge_transform_and_lower``'s + decomposition pass cannot proxy (``... is not tracked with proxy``). The range facts these + asserts encode survive on ``exported_program.range_constraints`` (further capped by + ``_fix_range_constraints``), so dropping the nodes (and the now-dead symint feeders) is safe. + """ + for module in exported_program.graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in list(module.graph.nodes): + if node.op == "call_function" and node.target in ( + torch.ops.aten._assert_tensor_metadata.default, + torch.ops.aten._assert_scalar.default, + ): + module.graph.erase_node(node) + module.graph.eliminate_dead_code() + module.recompile() + + @register_fx_program_fix("executorch") def _fix_missing_placeholder_vals(exported_program: ExportedProgram) -> None: """Ensure parameter/buffer/lifted-constant placeholders have a tensor ``meta["val"]``. diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index 42d499e0fb59..dd54539a3fd8 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -120,7 +120,7 @@ def export( with torch.no_grad(), patch_model_outputs(model) as (inputs_names, outputs_names), apply_patches("openvino"): exported_program: ExportedProgram = super().export(model, sample_inputs, config=config) - _drop_metadata_asserts(exported_program.graph_module) + _drop_runtime_asserts(exported_program.graph_module) # Run OV's own decomposition pass up front and decode the RESULT — handing the # ``ExportedProgram`` to ``convert_model`` would re-run it internally, regenerating node # names and discarding every fix applied below. @@ -565,20 +565,28 @@ def _pin_state_update_shapes(ov_model: openvino.Model) -> None: _OV_NAME_OK = re.compile(r"_\d+$") -def _drop_metadata_asserts(graph_module) -> None: - """Drop ``_assert_tensor_metadata`` nodes before the decomposition replay. +def _drop_runtime_asserts(graph_module) -> None: + """Drop ``_assert_tensor_metadata`` / ``_assert_scalar`` runtime asserts before the replay. - These asserts re-check trace-time dtypes/devices and fail the replay once other stages - have legitimately changed them. Scalar asserts are deliberately KEPT at this point — they - carry the ``torch._check`` range facts (``u >= 0``) that unbacked-symint guards inside the - replay rely on (splinter) — and are removed post-decomposition by ``_fix_drop_assert_ops``. + ``_assert_tensor_metadata`` re-checks trace-time dtypes/devices and fails the replay once + other stages have legitimately changed them. ``_assert_scalar`` lowers a ``torch._check`` + on an unbacked symint (e.g. the image-token count in ``get_placeholder_mask``) into a + ``cast_symbool_to_symint`` + ``eq`` chain whose ``Piecewise`` result OV's ``_ModuleStackTracer`` + cannot proxy, crashing the replay (``... is not tracked with proxy``). The range facts these + asserts encode survive on ``exported_program.range_constraints``, so dropping the nodes (and + the now-dead symint feeders via ``eliminate_dead_code``) is safe. ``_fix_drop_assert_ops`` + still removes any that reappear post-decomposition. """ for module in graph_module.modules(): if not isinstance(module, torch.fx.GraphModule): continue for node in list(module.graph.nodes): - if node.op == "call_function" and node.target is torch.ops.aten._assert_tensor_metadata.default: + if node.op == "call_function" and node.target in ( + torch.ops.aten._assert_tensor_metadata.default, + torch.ops.aten._assert_scalar.default, + ): module.graph.erase_node(node) + module.graph.eliminate_dead_code() module.recompile() @@ -906,17 +914,28 @@ def _fix_scatter_reduce(gm, node): k_shape = [1] * (ndim + 1) k_shape[d] = -1 with gm.graph.inserting_before(node): + # ``k_size`` is symbolic under dynamic shapes (e.g. BLT's ``max_num_patches``); baking + # the ``SymInt`` as an ``arange`` literal makes OV decode it as a malformed inlined + # constant. Feed the dimension through a ``sym_size`` node so it stays a real Range input. + arange_size = k_size if isinstance(k_size, int) else gm.graph.call_function( + torch.ops.aten.sym_size.int, args=(self_arg, d) + ) arange = gm.graph.call_function( - torch.ops.aten.arange.default, args=(k_size,), kwargs={"device": self_val.device} + torch.ops.aten.arange.default, args=(arange_size,), kwargs={"device": self_val.device} ) k_range = gm.graph.call_function(torch.ops.aten.view.default, args=(arange, k_shape)) index_unsq = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(index, d)) mask = gm.graph.call_function(torch.ops.aten.eq.Tensor, args=(index_unsq, k_range)) src_unsq = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(src, d)) - masked = gm.graph.call_function(torch.ops.aten.where.ScalarOther, args=(mask, src_unsq, min_value)) + # OV's frontend has no ``where.ScalarOther`` translation, so materialise the scalar + # branches as 0-dim tensors and use ``where.self`` (broadcasts the same way). + scalar_kwargs = {"dtype": src_val.dtype, "device": src_val.device} + min_tensor = gm.graph.call_function(torch.ops.aten.scalar_tensor.default, args=(min_value,), kwargs=scalar_kwargs) + masked = gm.graph.call_function(torch.ops.aten.where.self, args=(mask, src_unsq, min_tensor)) maxes = gm.graph.call_function(torch.ops.aten.amax.default, args=(masked, [d + 1])) any_match = gm.graph.call_function(torch.ops.aten.any.dim, args=(mask, d + 1)) - result = gm.graph.call_function(torch.ops.aten.where.ScalarOther, args=(any_match, maxes, 0.0)) + zero_tensor = gm.graph.call_function(torch.ops.aten.scalar_tensor.default, args=(0.0,), kwargs=scalar_kwargs) + result = gm.graph.call_function(torch.ops.aten.where.self, args=(any_match, maxes, zero_tensor)) result.meta.update(node.meta) node.replace_all_uses_with(result) gm.graph.erase_node(node) From c60e49eaa0c1a9624d87c0181cb6b4f9ae6a956d Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 09:23:38 +0200 Subject: [PATCH 04/14] style --- src/transformers/exporters/exporter_openvino.py | 14 ++++++++++---- tests/exporters/test_export.py | 9 ++++++--- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index dd54539a3fd8..54b7e5771372 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -917,8 +917,10 @@ def _fix_scatter_reduce(gm, node): # ``k_size`` is symbolic under dynamic shapes (e.g. BLT's ``max_num_patches``); baking # the ``SymInt`` as an ``arange`` literal makes OV decode it as a malformed inlined # constant. Feed the dimension through a ``sym_size`` node so it stays a real Range input. - arange_size = k_size if isinstance(k_size, int) else gm.graph.call_function( - torch.ops.aten.sym_size.int, args=(self_arg, d) + arange_size = ( + k_size + if isinstance(k_size, int) + else gm.graph.call_function(torch.ops.aten.sym_size.int, args=(self_arg, d)) ) arange = gm.graph.call_function( torch.ops.aten.arange.default, args=(arange_size,), kwargs={"device": self_val.device} @@ -930,11 +932,15 @@ def _fix_scatter_reduce(gm, node): # OV's frontend has no ``where.ScalarOther`` translation, so materialise the scalar # branches as 0-dim tensors and use ``where.self`` (broadcasts the same way). scalar_kwargs = {"dtype": src_val.dtype, "device": src_val.device} - min_tensor = gm.graph.call_function(torch.ops.aten.scalar_tensor.default, args=(min_value,), kwargs=scalar_kwargs) + min_tensor = gm.graph.call_function( + torch.ops.aten.scalar_tensor.default, args=(min_value,), kwargs=scalar_kwargs + ) masked = gm.graph.call_function(torch.ops.aten.where.self, args=(mask, src_unsq, min_tensor)) maxes = gm.graph.call_function(torch.ops.aten.amax.default, args=(masked, [d + 1])) any_match = gm.graph.call_function(torch.ops.aten.any.dim, args=(mask, d + 1)) - zero_tensor = gm.graph.call_function(torch.ops.aten.scalar_tensor.default, args=(0.0,), kwargs=scalar_kwargs) + zero_tensor = gm.graph.call_function( + torch.ops.aten.scalar_tensor.default, args=(0.0,), kwargs=scalar_kwargs + ) result = gm.graph.call_function(torch.ops.aten.where.self, args=(any_match, maxes, zero_tensor)) result.meta.update(node.meta) node.replace_all_uses_with(result) diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 8a66deb1b8ab..f07b6e2e58a4 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -97,15 +97,16 @@ ), "Sam2Model": "Same Hiera-backbone dynamic-shape budget overrun as `Sam2VisionModel`.", }, - # ExecuTorch — lowering failures grouped by root cause; see the first entry of each - # `Same ... as` chain for the full description. + # ExecuTorch, every variant. "executorch": { "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", }, + # ExecuTorch, generate path only. "executorch.generate": {}, + # ExecuTorch, dynamic-shape only. "executorch.dynamic": { "BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."), "BigBirdForPreTraining": "Same `timeout` failure as `BigBirdModel`.", @@ -141,8 +142,10 @@ "Wav2Vec2BertForAudioFrameClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", "Wav2Vec2BertForXVector": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", }, - # OpenVINO, every variant (currently empty — populated as we sweep). + # OpenVINO, every variant. "openvino": {}, + # OpenVINO, generate path only. + "openvino.generate": {}, # OpenVINO, dynamic-shape only. "openvino.dynamic": {}, } From 6c9e7177b792d44fcea3c9a44452b1a4913b7fb7 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 11:44:34 +0200 Subject: [PATCH 05/14] more fixes and targeted skips --- src/transformers/exporters/exporter_dynamo.py | 16 +++ .../exporters/exporter_executorch.py | 52 ++++++- .../exporters/exporter_openvino.py | 66 ++++++++- .../models/chameleon/modeling_chameleon.py | 5 +- .../mask2former/modeling_mask2former.py | 17 +-- .../models/videomae/modeling_videomae.py | 7 +- tests/exporters/test_export.py | 128 +++++++++--------- 7 files changed, 201 insertions(+), 90 deletions(-) diff --git a/src/transformers/exporters/exporter_dynamo.py b/src/transformers/exporters/exporter_dynamo.py index 152387c39b06..a14a4f7ef128 100644 --- a/src/transformers/exporters/exporter_dynamo.py +++ b/src/transformers/exporters/exporter_dynamo.py @@ -449,6 +449,22 @@ def _flatten_to_context(obj: Any, tensors: list) -> Any: if isinstance(obj, torch.layout): return {"_t": "layout", "n": str(obj).removeprefix("torch.")} if isinstance(obj, (torch.SymInt, torch.SymFloat, torch.SymBool)): + # A Sym* that has already specialized to a concrete constant is not a genuine dynamic + # graph output — bake it as a plain scalar instead of a leaf. Leaving it as a leaf makes + # flatten non-deterministic: the same field can be a constant-valued SymInt at one trace + # point (the dynamo out_spec capture) and an already-materialized python scalar at another + # (the aot-decomposition retrace), so the leaf count flips and `treespec.unflatten` fails + # with an off-by-one. deepseek_v4 hits this via two sibling cache layers + # (DeepseekV4HCACache / DeepseekV4CSACache) sharing a `cumulative_length` counter that one + # path leaves as a constant SymInt and the other as a python int. + if isinstance(obj, torch.SymBool): + const = obj.node.maybe_as_bool() + elif isinstance(obj, torch.SymFloat): + const = obj.node.maybe_as_float() + else: + const = obj.node.maybe_as_int() + if const is not None: + return const idx = len(tensors) tensors.append(obj) return {"_t": "sym", "i": idx} diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 721e021c6ba1..90d3db41cd33 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -356,14 +356,35 @@ def _patch_scaled_dot_product_attention(original): ``(L, N, H, E)`` buffer — invalidating those recorded views (``Cannot view a tensor with shape/strides``). ``clone`` records unconditionally and re-executes correctly under either layout, normalizing the strides the rest of the graph was recorded against. + + The eager fallback also fires on **any** device when ``attn_mask`` has a data-dependent + (unbacked) batch dim — the Idefics2/3 / SmolVLM vision tower drops padding images via + ``pixel_values[real_images_inds]`` (a boolean index → unbacked ``u0`` image count), so the + vision attention mask carries batch ``u0``. ``to_edge_transform_and_lower`` decomposes the + surviving ``aten.scaled_dot_product_attention`` node through the SDPA math CIA kernel, which + guards ``Eq(u0, 1)`` on the mask's batch (broadcast-vs-not) and raises + ``GuardOnDataDependentSymNode``. The manual matmul+softmax path masks against ``attn_weight`` + (both batch ``u0``) with plain broadcasting, so no ``Eq(u0, 1)`` guard is needed and no SDPA + node survives to be re-decomposed. """ + from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols + + def _has_unbacked_batch(t): + # True when ``t``'s batch dim is a data-dependent (unbacked, ``u*``) SymInt. + if t is None or t.ndim == 0: + return False + batch = t.shape[0] + return isinstance(batch, torch.SymInt) and bool(free_unbacked_symbols(batch.node.expr)) def patch(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, **kwargs): - needs_eager_attention = query.device.type == "cuda" and ( - kwargs.get("enable_gqa", False) - or query.shape[-1] != value.shape[-1] - or (attn_mask is not None and attn_mask.is_floating_point()) - ) + needs_eager_attention = ( + query.device.type == "cuda" + and ( + kwargs.get("enable_gqa", False) + or query.shape[-1] != value.shape[-1] + or (attn_mask is not None and attn_mask.is_floating_point()) + ) + ) or (attn_mask is not None and _has_unbacked_batch(attn_mask)) if needs_eager_attention: scale_factor = scale if scale is not None else math.sqrt(query.shape[-1]) ** -1 if key.shape[1] != query.shape[1]: @@ -376,7 +397,10 @@ def patch(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, sca causal_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril() attn_weight = attn_weight.masked_fill(~causal_mask, float("-inf")) if attn_mask is not None: - attn_weight = attn_weight + attn_mask + if attn_mask.dtype == torch.bool: + attn_weight = attn_weight.masked_fill(~attn_mask, float("-inf")) + else: + attn_weight = attn_weight + attn_mask attn_weight = torch.nn.functional.softmax(attn_weight, dim=-1) return torch.matmul(attn_weight, value) with sdpa_kernel(SDPBackend.MATH): @@ -1188,11 +1212,17 @@ def _drop_runtime_asserts(exported_program: ExportedProgram) -> None: """Drop ``_assert_scalar`` / ``_assert_tensor_metadata`` runtime asserts before lowering. ``_assert_scalar`` lowers a ``torch._check`` on an unbacked symint (e.g. the image-token - count in ``get_placeholder_mask``) into a ``cast_symbool_to_symint`` + ``eq`` chain whose + count in ``get_placeholder_mask``, or SmolVLM's ``inputs_merger`` registering ``Eq(u2, 1)`` + off the unbacked real-image count) into a ``cast_symbool_to_symint`` + ``eq`` chain whose ``Piecewise`` result the ``_ModuleStackTracer`` used by ``to_edge_transform_and_lower``'s decomposition pass cannot proxy (``... is not tracked with proxy``). The range facts these asserts encode survive on ``exported_program.range_constraints`` (further capped by ``_fix_range_constraints``), so dropping the nodes (and the now-dead symint feeders) is safe. + + Erasing the nodes is not sufficient on its own: ``to_edge``'s internal re-export re-runs + ``insert_deferred_runtime_asserts``, which reads ``shape_env.deferred_runtime_asserts`` to + decide what to regenerate — so that registry has to be cleared too, else the equality assert + reappears and trips the tracer again. """ for module in exported_program.graph_module.modules(): if not isinstance(module, torch.fx.GraphModule): @@ -1206,6 +1236,14 @@ def _drop_runtime_asserts(exported_program: ExportedProgram) -> None: module.graph.eliminate_dead_code() module.recompile() + for node in exported_program.graph_module.graph.nodes: + val = node.meta.get("val") + if isinstance(val, torch.Tensor) and getattr(val, "fake_mode", None) is not None: + shape_env = val.fake_mode.shape_env + if shape_env is not None: + shape_env.deferred_runtime_asserts.clear() + break # all nodes share the same shape_env + @register_fx_program_fix("executorch") def _fix_missing_placeholder_vals(exported_program: ExportedProgram) -> None: diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index 54b7e5771372..a491c9950faf 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -607,14 +607,17 @@ def _run_openvino_decompositions(exported_program: ExportedProgram) -> ExportedP def _deduplicate_output_args(graph_module) -> None: - """Give repeated graph outputs their own node via a real ``+ 0`` op. + """Give repeated graph outputs their own node via a fold-resistant self-identity op. Two Results sharing one OV tensor crash the translate session's ``is_number`` check: the results-cleanup pass erases the shared tensor's numeric id on the first visit and fails decoding the debug alias on the second. Repeats arise when decomposition collapses the distinction between two output nodes (and ``aten.clone`` is no protection — OV folds it to - identity); ``+ 0`` is the cheapest op that survives translation as a distinct tensor - (constant-folded away afterwards by OV's optimization passes). + identity). The copy must survive OV's neutral-constant elimination: ``add(x, 0)`` gets folded + back to ``x`` (so a duplicate output re-aliases the original — e.g. moshi's ``depth_past_key_values`` + ports collapsing onto the main-cache state buffer and losing their names), whereas ``maximum(x, x)`` + is a self-identity with no neutral constant and survives as a distinct tensor (mirroring the + ``logical_and(x, x)`` used for the bool branch). """ # Ops OV translates as pass-through — their output IS their input's tensor, so an output # arg behind one of these still aliases the underlying node. @@ -634,7 +637,7 @@ def dedup(arg): if val is not None and val.dtype == torch.bool: copy = graph_module.graph.call_function(torch.ops.aten.logical_and.default, args=(source, source)) else: - copy = graph_module.graph.call_function(torch.ops.aten.add.Tensor, args=(source, 0)) + copy = graph_module.graph.call_function(torch.ops.aten.maximum.default, args=(source, source)) copy.meta.update(source.meta) seen.add(source) return copy @@ -1037,6 +1040,42 @@ def _fix_empty_expand(gm, node): return True +@register_fx_node_fix("openvino") +def _fix_view_inferred_dim(gm, node): + """Replace the inferred ``-1`` in an ``aten.view`` target that also carries a symbolic dim. + + OV lowers ``aten.view`` to a ``Reshape`` and infers the ``-1`` dimension from the input's + element count. When another target dim is a runtime ``sym_size`` expression, OV's shape + inference can't reconcile the dynamic dim with the inferred ``-1`` and mis-resolves it — + edgetam/sam3_tracker's mask decoder does ``x.view(pixel_values.shape[0], -1, 8, 8)`` on a + ``[batch, 32, spatial]`` tensor and OV folds the ``-1`` to ``1`` while shifting the other + axes, so the runtime Reshape sees ``(64, 1, 8, 8)`` instead of ``(2, 32, 8, 8)`` and the + pattern product no longer matches the input. Substituting the ``-1`` with its concrete size + from the node's traced output — a static int for every graph that hits this — removes the + inference entirely and leaves OV a fully-determined pattern. + """ + if node.target not in (torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default): + return False + shape = node.args[1] + if not isinstance(shape, (list, tuple)): + return False + minus_one = [i for i, dim in enumerate(shape) if isinstance(dim, int) and dim == -1] + has_symbolic = any(not isinstance(dim, int) for dim in shape) + if len(minus_one) != 1 or not has_symbolic: + return False + out_val = node.meta.get("val") + if out_val is None: + return False + index = minus_one[0] + resolved = out_val.shape[index] + if not isinstance(resolved, int): + return False + new_shape = list(shape) + new_shape[index] = resolved + node.args = (node.args[0], new_shape) + tuple(node.args[2:]) + return True + + # ── Torch patches ─────────────────────────────────────────────────────────── # Each `_patch_*(original)` factory is registered via `@register_patch("openvino", path)` # and reversibly swaps a `torch` op the OV frontend can't lower with a decomposed @@ -1908,6 +1947,23 @@ def _convert_sym_floordiv(context): return [ov_ops.convert(ov_ops.floor(ov_ops.divide(a, b)), "i64").output(0)] +def _convert_sym_truediv(context): + """``a / b`` over SymInts → **float** division, matching Python's ``truediv``. + + OV's ``Divide`` on two integer operands does integer (truncating) division, but Python's + ``/`` always returns a float. granite_speech's chunked-attention reshape computes the merged + batch dim as ``batch * ceil(seq / chunk)`` — traced as ``ceil(truediv(sym_size, 200))``. + With integer operands ``200 / 200`` … ``128 / 200`` truncated to ``0``, so ``ceil(0) == 0`` + and the reshape got a ``0`` batch dim (pattern ``(0, 200, 2, 16)`` vs input ``(2, 200, 32)``). + Promoting integer operands to ``f32`` restores true division so the ``ceil`` rounds up.""" + a, b = context.get_input(0), context.get_input(1) + if a.get_element_type().is_integral(): + a = ov_ops.convert(a, "f32") + if b.get_element_type().is_integral(): + b = ov_ops.convert(b, "f32") + return [ov_ops.divide(a, b).output(0)] + + _OV_CONVERSION_EXTENSIONS: list[Any] = [] if is_openvino_available(): _OV_CONVERSION_EXTENSIONS.extend( @@ -1933,7 +1989,7 @@ def _convert_sym_floordiv(context): ConversionExtension("", _convert_sym_binop(ov_ops.add)), ConversionExtension("", _convert_sym_binop(ov_ops.subtract)), ConversionExtension("", _convert_sym_binop(ov_ops.multiply)), - ConversionExtension("", _convert_sym_binop(ov_ops.divide)), + ConversionExtension("", _convert_sym_truediv), ConversionExtension("", _convert_sym_floordiv), ConversionExtension("", _convert_sym_binop(ov_ops.floor_mod)), ConversionExtension("", _convert_sym_binop(ov_ops.power)), diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index 91132ef2f47e..d57cf3eb4125 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -771,7 +771,10 @@ def _build_img2bpe_mapping_tensor(self, device: torch.device) -> torch.Tensor: # assignment loop that trips FakeTensor propagation during export. keys = torch.tensor(list(self.img2bpe.keys()), dtype=torch.long) values = torch.tensor(list(self.img2bpe.values()), dtype=torch.int) - mapping = torch.zeros(int(keys.max().item()) + 1, dtype=torch.int) + # `max()` over the Python dict keys (a compile-time constant) instead of + # `keys.max().item()`: `.item()` on a tensor yields a data-dependent unbacked symint + # under `torch.export` that `torch.zeros(...)` cannot specialize. + mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int) mapping.scatter_(0, keys, values) self._img2bpe_mapping_cache[device] = mapping.to(device) return self._img2bpe_mapping_cache[device] diff --git a/src/transformers/models/mask2former/modeling_mask2former.py b/src/transformers/models/mask2former/modeling_mask2former.py index ab92b3c955c3..e9623bb9acb5 100644 --- a/src/transformers/models/mask2former/modeling_mask2former.py +++ b/src/transformers/models/mask2former/modeling_mask2former.py @@ -978,19 +978,10 @@ def forward( ) # batch_size, num_queries, n_heads, n_levels, n_points, 2 if reference_points.shape[-1] == 2: - # ``torch.tensor([[w, h], ...], device=cuda)`` where dims are SymInts materialises - # on CPU and moves — trips FakeTensor device propagation during ``torch.export``. - # Build via ``torch.stack`` on device-side scalars instead. - offset_normalizer = torch.stack( - [ - torch.stack( - [ - torch.as_tensor(shape[1], dtype=torch.long, device=reference_points.device), - torch.as_tensor(shape[0], dtype=torch.long, device=reference_points.device), - ] - ) - for shape in spatial_shapes_list - ] + offset_normalizer = torch.tensor( + [[shape[1], shape[0]] for shape in spatial_shapes_list], + dtype=torch.long, + device=reference_points.device, ) sampling_locations = ( reference_points[:, :, None, :, None, :] diff --git a/src/transformers/models/videomae/modeling_videomae.py b/src/transformers/models/videomae/modeling_videomae.py index c85f95ff1e95..a6ad7d85f8d5 100755 --- a/src/transformers/models/videomae/modeling_videomae.py +++ b/src/transformers/models/videomae/modeling_videomae.py @@ -494,7 +494,12 @@ def forward(self, hidden_states: torch.Tensor, return_token_num: int): for layer_module in self.decoder_layers: hidden_states = layer_module(hidden_states) - hidden_states = hidden_states[:, -return_token_num:] + # Equivalent to `hidden_states[:, -return_token_num:]`, but with a non-negative start + # index. `return_token_num` is the number of masked patches (produced by boolean + # indexing → a data-dependent unbacked symint); slicing from the end forces `torch.export` + # to guard on the sign of that symint (`-(u//13) < 0`), which is unresolvable and breaks + # ExecuTorch's `slice_copy` lowering. A non-negative start avoids the guard. + hidden_states = hidden_states[:, hidden_states.shape[1] - return_token_num :] # predictor projection hidden_states = self.norm(hidden_states) diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index f07b6e2e58a4..497323b15a89 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -46,14 +46,17 @@ # ──────────────────────────── skip lists ──────────────────────────── # # A single mapping ``EXPORT_SKIPS[scope][model_class_name] = reason`` drives every skip. -# ``scope`` is a dotted path that narrows from broad (``"all"`` — every backend, every variant) -# to specific (``"onnx.generate"``, ``"onnx.dynamic"``, ``"openvino"``, …). At test time -# ``_should_skip`` walks the scopes that match the current ``(backend, generate, dynamic)`` -# triple and returns ``True`` as soon as the model is found in any of them. Reasons live next -# to the model name so the "why" travels with the entry. +# A ``scope`` key is a **dotted set of tags** drawn from the active context — the backend +# (``"onnx"`` / ``"openvino"`` / ``"executorch"``), ``"generate"``, and ``"dynamic"``. It applies +# when *all* its tags are active, so tag order doesn't matter and any combination composes: +# ``"all"`` (no tags) always applies, ``"dynamic"`` matches any dynamic export, ``"onnx.dynamic"`` +# matches dynamic ONNX, ``"openvino.generate.dynamic"`` matches only OpenVINO generate-under-dynamic. +# ``_scope_applies`` does the subset check; ``_should_skip`` returns ``True`` as soon as the model +# is found under any applicable scope. Reasons live next to the model name so the "why" travels +# with the entry. # -# Adding a new skip: pick the most specific scope that applies and add a ``"Name": "reason"`` -# entry. Add a new scope key if the existing ones don't fit. +# Adding a new skip: pick the tightest tag-set that covers the failure and add a ``"Name": "reason"`` +# entry under that dotted key (creating the key if needed — no matcher change required). EXPORT_SKIPS: dict[str, dict[str, str]] = { @@ -96,6 +99,25 @@ "× 3 Q-pool stage transitions on symbolic H/W). ONNX + ORT push past 1000s timeout." ), "Sam2Model": "Same Hiera-backbone dynamic-shape budget overrun as `Sam2VisionModel`.", + "SwinModel": ( + "Shifted-window attention on symbolic H/W: `torch.export` + onnxscript exceed the " + "1000s test timeout under dynamic shapes (static exports fine)." + ), + "SwinBackbone": "Same shifted-window `timeout` as `SwinModel`.", + "SwinForImageClassification": "Same shifted-window `timeout` as `SwinModel`.", + "SwinForMaskedImageModeling": "Same shifted-window `timeout` as `SwinModel`.", + "Swinv2Model": "Same shifted-window `timeout` as `SwinModel`.", + "Swinv2ForImageClassification": "Same shifted-window `timeout` as `SwinModel`.", + "Swinv2ForMaskedImageModeling": "Same shifted-window `timeout` as `SwinModel`.", + "Swinv2Backbone": "Same shifted-window `timeout` as `SwinModel`.", + "DonutSwinModel": "Same shifted-window `timeout` as `SwinModel`.", + "DonutSwinForImageClassification": "Same shifted-window `timeout` as `SwinModel`.", + "MaskFormerModel": "Same shifted-window (Swin backbone) `timeout` as `SwinModel`.", + "MaskFormerForInstanceSegmentation": "Same `timeout` as `MaskFormerModel`.", + "Mask2FormerModel": ( + "Deformable-attention pixel decoder exceeds the 1000s test timeout under dynamic shapes." + ), + "Mask2FormerForUniversalSegmentation": "Same `timeout` as `Mask2FormerModel`.", }, # ExecuTorch, every variant. "executorch": { @@ -105,7 +127,12 @@ "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", }, # ExecuTorch, generate path only. - "executorch.generate": {}, + "executorch.generate": { + "MiniMaxM3SparseForConditionalGeneration": ( + "`flatc` schema compilation fails when serializing the ExecuTorch program for the MoE " + "decoder generate graph." + ), + }, # ExecuTorch, dynamic-shape only. "executorch.dynamic": { "BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."), @@ -116,31 +143,16 @@ "BigBirdForQuestionAnswering": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForSequenceClassification": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForTokenClassification": "Same `timeout` failure as `BigBirdModel`.", - "MaskFormerModel": "Same `timeout` failure as `BigBirdModel`.", - "MaskFormerForInstanceSegmentation": "Same `timeout` failure as `BigBirdModel`.", - "MaskFormerSwinModel": ( - "ExecuTorch memory planning overflows under unbounded dynamic shapes (`mem_offset does not fit in 64 bits`)." - ), - "MaskFormerSwinBackbone": "Same `overflow` failure as `MaskFormerSwinModel`.", - "MllamaModel": "Same `overflow` failure as `MaskFormerSwinModel`.", - "MllamaForConditionalGeneration": "Same `overflow` failure as `MaskFormerSwinModel`.", "Sam2Model": ( "The dropped-weight `KeyError` is fixed, but `torch.export` of the Hiera vision " "backbone under dynamic shapes then exceeds the 1000s test timeout (same " "Hiera-backbone dynamic-shape overrun as `Sam2VisionModel`)." ), "Sam2VisionModel": "Same `timeout` failure as `BigBirdModel`.", - "Swin2SRModel": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Swin2SRForImageSuperResolution": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Swinv2Model": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Swinv2ForImageClassification": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Swinv2ForMaskedImageModeling": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Swinv2Backbone": "Same `overflow` failure as `MaskFormerSwinModel`.", - "Wav2Vec2BertForCTC": ("`flatc` schema compilation fails when serializing the program."), - "Wav2Vec2BertModel": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", - "Wav2Vec2BertForSequenceClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", - "Wav2Vec2BertForAudioFrameClassification": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", - "Wav2Vec2BertForXVector": "Same `flatc` failure as `Wav2Vec2BertForCTC`.", + "Qwen3NextModel": ("Lowering exceeds the 1000s test timeout under dynamic shapes."), + "Qwen3NextForCausalLM": "Same `timeout` failure as `Qwen3NextModel`.", + "DiffusionGemmaModel": ("Lowering exceeds the 1000s test timeout under dynamic shapes."), + "DiffusionGemmaForBlockDiffusion": "Same `timeout` failure as `DiffusionGemmaModel`.", }, # OpenVINO, every variant. "openvino": {}, @@ -148,17 +160,13 @@ "openvino.generate": {}, # OpenVINO, dynamic-shape only. "openvino.dynamic": {}, -} - - -# ──────────────────────────── ONNX optimization toggles ──────────────────────────── -# Not "skips" — these select whether `onnxscript` optimisation runs for a given model. -# Same scope-keyed shape as ``EXPORT_SKIPS`` for symmetry. - - -ONNX_DISABLE_OPTIMIZE: dict[str, dict[str, str]] = { - # Disable for dynamic-shape only — static benefits from optimisation. - "dynamic": {}, + # OpenVINO, generate + dynamic-shape only. + "openvino.generate.dynamic": { + "MiniMaxM3SparseForConditionalGeneration": ( + "OpenVINO CPU runtime error (`infer_request.cpp`) on the MoE decoder generate graph " + "under dynamic shapes; static and non-generate variants export fine." + ), + }, } @@ -264,15 +272,16 @@ def _state_path(state): return outputs -def _onnx_optimize_enabled(model_class, dynamic: bool) -> bool: - """Return whether onnxscript optimisation should run for this model under this shape mode. +def _scope_applies(mapping: dict[str, dict[str, str]], active: set[str], name: str) -> bool: + """Whether ``name`` is listed under any scope in ``mapping`` applicable to ``active``. - Mirrors ``_should_skip``'s scope walk on ``ONNX_DISABLE_OPTIMIZE`` — ``"all"`` always - applies; ``"dynamic"`` adds the dynamic-only entries. + A dotted scope key is a set of required tags; it applies when every tag is in ``active`` + (``"all"`` = no tags = always). Tag order is irrelevant and any combination composes. """ - name = model_class.__name__ - scopes = ["all"] + (["dynamic"] if dynamic else []) - return not any(name in ONNX_DISABLE_OPTIMIZE.get(scope, {}) for scope in scopes) + return any( + (set() if scope == "all" else set(scope.split("."))) <= active and name in entries + for scope, entries in mapping.items() + ) # ──────────────────────────── mixins ──────────────────────────── @@ -310,22 +319,17 @@ def _skip_if_not_exportable(self): def _should_skip(self, model_class, generate=False, dynamic=False, backend=None): """Return True if this model class should be skipped for export tests. - Walks the scopes in ``EXPORT_SKIPS`` from broad to specific that match the current - ``(backend, generate, dynamic)`` triple — ``"all"`` always applies, ``"generate"`` only - for generate tests, ``""`` for that backend, and ``"."`` for - the more-specific intersections. + Builds the active tag-set from ``(backend, generate, dynamic)`` and returns True if + ``EXPORT_SKIPS`` lists the model under any applicable scope (see ``_scope_applies``). """ - name = model_class.__name__ - scopes = ["all"] - if generate: - scopes.append("generate") + active = set() if backend: - scopes.append(backend) - if generate: - scopes.append(f"{backend}.generate") - if dynamic: - scopes.append(f"{backend}.dynamic") - return any(name in EXPORT_SKIPS.get(scope, {}) for scope in scopes) + active.add(backend) + if generate: + active.add("generate") + if dynamic: + active.add("dynamic") + return _scope_applies(EXPORT_SKIPS, active, model_class.__name__) def _prepare_export_model_and_inputs(self, model_class): """Create model and forward inputs ready for export. @@ -424,9 +428,8 @@ def test_onnx_export(self, dynamic): if self._should_skip(model_class, dynamic=dynamic, backend="onnx"): continue - optimize = _onnx_optimize_enabled(model_class, dynamic) exporter = OnnxExporter() - config = OnnxConfig(dynamic=dynamic, optimize=optimize) + config = OnnxConfig(dynamic=dynamic) components = self._prepare_export_model_and_inputs(model_class) eager_outputs = self._collect_eager_outputs(components) @@ -571,9 +574,8 @@ def test_onnx_export_generate(self, dynamic): if self._should_skip(model_class, generate=True, dynamic=dynamic, backend="onnx"): continue - optimize = _onnx_optimize_enabled(model_class, dynamic) exporter = OnnxExporter() - config = OnnxConfig(dynamic=dynamic, optimize=optimize) + config = OnnxConfig(dynamic=dynamic) components = self._prepare_export_generate_model_and_inputs(model_class) eager_outputs = self._collect_eager_outputs(components) From 2ce279d9f2a816df75d06b388ec0f0039802cdd9 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 12:34:34 +0200 Subject: [PATCH 06/14] fixes, skips and reverts --- .../exporters/exporter_executorch.py | 2 +- .../exporters/exporter_openvino.py | 63 ++++++- .../models/chameleon/modeling_chameleon.py | 36 ++-- .../models/phimoe/modeling_phimoe.py | 33 ++-- .../modeling_pp_doclayout_v3.py | 7 +- .../models/videomae/modeling_videomae.py | 168 ++++++++---------- tests/exporters/test_export.py | 54 +++--- 7 files changed, 195 insertions(+), 168 deletions(-) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 90d3db41cd33..76f12554686f 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -336,7 +336,7 @@ def _patch_scaled_dot_product_attention(original): ``sdpa_kernel(MATH)`` forces the decomposable SDPA variant on any device — without it, CUDA traces pick ``_scaled_dot_product_efficient_attention``, which XNNPACK's edge-dialect - verifier rejects as non-core-ATen. Same shape of fix as the dynamo-path ``_patch_sdpa``, + verifier rejects as non-core-ATen. Same shape of fix as the Dynamo-path ``_patch_sdpa``, but unconditional here since the CUDA fused kernel is never lowerable by ExecuTorch's xnnpack backend. No-op on CPU (MATH is already the default), so this is safe everywhere. diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index a491c9950faf..ec9652964f12 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -1076,6 +1076,57 @@ def _fix_view_inferred_dim(gm, node): return True +@register_fx_node_fix("openvino") +def _fix_index_put_none_indices(gm, node): + """Rewrite ``aten.index_put`` with ``None`` index entries into a broadcast ``where``. + + ``x[:, :, idx] = value`` traces as ``index_put(x, [None, None, idx], value)``; OV's frontend + turns each ``None`` into a ``torch::None`` constant it can't translate (chameleon masks image + tokens out of its logits this way). For a single 1-D index tensor on one dim (the rest ``None``) + and a scalar value, this is equivalent to marking that dim's ``idx`` positions with a boolean + mask and ``where``-ing the value in — built from ``arange``/``eq``/``any``/``where``, all of + which translate cleanly. Non-scalar values or multi-index puts fall through unchanged. + """ + if node.target not in (torch.ops.aten.index_put.default, torch.ops.aten.index_put_.default): + return False + if len(node.args) < 3: + return False + self_arg, indices, values = node.args[0], node.args[1], node.args[2] + accumulate = node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + if accumulate or not isinstance(indices, (list, tuple)): + return False + non_none = [(dim, ix) for dim, ix in enumerate(indices) if ix is not None] + # Only the "some `None`s + exactly one 1-D index tensor" pattern; skip fully-explicit puts + # (OV lowers those) and multi-index puts. + if len(non_none) != 1 or len(indices) == len(non_none): + return False + dim, idx = non_none[0] + self_val = self_arg.meta.get("val") + idx_val = idx.meta.get("val") if hasattr(idx, "meta") else None + values_val = values.meta.get("val") if hasattr(values, "meta") else None + if self_val is None or idx_val is None or idx_val.ndim != 1: + return False + if values_val is None or values_val.numel() != 1: # scalar / broadcast value only + return False + size = self_val.shape[dim] + if not isinstance(size, int): + return False + with gm.graph.inserting_before(node): + iota = gm.graph.call_function(torch.ops.aten.arange.default, args=(size,), kwargs={"device": self_val.device}) + iota_u = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(iota, 1)) + idx_u = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(idx, 0)) + eq = gm.graph.call_function(torch.ops.aten.eq.Tensor, args=(iota_u, idx_u)) + mask = gm.graph.call_function(torch.ops.aten.any.dim, args=(eq, 1)) + broadcast_shape = [1] * self_val.ndim + broadcast_shape[dim] = size + mask = gm.graph.call_function(torch.ops.aten.view.default, args=(mask, broadcast_shape)) + result = gm.graph.call_function(torch.ops.aten.where.self, args=(mask, values, self_arg)) + result.meta.update(node.meta) + node.replace_all_uses_with(result) + gm.graph.erase_node(node) + return True + + # ── Torch patches ─────────────────────────────────────────────────────────── # Each `_patch_*(original)` factory is registered via `@register_patch("openvino", path)` # and reversibly swaps a `torch` op the OV frontend can't lower with a decomposed @@ -1942,8 +1993,18 @@ def _convert(context): def _convert_sym_floordiv(context): """``a // b`` over SymInts → ``floor(a / b)``, cast to i64. Used by patch/window-size computations (focalnet, donut_swin). The i64 cast keeps the result shape-op-friendly — - downstream ``SequenceMark → Concat`` requires a uniform int dtype.""" + downstream ``SequenceMark → Concat`` requires a uniform int dtype. + + Integer operands must be promoted to ``f32`` first: OV's ``Divide`` on two integers truncates + toward zero, so a subsequent ``floor`` is a no-op and the result is wrong for negative operands + (``-200 // 64`` gives ``-3`` instead of ``-4``). This breaks the ceil-div idiom ``-(-x // n)`` + on a symbolic ``x`` — e.g. minimax_m3_vl's ``num_key_blocks``, which then comes out one too small + and sends a downstream ``scatter`` out of bounds. Dividing in float restores true floor division.""" a, b = context.get_input(0), context.get_input(1) + if a.get_element_type().is_integral(): + a = ov_ops.convert(a, "f32") + if b.get_element_type().is_integral(): + b = ov_ops.convert(b, "f32") return [ov_ops.convert(ov_ops.floor(ov_ops.divide(a, b)), "i64").output(0)] diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index d57cf3eb4125..52cf77ec9218 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -761,26 +761,17 @@ def img2bpe(self): def bpe2img_search_tensors(self): return torch.tensor(sorted(self.bpe2img.keys())), torch.tensor(sorted(self.bpe2img.values())) - def _build_img2bpe_mapping_tensor(self, device: torch.device) -> torch.Tensor: - """Build a device-local ``(max_img_id + 1,)`` int lookup table mapping VQ image ids to - BPE token ids. Cached per-device so we don't rebuild on every call.""" - if not hasattr(self, "_img2bpe_mapping_cache"): - self._img2bpe_mapping_cache: dict[torch.device, torch.Tensor] = {} - if device not in self._img2bpe_mapping_cache: - # Build on CPU via ``scatter`` then move — avoids a Python-level element-wise - # assignment loop that trips FakeTensor propagation during export. - keys = torch.tensor(list(self.img2bpe.keys()), dtype=torch.long) - values = torch.tensor(list(self.img2bpe.values()), dtype=torch.int) - # `max()` over the Python dict keys (a compile-time constant) instead of - # `keys.max().item()`: `.item()` on a tensor yields a data-dependent unbacked symint - # under `torch.export` that `torch.zeros(...)` cannot specialize. - mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int) - mapping.scatter_(0, keys, values) - self._img2bpe_mapping_cache[device] = mapping.to(device) - return self._img2bpe_mapping_cache[device] + @cached_property + def img2bpe_mapping_tensor(self): + mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int) + for k, v in self.img2bpe.items(): + mapping[k] = v + return mapping def convert_img2bpe(self, img_batch: torch.Tensor) -> torch.Tensor: - return self._build_img2bpe_mapping_tensor(img_batch.device)[img_batch] + device = img_batch.device + img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")] + return img_tokens.to(device) @auto_docstring @@ -1086,14 +1077,9 @@ def forward( slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) - # Disallow image tokens which does not include special begin-image and end-image tokens. - # ``logits[:, :, image_tokens] = ...`` traces as ``aten.index_put`` with ``[None, None, tensor]`` - # indices — the ``None`` values become ``torch::None`` under torch.export, which OpenVINO's - # frontend can't lower. A vocab-sized bool mask + ``masked_fill`` is equivalent and traces cleanly. + # Disallow image tokens which does not include special begin-image and end-image tokens image_tokens = self.model.vocabulary_mapping.image_tokens - image_token_mask = torch.zeros(logits.shape[-1], dtype=torch.bool, device=logits.device) - image_token_mask[image_tokens] = True - logits = logits.masked_fill(image_token_mask, torch.finfo(logits.dtype).min) + logits[:, :, image_tokens] = torch.finfo(logits.dtype).min loss = None if labels is not None: diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 25f180050c38..eda3f270d5e4 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -56,17 +56,10 @@ def __init__(self, config: PhimoeConfig, device=None): self.rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - # Precompute both LongRoPE variants at init so forward stays traceable: at forward we - # blend by the runtime ``max(position_ids) + 1 > original_max_position_embeddings`` mask - # via ``torch.where``, avoiding a data-dependent Python-``bool`` on the seq_len tensor. inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) - if self.rope_type != "default": - inv_freq_long, _ = self.rope_init_fn( - self.config, device, seq_len=self.config.rope_parameters["original_max_position_embeddings"] + 1 - ) - self.register_buffer("inv_freq_long", inv_freq_long, persistent=False) @staticmethod def compute_default_rope_parameters( @@ -106,22 +99,16 @@ def forward(self, x, position_ids=None, layer_type=None): f"{self.__class__.__name__} does not support layer types, but got `layer_type={layer_type}`" ) - if self.rope_type != "default": - threshold = self.config.rope_parameters["original_max_position_embeddings"] - # ``max(position_ids) + 1`` (not ``shape[-1]``) so decode / sliding-window paths that - # ship a short slice of position_ids still see the true reached max — matters for - # LongRoPE scale selection. Blending via ``torch.where`` avoids the ``Python bool`` - # on a tensor that trips ``GuardOnDataDependentSymNode`` under ``torch.export``. - is_long_context = (torch.max(position_ids) + 1) > threshold - inv_freq = torch.where(is_long_context, self.inv_freq_long, self.inv_freq) - mscale = torch.where( - is_long_context, - torch.tensor(self.config.rope_parameters["long_mscale"], device=x.device, dtype=torch.float32), - torch.tensor(self.config.rope_parameters["short_mscale"], device=x.device, dtype=torch.float32), + mscale = None + seq_len = torch.max(position_ids) + 1 + if self.config.rope_parameters["rope_type"] != "default" and seq_len: + mscale = ( + self.config.rope_parameters["long_mscale"] + if seq_len > self.config.rope_parameters["original_max_position_embeddings"] + else self.config.rope_parameters["short_mscale"] ) - else: - inv_freq = self.inv_freq - mscale = self.attention_scaling + inv_freq, attention_scaling = self.rope_init_fn(self.config, x.device, seq_len) + mscale = attention_scaling if mscale is None else mscale inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() diff --git a/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py b/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py index 944bf34415c0..bc1cf002f198 100644 --- a/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py +++ b/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py @@ -1555,12 +1555,7 @@ def mask_to_box_coordinate(mask, dtype): is_mask_non_empty = torch.any(mask, dim=(-2, -1)).unsqueeze(-1) unnormalized_bbox = unnormalized_bbox * is_mask_non_empty - # ``torch.tensor([w, h, w, h])`` where ``w`` / ``h`` are SymInts materialises on CPU and - # then transfers, tripping FakeTensor device propagation during export. Build via - # ``torch.stack`` on already-device-side scalars instead. - width_t = torch.as_tensor(width, device=mask.device, dtype=dtype) - height_t = torch.as_tensor(height, device=mask.device, dtype=dtype) - norm_tensor = torch.stack([width_t, height_t, width_t, height_t]) + norm_tensor = torch.tensor([width, height, width, height], device=mask.device, dtype=dtype) normalized_bbox_xyxy = unnormalized_bbox / norm_tensor x_min_norm, y_min_norm, x_max_norm, y_max_norm = normalized_bbox_xyxy.unbind(dim=-1) diff --git a/src/transformers/models/videomae/modeling_videomae.py b/src/transformers/models/videomae/modeling_videomae.py index a6ad7d85f8d5..35bcce23c973 100755 --- a/src/transformers/models/videomae/modeling_videomae.py +++ b/src/transformers/models/videomae/modeling_videomae.py @@ -63,7 +63,7 @@ class VideoMAEDecoderOutput(ModelOutput): @dataclass class VideoMAEForPreTrainingOutput(ModelOutput): r""" - loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): + loss (`torch.FloatTensor` of shape `(1,)`): Pixel reconstruction loss. logits (`torch.FloatTensor` of shape `(batch_size, patch_size ** 2 * num_channels)`): Pixel reconstruction logits. @@ -494,12 +494,7 @@ def forward(self, hidden_states: torch.Tensor, return_token_num: int): for layer_module in self.decoder_layers: hidden_states = layer_module(hidden_states) - # Equivalent to `hidden_states[:, -return_token_num:]`, but with a non-negative start - # index. `return_token_num` is the number of masked patches (produced by boolean - # indexing → a data-dependent unbacked symint); slicing from the end forces `torch.export` - # to guard on the sign of that symint (`-(u//13) < 0`), which is unresolvable and breaks - # ExecuTorch's `slice_copy` lowering. A non-negative start avoids the guard. - hidden_states = hidden_states[:, hidden_states.shape[1] - return_token_num :] + hidden_states = hidden_states[:, -return_token_num:] # predictor projection hidden_states = self.norm(hidden_states) @@ -537,7 +532,6 @@ def forward( self, pixel_values: torch.FloatTensor, bool_masked_pos: torch.BoolTensor, - return_loss: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> VideoMAEForPreTrainingOutput: r""" @@ -545,9 +539,6 @@ def forward( Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Each video in the batch must have the same number of masked patches. Sequence length is `(num_frames // tubelet_size) * (image_size // patch_size) ** 2`. - return_loss (`bool`, *optional*): - Whether to compute and return the pixel reconstruction loss. Defaults to `self.config.return_loss` if - that attribute is set, and `True` otherwise. Examples: ```python @@ -595,89 +586,82 @@ def forward( logits = decoder_outputs.logits loss = None - if return_loss is None: - return_loss = getattr(self.config, "return_loss", True) - if return_loss: - with torch.no_grad(): - # calculate the labels to be predicted + with torch.no_grad(): + # calculate the labels to be predicted + if self.config.num_channels != 3: + # Can't unnormalize with default means/stds + frames = pixel_values + else: + # first, unnormalize the frames + device = pixel_values.device + dtype = pixel_values.dtype + mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None] + std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None] + frames = pixel_values * std + mean # in [0, 1] + + batch_size, time, num_channels, height, width = frames.shape + tubelet_size, patch_size = self.config.tubelet_size, self.config.patch_size + if self.config.norm_pix_loss: + # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) + frames = frames.view( + batch_size, + time // tubelet_size, + tubelet_size, + num_channels, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + # step 2: move dimensions to concatenate: + frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() + # step 3: concatenate: + frames = frames.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size, + num_channels, + ) + # step 4: normalize. The authors find that the mean is about 0.48 and standard deviation is about 0.08. + frames_norm = (frames - frames.mean(dim=-2, keepdim=True)) / ( + frames.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6 + ) + # step 5: reshape to (batch_size, T//ts * H//ps * W//ps, ts * ps * ps * C) + videos_patch = frames_norm.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size * num_channels, + ) + else: if self.config.num_channels != 3: - # Can't unnormalize with default means/stds - frames = pixel_values - else: - # first, unnormalize the frames - device = pixel_values.device - dtype = pixel_values.dtype - mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[ - None, None, :, None, None - ] - std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[ - None, None, :, None, None - ] - frames = pixel_values * std + mean # in [0, 1] - - batch_size, time, num_channels, height, width = frames.shape - tubelet_size, patch_size = self.config.tubelet_size, self.config.patch_size - if self.config.norm_pix_loss: - # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) - frames = frames.view( - batch_size, - time // tubelet_size, - tubelet_size, - num_channels, - height // patch_size, - patch_size, - width // patch_size, - patch_size, + raise ValueError( + "Can't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False." ) - # step 2: move dimensions to concatenate: - frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() - # step 3: concatenate: - frames = frames.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size, - num_channels, - ) - # step 4: normalize. The authors find that the mean is about 0.48 and standard deviation is about 0.08. - frames_norm = (frames - frames.mean(dim=-2, keepdim=True)) / ( - frames.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6 - ) - # step 5: reshape to (batch_size, T//ts * H//ps * W//ps, ts * ps * ps * C) - videos_patch = frames_norm.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size * num_channels, - ) - else: - if self.config.num_channels != 3: - raise ValueError( - "Can't unnormalize non-RGB images. Consider setting config.norm_pix_loss to False." - ) - # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) - frames = frames.view( - batch_size, - time // tubelet_size, - tubelet_size, - num_channels, - height // patch_size, - patch_size, - width // patch_size, - patch_size, - ) - # step 2: move dimensions to concatenate: (batch_size, T//ts, H//ps, W//ps, ts, ps, ps, C) - frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() - # step 3: concatenate - videos_patch = frames.view( - batch_size, - time // tubelet_size * height // patch_size * width // patch_size, - tubelet_size * patch_size * patch_size * num_channels, - ) - - batch_size, _, num_channels = videos_patch.shape - labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels) - - loss_fct = MSELoss() - loss = loss_fct(logits, labels) + # step 1: split up dimensions (time by tubelet_size, height by patch_size, width by patch_size) + frames = frames.view( + batch_size, + time // tubelet_size, + tubelet_size, + num_channels, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + # step 2: move dimensions to concatenate: (batch_size, T//ts, H//ps, W//ps, ts, ps, ps, C) + frames = frames.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous() + # step 3: concatenate + videos_patch = frames.view( + batch_size, + time // tubelet_size * height // patch_size * width // patch_size, + tubelet_size * patch_size * patch_size * num_channels, + ) + + batch_size, _, num_channels = videos_patch.shape + labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels) + + loss_fct = MSELoss() + loss = loss_fct(logits, labels) return VideoMAEForPreTrainingOutput( loss=loss, diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 497323b15a89..15ced7f28a2b 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -61,7 +61,23 @@ EXPORT_SKIPS: dict[str, dict[str, str]] = { # Every backend, every variant. - "all": {}, + "all": { + "VideoMAEForPreTraining": ( + "The reconstruction-loss path indexes `videos_patch[bool_masked_pos]` (a boolean " + "index → data-dependent unbacked count) and compares its shape against `logits`, which " + "`torch.export` can't guard (`Eq(u2 // 13, u3)`). Only the pretraining head hits this; " + "`VideoMAEModel` / `VideoMAEForVideoClassification` export fine." + ), + }, + # Any backend (incl. plain `torch.export`), dynamic-shape variant only. + "dynamic": { + "Sam2Model": ( + "`torch.export` of the Hiera vision backbone under dynamic shapes exceeds the 1000s " + "test timeout on every backend (Dynamo/ONNX/OpenVINO/ExecuTorch) — 12 attention blocks " + "× 3 Q-pool stage transitions on symbolic H/W. Static exports fine." + ), + "Sam2VisionModel": "Same Hiera-backbone dynamic-shape `timeout` as `Sam2Model`.", + }, # Every backend, generate path only. "generate": { "RecurrentGemmaForCausalLM": ( @@ -93,12 +109,6 @@ "GroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.", "MMGroundingDinoModel": "Same as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.", - "Sam2VisionModel": ( - "`torch.export` of the Hiera vision backbone under dynamic shapes takes ~7.5 min " - "even after simplifying `window_partition`/`window_unpartition` (12 attention blocks " - "× 3 Q-pool stage transitions on symbolic H/W). ONNX + ORT push past 1000s timeout." - ), - "Sam2Model": "Same Hiera-backbone dynamic-shape budget overrun as `Sam2VisionModel`.", "SwinModel": ( "Shifted-window attention on symbolic H/W: `torch.export` + onnxscript exceed the " "1000s test timeout under dynamic shapes (static exports fine)." @@ -143,12 +153,6 @@ "BigBirdForQuestionAnswering": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForSequenceClassification": "Same `timeout` failure as `BigBirdModel`.", "BigBirdForTokenClassification": "Same `timeout` failure as `BigBirdModel`.", - "Sam2Model": ( - "The dropped-weight `KeyError` is fixed, but `torch.export` of the Hiera vision " - "backbone under dynamic shapes then exceeds the 1000s test timeout (same " - "Hiera-backbone dynamic-shape overrun as `Sam2VisionModel`)." - ), - "Sam2VisionModel": "Same `timeout` failure as `BigBirdModel`.", "Qwen3NextModel": ("Lowering exceeds the 1000s test timeout under dynamic shapes."), "Qwen3NextForCausalLM": "Same `timeout` failure as `Qwen3NextModel`.", "DiffusionGemmaModel": ("Lowering exceeds the 1000s test timeout under dynamic shapes."), @@ -159,13 +163,23 @@ # OpenVINO, generate path only. "openvino.generate": {}, # OpenVINO, dynamic-shape only. - "openvino.dynamic": {}, - # OpenVINO, generate + dynamic-shape only. - "openvino.generate.dynamic": { - "MiniMaxM3SparseForConditionalGeneration": ( - "OpenVINO CPU runtime error (`infer_request.cpp`) on the MoE decoder generate graph " - "under dynamic shapes; static and non-generate variants export fine." - ), + "openvino.dynamic": { + "BigBirdModel": ("OpenVINO conversion exceeds the 1000s test timeout under dynamic shapes."), + "BigBirdForPreTraining": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForMaskedLM": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForCausalLM": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForMultipleChoice": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForQuestionAnswering": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForSequenceClassification": "Same `timeout` failure as `BigBirdModel`.", + "BigBirdForTokenClassification": "Same `timeout` failure as `BigBirdModel`.", + "MaskFormerModel": "Shifted-window (Swin) backbone exceeds the 1000s test timeout under dynamic shapes.", + "MaskFormerForInstanceSegmentation": "Same `timeout` as `MaskFormerModel`.", + "Mask2FormerModel": "Deformable-attention pixel decoder exceeds the 1000s test timeout under dynamic shapes.", + "Mask2FormerForUniversalSegmentation": "Same `timeout` as `Mask2FormerModel`.", + "GroundingDinoModel": ("Deformable-attention encoder exceeds the 1000s test timeout under dynamic shapes."), + "GroundingDinoForObjectDetection": "Same `timeout` as `GroundingDinoModel`.", + "MMGroundingDinoModel": "Same `timeout` as `GroundingDinoModel`.", + "MMGroundingDinoForObjectDetection": "Same `timeout` as `GroundingDinoModel`.", }, } From 5310517d334128e9d0710687e8518487cd6dd559 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 13:09:10 +0200 Subject: [PATCH 07/14] more reverts --- src/transformers/masking_utils.py | 11 ++----- .../models/phimoe/modular_phimoe.py | 33 ++++++------------- .../modular_pp_doclayout_v3.py | 7 +--- src/transformers/pytorch_utils.py | 4 +-- 4 files changed, 15 insertions(+), 40 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index c508a65a1efe..130b1a1dec66 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -533,11 +533,8 @@ def sdpa_mask( batch_arange = torch.arange(batch_size, device=device) head_arange = torch.arange(1, device=device) - # Fold the offsets into the arange bounds rather than adding them afterwards — a python - # scalar in `aten.add` gets materialized as an int64 constant by ExecuTorch's edge passes, - # whose buffer the XNNPACK lowering then downcasts without updating the spec. - q_arange = torch.arange(q_offset, q_offset + q_length, device=device) - kv_arange = torch.arange(kv_offset, kv_offset + kv_length, device=device) + q_arange = torch.arange(q_length, device=device) + q_offset + kv_arange = torch.arange(kv_length, device=device) + kv_offset # Actual mask creation # Option 1: Fast non-vmap mask creation (default) @@ -635,9 +632,7 @@ def eager_mask( ) # only bidirectional masks can be skipped, otherwise we convert bool -> float if mask is not None: - # A python-float `other` would be lifted as a float64 constant by `torch.export`, - # whose serialized buffer then mismatches the spec — keep both branches in `dtype`. - min_dtype = torch.tensor(torch.finfo(dtype).min, device=mask.device, dtype=dtype) + min_dtype = torch.finfo(dtype).min # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type) mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype) return mask diff --git a/src/transformers/models/phimoe/modular_phimoe.py b/src/transformers/models/phimoe/modular_phimoe.py index 0913791c8396..a4c0a081d89c 100644 --- a/src/transformers/models/phimoe/modular_phimoe.py +++ b/src/transformers/models/phimoe/modular_phimoe.py @@ -49,17 +49,10 @@ def __init__(self, config: PhimoeConfig, device=None): self.rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - # Precompute both LongRoPE variants at init so forward stays traceable: at forward we - # blend by the runtime ``max(position_ids) + 1 > original_max_position_embeddings`` mask - # via ``torch.where``, avoiding a data-dependent Python-``bool`` on the seq_len tensor. inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) - if self.rope_type != "default": - inv_freq_long, _ = self.rope_init_fn( - self.config, device, seq_len=self.config.rope_parameters["original_max_position_embeddings"] + 1 - ) - self.register_buffer("inv_freq_long", inv_freq_long, persistent=False) def forward(self, x, position_ids=None, layer_type=None): if layer_type is not None: @@ -67,22 +60,16 @@ def forward(self, x, position_ids=None, layer_type=None): f"{self.__class__.__name__} does not support layer types, but got `layer_type={layer_type}`" ) - if self.rope_type != "default": - threshold = self.config.rope_parameters["original_max_position_embeddings"] - # ``max(position_ids) + 1`` (not ``shape[-1]``) so decode / sliding-window paths that - # ship a short slice of position_ids still see the true reached max — matters for - # LongRoPE scale selection. Blending via ``torch.where`` avoids the ``Python bool`` - # on a tensor that trips ``GuardOnDataDependentSymNode`` under ``torch.export``. - is_long_context = (torch.max(position_ids) + 1) > threshold - inv_freq = torch.where(is_long_context, self.inv_freq_long, self.inv_freq) - mscale = torch.where( - is_long_context, - torch.tensor(self.config.rope_parameters["long_mscale"], device=x.device, dtype=torch.float32), - torch.tensor(self.config.rope_parameters["short_mscale"], device=x.device, dtype=torch.float32), + mscale = None + seq_len = torch.max(position_ids) + 1 + if self.config.rope_parameters["rope_type"] != "default" and seq_len: + mscale = ( + self.config.rope_parameters["long_mscale"] + if seq_len > self.config.rope_parameters["original_max_position_embeddings"] + else self.config.rope_parameters["short_mscale"] ) - else: - inv_freq = self.inv_freq - mscale = self.attention_scaling + inv_freq, attention_scaling = self.rope_init_fn(self.config, x.device, seq_len) + mscale = attention_scaling if mscale is None else mscale inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() diff --git a/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py b/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py index 222dfe4d4c64..acdb3691f58e 100644 --- a/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py +++ b/src/transformers/models/pp_doclayout_v3/modular_pp_doclayout_v3.py @@ -599,12 +599,7 @@ def mask_to_box_coordinate(mask, dtype): is_mask_non_empty = torch.any(mask, dim=(-2, -1)).unsqueeze(-1) unnormalized_bbox = unnormalized_bbox * is_mask_non_empty - # ``torch.tensor([w, h, w, h])`` where ``w`` / ``h`` are SymInts materialises on CPU and - # then transfers, tripping FakeTensor device propagation during export. Build via - # ``torch.stack`` on already-device-side scalars instead. - width_t = torch.as_tensor(width, device=mask.device, dtype=dtype) - height_t = torch.as_tensor(height, device=mask.device, dtype=dtype) - norm_tensor = torch.stack([width_t, height_t, width_t, height_t]) + norm_tensor = torch.tensor([width, height, width, height], device=mask.device, dtype=dtype) normalized_bbox_xyxy = unnormalized_bbox / norm_tensor x_min_norm, y_min_norm, x_max_norm, y_max_norm = normalized_bbox_xyxy.unbind(dim=-1) diff --git a/src/transformers/pytorch_utils.py b/src/transformers/pytorch_utils.py index ccaff43c8520..444c8bc457bb 100644 --- a/src/transformers/pytorch_utils.py +++ b/src/transformers/pytorch_utils.py @@ -118,9 +118,7 @@ def __repr__(self) -> str: def forward(self, x): size_out = x.size()[:-1] + (self.nf,) - # reshape, not view: export passes (e.g. ExecuTorch's dim-order pass) can hand this a - # non-contiguous input, which view rejects; reshape is free when contiguous. - x = torch.addmm(self.bias, x.reshape(-1, x.size(-1)), self.weight) + x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight) x = x.view(size_out) return x From 97c19f96c4428e9927c215921cddece1ea2b2a3a Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 13:35:43 +0200 Subject: [PATCH 08/14] reverts --- .../exporters/exporter_executorch.py | 16 +------ .../exporters/exporter_openvino.py | 42 +++++++++++++++++++ .../modeling_gpt_neox_japanese.py | 4 +- .../models/longt5/modeling_longt5.py | 27 ++++++------ .../models/t5gemma2/modeling_t5gemma2.py | 5 +-- .../models/t5gemma2/modular_t5gemma2.py | 5 +-- 6 files changed, 59 insertions(+), 40 deletions(-) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 76f12554686f..563f3c84db09 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -1212,17 +1212,11 @@ def _drop_runtime_asserts(exported_program: ExportedProgram) -> None: """Drop ``_assert_scalar`` / ``_assert_tensor_metadata`` runtime asserts before lowering. ``_assert_scalar`` lowers a ``torch._check`` on an unbacked symint (e.g. the image-token - count in ``get_placeholder_mask``, or SmolVLM's ``inputs_merger`` registering ``Eq(u2, 1)`` - off the unbacked real-image count) into a ``cast_symbool_to_symint`` + ``eq`` chain whose + count in ``get_placeholder_mask``) into a ``cast_symbool_to_symint`` + ``eq`` chain whose ``Piecewise`` result the ``_ModuleStackTracer`` used by ``to_edge_transform_and_lower``'s decomposition pass cannot proxy (``... is not tracked with proxy``). The range facts these asserts encode survive on ``exported_program.range_constraints`` (further capped by ``_fix_range_constraints``), so dropping the nodes (and the now-dead symint feeders) is safe. - - Erasing the nodes is not sufficient on its own: ``to_edge``'s internal re-export re-runs - ``insert_deferred_runtime_asserts``, which reads ``shape_env.deferred_runtime_asserts`` to - decide what to regenerate — so that registry has to be cleared too, else the equality assert - reappears and trips the tracer again. """ for module in exported_program.graph_module.modules(): if not isinstance(module, torch.fx.GraphModule): @@ -1236,14 +1230,6 @@ def _drop_runtime_asserts(exported_program: ExportedProgram) -> None: module.graph.eliminate_dead_code() module.recompile() - for node in exported_program.graph_module.graph.nodes: - val = node.meta.get("val") - if isinstance(val, torch.Tensor) and getattr(val, "fake_mode", None) is not None: - shape_env = val.fake_mode.shape_env - if shape_env is not None: - shape_env.deferred_runtime_asserts.clear() - break # all nodes share the same shape_env - @register_fx_program_fix("executorch") def _fix_missing_placeholder_vals(exported_program: ExportedProgram) -> None: diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index ec9652964f12..dbd932285b20 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -1127,6 +1127,48 @@ def _fix_index_put_none_indices(gm, node): return True +@register_fx_node_fix("openvino") +def _fix_index_put_bool_mask(gm, node): + """Rewrite ``aten.index_put`` with a single boolean-mask index into a broadcast ``where``. + + ``x[bool_mask] = values`` (e.g. t5gemma2 swapping in an end-of-image embedding where + ``input_ids == eoi_token``) traces as ``index_put(x, [bool_mask], values)``; OV lowers the + boolean advanced index through a ``nonzero``-style dynamic gather its frontend can't convert + (``SequenceMark`` OpConversionFailure). When ``bool_mask`` indexes ``x``'s leading dims and + ``values`` broadcasts over the trailing ones, this equals ``where(mask[..., None], values, x)`` + — pure elementwise, no dynamic indexing. Flattened per-row values (which ``where`` can't + express) are left untouched. + """ + if node.target not in (torch.ops.aten.index_put.default, torch.ops.aten.index_put_.default): + return False + if len(node.args) < 3: + return False + self_arg, indices, values = node.args[0], node.args[1], node.args[2] + accumulate = node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + if accumulate or not isinstance(indices, (list, tuple)) or len(indices) != 1 or indices[0] is None: + return False + mask = indices[0] + self_val = self_arg.meta.get("val") + mask_val = mask.meta.get("val") if hasattr(mask, "meta") else None + values_val = values.meta.get("val") if hasattr(values, "meta") else None + if self_val is None or mask_val is None or getattr(mask_val, "dtype", None) != torch.bool: + return False + # The mask must cover the leading dims and the value must fit the trailing (non-mask) dims — + # otherwise ``values`` is a flattened selected-rows tensor that a broadcast ``where`` can't + # reproduce. + if mask_val.ndim > self_val.ndim or values_val is None or values_val.ndim > self_val.ndim - mask_val.ndim: + return False + with gm.graph.inserting_before(node): + broadcast_mask = mask + for _ in range(self_val.ndim - mask_val.ndim): + broadcast_mask = gm.graph.call_function(torch.ops.aten.unsqueeze.default, args=(broadcast_mask, -1)) + result = gm.graph.call_function(torch.ops.aten.where.self, args=(broadcast_mask, values, self_arg)) + result.meta.update(node.meta) + node.replace_all_uses_with(result) + gm.graph.erase_node(node) + return True + + # ── Torch patches ─────────────────────────────────────────────────────────── # Each `_patch_*(original)` factory is registered via `@register_patch("openvino", path)` # and reversibly swaps a `torch` op the OV frontend can't lower with a decomposed diff --git a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py index c6ab35493bdd..34fba03f5687 100755 --- a/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py +++ b/src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py @@ -258,8 +258,8 @@ def _attn(self, query, key, value, attention_mask=None): batch_size, num_attention_heads, query_length, attn_head_size = query.size() key_length = key.size(-2) - query = query.reshape(batch_size * num_attention_heads, query_length, attn_head_size) - key = key.reshape(batch_size * num_attention_heads, key_length, attn_head_size) + query = query.view(batch_size * num_attention_heads, query_length, attn_head_size) + key = key.view(batch_size * num_attention_heads, key_length, attn_head_size) # [batch_size * num_heads, q_length, kv_length] attn_scores = torch.zeros( diff --git a/src/transformers/models/longt5/modeling_longt5.py b/src/transformers/models/longt5/modeling_longt5.py index 83bed253e88f..b3e082789e9c 100644 --- a/src/transformers/models/longt5/modeling_longt5.py +++ b/src/transformers/models/longt5/modeling_longt5.py @@ -70,10 +70,9 @@ def _split_into_blocks(x: torch.Tensor, block_len: int, dim: int) -> torch.Tenso """Split an input tensor into blocks of a given `block_len` along the given `dim`. If the dimension length is not a multiple of `block_len`, it will be padded first with selected `pad_value`. """ - # Pad tensor to multiple of block_len — a no-op when the length is already a multiple. - # Padding unconditionally keeps the shape arithmetic branch-free, so ``torch.export`` - # traces a single graph that holds for every sequence length. - x = _pad_to_multiple(x, block_len, dim, pad_value=0) + # pad tensor to multiple of block_len + if x.shape[dim] % block_len != 0: + x = _pad_to_multiple(x, block_len, dim, pad_value=0) num_blocks = x.shape[dim] // block_len output_shape = x.shape[:dim] + (num_blocks, block_len) + x.shape[(dim + 1) :] # If 0 is in output_shape, we cannot apply reshape because of incompatibility with ONNX conversion @@ -107,9 +106,9 @@ def _concatenate_3_blocks(x: torch.Tensor, block_dim: int, sequence_dim: int, pa return torch.cat(blocks_list, dim=sequence_dim) -def _make_3block_relative_position_ids(block_len: int, device: torch.device | None = None) -> torch.Tensor: +def _make_3block_relative_position_ids(block_len: int) -> torch.Tensor: """Makes 3-blocked relative position ids for local attention.""" - position_ids = torch.arange(3 * block_len, dtype=torch.int32, device=device) + position_ids = torch.arange(3 * block_len, dtype=torch.int32) center_position_ids = position_ids[block_len:-block_len] # [block_len, 3 * block_len] relative_position_ids = position_ids.unsqueeze(0) - center_position_ids.unsqueeze(1) @@ -118,9 +117,10 @@ def _make_3block_relative_position_ids(block_len: int, device: torch.device | No def _mask_local_attention_mask(local_attention_mask: torch.Tensor, block_len: int) -> torch.Tensor: """Mask local attention mask to enforce that tokens are not allowed to attend tokens farther than ``local_radius.""" - relative_position_ids = _make_3block_relative_position_ids(block_len, device=local_attention_mask.device) + relative_position_ids = _make_3block_relative_position_ids(block_len) locality_mask = torch.abs(relative_position_ids) < block_len locality_mask = locality_mask[None, None, :, :] + locality_mask = locality_mask.to(local_attention_mask.device) return torch.logical_and(local_attention_mask, locality_mask) @@ -156,7 +156,8 @@ def _make_global_fixed_block_ids( batch_size, seq_len = attention_mask.shape[:2] def handle_orphan_tokens(block_ids: torch.Tensor) -> torch.Tensor: - block_ends = (torch.arange(seq_len, device=block_ids.device) % global_block_size) == global_block_size - 1 + block_ends = (torch.arange(seq_len) % global_block_size) == global_block_size - 1 + block_ends = block_ends.to(block_ids.device) true_block_ends = torch.logical_and(block_ends, block_ids >= 0) full_blocks = true_block_ends.sum(-1).unsqueeze(-1).type(block_ids.dtype) - 1 block_ids = torch.where(block_ids < full_blocks, block_ids, full_blocks) @@ -182,12 +183,8 @@ def handle_orphan_tokens(block_ids: torch.Tensor) -> torch.Tensor: _sequence_block_ids_max = torch.zeros( batch_size, 0, dtype=global_block_ids.dtype, device=global_block_ids.device ) - # Equivalent to ``cumsum(ones(...), dim=-1) - 1``: [0, 1, …, num_globals-1] per batch. - global_segment_ids = ( - torch.arange(num_globals, dtype=attention_mask.dtype, device=attention_mask.device) - .unsqueeze(0) - .expand(batch_size, -1) - ) + global_segment_ids = torch.cumsum(torch.ones(batch_size, num_globals), dim=-1) - 1 + global_segment_ids = global_segment_ids.to(attention_mask.device) global_segment_ids = torch.where(global_segment_ids <= _sequence_block_ids_max, 1, 0) return global_block_ids.type(torch.int), global_segment_ids.type(torch.int) @@ -811,7 +808,7 @@ def unshape(states): # global_seq_len := seq_len // self.global_block_size # shapes: (batch_size, seq_len) & (batch_size, global_seq_len) block_ids, global_segment_ids = _make_global_fixed_block_ids( - mask if mask is not None else torch.ones(hidden_states.shape[:-1], device=hidden_states.device), + mask if mask is not None else torch.ones(hidden_states.shape[:-1]), self.global_block_size, ) # Create global inputs diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index bdd0edbc0ad6..0085b9a31e19 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -641,10 +641,7 @@ def __init__( def forward(self, input_ids: torch.Tensor): input_embeddings = super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) - # ``embeddings[bool_mask] = vector`` traces as ``aten.index_put`` with a mask list, - # which several export backends can't lower — ``torch.where`` is equivalent. - eoi_mask = (input_ids == self.eoi_token_index).unsqueeze(-1) - input_embeddings = torch.where(eoi_mask, self.eoi_embedding.to(input_embeddings.dtype), input_embeddings) + input_embeddings[input_ids == self.eoi_token_index] = self.eoi_embedding.to(input_embeddings.dtype) return input_embeddings diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index 108d99ad8336..c59546c72520 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -457,10 +457,7 @@ def __init__( def forward(self, input_ids: torch.Tensor): input_embeddings = super().forward(input_ids) * self.embed_scale.to(self.weight.dtype) - # ``embeddings[bool_mask] = vector`` traces as ``aten.index_put`` with a mask list, - # which several export backends can't lower — ``torch.where`` is equivalent. - eoi_mask = (input_ids == self.eoi_token_index).unsqueeze(-1) - input_embeddings = torch.where(eoi_mask, self.eoi_embedding.to(input_embeddings.dtype), input_embeddings) + input_embeddings[input_ids == self.eoi_token_index] = self.eoi_embedding.to(input_embeddings.dtype) return input_embeddings From b45c4e9b45c5713fdf4992ba2f177f582aa57fe9 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 15:04:07 +0200 Subject: [PATCH 09/14] claude review --- src/transformers/exporters/auto.py | 2 +- src/transformers/exporters/configs.py | 3 +- .../exporters/exporter_executorch.py | 14 ++- src/transformers/exporters/exporter_onnx.py | 51 ++++++++-- .../exporters/exporter_openvino.py | 93 ++++++++----------- src/transformers/exporters/utils.py | 18 +++- 6 files changed, 114 insertions(+), 67 deletions(-) diff --git a/src/transformers/exporters/auto.py b/src/transformers/exporters/auto.py index adef5a2928dc..c2d44dcd6278 100644 --- a/src/transformers/exporters/auto.py +++ b/src/transformers/exporters/auto.py @@ -72,7 +72,7 @@ def from_dict(cls, export_config_dict: dict): class AutoHfExporter: """ - The Auto-HF expoerter class that takes care of automatically instantiating to the correct + The Auto-HF exporter class that takes care of automatically instantiating to the correct `HfExporter` given the `ExportConfig`. """ diff --git a/src/transformers/exporters/configs.py b/src/transformers/exporters/configs.py index 982751aa2aa1..e05d79cbae64 100644 --- a/src/transformers/exporters/configs.py +++ b/src/transformers/exporters/configs.py @@ -100,8 +100,8 @@ class DynamoConfig(ExportConfigMixin): """ export_format: ExportFormat = ExportFormat.DYNAMO - dynamic: bool = False + dynamic: bool = False strict: bool = False dynamic_shapes: dict[str, Any] | None = None prefer_deferred_runtime_asserts_over_guards: bool = False @@ -142,7 +142,6 @@ class OnnxConfig(DynamoConfig): export_format: ExportFormat = ExportFormat.ONNX output_path: str | PathLike | None = None - dynamic_shapes: dict[str, Any] | None = None opset_version: int | None = None external_data: bool = True optimize: bool = True diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 563f3c84db09..ef915d150bae 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -446,7 +446,12 @@ def _patch_expand(original): def patch(self, *sizes): if len(sizes) == 1 and isinstance(sizes[0], (list, tuple, torch.Size)): sizes = tuple(sizes[0]) - return original(self, *sizes).clone(memory_format=torch.contiguous_format) + result = original(self, *sizes) + # Only materialise when ``expand`` actually introduced a stride-0 (broadcast) dim; a + # no-broadcast expand is a plain view ExecuTorch's memory planner accepts as-is. + if 0 in result.stride(): + return result.clone(memory_format=torch.contiguous_format) + return result return patch @@ -1372,10 +1377,17 @@ def _fix_sym_pow_as_mul(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool: mul_scalar = _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS.get(operator.mul) if mul_scalar is None: return False + base_val = base.meta.get("val") if isinstance(base, torch.fx.Node) else base with gm.graph.inserting_before(node): running = base + running_val = base_val for _ in range(exp - 1): running = gm.graph.call_function(mul_scalar, (running, base)) + # Propagate the symbolic value so downstream passes / the emitter see a ``meta["val"]`` + # on the synthesised products (matches the original ``pow`` node's value). + if base_val is not None and running_val is not None: + running_val = running_val * base_val + running.meta["val"] = running_val node.replace_all_uses_with(running) gm.graph.erase_node(node) return True diff --git a/src/transformers/exporters/exporter_onnx.py b/src/transformers/exporters/exporter_onnx.py index b31b46f282c3..0f97cea51997 100644 --- a/src/transformers/exporters/exporter_onnx.py +++ b/src/transformers/exporters/exporter_onnx.py @@ -74,6 +74,27 @@ from onnxscript.function_libs.torch_lib.ops.core import aten_index_put from onnxscript.onnx_opset import opset18 as op + # torch.dtype -> onnx_ir.DataType, mirroring torch.onnx's private _TORCH_DTYPE_TO_ONNX so we + # don't depend on that path. Only the dtypes a `Cast` target can realistically be; exotic + # float8/float4 variants (never emitted as an ``out_dtype``) are omitted. + _TORCH_DTYPE_TO_ONNX: dict[torch.dtype, onnx_ir.DataType] = { + torch.float32: onnx_ir.DataType.FLOAT, + torch.float64: onnx_ir.DataType.DOUBLE, + torch.float16: onnx_ir.DataType.FLOAT16, + torch.bfloat16: onnx_ir.DataType.BFLOAT16, + torch.bool: onnx_ir.DataType.BOOL, + torch.int8: onnx_ir.DataType.INT8, + torch.int16: onnx_ir.DataType.INT16, + torch.int32: onnx_ir.DataType.INT32, + torch.int64: onnx_ir.DataType.INT64, + torch.uint8: onnx_ir.DataType.UINT8, + torch.uint16: onnx_ir.DataType.UINT16, + torch.uint32: onnx_ir.DataType.UINT32, + torch.uint64: onnx_ir.DataType.UINT64, + torch.complex64: onnx_ir.DataType.COMPLEX64, + torch.complex128: onnx_ir.DataType.COMPLEX128, + } + if TYPE_CHECKING: from ..modeling_utils import PreTrainedModel @@ -125,6 +146,7 @@ def export( output_names=outputs_names, kwargs=copy.deepcopy(dict(sample_inputs)), custom_translation_table=_ONNX_TRANSLATION_TABLE, + keep_initializers_as_inputs=config.keep_initializers_as_inputs, opset_version=config.opset_version, external_data=config.external_data, export_params=config.export_params, @@ -472,7 +494,8 @@ def patch(*args, dtype=None, **kwargs): if dtype is None: # find fill_value: positional arg or kwarg fill_value = kwargs.get("fill_value", args[1] if len(args) > 1 else None) - if isinstance(fill_value, int): + # `bool` is a subclass of `int` — exclude it so `torch.full(size, True)` stays bool. + if isinstance(fill_value, int) and not isinstance(fill_value, bool): dtype = torch.long return original(*args, dtype=dtype, **kwargs) @@ -723,8 +746,10 @@ def _fix_sort_stable(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool: if node.target is not torch.ops.aten.sort.stable: return False self_arg = node.args[0] - dim = node.args[2] if len(node.args) > 2 else -1 - descending = node.args[3] if len(node.args) > 3 else False + # `dim`/`descending` are keyword-only in `sort.stable` (schema: `sort.stable(self, *, stable, dim=-1, + # descending=False)`), so they arrive in `node.kwargs`, never `node.args`. + dim = node.kwargs.get("dim", -1) + descending = node.kwargs.get("descending", False) with gm.graph.inserting_before(node): new = gm.graph.call_function(torch.ops.aten.sort.default, args=(self_arg, dim, descending)) node.replace_all_uses_with(new) @@ -801,7 +826,9 @@ def _aten_index_put( is_bool = ( bool_mask is not None and getattr(getattr(bool_mask, "type", None), "dtype", None) == onnx_ir.DataType.BOOL ) - if not is_bool: + # The Where-based paths below overwrite; they can't express `self[mask] += values`. Delegate the + # accumulate case (and any non-bool-mask index) to torchlib, which handles both correctly. + if not is_bool or accumulate: return aten_index_put(self, indices, values, accumulate) for _ in range(len(self.shape) - len(bool_mask.shape)): bool_mask = op.Unsqueeze(bool_mask, op.Constant(value_ints=[-1])) @@ -835,6 +862,11 @@ def _aten_bincount(self: INT64, weights=None, minlength: int = 0) -> INT64: return op.ReduceSum(one_hot, op.Constant(value_ints=[0]), keepdims=0) +def _torch_dtype_to_onnx(dtype: torch.dtype) -> int: + """Map a ``torch.dtype`` to the ONNX ``op.Cast(to=...)`` TensorProto int.""" + return _TORCH_DTYPE_TO_ONNX[dtype].value + + def _aten_grouped_mm(mat_a: TReal, mat_b: TReal, offs: INT64, bias=None, out_dtype=None) -> TReal: """ONNX implementation of `aten._grouped_mm.default`. @@ -863,10 +895,17 @@ def _aten_grouped_mm(mat_a: TReal, mat_b: TReal, offs: INT64, bias=None, out_dty end = op.Slice(offs_i64, g_lo, g_hi, axes_0) # (1,) — offs[g] a_g = op.Slice(mat_a, prev_end, end, axes_0) # (n_g, K) w_g = op.Squeeze(op.Slice(mat_b, g_lo, g_hi, axes_0), axes_0) # (K, N) - outputs.append(op.MatMul(a_g, w_g)) # (n_g, N) + out_g = op.MatMul(a_g, w_g) # (n_g, N) + if bias is not None: + # per-group bias ``(G, N)`` → ``(N,)`` broadcasts over the group's rows + out_g = op.Add(out_g, op.Squeeze(op.Slice(bias, g_lo, g_hi, axes_0), axes_0)) + outputs.append(out_g) prev_end = end - return op.Concat(*outputs, axis=0) # (M, N) + result = op.Concat(*outputs, axis=0) # (M, N) + if out_dtype is not None: + result = op.Cast(result, to=_torch_dtype_to_onnx(out_dtype)) + return result def _aten_repeat_interleave_self_int(self, repeats, dim=None, output_size=None): diff --git a/src/transformers/exporters/exporter_openvino.py b/src/transformers/exporters/exporter_openvino.py index dbd932285b20..3a1f142d06c2 100644 --- a/src/transformers/exporters/exporter_openvino.py +++ b/src/transformers/exporters/exporter_openvino.py @@ -634,7 +634,7 @@ def dedup(arg): return arg with graph_module.graph.inserting_before(output_node): val = source.meta.get("val") - if val is not None and val.dtype == torch.bool: + if isinstance(val, torch.Tensor) and val.dtype == torch.bool: copy = graph_module.graph.call_function(torch.ops.aten.logical_and.default, args=(source, source)) else: copy = graph_module.graph.call_function(torch.ops.aten.maximum.default, args=(source, source)) @@ -720,16 +720,14 @@ def _fix_sym_min_max(gm, node): @register_fx_program_fix("openvino") def _fix_to_dtype_layout_in_subgraphs(exported_program): - """Rewrite ``aten.to.dtype_layout`` and ``aten.to.device`` in *submodule* FX graphs to - ``aten._to_copy`` (with just the dtype kwarg). - - OV's PyTorch frontend registers ``ConversionExtension`` handlers only for the top-level - graph; nested ``HigherOrderOp`` subgraphs (like the one wrapped by - ``wrap_with_set_grad_enabled`` in Chameleon's rotary path) still see the raw - ``aten.to.dtype_layout`` node, for which OV has no translator — resulting in a dangling - ``torch::None`` constant. Rewriting the FX target here lets OV's built-in - ``aten._to_copy.default`` handler take over (which we also override at the top level to - swallow complex-dtype casts).""" + """Rewrite every ``aten.to.{dtype,dtype_layout,device,other}`` node to ``aten._to_copy`` + (keeping only the dtype kwarg), in the top-level graph and every submodule graph. + + OV's PyTorch frontend has no ``aten.to.*`` translators at all — an unhandled variant falls + back to a dangling ``torch::None`` constant that fails conversion (e.g. the + ``wrap_with_set_grad_enabled`` HigherOrderOp subgraph in Chameleon's rotary path hits + ``aten.to.dtype_layout``). Rewriting the FX target here — before conversion — lets our + ``_convert_to_copy`` override handle every case (it also swallows complex-dtype casts).""" # Walk the top-level graph AND every submodule's graph (higher-order-op subgraphs). graphs = [exported_program.graph_module] graphs.extend(m for _, m in exported_program.graph_module.named_children() if hasattr(m, "graph")) @@ -901,11 +899,11 @@ def _fix_scatter_reduce(gm, node): return True if reduce == "amax" and include_self is False: - # ``self`` is initialised to zeros; each source element competes for the max at - # ``index[j]``; empty positions stay zero. Decompose to a broadcast comparison + amax: - # build a one-hot mask ``(index.unsqueeze(dim) == arange(K))``, then take the elementwise - # max of ``src`` where the mask is set (``-inf`` elsewhere), then zero any positions with - # no scatter. + # ``amax`` with ``include_self=False``: each source element competes for the max at + # ``index[j]``; positions no source scatters to keep ``self``'s original value. Decompose to + # a broadcast comparison + amax: build a one-hot mask ``(index.unsqueeze(dim) == arange(K))``, + # take the elementwise max of ``src`` where the mask is set (``-inf`` elsewhere), then fall + # back to ``self`` for positions with no scatter. self_val = self_arg.meta.get("val") src_val = src.meta.get("val") if self_val is None or src_val is None or not src_val.dtype.is_floating_point: @@ -941,10 +939,7 @@ def _fix_scatter_reduce(gm, node): masked = gm.graph.call_function(torch.ops.aten.where.self, args=(mask, src_unsq, min_tensor)) maxes = gm.graph.call_function(torch.ops.aten.amax.default, args=(masked, [d + 1])) any_match = gm.graph.call_function(torch.ops.aten.any.dim, args=(mask, d + 1)) - zero_tensor = gm.graph.call_function( - torch.ops.aten.scalar_tensor.default, args=(0.0,), kwargs=scalar_kwargs - ) - result = gm.graph.call_function(torch.ops.aten.where.self, args=(any_match, maxes, zero_tensor)) + result = gm.graph.call_function(torch.ops.aten.where.self, args=(any_match, maxes, self_arg)) result.meta.update(node.meta) node.replace_all_uses_with(result) gm.graph.erase_node(node) @@ -1780,6 +1775,14 @@ def _convert_index_add(context): dim = int(context.get_values_from_const_input(1)) index = context.get_input(2) source = context.get_input(3) + # ``index_add`` is ``self[index] += alpha * source``; fold a non-default ``alpha`` (FX input 4) + # into ``source`` before the scatter-add. + if context.get_input_size() > 4 and context.get_input(4).get_node().get_type_name() == "Constant": + alpha = context.get_values_from_const_input(4) + if alpha != 1: + source = ov_ops.multiply( + source, ov_ops.convert(ov_ops.constant(np.array(alpha)), source.get_element_type()) + ) # Broadcast 1-D index to source's rank/shape along ``dim`` so ScatterElementsUpdate # can consume element-wise ``source`` values. src_shape = ov_ops.shape_of(source, output_type="i64") @@ -1882,28 +1885,6 @@ def _convert_layer_norm(context): return [shifted.output(0)] -def _convert_aten_to(context): - """Convert ``aten.to.{dtype,device,dtype_layout,other}`` — emit a real ``Convert`` when the - target dtype is present, else identity. - - OV's frontend has no ``aten.to.*`` translations at all (only ``aten._to_copy.default``); - every unhandled variant falls back to a ``torch::None`` constant that fails conversion - (chameleon's rotary sub-module hits ``aten.to.dtype_layout``). ``layout`` / ``device`` - kwargs are silently dropped — OV exports are inherently device-neutral.""" - data = context.get_input(0) - if not context.has_attribute("dtype"): - # ``aten.to.device`` — just device move, no dtype. Emit identity. - return [data] - try: - dtype = context.get_attribute("dtype") - except Exception: - # Complex dtypes throw. Skip (identity) — see ``_convert_to_copy``. - return [data] - if dtype is None: - return [data] - return [ov_ops.convert(data, dtype).output(0)] - - def _convert_to_copy(context): """Convert ``aten._to_copy(self, dtype=..., ...)`` to an OV ``Convert``. @@ -1951,23 +1932,35 @@ def _convert_sdpa(context): OV's ``opset13::ScaledDotProductAttention`` rejects int-typed masks. Under CUDA export ``aten.expand`` promotes bool masks to ``i64`` during OV translation, so we insert a - ``Convert(→ boolean)`` on the mask input before instantiating the op. Q/K/V/scale pass - through unchanged.""" + ``Convert(→ boolean)`` on the mask input before instantiating the op. + + The ``scale`` arg (FX input 6) is threaded through when present: it defaults to + ``head_dim**-0.5`` in both aten and OV, but models like Gemma2/Gemma3 pass an explicit + ``query_pre_attn_scalar**-0.5`` that differs — dropping it would silently change the attention + temperature. Q/K/V pass through unchanged.""" q, k, v = context.get_input(0), context.get_input(1), context.get_input(2) # A ``None`` FX arg reaches the extension as an unconverted ``PtFrameworkNode``. mask = None if context.get_input_size() > 3: candidate = context.get_input(3) if candidate.get_node().get_type_name() != "PtFrameworkNode": - mask = ov_ops.convert(candidate, "boolean") + # A bool mask that `aten.expand` promoted to ``i64`` under CUDA export is cast back to + # boolean; a float *additive* mask (``0`` attend / ``-inf`` masked) passes through unchanged + # — OV's SDPA adds it to the scores, whereas casting it to bool would invert/destroy it. + mask = candidate if candidate.get_element_type().is_real() else ov_ops.convert(candidate, "boolean") is_causal = False if context.get_input_size() > 5: # ``is_causal`` is a positional FX input (arg 5), not a node attribute. if context.get_input(5).get_node().get_type_name() == "Constant": is_causal = bool(context.get_values_from_const_input(5)) - if mask is None: - return [ov_ops.scaled_dot_product_attention(q, k, v, causal=is_causal).output(0)] - return [ov_ops.scaled_dot_product_attention(q, k, v, mask, causal=is_causal).output(0)] + # ``scale`` is FX input 6; a ``None`` (default) arrives as a non-``Constant`` and is skipped so + # OV falls back to its own ``head_dim**-0.5`` default (identical to aten's). + kwargs = {"causal": is_causal} + if mask is not None: + kwargs["attention_mask"] = mask + if context.get_input_size() > 6 and context.get_input(6).get_node().get_type_name() == "Constant": + kwargs["scale"] = context.get_input(6) + return [ov_ops.scaled_dot_product_attention(q, k, v, **kwargs).output(0)] def _convert_complex(context): @@ -2081,10 +2074,6 @@ def _convert_sym_truediv(context): ConversionExtension("aten._fft_c2c.default", _convert_fft_c2c), ConversionExtension("aten._conj.default", _convert_conj), ConversionExtension("aten._to_copy.default", _convert_to_copy), - ConversionExtension("aten.to.dtype", _convert_aten_to), - ConversionExtension("aten.to.dtype_layout", _convert_aten_to), - ConversionExtension("aten.to.device", _convert_aten_to), - ConversionExtension("aten.to.other", _convert_aten_to), ConversionExtension("aten.layer_norm.default", _convert_layer_norm), ConversionExtension("aten.scaled_dot_product_attention.default", _convert_sdpa), ConversionExtension("aten.bitwise_not.default", _convert_bitwise_not), diff --git a/src/transformers/exporters/utils.py b/src/transformers/exporters/utils.py index 0ffd831d23e8..9d721a980af2 100644 --- a/src/transformers/exporters/utils.py +++ b/src/transformers/exporters/utils.py @@ -232,7 +232,7 @@ def _map_leaf_tensors(obj: Any, fn: callable) -> Any: Mutates dicts and `__dict__`-bearing objects in place (preserving identity — callers rely on this so downstream pops/mutations propagate back to the original mapping); - rebuilds lists/tuples/sets/frozensets (immutable or order-sensitive containers). + rebuilds lists/tuples/sets (immutable or order-sensitive containers). Skips non-traversable leaf types (enum, SymInt, etc.). """ if isinstance(obj, _LEAF_SKIP_TYPES): @@ -444,10 +444,18 @@ def _prepare_grid_thw_vision_inputs(model: torch.nn.Module, inputs: dict[str, An spatial_merge_size = inputs.get("merge_sizes", 1) inputs["cu_seqlens"] = get_vision_cu_seqlens(grid_thw) - # 3-axis (t, h, w) rotary encoders expose an ``axis_dim`` attr on their rotary_emb - # (minimax_m3_vl); default 2-axis (h, w) covers qwen2_5_vl / qwen3_vl / glm4v / paddleocr_vl. - include_temporal = _find_submodule_attr(model, "axis_dim") is not None - inputs["position_ids"] = get_vision_position_ids(grid_thw, spatial_merge_size, include_temporal=include_temporal) + + # Only set the vision `position_ids` when nothing already has: a full (non-decomposed) multimodal + # export precomputes the LLM's `position_ids` via `get_rope_index` first, and the vision ids + # (different shape/meaning) must not overwrite it. On the decomposed vision sub-model nothing sets + # it first, so this still fires there. + if inputs.get("position_ids") is None: + # 3-axis (t, h, w) rotary encoders expose an ``axis_dim`` attr on their rotary_emb + # (minimax_m3_vl); default 2-axis (h, w) covers qwen2_5_vl / qwen3_vl / glm4v / paddleocr_vl. + include_temporal = _find_submodule_attr(model, "axis_dim") is not None + inputs["position_ids"] = get_vision_position_ids( + grid_thw, spatial_merge_size, include_temporal=include_temporal + ) window_size = _find_submodule_attr(model, "window_size") patch_size = _find_submodule_attr(model, "patch_size") From 7f0d32ed9786e102cd4e9c8543a20b5516531fbd Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 6 Jul 2026 21:47:11 +0200 Subject: [PATCH 10/14] Add Kimi-K2.5 exportability and standardize packed-vision attention Exporter support for Kimi-K2.5 (vision/audio attention registered with the reshaped vision-attention patch) plus: - Kimi vision rotary now emits the standard packed (seq, head_dim) cos/sin instead of the batched (1, seq, head_dim) LLM form, matching every other vision encoder (no exporter-side normalization needed). - ExecuTorch: materialise non-contiguous reshapes via a backend-local _patch_reshape rather than an unconditional clone in the shared vision patch, so ONNX/OpenVINO/dynamo graphs don't carry the copy. - is_multimodal short-circuits to False for non-PreTrainedModel inputs. - vision_utils.get_vision_cu_seqlens gains a merge_temporal option. Co-Authored-By: Claude Opus 4.8 --- src/transformers/exporters/exporter_dynamo.py | 6 +- .../exporters/exporter_executorch.py | 19 ++ src/transformers/exporters/utils.py | 46 +++- .../models/kimi_k25/modeling_kimi_k25.py | 225 ++++++++++-------- .../models/kimi_k25/modular_kimi_k25.py | 225 ++++++++++-------- src/transformers/vision_utils.py | 20 +- tests/exporters/test_export.py | 2 + .../models/kimi_k25/test_modeling_kimi_k25.py | 1 - 8 files changed, 332 insertions(+), 212 deletions(-) diff --git a/src/transformers/exporters/exporter_dynamo.py b/src/transformers/exporters/exporter_dynamo.py index 9dee78ad15bf..0e2c9643dde2 100644 --- a/src/transformers/exporters/exporter_dynamo.py +++ b/src/transformers/exporters/exporter_dynamo.py @@ -350,8 +350,8 @@ def _to_batched(t): enable_gqa=getattr(self, "num_key_value_heads", self.num_heads) != self.num_heads, ) - # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim) - attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1).contiguous() + # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim). + attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1) out_proj = self.proj if hasattr(self, "proj") else self.out_proj attn_output = out_proj(attn_output) @@ -381,6 +381,8 @@ def _to_batched(t): "transformers.models.glm_image.modeling_glm_image.GlmImageVisionAttention.forward", # Separate `.q` / `.k` / `.v` + single rotary tensor + `.proj` "transformers.models.qwen2_5_omni.modeling_qwen2_5_omni.Qwen2_5OmniVisionAttention.forward", + # Separate `q_proj`/`k_proj`/`v_proj` + `(cos, sin)` rotary + `.proj` (single return) + "transformers.models.kimi_k25.modeling_kimi_k25.Kimi_K25VisionAttention.forward", # Separate `_proj` + `(cos, sin)` rotary + `.out_proj` (tuple return) "transformers.models.video_llama_3.modeling_video_llama_3.VideoLlama3VisionAttention.forward", "transformers.models.paddleocr_vl.modeling_paddleocr_vl.PaddleOCRVisionAttention.forward", diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index bc6513555098..a42f5892bb04 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -418,6 +418,25 @@ def patch(self, *sizes): return patch +@register_patch("executorch", "torch.reshape", "torch.Tensor.reshape") +def _patch_reshape(original): + """Materialise a non-contiguous input before ``reshape``. + + ExecuTorch's edge-lowering reshape reference refuses a non-contiguous input (e.g. the + ``transpose(1, 2).reshape(...)`` in the packed vision-attention forward). A plain + ``.contiguous()`` gets folded away by functionalization, but a ``.clone()`` survives. Eager + ``reshape`` already copies a non-contiguous tensor, so this adds no extra work — it just moves + the copy where ExecuTorch's lowering needs it. + """ + + def patch(input, *shape): + if not input.is_contiguous(): + input = input.clone() + return original(input, *shape) + + return patch + + # ── Stage 3: ExecuTorch patches ─────────────────────────────────────────────── # Reversible swaps of ExecuTorch internals (passes, verifiers, op dicts) that crash # on legitimate dynamic-shape patterns: `SpecPropPass.update_placeholder_tensor_specs`, diff --git a/src/transformers/exporters/utils.py b/src/transformers/exporters/utils.py index ef876702daa5..98bf094f733d 100644 --- a/src/transformers/exporters/utils.py +++ b/src/transformers/exporters/utils.py @@ -432,9 +432,10 @@ def _prepare_grid_thw_vision_inputs(model: torch.nn.Module, inputs: dict[str, An `window_index`/`cu_window_seqlens` (XNet-style window attn) and `bilinear_indices`/`bilinear_weights` (interpolation-based merging). - Optional helpers are gated by the presence of their config attribute on the encoder - (`window_size`+`patch_size` for window attention, `num_grid_per_side` for bilinear), - so a model that doesn't use that feature won't get its kwarg injected. + Optional helpers are gated by a submodule attribute (`window_size`+`patch_size` for window + attention, `num_grid_per_side` for bilinear) or, for model-specific ones, by the encoder's + modeling module defining the helper (`get_vision_frame_index` / `get_vision_temporal_merge_index` + for kimi_k25) — so a model that doesn't use a feature won't get its kwarg injected. """ grid_thw = inputs["grid_thw"] spatial_merge_size = _find_submodule_attr(model, "spatial_merge_size") @@ -443,7 +444,12 @@ def _prepare_grid_thw_vision_inputs(model: torch.nn.Module, inputs: dict[str, An # none (its encoder hard-codes `1` because spatial merging happens in the projector). spatial_merge_size = inputs.get("merge_sizes", 1) - inputs["cu_seqlens"] = get_vision_cu_seqlens(grid_thw) + # kimi_k25-style encoders define their own per-frame / temporal-merge precompute helpers in their + # modeling module (resolved below) and attend over the whole clip, so `cu_seqlens` is per-clip + # (matching the encoder's util call). Other grid_thw encoders lack these and stay per-frame. + module = sys.modules[type(model).__module__] + temporal_encoder = hasattr(module, "get_vision_frame_index") + inputs["cu_seqlens"] = get_vision_cu_seqlens(grid_thw, merge_temporal=temporal_encoder) # 3-axis (t, h, w) rotary encoders expose an ``axis_dim`` attr on their rotary_emb # (minimax_m3_vl); default 2-axis (h, w) covers qwen2_5_vl / qwen3_vl / glm4v / paddleocr_vl. include_temporal = _find_submodule_attr(model, "axis_dim") is not None @@ -458,9 +464,27 @@ def _prepare_grid_thw_vision_inputs(model: torch.nn.Module, inputs: dict[str, An num_grid_per_side = _find_submodule_attr(model, "num_grid_per_side") if num_grid_per_side is not None: - inputs["bilinear_indices"], inputs["bilinear_weights"] = get_vision_bilinear_indices_and_weights( - grid_thw, num_grid_per_side, spatial_merge_size + if hasattr(module, "get_vision_bicubic_indices_and_weights"): + # kimi_k25 resamples its learned grid bicubically (helper defined in its own module). + inputs["bicubic_indices"], inputs["bicubic_weights"] = module.get_vision_bicubic_indices_and_weights( + grid_thw, num_grid_per_side + ) + else: + inputs["bilinear_indices"], inputs["bilinear_weights"] = get_vision_bilinear_indices_and_weights( + grid_thw, num_grid_per_side, spatial_merge_size + ) + + # Per-frame additive position table (kimi_k25): gathered by frame index instead of a per-clip loop. + if temporal_encoder: + inputs["frame_index"] = module.get_vision_frame_index(grid_thw) + + # Temporal-pooling spatial merger (kimi_k25): one gather index replaces its per-clip merge loop. + if hasattr(module, "get_vision_temporal_merge_index"): + merge_kernel_size = _find_submodule_attr(model, "merge_kernel_size") + kernel_height, kernel_width = ( + merge_kernel_size if not isinstance(merge_kernel_size, int) else (merge_kernel_size, merge_kernel_size) ) + inputs["temporal_merge_index"] = module.get_vision_temporal_merge_index(grid_thw, kernel_height, kernel_width) @register_export_input_preparer("target_sizes") @@ -684,9 +708,13 @@ def _find_multimodal_submodules(model: PreTrainedModel) -> dict[str, torch.nn.Mo return found -def is_multimodal(model: PreTrainedModel) -> bool: - """Returns `True` if the model is multi-modal with modal encoders and a language model.""" - return bool(_find_multimodal_submodules(model)) +def is_multimodal(model: PreTrainedModel | torch.nn.Module) -> bool: + """Returns `True` if the model is multi-modal with modal encoders and a language model. + + A non-`PreTrainedModel` (e.g. a bare `nn.Module`) has no canonical `get_encoder`/`get_decoder` + accessors and is trivially not multi-modal, so it short-circuits to `False`. + """ + return isinstance(model, PreTrainedModel) and bool(_find_multimodal_submodules(model)) def decompose_multimodal(model: PreTrainedModel, inputs: dict[str, Any]) -> dict[str, tuple[torch.nn.Module, dict]]: diff --git a/src/transformers/models/kimi_k25/modeling_kimi_k25.py b/src/transformers/models/kimi_k25/modeling_kimi_k25.py index 72d5bba859d9..0f77136adefd 100644 --- a/src/transformers/models/kimi_k25/modeling_kimi_k25.py +++ b/src/transformers/models/kimi_k25/modeling_kimi_k25.py @@ -43,7 +43,7 @@ from ...utils.deprecation import deprecate_kwarg from ...utils.generic import is_flash_attention_requested, maybe_autocast from ...utils.output_capturing import capture_outputs -from ...vision_utils import get_vision_position_ids +from ...vision_utils import get_vision_cu_seqlens, get_vision_position_ids from ..auto import AutoModel from .configuration_kimi_k25 import Kimi_K25Config, Kimi_K25VisionConfig @@ -99,6 +99,74 @@ class Kimi_K25CausalLMOutputWithPast(ModelOutput): image_hidden_states: torch.FloatTensor | None = None +def get_vision_bicubic_indices_and_weights( + grid_thw: torch.Tensor, num_grid_per_side: int, kwargs: dict | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-patch 16-tap bicubic gather indices/weights for resampling a square learned + `(num_grid_per_side, num_grid_per_side)` position-embedding table to each image's `(h, w)`, + or pop `"bicubic_indices"`/`"bicubic_weights"` from `kwargs`. + + Reproduces `F.interpolate(mode="bicubic", align_corners=False)` (Keys cubic kernel, `a=-0.75`) + as `(total_patches, 16)` indices + weights, consumed by a single fused `F.embedding_bag`. Fully + vectorised over packed patches (ragged `(h, w)` handled with `repeat_interleave`, no per-image + loop), so it traces and supports dynamic shapes like the other grid_thw precompute helpers. + """ + if kwargs is not None: + bicubic_indices = kwargs.pop("bicubic_indices", None) + bicubic_weights = kwargs.pop("bicubic_weights", None) + if bicubic_indices is not None and bicubic_weights is not None: + return bicubic_indices, bicubic_weights + + a = -0.75 + side = num_grid_per_side + device = grid_thw.device + offsets = torch.arange(-1, 3, device=device) # the 4 bicubic taps: floor-1 .. floor+2 + + def cubic_weights(distance): + # Keys convolution kernel (a=-0.75): near lobe for |d| <= 1, far lobe for 1 < |d| < 2. + near = ((a + 2) * distance - (a + 3)) * distance * distance + 1 + far = ((a * distance - 5 * a) * distance + 8 * a) * distance - 4 * a + return torch.where(distance <= 1, near, far) + + def axis_taps_weights(index, size): + src = (index + 0.5) * side / size - 0.5 # source coordinate, align_corners=False + floor = torch.floor(src) + taps = (floor.long()[:, None] + offsets).clamp(0, side - 1) # (total, 4) + return taps, cubic_weights((src[:, None] - floor[:, None] - offsets).abs()) + + # Per-patch (row, col) within its image, derived from packed offsets — no per-image loop. + counts = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] + heights = torch.repeat_interleave(grid_thw[:, 1], counts) + widths = torch.repeat_interleave(grid_thw[:, 2], counts) + starts = torch.repeat_interleave(F.pad(counts.cumsum(0)[:-1], (1, 0)), counts) + within = (torch.arange(counts.sum(), device=device) - starts) % (heights * widths) + h_taps, h_weights = axis_taps_weights(within // widths, heights) + w_taps, w_weights = axis_taps_weights(within % widths, widths) + # 2D separable: outer of the 4 h-taps × 4 w-taps → 16 taps per patch. + bicubic_indices = (h_taps[:, :, None] * side + w_taps[:, None, :]).reshape(-1, 16) + bicubic_weights = (h_weights[:, :, None] * w_weights[:, None, :]).reshape(-1, 16) + return bicubic_indices, bicubic_weights + + +def get_vision_frame_index(grid_thw: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor: + """Per-patch index into a temporal embedding table whose row `0` is a zero pad, or pop + `"frame_index"` from `kwargs`. + + Single-frame clips (`t == 1`, images) map every patch to `0` (no temporal term); frame `f` of a + multi-frame clip maps to `f + 1`. Precomputable, so the encoder avoids a per-clip `if t > 1` loop. + """ + if kwargs is not None and (frame_index := kwargs.pop("frame_index", None)) is not None: + return frame_index + device = grid_thw.device + parts = [] + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + # t == 1 → [0] (padded row 0 = zero); t > 1 → [1..t] → time_emb[0..t-1] + frames = torch.arange(t, device=device) + int(t > 1) + parts.append(frames.repeat_interleave(h * w)) + return torch.cat(parts) + + class Kimi_K25VisionPositionEmbeddings(nn.Module): def __init__(self, config): super().__init__() @@ -108,6 +176,9 @@ def __init__(self, config): self.position_embeddings = nn.Parameter( torch.zeros(config.pos_emb_height, config.pos_emb_width, config.hidden_size) ) + # Side of the (square) learned grid; the exporter's input preparer reads it to precompute + # the bicubic gather indices. + self.num_grid_per_side = config.pos_emb_height # Time-axis pos_emb are an additive sinusoidal table, i.e. add pos to hiddens rather than rotating time_position_embeddings = self.compute_pos_embed() @@ -120,35 +191,22 @@ def compute_pos_embed(self): pos_embed = torch.cat([freqs.sin(), freqs.cos()], dim=1) # (M, D) return pos_embed.unsqueeze(1) - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - pos_embs = [] - for t, h, w in grid_thw.tolist(): - if t > self.num_frames: - raise ValueError( - f"Got an input with {t} frames. Number of frames should be less than config.pos_emb_time=({self.num_frames})" - ) - - # Apply learned positions on h/w grids with optional interpolation for bigger images - if (h, w) == self.position_embeddings.shape[:-1]: - position_embeddings = self.position_embeddings.flatten(0, 1) - else: - position_embeddings = self.position_embeddings.permute(2, 0, 1).unsqueeze(0) - position_embeddings = F.interpolate( - position_embeddings, - size=(h, w), - mode="bicubic", - ) - position_embeddings = position_embeddings.squeeze(0).permute(1, 2, 0).flatten(0, 1) - - position_embeddings = position_embeddings.unsqueeze(0) # Add T axis - # Add sinusoidal positions for time grid if processing videos - if t > 1: - position_embeddings = position_embeddings.repeat(t, 1, 1) - position_embeddings = position_embeddings + self.time_position_embeddings[0:t] - - pos_embs.append(position_embeddings.flatten(0, 1)) - hidden_states = hidden_states + torch.cat(pos_embs, dim=0).to(hidden_states.dtype) - return hidden_states + def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + # Spatial: bicubically resample the learned grid to each image's (h, w) as a fused weighted + # gather (`embedding_bag`), equivalent to a per-image `F.interpolate(mode="bicubic")` but a + # single traceable op over all patches — and faster. + table = self.position_embeddings.flatten(0, 1) + bicubic_indices, bicubic_weights = get_vision_bicubic_indices_and_weights( + grid_thw, self.num_grid_per_side, kwargs=kwargs + ) + pos = F.embedding_bag(bicubic_indices, table, per_sample_weights=bicubic_weights.to(table.dtype), mode="sum") + # Temporal: add a per-frame sinusoid. Row 0 of the table is a zero pad, so single-frame clips + # (frame index 0) get none. + time_table = torch.cat( + [self.time_position_embeddings.new_zeros(1, self.dim), self.time_position_embeddings.squeeze(1)] + ) + pos = pos + time_table[get_vision_frame_index(grid_thw, kwargs=kwargs)] + return hidden_states + pos.to(hidden_states.dtype) class Kimi_K25VisionPatchEmbed(nn.Module): @@ -160,9 +218,9 @@ def __init__(self, config): self.proj = nn.Conv2d(3, config.hidden_size, kernel_size=patch_size, stride=patch_size) self.pos_emb = Kimi_K25VisionPositionEmbeddings(config) - def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: hidden_states = self.proj(pixel_values).view(pixel_values.size(0), -1) - hidden_states = self.pos_emb(hidden_states, grid_thw) + hidden_states = self.pos_emb(hidden_states, grid_thw, **kwargs) return hidden_states @@ -223,17 +281,14 @@ def compute_default_rope_parameters( @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): - position_ids_expanded = position_ids.permute(1, 2, 0)[..., None].float() # shape (bs, positions, 2, 1) + position_ids_expanded = position_ids.transpose(0, 1)[..., None].float() # (positions, 2, 1) inv_freq_expanded = ( - self.inv_freq[None, None, None, :] - .float() - .expand(position_ids_expanded.shape[0], position_ids_expanded.shape[1], 2, -1) - .to(x.device) - ) # shape (bs, positions, 2, freq_dim) + self.inv_freq[None, None, :].float().expand(position_ids_expanded.shape[0], 2, -1).to(x.device) + ) # (positions, 2, freq_dim) device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with maybe_autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(2, 3).flatten(2) + freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(1, 2).flatten(1) emb = torch.cat([freqs, freqs], dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling @@ -445,6 +500,33 @@ def _init_weights(self, module): init.trunc_normal_(module.position_embeddings, mean=0.0) +def get_vision_temporal_merge_index( + grid_thw: torch.Tensor, kernel_height: int, kernel_width: int, kwargs: dict | None = None +) -> torch.Tensor: + """Gather index regrouping a flat patch sequence into `(total_merged, t, kernel_height * + kernel_width)` for the temporal-pooling merger, or pop `"temporal_merge_index"` from `kwargs`. + + Row `m` collects the `t` frames × `kernel_height*kernel_width` source patches that pool into + merged token `m`; the caller means over the frame axis. Precomputable, so the encoder avoids a + per-clip `grid_thw.tolist()` loop. + """ + if kwargs is not None and (index := kwargs.pop("temporal_merge_index", None)) is not None: + return index + device = grid_thw.device + running, rows = 0, [] + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + new_h, new_w = h // kernel_height, w // kernel_width + base = torch.arange(running, running + t * h * w, device=device).view( + t, new_h, kernel_height, new_w, kernel_width + ) + # (t, new_h, new_w, kh, kw) → (new_h*new_w, t, kh*kw): frame axis kept for the caller's mean. + base = base.permute(1, 3, 0, 2, 4).reshape(new_h * new_w, t, kernel_height * kernel_width) + rows.append(base) + running += t * h * w + return torch.cat(rows, dim=0) + + class Kimi_K25VisionModel(Kimi_K25PreTrainedModel): config: Kimi_K25VisionConfig input_modalities = ("image", "video") @@ -463,51 +545,6 @@ def __init__(self, config: Kimi_K25VisionConfig): self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-05) self.post_init() - def temporal_patch_merger( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor, - ) -> list[torch.Tensor]: - r""" - Merges temporal frames by spatially pooling patch embeddings across time. - - For each video clip defined by `grid_thw`, the method reshapes the flat patch sequence - into a `(T, H, W)` grid, averages over the temporal dimension, then rearranges spatial - patches into groups of `kernel_height * kernel_width` — matching the merged-token layout - expected by downstream layers. - - Args: - hidden_states (`torch.Tensor` of shape `(total_patches, hidden_dim)`): - Concatenated patch embeddings for all clips in the batch. `total_patches` equals - the sum of `t * h * w` over all entries in `grid_thw`. - grid_thw (`torch.Tensor` of shape `(batch_size, 3)`): - Temporal and spatial grid dimensions for each clip, where each row is - `(num_frames, grid_height, grid_width)`. `grid_height` and `grid_width` must be - divisible by `kernel_height` and `kernel_width` respectively. - - Returns: - `torch.Tensor` of shape `(total_merged_patches, kernel_height * kernel_width, hidden_dim)`: - Temporally pooled patch embeddings. `total_merged_patches` equals the sum of - `(h // kernel_height) * (w // kernel_width)` over all clips. - """ - hidden_dim = hidden_states.size(-1) - kernel_height, kernel_width = self.merge_kernel_size - - outputs = [] - running_length = 0 - for t, h, w in grid_thw.tolist(): - # Get the current sequence - seq = hidden_states[running_length : running_length + t * h * w] - # Reshape along self.merge_kernel_size and concat to the last dimension - new_height, new_width = h // kernel_height, w // kernel_width - reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, hidden_dim) - reshaped_seq = reshaped_seq.transpose(2, 3).mean(dim=0) # temporal pooling - padded_seq = reshaped_seq.reshape(new_height * new_width, kernel_height * kernel_width, -1) - outputs.append(padded_seq) - running_length += t * h * w - - return torch.cat(outputs, dim=0) - @capture_outputs @auto_docstring def forward( @@ -520,20 +557,13 @@ def forward( grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ - hidden_states = self.patch_embed(pixel_values, grid_thw=grid_thw) - position_ids = get_vision_position_ids(grid_thw, spatial_merge_size=1) - position_ids = position_ids.transpose(0, 1).flip(0)[:, None, :] + hidden_states = self.patch_embed(pixel_values, grid_thw=grid_thw, **kwargs) + position_ids = get_vision_position_ids(grid_thw, spatial_merge_size=1, kwargs=kwargs) + position_ids = position_ids.transpose(0, 1).flip(0) # (2, positions) position_embeddings = self.rotary_emb(hidden_states, position_ids) - lengths = torch.cat( - ( - torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device), - grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2], - ) - ) - - max_seqlen = lengths.max() - cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32) + cu_seqlens = get_vision_cu_seqlens(grid_thw, merge_temporal=True, kwargs=kwargs) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() for block in self.layers: hidden_states = block( @@ -545,7 +575,8 @@ def forward( ) hidden_states = self.final_layernorm(hidden_states) - pooled_hidden_states = self.temporal_patch_merger(hidden_states, grid_thw) + merge_index = get_vision_temporal_merge_index(grid_thw, *self.merge_kernel_size, kwargs=kwargs) + pooled_hidden_states = hidden_states[merge_index].mean(dim=1) return BaseModelOutputWithPooling( last_hidden_state=hidden_states, diff --git a/src/transformers/models/kimi_k25/modular_kimi_k25.py b/src/transformers/models/kimi_k25/modular_kimi_k25.py index 64a1dc24896b..fcb66df8710c 100644 --- a/src/transformers/models/kimi_k25/modular_kimi_k25.py +++ b/src/transformers/models/kimi_k25/modular_kimi_k25.py @@ -35,7 +35,7 @@ ) from ...utils.generic import is_flash_attention_requested, maybe_autocast from ...utils.output_capturing import capture_outputs -from ...vision_utils import get_vision_position_ids +from ...vision_utils import get_vision_cu_seqlens, get_vision_position_ids from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel from ..gemma4.modeling_gemma4 import Gemma4VisionRotaryEmbedding from ..glm4v.modeling_glm4v import Glm4vForConditionalGeneration @@ -54,6 +54,101 @@ logger = logging.get_logger(__name__) +def get_vision_bicubic_indices_and_weights( + grid_thw: torch.Tensor, num_grid_per_side: int, kwargs: dict | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-patch 16-tap bicubic gather indices/weights for resampling a square learned + `(num_grid_per_side, num_grid_per_side)` position-embedding table to each image's `(h, w)`, + or pop `"bicubic_indices"`/`"bicubic_weights"` from `kwargs`. + + Reproduces `F.interpolate(mode="bicubic", align_corners=False)` (Keys cubic kernel, `a=-0.75`) + as `(total_patches, 16)` indices + weights, consumed by a single fused `F.embedding_bag`. Fully + vectorised over packed patches (ragged `(h, w)` handled with `repeat_interleave`, no per-image + loop), so it traces and supports dynamic shapes like the other grid_thw precompute helpers. + """ + if kwargs is not None: + bicubic_indices = kwargs.pop("bicubic_indices", None) + bicubic_weights = kwargs.pop("bicubic_weights", None) + if bicubic_indices is not None and bicubic_weights is not None: + return bicubic_indices, bicubic_weights + + a = -0.75 + side = num_grid_per_side + device = grid_thw.device + offsets = torch.arange(-1, 3, device=device) # the 4 bicubic taps: floor-1 .. floor+2 + + def cubic_weights(distance): + # Keys convolution kernel (a=-0.75): near lobe for |d| <= 1, far lobe for 1 < |d| < 2. + near = ((a + 2) * distance - (a + 3)) * distance * distance + 1 + far = ((a * distance - 5 * a) * distance + 8 * a) * distance - 4 * a + return torch.where(distance <= 1, near, far) + + def axis_taps_weights(index, size): + src = (index + 0.5) * side / size - 0.5 # source coordinate, align_corners=False + floor = torch.floor(src) + taps = (floor.long()[:, None] + offsets).clamp(0, side - 1) # (total, 4) + return taps, cubic_weights((src[:, None] - floor[:, None] - offsets).abs()) + + # Per-patch (row, col) within its image, derived from packed offsets — no per-image loop. + counts = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] + heights = torch.repeat_interleave(grid_thw[:, 1], counts) + widths = torch.repeat_interleave(grid_thw[:, 2], counts) + starts = torch.repeat_interleave(F.pad(counts.cumsum(0)[:-1], (1, 0)), counts) + within = (torch.arange(counts.sum(), device=device) - starts) % (heights * widths) + h_taps, h_weights = axis_taps_weights(within // widths, heights) + w_taps, w_weights = axis_taps_weights(within % widths, widths) + # 2D separable: outer of the 4 h-taps × 4 w-taps → 16 taps per patch. + bicubic_indices = (h_taps[:, :, None] * side + w_taps[:, None, :]).reshape(-1, 16) + bicubic_weights = (h_weights[:, :, None] * w_weights[:, None, :]).reshape(-1, 16) + return bicubic_indices, bicubic_weights + + +def get_vision_frame_index(grid_thw: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor: + """Per-patch index into a temporal embedding table whose row `0` is a zero pad, or pop + `"frame_index"` from `kwargs`. + + Single-frame clips (`t == 1`, images) map every patch to `0` (no temporal term); frame `f` of a + multi-frame clip maps to `f + 1`. Precomputable, so the encoder avoids a per-clip `if t > 1` loop. + """ + if kwargs is not None and (frame_index := kwargs.pop("frame_index", None)) is not None: + return frame_index + device = grid_thw.device + parts = [] + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + # t == 1 → [0] (padded row 0 = zero); t > 1 → [1..t] → time_emb[0..t-1] + frames = torch.arange(t, device=device) + int(t > 1) + parts.append(frames.repeat_interleave(h * w)) + return torch.cat(parts) + + +def get_vision_temporal_merge_index( + grid_thw: torch.Tensor, kernel_height: int, kernel_width: int, kwargs: dict | None = None +) -> torch.Tensor: + """Gather index regrouping a flat patch sequence into `(total_merged, t, kernel_height * + kernel_width)` for the temporal-pooling merger, or pop `"temporal_merge_index"` from `kwargs`. + + Row `m` collects the `t` frames × `kernel_height*kernel_width` source patches that pool into + merged token `m`; the caller means over the frame axis. Precomputable, so the encoder avoids a + per-clip `grid_thw.tolist()` loop. + """ + if kwargs is not None and (index := kwargs.pop("temporal_merge_index", None)) is not None: + return index + device = grid_thw.device + running, rows = 0, [] + for t, h, w in grid_thw.tolist(): + t, h, w = int(t), int(h), int(w) + new_h, new_w = h // kernel_height, w // kernel_width + base = torch.arange(running, running + t * h * w, device=device).view( + t, new_h, kernel_height, new_w, kernel_width + ) + # (t, new_h, new_w, kh, kw) → (new_h*new_w, t, kh*kw): frame axis kept for the caller's mean. + base = base.permute(1, 3, 0, 2, 4).reshape(new_h * new_w, t, kernel_height * kernel_width) + rows.append(base) + running += t * h * w + return torch.cat(rows, dim=0) + + @auto_docstring(checkpoint="moonshotai/Kimi-K2.6") @strict class Kimi_K25VisionConfig(PreTrainedConfig): @@ -145,6 +240,9 @@ def __init__(self, config): self.position_embeddings = nn.Parameter( torch.zeros(config.pos_emb_height, config.pos_emb_width, config.hidden_size) ) + # Side of the (square) learned grid; the exporter's input preparer reads it to precompute + # the bicubic gather indices. + self.num_grid_per_side = config.pos_emb_height # Time-axis pos_emb are an additive sinusoidal table, i.e. add pos to hiddens rather than rotating time_position_embeddings = self.compute_pos_embed() @@ -157,35 +255,22 @@ def compute_pos_embed(self): pos_embed = torch.cat([freqs.sin(), freqs.cos()], dim=1) # (M, D) return pos_embed.unsqueeze(1) - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: - pos_embs = [] - for t, h, w in grid_thw.tolist(): - if t > self.num_frames: - raise ValueError( - f"Got an input with {t} frames. Number of frames should be less than config.pos_emb_time=({self.num_frames})" - ) - - # Apply learned positions on h/w grids with optional interpolation for bigger images - if (h, w) == self.position_embeddings.shape[:-1]: - position_embeddings = self.position_embeddings.flatten(0, 1) - else: - position_embeddings = self.position_embeddings.permute(2, 0, 1).unsqueeze(0) - position_embeddings = F.interpolate( - position_embeddings, - size=(h, w), - mode="bicubic", - ) - position_embeddings = position_embeddings.squeeze(0).permute(1, 2, 0).flatten(0, 1) - - position_embeddings = position_embeddings.unsqueeze(0) # Add T axis - # Add sinusoidal positions for time grid if processing videos - if t > 1: - position_embeddings = position_embeddings.repeat(t, 1, 1) - position_embeddings = position_embeddings + self.time_position_embeddings[0:t] - - pos_embs.append(position_embeddings.flatten(0, 1)) - hidden_states = hidden_states + torch.cat(pos_embs, dim=0).to(hidden_states.dtype) - return hidden_states + def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + # Spatial: bicubically resample the learned grid to each image's (h, w) as a fused weighted + # gather (`embedding_bag`), equivalent to a per-image `F.interpolate(mode="bicubic")` but a + # single traceable op over all patches — and faster. + table = self.position_embeddings.flatten(0, 1) + bicubic_indices, bicubic_weights = get_vision_bicubic_indices_and_weights( + grid_thw, self.num_grid_per_side, kwargs=kwargs + ) + pos = F.embedding_bag(bicubic_indices, table, per_sample_weights=bicubic_weights.to(table.dtype), mode="sum") + # Temporal: add a per-frame sinusoid. Row 0 of the table is a zero pad, so single-frame clips + # (frame index 0) get none. + time_table = torch.cat( + [self.time_position_embeddings.new_zeros(1, self.dim), self.time_position_embeddings.squeeze(1)] + ) + pos = pos + time_table[get_vision_frame_index(grid_thw, kwargs=kwargs)] + return hidden_states + pos.to(hidden_states.dtype) class Kimi_K25VisionPatchEmbed(nn.Module): @@ -197,9 +282,9 @@ def __init__(self, config): self.proj = nn.Conv2d(3, config.hidden_size, kernel_size=patch_size, stride=patch_size) self.pos_emb = Kimi_K25VisionPositionEmbeddings(config) - def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: hidden_states = self.proj(pixel_values).view(pixel_values.size(0), -1) - hidden_states = self.pos_emb(hidden_states, grid_thw) + hidden_states = self.pos_emb(hidden_states, grid_thw, **kwargs) return hidden_states @@ -207,17 +292,14 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.T # The difference is that gemma4 stacks H/W embeds on `dim`, while Kimi interleaves them class Kimi_K25VisionRotaryEmbedding(Gemma4VisionRotaryEmbedding): def forward(self, x, position_ids): - position_ids_expanded = position_ids.permute(1, 2, 0)[..., None].float() # shape (bs, positions, 2, 1) + position_ids_expanded = position_ids.transpose(0, 1)[..., None].float() # (positions, 2, 1) inv_freq_expanded = ( - self.inv_freq[None, None, None, :] - .float() - .expand(position_ids_expanded.shape[0], position_ids_expanded.shape[1], 2, -1) - .to(x.device) - ) # shape (bs, positions, 2, freq_dim) + self.inv_freq[None, None, :].float().expand(position_ids_expanded.shape[0], 2, -1).to(x.device) + ) # (positions, 2, freq_dim) device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with maybe_autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(2, 3).flatten(2) + freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(1, 2).flatten(1) emb = torch.cat([freqs, freqs], dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling @@ -349,51 +431,6 @@ def __init__(self, config: Kimi_K25VisionConfig): self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-05) self.post_init() - def temporal_patch_merger( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor, - ) -> list[torch.Tensor]: - r""" - Merges temporal frames by spatially pooling patch embeddings across time. - - For each video clip defined by `grid_thw`, the method reshapes the flat patch sequence - into a `(T, H, W)` grid, averages over the temporal dimension, then rearranges spatial - patches into groups of `kernel_height * kernel_width` — matching the merged-token layout - expected by downstream layers. - - Args: - hidden_states (`torch.Tensor` of shape `(total_patches, hidden_dim)`): - Concatenated patch embeddings for all clips in the batch. `total_patches` equals - the sum of `t * h * w` over all entries in `grid_thw`. - grid_thw (`torch.Tensor` of shape `(batch_size, 3)`): - Temporal and spatial grid dimensions for each clip, where each row is - `(num_frames, grid_height, grid_width)`. `grid_height` and `grid_width` must be - divisible by `kernel_height` and `kernel_width` respectively. - - Returns: - `torch.Tensor` of shape `(total_merged_patches, kernel_height * kernel_width, hidden_dim)`: - Temporally pooled patch embeddings. `total_merged_patches` equals the sum of - `(h // kernel_height) * (w // kernel_width)` over all clips. - """ - hidden_dim = hidden_states.size(-1) - kernel_height, kernel_width = self.merge_kernel_size - - outputs = [] - running_length = 0 - for t, h, w in grid_thw.tolist(): - # Get the current sequence - seq = hidden_states[running_length : running_length + t * h * w] - # Reshape along self.merge_kernel_size and concat to the last dimension - new_height, new_width = h // kernel_height, w // kernel_width - reshaped_seq = seq.view(t, new_height, kernel_height, new_width, kernel_width, hidden_dim) - reshaped_seq = reshaped_seq.transpose(2, 3).mean(dim=0) # temporal pooling - padded_seq = reshaped_seq.reshape(new_height * new_width, kernel_height * kernel_width, -1) - outputs.append(padded_seq) - running_length += t * h * w - - return torch.cat(outputs, dim=0) - @capture_outputs @auto_docstring def forward( @@ -406,20 +443,13 @@ def forward( grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. """ - hidden_states = self.patch_embed(pixel_values, grid_thw=grid_thw) - position_ids = get_vision_position_ids(grid_thw, spatial_merge_size=1) - position_ids = position_ids.transpose(0, 1).flip(0)[:, None, :] + hidden_states = self.patch_embed(pixel_values, grid_thw=grid_thw, **kwargs) + position_ids = get_vision_position_ids(grid_thw, spatial_merge_size=1, kwargs=kwargs) + position_ids = position_ids.transpose(0, 1).flip(0) # (2, positions) position_embeddings = self.rotary_emb(hidden_states, position_ids) - lengths = torch.cat( - ( - torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device), - grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2], - ) - ) - - max_seqlen = lengths.max() - cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32) + cu_seqlens = get_vision_cu_seqlens(grid_thw, merge_temporal=True, kwargs=kwargs) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() for block in self.layers: hidden_states = block( @@ -431,7 +461,8 @@ def forward( ) hidden_states = self.final_layernorm(hidden_states) - pooled_hidden_states = self.temporal_patch_merger(hidden_states, grid_thw) + merge_index = get_vision_temporal_merge_index(grid_thw, *self.merge_kernel_size, kwargs=kwargs) + pooled_hidden_states = hidden_states[merge_index].mean(dim=1) return BaseModelOutputWithPooling( last_hidden_state=hidden_states, diff --git a/src/transformers/vision_utils.py b/src/transformers/vision_utils.py index 6293cdc7b290..dd66ff72be14 100644 --- a/src/transformers/vision_utils.py +++ b/src/transformers/vision_utils.py @@ -32,22 +32,30 @@ import torch.nn.functional as F -def get_vision_cu_seqlens(grid_thw: torch.Tensor, kwargs: dict | None = None) -> torch.Tensor: +def get_vision_cu_seqlens( + grid_thw: torch.Tensor, merge_temporal: bool = False, kwargs: dict | None = None +) -> torch.Tensor: """Get cumulative sequence lengths from vision grid info, or pop from `kwargs` if precomputed. Args: grid_thw: `(num_images_or_videos, 3)` — temporal, height, width per entry. + merge_temporal: when `False` (default), each frame is its own attention segment (`h * w` + per frame, `t` segments per entry — the qwen2_vl / glm4v convention). When `True`, + the whole clip is a single segment (`t * h * w`), i.e. attention spans all frames + jointly (the kimi_k25 convention). kwargs: optional caller kwargs — if it contains `"cu_seqlens"` it is popped and returned. Returns: - `cu_seqlens`: `(total_patches + 1,)` int32 cumulative sequence boundaries. + `cu_seqlens`: `(num_segments + 1,)` int32 cumulative sequence boundaries. """ if kwargs is not None and (cu_seqlens := kwargs.pop("cu_seqlens", None)) is not None: return cu_seqlens - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( - dim=0, dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32 - ) - return F.pad(cu_seqlens, (1, 0), value=0) + dtype = grid_thw.dtype if torch.jit.is_tracing() else torch.int32 + if merge_temporal: + seqlens = grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] + else: + seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + return F.pad(seqlens.cumsum(dim=0, dtype=dtype), (1, 0), value=0) def get_vision_position_ids( diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 18affeac92e9..60d67cb01c48 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -181,6 +181,8 @@ "InstructBlipForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", "InstructBlipVideoForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.", "InstructBlipVideoModel": "Same `spec` failure as `BarkFineModel`.", + "Kimi_K25Model": "Same `spec` failure as `BarkFineModel`, in the DeepseekV3 language model (vision encoder exports fine).", + "Kimi_K25ForConditionalGeneration": "Same `spec` failure as `BarkFineModel`, in the DeepseekV3 language model (vision encoder exports fine).", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", "MiniMaxM3VLModel": ("Serialization rejects an i64 constant (`bad number for type int32`)."), diff --git a/tests/models/kimi_k25/test_modeling_kimi_k25.py b/tests/models/kimi_k25/test_modeling_kimi_k25.py index d44bcb1dbc60..0ec5b6f578ef 100644 --- a/tests/models/kimi_k25/test_modeling_kimi_k25.py +++ b/tests/models/kimi_k25/test_modeling_kimi_k25.py @@ -127,7 +127,6 @@ def get_additional_inputs(self, config, input_ids, pixel_values): @require_torch class Kimi_K25ModelTest(VLMModelTest, unittest.TestCase): model_tester_class = Kimi_K25VisionText2TextModelTester - test_torch_exportable = False # Kimi has images shaped as (bs*patch_len, dim) so we can't slice to batches in generate def prepare_config_and_inputs_for_generate(self, batch_size=2): From fe9066d1ea545700e4e4bc3b7f8369ae05e2d666 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 7 Jul 2026 08:19:15 +0200 Subject: [PATCH 11/14] more passing models --- docs/source/en/exporters.md | 20 +----------- src/transformers/exporters/exporter_onnx.py | 7 +++-- .../models/videomae/modeling_videomae.py | 6 +++- tests/exporters/test_export.py | 31 +++++-------------- 4 files changed, 19 insertions(+), 45 deletions(-) diff --git a/docs/source/en/exporters.md b/docs/source/en/exporters.md index 7aa4459938b6..901b3132e3d7 100644 --- a/docs/source/en/exporters.md +++ b/docs/source/en/exporters.md @@ -458,9 +458,7 @@ visible from the public `export()` API, but the most common things to know: - `grouped_mm` traces fine through `DynamoExporter` and is auto-translated for `OnnxExporter`; for `ExecutorchExporter` with the XNNPACK backend, the exporter swaps MoE experts to `batched_mm` because XNNPACK has no `_grouped_mm.out` kernel. -- A short list of models (`EXPORT_SKIP_MODEL_CLASSES`) is skipped from the export sweep when - the model itself is fundamentally non-exportable; each entry carries a TODO with the - model-side change needed. +- Not every architecture exports cleanly yet — a few hit data-dependent control flow that can't be vectorised, or exceed practical export time under dynamic shapes. When that happens the failure surfaces at `export()` time with a concrete error, not silently.
Export pipeline — internals (per-backend stages and how to extend) @@ -544,22 +542,6 @@ The split is intentional: translation, an ORT validation quirk, an FX decomposition that emits a dead op. Keep the workaround in the exporter and the modeling code stays clean. -### Known upstream workarounds - -A small number of model classes hit confirmed bugs in `onnxscript`'s graph optimizer -(constant folding crashing on `SplitToSequence`, FPN initialisers being dropped). For those, -ONNX optimisation is selectively disabled via -[`ONNX_DISABLE_OPTIMIZE_MODEL_CLASSES`](https://github.com/huggingface/transformers/blob/main/tests/exporters/test_utils.py) -in the test suite — each entry is annotated with the upstream issue it works around. This -list is **expected to shrink** as upstream bugs land; it is not an extension point for -arbitrary skipping, and new entries should reference a specific upstream bug. - -A second list, [`EXPORT_SKIP_MODEL_CLASSES`](https://github.com/huggingface/transformers/blob/main/tests/exporters/test_utils.py), -opts a handful of model classes out of the entire export sweep when the model itself is -fundamentally non-exportable as-is (data-dependent control flow that can't be vectorised, -modules treated as forward arguments, …). Same expectations: every entry carries a TODO -naming the underlying model change needed; the list should shrink, not grow. -
## API reference diff --git a/src/transformers/exporters/exporter_onnx.py b/src/transformers/exporters/exporter_onnx.py index 0f97cea51997..7af8349aa783 100644 --- a/src/transformers/exporters/exporter_onnx.py +++ b/src/transformers/exporters/exporter_onnx.py @@ -210,9 +210,12 @@ def patch(condition, x=None, y=None): if isinstance(x, torch.Tensor) and isinstance(y, torch.Tensor) and x.dtype != y.dtype: y = y.to(x.dtype) elif isinstance(x, torch.Tensor) and isinstance(y, (int, float, bool)): - y = torch.tensor(y, dtype=x.dtype, device=x.device) + # `full_like` (a traced op) rather than `torch.tensor(...)` (a fresh leaf constant): the + # latter, if materialised during `run_decompositions`' retrace, becomes an unregistered + # `_tensor_constant` → `alias` → `detach_` that trips aot's functional-graph assertion. + y = torch.full_like(x, y) elif isinstance(y, torch.Tensor) and isinstance(x, (int, float, bool)): - x = torch.tensor(x, dtype=y.dtype, device=y.device) + x = torch.full_like(y, x) if x is None and y is None: return original(condition) elif y is None: diff --git a/src/transformers/models/videomae/modeling_videomae.py b/src/transformers/models/videomae/modeling_videomae.py index 35bcce23c973..df033c6f892e 100755 --- a/src/transformers/models/videomae/modeling_videomae.py +++ b/src/transformers/models/videomae/modeling_videomae.py @@ -28,7 +28,7 @@ from ...modeling_outputs import BaseModelOutput, ImageClassifierOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging, torch_compilable_check from ...utils.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from ...utils.generic import can_return_tuple, merge_with_config_defaults from ...utils.output_capturing import capture_outputs @@ -661,6 +661,10 @@ def forward( labels = videos_patch[bool_masked_pos].reshape(batch_size, -1, num_channels) loss_fct = MSELoss() + torch_compilable_check( + logits.shape[1] == labels.shape[1], + "VideoMAE reconstruction logits and labels must cover the same number of masked patches.", + ) loss = loss_fct(logits, labels) return VideoMAEForPreTrainingOutput( diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 15ced7f28a2b..74e8126d611f 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -61,14 +61,7 @@ EXPORT_SKIPS: dict[str, dict[str, str]] = { # Every backend, every variant. - "all": { - "VideoMAEForPreTraining": ( - "The reconstruction-loss path indexes `videos_patch[bool_masked_pos]` (a boolean " - "index → data-dependent unbacked count) and compares its shape against `logits`, which " - "`torch.export` can't guard (`Eq(u2 // 13, u3)`). Only the pretraining head hits this; " - "`VideoMAEModel` / `VideoMAEForVideoClassification` export fine." - ), - }, + "all": {}, # Any backend (incl. plain `torch.export`), dynamic-shape variant only. "dynamic": { "Sam2Model": ( @@ -90,25 +83,11 @@ ), }, # ONNX, every variant. - "onnx": { - "CHMv2ForDepthEstimation": ( - "`run_decompositions` retraces through aot_autograd which emits a `detach_(alias(...))` " - "pair the functional-graph assertion rejects (independent of any source `.detach()` — " - "verified). Torch export works. TODO: file upstream `torch.export` issue." - ), - }, + "onnx": {}, # ONNX, generate path only. "onnx.generate": {}, # ONNX, dynamic-shape only. "onnx.dynamic": { - "GroundingDinoModel": ( - "Same `detach_(alias(...))` retrace bug as CHMv2, but only triggered under dynamic " - "shapes — `aot_autograd`'s decomposition pipeline emits the detach itself (verified " - "by guarding all three modeling-side detaches with `if self.training`). Static works." - ), - "GroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.", - "MMGroundingDinoModel": "Same as `GroundingDinoModel`.", - "MMGroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.", "SwinModel": ( "Shifted-window attention on symbolic H/W: `torch.export` + onnxscript exceed the " "1000s test timeout under dynamic shapes (static exports fine)." @@ -131,6 +110,12 @@ }, # ExecuTorch, every variant. "executorch": { + "VideoMAEForPreTraining": ( + "Torch/ONNX/OpenVINO export fine (a `torch._check` states the logits/labels masked-token " + "count invariant), but ExecuTorch edge-lowering still fails on the decoder's data-dependent " + "negative slice `hidden_states[:, -return_token_num:]` — `aten.slice_copy`'s meta raises " + "`GuardOnDataDependentSymNode` on the symbolic masked-token count." + ), "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", From 30a151c94a5cdafd6d27b9677ecb0a40ea52c422 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 7 Jul 2026 09:56:28 +0200 Subject: [PATCH 12/14] more fixes --- src/transformers/cache_utils.py | 6 +- src/transformers/exporters/exporter_dynamo.py | 2 - .../exporters/exporter_executorch.py | 102 +++++++++++++++--- .../configuration_recurrent_gemma.py | 4 + .../modeling_recurrent_gemma.py | 99 ++++++----------- tests/exporters/test_export.py | 24 +---- 6 files changed, 128 insertions(+), 109 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 568cfce3abaa..0fa4145e00c6 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -951,8 +951,10 @@ def lazy_initialization( self.is_conv_states_initialized = True if recurrent_states is not None: - # The shape is always static, so we init as such - self.recurrent_states = torch.zeros_like(recurrent_states, dtype=self.dtype, device=self.device) + # The shape is always static, so we init as such. Preserve the recurrent tensor's OWN + # dtype (not `self.dtype`, which follows the conv state): models like recurrent_gemma + # accumulate the recurrence in float32 while the conv state stays in the compute dtype. + self.recurrent_states = torch.zeros_like(recurrent_states) # Mark as static address to be able to use cudagraphs if not is_torchdynamo_compiling(): torch._dynamo.mark_static_address(self.recurrent_states) diff --git a/src/transformers/exporters/exporter_dynamo.py b/src/transformers/exporters/exporter_dynamo.py index a14a4f7ef128..00fa9ec19694 100644 --- a/src/transformers/exporters/exporter_dynamo.py +++ b/src/transformers/exporters/exporter_dynamo.py @@ -675,8 +675,6 @@ def get_auto_dynamic_shapes(inputs: Any) -> Any: "_prefill_len", # glm_image (m-rope prefill length) "cached_rotary_positional_embedding", # wav2vec2_bert, seamless_m4t, clvp "cached_sequence_length", # wav2vec2_bert, seamless_m4t, clvp - "conv1d_state", # recurrent_gemma (conv state on RecurrentGemmaRecurrentBlock) - "recurrent_states", # recurrent_gemma (RG-LRU state on RecurrentGemmaRglru) ) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index ef915d150bae..688718880c44 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -38,6 +38,7 @@ import math import operator +import re from collections.abc import MutableMapping from typing import Any @@ -60,7 +61,14 @@ if is_torch_available(): import torch from torch.export import ExportedProgram - from torch.fx.experimental.symbolic_shapes import guard_or_true + from torch.fx.experimental.symbolic_shapes import ( + free_symbols, + free_unbacked_symbols, + guard_or_false, + guard_or_true, + statically_known_true, + ) + from torch.fx.passes.infra.pass_base import PassResult from torch.nn.attention import SDPBackend, sdpa_kernel from torch.utils._sympy.numbers import IntInfinity from torch.utils._sympy.value_ranges import ValueRanges @@ -73,9 +81,19 @@ from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner + from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( # type: ignore[import-not-found] + XNNStaticReshape, + XNode, + ) + from executorch.backends.xnnpack.utils.utils import get_input_node from executorch.exir.capture._config import EdgeCompileConfig + from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.passes.executorch_prim_ops_registry import _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS + from executorch.exir.passes.replace_view_copy_with_view_pass import _VIEW_OP, _is_view_copy, _ViewSpec + from executorch.exir.passes.spec_prop_pass import _is_mutable_buffer from executorch.exir.program import EdgeProgramManager, ExecutorchProgramManager, to_edge_transform_and_lower + from executorch.exir.sym_util import eval_expr + from executorch.exir.tensor import determine_tensor_dynanism logger = logging.get_logger(__name__) @@ -367,7 +385,6 @@ def _patch_scaled_dot_product_attention(original): (both batch ``u0``) with plain broadcasting, so no ``Eq(u0, 1)`` guard is needed and no SDPA node survives to be re-decomposed. """ - from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols def _has_unbacked_batch(t): # True when ``t``'s batch dim is a data-dependent (unbacked, ``u*``) SymInt. @@ -489,7 +506,6 @@ def _patch_eval_upper_bound(original): — the same trace-proportional heuristic ``_fix_range_constraints`` applies to the per-symbol ranges — so planned buffers stay proportional to the sampled inputs. """ - from executorch.exir.sym_util import eval_expr def patch(maybe_symint): if isinstance(maybe_symint, int): @@ -514,7 +530,6 @@ def _patch_remove_empty_tensors_from_cat(_original): either way at trace time. Using ``guard_or_true`` keeps unbacked-shape inputs conservatively (the pass is purely an optimisation). """ - from executorch.exir.dialects._ops import ops as exir_ops def patch(self, graph_module, cat_node): pruned = [arg for arg in cat_node.args[0] if guard_or_true(arg.meta["val"].numel() != 0)] @@ -575,7 +590,6 @@ def _patch_dim_order_from_stride(_original): so the sort still produces *a* dim order when the comparison is unbacked — the exact order on unbacked dims doesn't affect correctness, just memory layout. """ - from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true def patch(stride): for s in stride: @@ -609,7 +623,6 @@ def _patch_update_placeholder_tensor_specs(_original): (``val`` is ``None``) and the assignment raises ``AttributeError``. Skip ``None`` specs so user inputs aren't mis-marked const. """ - from executorch.exir.passes.spec_prop_pass import _is_mutable_buffer def patch(self, exported_program, graph_module): sig = exported_program.graph_signature @@ -678,8 +691,6 @@ def _view_replaceable_nodes(graph_module): base``. Those nodes must stay ``view_copy`` (an out-variant copy op, always correct) — only the storage-sharing optimisation is skipped for them. """ - from executorch.exir.passes.replace_view_copy_with_view_pass import _is_view_copy - from executorch.exir.tensor import determine_tensor_dynanism for node in graph_module.graph.nodes: if _is_view_copy(node) and all(user.op != "output" for user in node.users): @@ -697,8 +708,6 @@ def _view_replaceable_nodes(graph_module): def _patch_replace_view_copy_with_view_call(_original): """Replacement for ``ReplaceViewCopyWithViewPass.call`` that only replaces ``view_copy`` nodes whose shape dynamism matches their base's — see ``_view_replaceable_nodes``.""" - from executorch.exir.passes.replace_view_copy_with_view_pass import _VIEW_OP, _ViewSpec - from torch.fx.passes.infra.pass_base import PassResult def patch(self, graph_module): n_replaced = 0 @@ -781,12 +790,6 @@ def _make_squeeze_define_node(original): where both batch and time are dynamic). Replace the strict check with a no-op when the dynamic-dim count is preserved across the squeeze. """ - from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( # type: ignore[import-not-found] - XNNStaticReshape, - XNode, - ) - from executorch.backends.xnnpack.utils.utils import get_input_node - from torch.fx.experimental.symbolic_shapes import free_symbols def patch(self, node, xnn_graph, vals_to_ids, debug_handle): self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids) @@ -1115,6 +1118,37 @@ def patch(original_program, call_delegate_node, input_specs_to_delete, output_sp return patch +@register_patch( + "executorch", + "executorch.backends.xnnpack.serialization.xnnpack_graph_serialize._flatc_compile", +) +def _patch_flatc_compile_nonfinite(original): + """Rewrite non-finite float literals in the XNNPACK delegate JSON before ``flatc``. + + XNNPACK serializes its delegate graph via ``json.dumps``, which emits non-finite floats as the + bare tokens ``-Infinity`` / ``Infinity`` / ``NaN`` — not part of the flatbuffers JSON grammar, so + ``flatc`` fails with ``cannot parse value starting with: -``. MiniMaxM3's lightning-indexer block + padding (``F.pad(scores, ..., value=float("-inf"))``) lowers to a ``constant_pad_nd`` whose + ``-inf`` ``padding_value`` hits this. Swap the tokens for flatbuffers' own ``-inf`` / ``inf`` / + ``nan`` (parsed to the identical IEEE value) so the exact ``-inf`` semantics are preserved. + """ + + def patch(output_dir, schema_path, json_path): + with open(json_path) as f: + data = f.read() + # Lookbehind/lookahead on JSON delimiters so only bare numeric literals match (quoted + # strings are bounded by `"` and never touched). + fixed = re.sub(r"(?<=[:\[,\s])-Infinity(?=[,\]}\s])", "-inf", data) + fixed = re.sub(r"(?<=[:\[,\s])Infinity(?=[,\]}\s])", "inf", fixed) + fixed = re.sub(r"(?<=[:\[,\s])NaN(?=[,\]}\s])", "nan", fixed) + if fixed != data: + with open(json_path, "w") as f: + f.write(fixed) + return original(output_dir, schema_path, json_path) + + return patch + + @register_patch("executorch", "executorch.backends.xnnpack.operators.node_visitor._node_visitor_dict") def _patch_squeeze_node_visitors(original): """Swap the squeeze/unsqueeze visitor entries in ``_node_visitor_dict`` with subclasses @@ -1391,3 +1425,39 @@ def _fix_sym_pow_as_mul(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool: node.replace_all_uses_with(running) gm.graph.erase_node(node) return True + + +@register_fx_node_fix("executorch") +def _fix_negative_slice_start(gm: torch.fx.GraphModule, node: torch.fx.Node) -> bool: + """Rewrite a data-dependent negative slice start into its positive ``dim_size + start`` form. + + A negative slice on an unbacked length (VideoMAE's decoder keeps only the masked tokens via + ``hidden_states[:, -return_token_num:]``, ``return_token_num`` being the symbolic masked-patch + count) records ``start = -(u // 2)``. ``to_edge_transform_and_lower`` re-runs ``slice_forward``'s + meta, whose ``if start_val < 0`` guard can't be decided on a size-like symbol + (``GuardOnDataDependentSymNode``). Replace the start with ``sym_size(input, dim) + start`` — for a + tail slice this is the (size-like, hence provably ``>= 0``) number of leading elements, so the + guard is statically false. Drop the stale ``unbacked_bindings``: the output length is now a + computable expression, not the fresh unbacked symbol ``run_decompositions`` recorded. + """ + + if node.target not in (torch.ops.aten.slice.Tensor, torch.ops.aten.slice_copy.Tensor) or len(node.args) < 3: + return False + start = node.args[2] + if not isinstance(start, torch.fx.Node): + return False + start_val = start.meta.get("val") + if not isinstance(start_val, torch.SymInt) or statically_known_true(start_val >= 0): + return False + input_node, dim = node.args[0], node.args[1] + input_val = input_node.meta.get("val") + if not isinstance(input_val, torch.Tensor): + return False + with gm.graph.inserting_before(node): + size_node = gm.graph.call_function(torch.ops.aten.sym_size.int, (input_node, dim)) + size_node.meta["val"] = input_val.shape[dim] + add_node = gm.graph.call_function(operator.add, (size_node, start)) + add_node.meta["val"] = input_val.shape[dim] + start_val + node.args = (*node.args[:2], add_node, *node.args[3:]) + node.meta.pop("unbacked_bindings", None) + return True diff --git a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py index 56610ebdee47..5be09a7bd0a1 100644 --- a/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/configuration_recurrent_gemma.py @@ -87,6 +87,10 @@ def __post_init__(self, **kwargs): self.num_key_value_heads if self.num_key_value_heads is not None else self.num_attention_heads ) self.final_w_init_variance_scale = 2.0 / self.num_hidden_layers + self.layer_types = [ + "linear_attention" if block_type == "recurrent" else "sliding_attention" + for block_type in self.layers_block_type + ] kwargs.setdefault("partial_rotary_factor", 0.5) # assign default for BC super().__post_init__(**kwargs) diff --git a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py index fefcdef463cf..2548588e3ad5 100644 --- a/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py +++ b/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py @@ -25,7 +25,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...masking_utils import create_sliding_window_causal_mask +from ...masking_utils import create_recurrent_attention_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, CausalLMOutput from ...modeling_rope_utils import dynamic_rope_update @@ -317,12 +317,12 @@ def __init__(self, config): torch.empty([self.num_attention_heads, self.block_width, self.block_width]) ) self.recurrent_gate_bias = nn.Parameter(torch.empty([self.num_attention_heads, self.block_width])) - self.recurrent_states = None def forward( self, activations: torch.Tensor, position_ids: torch.Tensor, + recurrent_states: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: batch_size, seq_len, lru_width = activations.shape reset = position_ids[:, :, None] == 0 @@ -354,10 +354,9 @@ def forward( hidden_states=normalized_x, recurrent_gate=recurrent_gate, reset=reset, - recurrent_states=self.recurrent_states, + recurrent_states=recurrent_states, ) - self.recurrent_states = recurrent_states - return hidden_states + return hidden_states, recurrent_states # TODO refactor def _rnn_scan( @@ -415,6 +414,7 @@ class RecurrentGemmaRecurrentBlock(nn.Module): def __init__(self, config: RecurrentGemmaConfig, layer_idx: int): super().__init__() + self.layer_idx = layer_idx self.lru_width = config.lru_width self.hidden_size = config.hidden_size self.linear_y = nn.Linear(in_features=config.hidden_size, out_features=config.lru_width) @@ -431,18 +431,15 @@ def __init__(self, config: RecurrentGemmaConfig, layer_idx: int): self.rg_lru = RecurrentGemmaRglru(config) self.act_fn = ACT2FN[config.hidden_activation] - self.conv1d_state = None - def forward( self, input_states: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, - use_cache: bool = True, + past_key_values: Cache | None = None, **kwargs, ) -> tuple[torch.Tensor, None]: _, seq_len, _ = input_states.shape - batch_size = input_states.shape[0] y_branch = self.linear_y(input_states) y_branch = self.act_fn(y_branch) @@ -450,42 +447,31 @@ def forward( x_branch = self.linear_x(input_states) x_branch = x_branch.transpose(1, 2) - if use_cache: - # Check if cache needs initialization (None or batch size mismatch) - if self.conv1d_state is None or self.conv1d_state.shape[0] != batch_size: - self.conv1d_state = torch.zeros( - (batch_size, self.hidden_size, self.conv1d_width - 1), - device=input_states.device, - dtype=input_states.dtype, - ) - self.rg_lru.recurrent_states = torch.zeros( - (batch_size, self.lru_width), device=input_states.device, dtype=torch.float32 - ) + use_cache = past_key_values is not None + use_precomputed_states = use_cache and past_key_values.has_previous_state(self.layer_idx) - if position_ids.shape[1] != 1: # prefill - self.conv1d_state = nn.functional.pad(x_branch, (self.conv1d_width - x_branch.shape[-1] - 1, 0)) + if use_cache: + if not use_precomputed_states: # prefill + conv_state = nn.functional.pad(x_branch, (self.conv1d_width - x_branch.shape[-1] - 1, 0)) + past_key_values.update_conv_state(conv_state, self.layer_idx) x_branch = self.conv_1d(x_branch)[..., :seq_len] else: # decoding - conv_state = torch.cat((self.conv1d_state, x_branch), -1) + conv_state = torch.cat((past_key_values.layers[self.layer_idx].conv_states, x_branch), -1) x_branch = torch.sum(conv_state * self.conv_1d.weight[:, 0, :], dim=-1) + self.conv_1d.bias x_branch = x_branch.unsqueeze(-1) - self.conv1d_state = conv_state[:, :, 1:] + past_key_values.update_conv_state(conv_state[:, :, 1:], self.layer_idx) else: - self.conv1d_state = None - self.rg_lru.recurrent_states = None x_branch = self.conv_1d(x_branch)[..., :seq_len] - x_branch = self.rg_lru(x_branch.transpose(1, 2), position_ids) + recurrent_states = past_key_values.layers[self.layer_idx].recurrent_states if use_precomputed_states else None + x_branch, recurrent_states = self.rg_lru(x_branch.transpose(1, 2), position_ids, recurrent_states) + if use_cache: + past_key_values.update_recurrent_state(recurrent_states, self.layer_idx) hidden_states = x_branch * y_branch hidden_states = self.linear_out(hidden_states) return hidden_states, None - def _setup_cache(self, batch, device, dtype): - # recurrent_states always computed in full precision - self.rg_lru.recurrent_states = torch.zeros((batch, self.lru_width), device=device, dtype=torch.float32) - self.conv1d_state = torch.zeros((batch, self.hidden_size, self.conv1d_width - 1), device=device, dtype=dtype) - TEMPORAL_BLOCK_CLASSES = {"recurrent": RecurrentGemmaRecurrentBlock, "attention": RecurrentGemmaAttention} @@ -521,7 +507,6 @@ def forward( activations: torch.Tensor, position_ids: torch.Tensor, attention_mask: torch.Tensor, - use_cache: bool | None = None, past_key_values: Cache | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: @@ -532,7 +517,6 @@ def forward( inputs_normalized, position_ids=position_ids, attention_mask=attention_mask, - use_cache=use_cache, past_key_values=past_key_values, **kwargs, ) @@ -618,20 +602,6 @@ def _init_weights(self, module): init.copy_(module.inv_freq, buffer_value) init.copy_(module.original_inv_freq, buffer_value) - def _setup_cache(self, config, batch, device, dtype): - layers = getattr(self, "model", self).layers - for layer in layers: - if hasattr(layer.temporal_block, "_setup_cache"): - layer.temporal_block._setup_cache(batch, device, dtype) - - -def _get_seq_length(self, layer_idx: int = 0) -> int: - return self.layers[self.first_attention_layer].get_seq_length() - - -def _get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: - return self.layers[self.first_attention_layer].get_mask_sizes(query_length) - @auto_docstring class RecurrentGemmaModel(RecurrentGemmaPreTrainedModel): @@ -675,36 +645,31 @@ def forward( hidden_states = inputs_embeds if use_cache and past_key_values is None: - self._setup_cache(self.config, hidden_states.shape[0], hidden_states.device, hidden_states.dtype) past_key_values = DynamicCache(config=self.config) - # Hack because the mamba layer indices will stay empty in `past_key_values`, and we want `get_seq_length` and - # `get_mask_sizes` to use the first attention layer by default for the mask function to create correct masks - if past_key_values is not None: - past_key_values.first_attention_layer = self.config.layers_block_type.index("attention") - # bound new methods to this instance only - past_key_values.get_seq_length = _get_seq_length.__get__(past_key_values) - past_key_values.get_mask_sizes = _get_mask_sizes.__get__(past_key_values) - if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) - causal_mask = create_sliding_window_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + "linear_attention": create_recurrent_attention_mask(**mask_kwargs), + } hidden_states = hidden_states * self.normalizer.type(hidden_states.dtype) for i, residual_block in enumerate(self.layers): - hidden_states = residual_block( - hidden_states, position_ids, causal_mask, use_cache, past_key_values, **kwargs - ) + causal_mask = causal_mask_mapping[self.config.layer_types[i]] + hidden_states = residual_block(hidden_states, position_ids, causal_mask, past_key_values, **kwargs) hidden_states = self.final_norm(hidden_states) diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 74e8126d611f..28214d3b675c 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -72,16 +72,7 @@ "Sam2VisionModel": "Same Hiera-backbone dynamic-shape `timeout` as `Sam2Model`.", }, # Every backend, generate path only. - "generate": { - "RecurrentGemmaForCausalLM": ( - "Keeps RG-LRU/conv state as plain module attributes (not a `Cache` passed via " - "`past_key_values`), so the state cannot be a graph input/output — the exported " - "decode step computes from zero-initialized state and its logits diverge from eager " - "(prefill exports and matches). " - "TODO: migrate the state to `LinearAttentionLayer` entries in `past_key_values` " - "(the qwen3_next / Mamba pattern)." - ), - }, + "generate": {}, # ONNX, every variant. "onnx": {}, # ONNX, generate path only. @@ -110,24 +101,13 @@ }, # ExecuTorch, every variant. "executorch": { - "VideoMAEForPreTraining": ( - "Torch/ONNX/OpenVINO export fine (a `torch._check` states the logits/labels masked-token " - "count invariant), but ExecuTorch edge-lowering still fails on the decoder's data-dependent " - "negative slice `hidden_states[:, -return_token_num:]` — `aten.slice_copy`'s meta raises " - "`GuardOnDataDependentSymNode` on the symbolic masked-token count." - ), "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", "MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", }, # ExecuTorch, generate path only. - "executorch.generate": { - "MiniMaxM3SparseForConditionalGeneration": ( - "`flatc` schema compilation fails when serializing the ExecuTorch program for the MoE " - "decoder generate graph." - ), - }, + "executorch.generate": {}, # ExecuTorch, dynamic-shape only. "executorch.dynamic": { "BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."), From 659cb5e425e4f7a2c02de5fc2a37c0c366584bfe Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 9 Jul 2026 16:44:26 +0200 Subject: [PATCH 13/14] more fixes --- src/transformers/exporters/utils.py | 17 +++++++++++++++++ src/transformers/models/hiera/modeling_hiera.py | 13 +++++++++++-- .../models/pixtral/modeling_pixtral.py | 7 ++++--- tests/exporters/test_export.py | 12 +++++++++--- tests/models/emu3/test_modeling_emu3.py | 2 -- tests/models/hiera/test_modeling_hiera.py | 1 - .../lighton_ocr/test_modeling_lighton_ocr.py | 1 - tests/models/mistral3/test_modeling_mistral3.py | 1 - tests/models/pixtral/test_modeling_pixtral.py | 1 - 9 files changed, 41 insertions(+), 14 deletions(-) diff --git a/src/transformers/exporters/utils.py b/src/transformers/exporters/utils.py index 1242f70398b7..6a80d55c2722 100644 --- a/src/transformers/exporters/utils.py +++ b/src/transformers/exporters/utils.py @@ -514,6 +514,23 @@ def _prepare_navit_vision_inputs(model: torch.nn.Module, inputs: dict[str, Any]) inputs["merged_shape"] = get_vision_merged_shape(target_sizes, window_kernel_size) +@register_export_input_preparer("image_sizes") +def _prepare_image_sizes_as_ints(model: torch.nn.Module, inputs: dict[str, Any]) -> None: + """Replace a tensor `image_sizes` with a python list of `(h, w)` int-tuples (the `.tolist()` runs here, + outside the traced graph). + + `image_sizes` is per-image geometry, and encoders crop/split each image by it — e.g. + `image_sizes[i] // patch_size` (Pixtral) or `int(image_sizes[i] / factor)` (Emu3 VQVAE). As a tensor + those bounds become unbacked symints under `torch.export`; as python ints they stay static (matching + each encoder's own `image_sizes is None` fallback, which already builds int-tuples). Models that route + `image_sizes` around the traced graph (e.g. LLaVA-NeXT resolves anyres before tracing) never hit this. + """ + image_sizes = inputs["image_sizes"] + if not torch.is_tensor(image_sizes): + return + inputs["image_sizes"] = [tuple(int(v) for v in row) for row in image_sizes.tolist()] + + @register_export_input_preparer("input_features", "feature_lens") def _prepare_omni_audio_inputs(model: torch.nn.Module, inputs: dict[str, Any]) -> None: """Replace `input_features`/`feature_lens` with precomputed `padded_feature`, `chunk_lengths`, diff --git a/src/transformers/models/hiera/modeling_hiera.py b/src/transformers/models/hiera/modeling_hiera.py index 6bc352a97e1a..407033368c39 100644 --- a/src/transformers/models/hiera/modeling_hiera.py +++ b/src/transformers/models/hiera/modeling_hiera.py @@ -351,18 +351,27 @@ def forward( batch_size, seq_len, _ = hidden_states.shape num_windows = 1 + tokens_per_window = seq_len if self.use_mask_unit_attn: num_windows = seq_len // (self.query_stride * self.window_size) + tokens_per_window = self.query_stride * self.window_size qkv = self.qkv(hidden_states) - qkv = qkv.reshape(batch_size, -1, num_windows, 3, self.num_heads, self.head_dim) + qkv = qkv.reshape(batch_size, tokens_per_window, num_windows, 3, self.num_heads, self.head_dim) qkv = qkv.permute(3, 0, 4, 2, 1, 5) query, key, value = qkv.unbind(0) if self.query_stride > 1: # Refer to unroll to see how this performs a maxpool-Nd - query = query.view(batch_size, self.num_heads, num_windows, self.query_stride, -1, self.head_dim) + query = query.view( + batch_size, + self.num_heads, + num_windows, + self.query_stride, + tokens_per_window // self.query_stride, + self.head_dim, + ) query = query.max(dim=3).values attn_weights = (query * self.scale) @ key.transpose(-1, -2) diff --git a/src/transformers/models/pixtral/modeling_pixtral.py b/src/transformers/models/pixtral/modeling_pixtral.py index e7c92bfe4df7..6ae8c2d5baee 100644 --- a/src/transformers/models/pixtral/modeling_pixtral.py +++ b/src/transformers/models/pixtral/modeling_pixtral.py @@ -400,10 +400,11 @@ def generate_block_attention_mask(patch_embeds_list, tensor): d_min = torch.finfo(dtype).min causal_mask = torch.full((seq_len, seq_len), fill_value=d_min, dtype=dtype, device=device) - block_end_idx = torch.tensor(patch_embeds_list).cumsum(-1) - block_start_idx = torch.tensor([0] + patch_embeds_list[:-1]).cumsum(-1) - for start, end in zip(block_start_idx, block_end_idx): + start = 0 + for num_patches in patch_embeds_list: + end = start + num_patches causal_mask[start:end, start:end] = 0 + start = end causal_mask = causal_mask[None, None, :, :].expand(tensor.shape[0], 1, -1, -1) return causal_mask diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 0123c419f691..137ce1d53e59 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -71,6 +71,13 @@ "× 3 Q-pool stage transitions on symbolic H/W. Static exports fine." ), "Sam2VisionModel": "Same Hiera-backbone dynamic-shape `timeout` as `Sam2Model`.", + "HieraModel": ( + "Hiera mask-unit window `reroll` produces nested symbolic floordivs that `torch.export` " + "can't guard under dynamic shapes (same backbone family as `Sam2Model`). Static exports fine." + ), + "HieraBackbone": "Same Hiera `reroll` dynamic-shape failure as `HieraModel`.", + "HieraForImageClassification": "Same Hiera `reroll` dynamic-shape failure as `HieraModel`.", + "HieraForPreTraining": "Same Hiera `reroll` dynamic-shape failure as `HieraModel`.", }, # Every backend, generate path only. "generate": {}, @@ -105,8 +112,7 @@ "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", "Kimi_K25Model": ( - "ExecuTorch `to_edge` spec failure in the DeepseekV3 language model; the vision encoder " - "exports fine." + "ExecuTorch `to_edge` spec failure in the DeepseekV3 language model; the vision encoder exports fine." ), "Kimi_K25ForConditionalGeneration": "Same ExecuTorch `spec` failure as `Kimi_K25Model` (DeepseekV3 LM).", "MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.", @@ -386,7 +392,7 @@ def test_torch_export(self, dynamic, atol=1e-4, rtol=1e-4): config = DynamoConfig(dynamic=dynamic) for model_class in self.all_model_classes: - if self._should_skip(model_class): + if self._should_skip(model_class, dynamic=dynamic): continue components = self._prepare_export_model_and_inputs(model_class) diff --git a/tests/models/emu3/test_modeling_emu3.py b/tests/models/emu3/test_modeling_emu3.py index 31d2a936e391..80e27bd15f23 100644 --- a/tests/models/emu3/test_modeling_emu3.py +++ b/tests/models/emu3/test_modeling_emu3.py @@ -293,8 +293,6 @@ class Emu3Vision2TextModelTest(ModelTesterMixin, GenerationTesterMixin, Pipeline ) skip_test_image_features_output_shape = True # Emu3 uses index -3 for hidden_size instead of -1 - test_torch_exportable = False # data-dependent control flow in vision/segmentation head - def setUp(self): self.model_tester = Emu3Vision2TextModelTester(self) self.config_tester = ConfigTester(self, config_class=Emu3Config, has_text_modality=False, hidden_size=32) diff --git a/tests/models/hiera/test_modeling_hiera.py b/tests/models/hiera/test_modeling_hiera.py index 3b8195b40bea..20b4ddfb360d 100644 --- a/tests/models/hiera/test_modeling_hiera.py +++ b/tests/models/hiera/test_modeling_hiera.py @@ -245,7 +245,6 @@ class HieraModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): ) test_resize_embeddings = False - test_torch_exportable = False # massive symbolic expression from multi-stage pooling def setUp(self): self.model_tester = HieraModelTester(self) diff --git a/tests/models/lighton_ocr/test_modeling_lighton_ocr.py b/tests/models/lighton_ocr/test_modeling_lighton_ocr.py index e8594c979399..0bf3bae4196b 100644 --- a/tests/models/lighton_ocr/test_modeling_lighton_ocr.py +++ b/tests/models/lighton_ocr/test_modeling_lighton_ocr.py @@ -229,7 +229,6 @@ class LightOnOcrForConditionalGenerationModelTest(ModelTesterMixin, GenerationTe skip_test_image_features_output_shape = True _is_composite = True - test_torch_exportable = False # data-dependent multimodal placeholder mask def setUp(self): self.model_tester = LightOnOcrVisionText2TextModelTester(self) diff --git a/tests/models/mistral3/test_modeling_mistral3.py b/tests/models/mistral3/test_modeling_mistral3.py index 457dc6b8ea8b..3159ee012c29 100644 --- a/tests/models/mistral3/test_modeling_mistral3.py +++ b/tests/models/mistral3/test_modeling_mistral3.py @@ -177,7 +177,6 @@ class Mistral3ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterM # Mistral3 merges batch_size and num_patches in index 1, with index 0 hardcoded to 1 skip_test_image_features_output_shape = True _is_composite = True - test_torch_exportable = False # data-dependent multimodal placeholder mask def setUp(self): self.model_tester = Mistral3VisionText2TextModelTester(self) diff --git a/tests/models/pixtral/test_modeling_pixtral.py b/tests/models/pixtral/test_modeling_pixtral.py index 99be7747b6f2..863407336092 100644 --- a/tests/models/pixtral/test_modeling_pixtral.py +++ b/tests/models/pixtral/test_modeling_pixtral.py @@ -112,7 +112,6 @@ class PixtralVisionModelModelTest(ModelTesterMixin, unittest.TestCase): additional_model_inputs = ["image_sizes"] test_resize_embeddings = False - test_torch_exportable = False # data-dependent vision placeholder mask def setUp(self): self.model_tester = PixtralVisionModelTester(self) From 8c54bf29aa82898540944087be8ad6695fcfa1df Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Fri, 10 Jul 2026 10:42:42 +0200 Subject: [PATCH 14/14] more fixes --- .../exporters/exporter_executorch.py | 8 +- .../models/funnel/modeling_funnel.py | 26 +++--- .../models/hunyuan_vl/modeling_hunyuan_vl.py | 17 ++-- .../models/hunyuan_vl/modular_hunyuan_vl.py | 17 ++-- .../models/oneformer/modeling_oneformer.py | 31 +++---- .../models/tapas/modeling_tapas.py | 80 +++++++------------ tests/exporters/test_export.py | 39 ++++++++- tests/models/funnel/test_modeling_funnel.py | 4 +- .../hunyuan_vl/test_modeling_hunyuan_vl.py | 2 +- .../oneformer/test_modeling_oneformer.py | 2 +- tests/models/tapas/test_modeling_tapas.py | 2 +- 11 files changed, 123 insertions(+), 105 deletions(-) diff --git a/src/transformers/exporters/exporter_executorch.py b/src/transformers/exporters/exporter_executorch.py index 8284c7489b27..793d323256b7 100644 --- a/src/transformers/exporters/exporter_executorch.py +++ b/src/transformers/exporters/exporter_executorch.py @@ -460,10 +460,10 @@ def _patch_expand(original): the broadcast so the captured tensor has standard strides downstream. """ - def patch(self, *sizes): - if len(sizes) == 1 and isinstance(sizes[0], (list, tuple, torch.Size)): - sizes = tuple(sizes[0]) - result = original(self, *sizes) + def patch(self, *sizes, **kwargs): + # Forward whatever form the caller used — positional ``expand(*sizes)``, a single + # list/tuple, or the keyword form ``expand(size=...)`` — straight to the original. + result = original(self, *sizes, **kwargs) # Only materialise when ``expand`` actually introduced a stride-0 (broadcast) dim; a # no-broadcast expand is a plain view ExecuTorch's memory planner accepts as-is. if 0 in result.stride(): diff --git a/src/transformers/models/funnel/modeling_funnel.py b/src/transformers/models/funnel/modeling_funnel.py index 3eda3d4856ae..7e0f3ade51eb 100644 --- a/src/transformers/models/funnel/modeling_funnel.py +++ b/src/transformers/models/funnel/modeling_funnel.py @@ -148,7 +148,10 @@ def get_position_embeds( cos_embed = self.cos_dropout(torch.cos(sinusoid)) pos_embed = torch.cat([sin_embed, cos_embed], dim=-1) - pos = torch.arange(0, seq_len, dtype=torch.int64, device=device).to(dtype) + # Positions are integer indices fully determined by `seq_len`; keeping them as Python + # ints (instead of tensors) makes the relative-position `arange` sizes static, which is + # required for tracing / export. + pos = list(range(seq_len)) pooled_pos = pos position_embeds_list = [] for block_index in range(0, self.config.num_blocks): @@ -165,7 +168,7 @@ def get_position_embeds( # construct rel_pos_id stride = 2 ** (block_index - 1) - rel_pos = self.relative_pos(pos, stride, pooled_pos, shift=2) + rel_pos = self.relative_pos(pos, stride, pooled_pos, shift=2, device=device) rel_pos = rel_pos[:, None] + zero_offset rel_pos = rel_pos.expand(rel_pos.size(0), d_model) position_embeds_pooling = torch.gather(pos_embed, 0, rel_pos) @@ -173,7 +176,7 @@ def get_position_embeds( # Second type pos = pooled_pos stride = 2**block_index - rel_pos = self.relative_pos(pos, stride) + rel_pos = self.relative_pos(pos, stride, device=device) rel_pos = rel_pos[:, None] + zero_offset rel_pos = rel_pos.expand(rel_pos.size(0), d_model) @@ -182,7 +185,7 @@ def get_position_embeds( position_embeds_list.append([position_embeds_no_pooling, position_embeds_pooling]) return position_embeds_list - def stride_pool_pos(self, pos_id: torch.Tensor, block_index: int): + def stride_pool_pos(self, pos_id: list[int], block_index: int) -> list[int]: """ Pool `pos_id` while keeping the cls token separate (if `config.separate_cls=True`). """ @@ -191,25 +194,30 @@ def stride_pool_pos(self, pos_id: torch.Tensor, block_index: int): # the previous block of the 1st real block. Since the 1st real # block always has position 1, the position of the previous block # will be at `1 - 2 ** block_index`. - cls_pos = pos_id.new_tensor([-(2**block_index) + 1]) + cls_pos = [-(2**block_index) + 1] pooled_pos_id = pos_id[1:-1] if self.config.truncate_seq else pos_id[1:] - return torch.cat([cls_pos, pooled_pos_id[::2]], 0) + return cls_pos + pooled_pos_id[::2] else: return pos_id[::2] - def relative_pos(self, pos: torch.Tensor, stride: int, pooled_pos=None, shift: int = 1) -> torch.Tensor: + def relative_pos( + self, pos: list[int], stride: int, pooled_pos: list[int] | None = None, shift: int = 1, device=None + ) -> torch.Tensor: """ Build the relative positional vector between `pos` and `pooled_pos`. + + `pos` and `pooled_pos` are lists of Python ints so the bounds (and hence the output size) are + statically known, which keeps the produced `arange` export-friendly. """ if pooled_pos is None: pooled_pos = pos ref_point = pooled_pos[0] - pos[0] - num_remove = shift * pooled_pos.shape[0] + num_remove = shift * len(pooled_pos) max_dist = ref_point + num_remove * stride min_dist = pooled_pos[0] - pos[-1] - return torch.arange(max_dist, min_dist - 1, -stride, dtype=torch.long, device=pos.device) + return torch.arange(max_dist, min_dist - 1, -stride, dtype=torch.long, device=device) def stride_pool( self, diff --git a/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py index 7cb5f3b2f44b..66d86696cd5e 100644 --- a/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py +++ b/src/transformers/models/hunyuan_vl/modeling_hunyuan_vl.py @@ -222,16 +222,12 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: list[list[int]]) -> torc embeddings = patch_embeds.flatten(-2).squeeze(-1) embeddings = embeddings.reshape(batch_size, sequence_len, -1).squeeze(0) - start = 0 - image_embeddings_list = [] + position_embeddings_list = [] for t, h, w in grid_thw: - end = start + t * h * w - image_embeddings = embeddings[start:end, :] - position_embedding = self.interpolate_pos_encoding(image_embeddings, h, w).squeeze(0).repeat(t, 1) - image_embeddings_list.append(image_embeddings + position_embedding) - start = end + position_embeddings_list.append(self.interpolate_pos_encoding(embeddings, h, w).squeeze(0).repeat(t, 1)) + position_embeddings = torch.concat(position_embeddings_list, dim=0) - return torch.concat(image_embeddings_list, dim=0).unsqueeze(0) + return (embeddings + position_embeddings).unsqueeze(0) class HunYuanVLVisionPatchMerger(nn.Module): @@ -260,7 +256,10 @@ def __init__(self, config: HunYuanVLVisionConfig): def forward(self, hidden_states: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: hidden_states = self.before_rms(hidden_states) dtype = hidden_states.dtype - hidden_states = hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, *size) + hidden_states = hidden_states.permute(0, 2, 1) + hidden_states = hidden_states.reshape(hidden_states.shape[0], hidden_states.shape[1], *size) + torch._check(hidden_states.shape[2] > 1) + torch._check(hidden_states.shape[3] > 1) hidden_states = self.proj_conv(hidden_states) hidden_states = self.proj_act(hidden_states) hidden_states = self.proj_out(hidden_states) diff --git a/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py b/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py index 557035e3d67e..7d8a8e574561 100644 --- a/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py +++ b/src/transformers/models/hunyuan_vl/modular_hunyuan_vl.py @@ -584,16 +584,12 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: list[list[int]]) -> torc embeddings = patch_embeds.flatten(-2).squeeze(-1) embeddings = embeddings.reshape(batch_size, sequence_len, -1).squeeze(0) - start = 0 - image_embeddings_list = [] + position_embeddings_list = [] for t, h, w in grid_thw: - end = start + t * h * w - image_embeddings = embeddings[start:end, :] - position_embedding = self.interpolate_pos_encoding(image_embeddings, h, w).squeeze(0).repeat(t, 1) - image_embeddings_list.append(image_embeddings + position_embedding) - start = end + position_embeddings_list.append(self.interpolate_pos_encoding(embeddings, h, w).squeeze(0).repeat(t, 1)) + position_embeddings = torch.concat(position_embeddings_list, dim=0) - return torch.concat(image_embeddings_list, dim=0).unsqueeze(0) + return (embeddings + position_embeddings).unsqueeze(0) class HunYuanVLVisionPatchMerger(nn.Module): @@ -622,7 +618,10 @@ def __init__(self, config: HunYuanVLVisionConfig): def forward(self, hidden_states: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: hidden_states = self.before_rms(hidden_states) dtype = hidden_states.dtype - hidden_states = hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, *size) + hidden_states = hidden_states.permute(0, 2, 1) + hidden_states = hidden_states.reshape(hidden_states.shape[0], hidden_states.shape[1], *size) + torch._check(hidden_states.shape[2] > 1) + torch._check(hidden_states.shape[3] > 1) hidden_states = self.proj_conv(hidden_states) hidden_states = self.proj_act(hidden_states) hidden_states = self.proj_out(hidden_states) diff --git a/src/transformers/models/oneformer/modeling_oneformer.py b/src/transformers/models/oneformer/modeling_oneformer.py index af1a464237f4..6a31d27d132e 100644 --- a/src/transformers/models/oneformer/modeling_oneformer.py +++ b/src/transformers/models/oneformer/modeling_oneformer.py @@ -979,6 +979,7 @@ def forward( position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, + spatial_shapes_list=None, level_start_index=None, output_attentions: bool = False, ): @@ -1022,7 +1023,7 @@ def forward( else: raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}") # PyTorch implementation - output = multi_scale_deformable_attention(value, spatial_shapes, sampling_locations, attention_weights) + output = multi_scale_deformable_attention(value, spatial_shapes_list, sampling_locations, attention_weights) output = self.output_proj(output) return output, attention_weights @@ -1056,6 +1057,7 @@ def forward( position_embeddings: torch.Tensor | None = None, reference_points=None, spatial_shapes=None, + spatial_shapes_list=None, level_start_index=None, output_attentions: bool = False, ): @@ -1088,6 +1090,7 @@ def forward( position_embeddings=position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, output_attentions=output_attentions, ) @@ -1139,12 +1142,12 @@ def __init__(self, config: OneFormerConfig): self.layers = nn.ModuleList([OneFormerPixelDecoderEncoderLayer(config) for _ in range(config.encoder_layers)]) @staticmethod - def get_reference_points(spatial_shapes, valid_ratios, device): + def get_reference_points(spatial_shapes_list, valid_ratios, device): """ Get reference points for each feature map. Used in decoder. Args: - spatial_shapes (`torch.LongTensor` of shape `(num_feature_levels, 2)`): + spatial_shapes_list (`list[tuple[int, int]]`): Spatial shapes of each feature map. valid_ratios (`torch.FloatTensor` of shape `(batch_size, num_feature_levels, 2)`): Valid ratios of each feature map. @@ -1154,7 +1157,7 @@ def get_reference_points(spatial_shapes, valid_ratios, device): `torch.FloatTensor` of shape `(batch_size, num_queries, num_feature_levels, 2)` """ reference_points_list = [] - for lvl, (height, width) in enumerate(spatial_shapes): + for lvl, (height, width) in enumerate(spatial_shapes_list): ref_y, ref_x = torch.meshgrid( torch.linspace(0.5, height - 0.5, height, dtype=valid_ratios.dtype, device=device), torch.linspace(0.5, width - 0.5, width, dtype=valid_ratios.dtype, device=device), @@ -1174,6 +1177,7 @@ def forward( attention_mask=None, position_embeddings=None, spatial_shapes=None, + spatial_shapes_list=None, level_start_index=None, valid_ratios=None, output_attentions=None, @@ -1213,7 +1217,7 @@ def forward( return_dict = return_dict if return_dict is not None else self.config.return_dict hidden_states = inputs_embeds - reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=inputs_embeds.device) + reference_points = self.get_reference_points(spatial_shapes_list, valid_ratios, device=inputs_embeds.device) encoder_states = () if output_hidden_states else None all_attentions = () if output_attentions else None @@ -1226,6 +1230,7 @@ def forward( position_embeddings=position_embeddings, reference_points=reference_points, spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, output_attentions=output_attentions, ) @@ -1369,11 +1374,11 @@ def forward( source_flatten = [] mask_flatten = [] lvl_pos_embed_flatten = [] - spatial_shapes = [] + spatial_shapes_list = [] for level, (source, mask, pos_embed) in enumerate(zip(sources, masks, position_embeddings_list)): batch_size, num_channels, height, width = source.shape spatial_shape = (height, width) - spatial_shapes.append(spatial_shape) + spatial_shapes_list.append(spatial_shape) source = source.flatten(2).transpose(1, 2) mask = mask.flatten(1) pos_embed = pos_embed.flatten(2).transpose(1, 2) @@ -1384,7 +1389,7 @@ def forward( source_flatten = torch.cat(source_flatten, 1) mask_flatten = torch.cat(mask_flatten, 1) lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1) - spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=source_flatten.device) + spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=source_flatten.device) level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1])) valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1) @@ -1396,6 +1401,7 @@ def forward( attention_mask=mask_flatten, position_embeddings=lvl_pos_embed_flatten, spatial_shapes=spatial_shapes, + spatial_shapes_list=spatial_shapes_list, level_start_index=level_start_index, valid_ratios=valid_ratios, output_attentions=output_attentions, @@ -1406,19 +1412,14 @@ def forward( y = encoder_outputs.last_hidden_state bs = y.shape[0] - split_size_or_sections = [None] * self.num_feature_levels - for i in range(self.num_feature_levels): - if i < self.num_feature_levels - 1: - split_size_or_sections[i] = level_start_index[i + 1] - level_start_index[i] - else: - split_size_or_sections[i] = y.shape[1] - level_start_index[i] + split_size_or_sections = [height * width for height, width in spatial_shapes_list] y = torch.split(y, split_size_or_sections, dim=1) out = [] multi_scale_features = [] num_cur_levels = 0 for i, z in enumerate(y): - out.append(z.transpose(1, 2).view(bs, -1, spatial_shapes[i][0], spatial_shapes[i][1])) + out.append(z.transpose(1, 2).view(bs, -1, spatial_shapes_list[i][0], spatial_shapes_list[i][1])) # append `out` with extra FPN levels # Reverse feature maps into top-down order (from low to high resolution) diff --git a/src/transformers/models/tapas/modeling_tapas.py b/src/transformers/models/tapas/modeling_tapas.py index 040fe7590371..55f199b25d9e 100644 --- a/src/transformers/models/tapas/modeling_tapas.py +++ b/src/transformers/models/tapas/modeling_tapas.py @@ -1009,7 +1009,7 @@ class for more info. if self.config.average_logits_per_cell: logits_per_cell, _ = reduce_mean(logits, cell_index) logits = gather(logits_per_cell, cell_index) - dist_per_token = torch.distributions.Bernoulli(logits=logits) + dist_per_token = torch.distributions.Bernoulli(logits=logits, validate_args=False) # Compute cell selection loss per example. selection_loss_per_example = None @@ -1027,7 +1027,7 @@ class for more info. selection_loss_per_example, logits = _single_column_cell_selection_loss( logits, column_logits, labels, cell_index, col_index, cell_mask ) - dist_per_token = torch.distributions.Bernoulli(logits=logits) + dist_per_token = torch.distributions.Bernoulli(logits=logits, validate_args=False) # Supervised cell selection if self.config.disable_per_token_loss: @@ -1275,7 +1275,7 @@ def __init__(self, indices, num_segments, batch_dims=0): index. """ self.indices = torch.as_tensor(indices, device=indices.device) - self.num_segments = torch.as_tensor(num_segments, device=indices.device) + self.num_segments = num_segments self.batch_dims = batch_dims def batch_shape(self): @@ -1376,11 +1376,11 @@ def flatten(index, name="segmented_flatten"): Returns: (`IndexMap`): The flattened IndexMap. """ - # first, get batch_size as scalar tensor - batch_size = torch.prod(torch.tensor(list(index.batch_shape()))) + # first, get batch_size as a python int (static, derived from the batch dimensions) + batch_size = math.prod(index.batch_shape()) # next, create offset as 1-D tensor of length batch_size, # and multiply element-wise by num segments (to offset different elements in the batch) e.g. if batch size is 2: [0, 64] - offset = torch.arange(start=0, end=batch_size, device=index.num_segments.device) * index.num_segments + offset = torch.arange(start=0, end=batch_size, device=index.indices.device) * index.num_segments offset = offset.view(index.batch_shape()) for _ in range(index.batch_dims, len(index.indices.size())): # typically range(1,2) offset = offset.unsqueeze(-1) @@ -1389,7 +1389,7 @@ def flatten(index, name="segmented_flatten"): return IndexMap(indices=indices.view(-1), num_segments=index.num_segments * batch_size, batch_dims=0) -def range_index_map(batch_shape, num_segments, name="range_index_map"): +def range_index_map(batch_shape, num_segments, device=None, name="range_index_map"): """ Constructs an index map equal to range(num_segments). @@ -1398,39 +1398,23 @@ def range_index_map(batch_shape, num_segments, name="range_index_map"): Batch shape num_segments (`int`): Number of segments + device (`torch.device`, *optional*): + Device on which to place the resulting indices. name (`str`, *optional*, defaults to 'range_index_map'): Name for the operation. Currently not used Returns: (`IndexMap`): IndexMap of shape batch_shape with elements equal to range(num_segments). """ - device = num_segments.device if torch.is_tensor(num_segments) else "cpu" - batch_shape = torch.as_tensor( - batch_shape, dtype=torch.long, device=device - ) # create a rank 1 tensor vector containing batch_shape (e.g. [2]) - assert len(batch_shape.size()) == 1 - num_segments = torch.as_tensor( - num_segments, device=device - ) # create a rank 0 tensor (scalar) containing num_segments (e.g. 64) - assert len(num_segments.size()) == 0 - - indices = torch.arange( - start=0, end=num_segments, device=num_segments.device - ) # create a rank 1 vector with num_segments elements - new_tensor = torch.cat( - [torch.ones_like(batch_shape, dtype=torch.long, device=num_segments.device), num_segments.unsqueeze(dim=0)], - dim=0, - ) - # new_tensor is just a vector of [1 64] for example (assuming only 1 batch dimension) - new_shape = [int(x) for x in new_tensor.tolist()] - indices = indices.view(new_shape) + batch_shape = tuple(batch_shape) # e.g. (2,) for a single batch dimension + num_segments = int(num_segments) - multiples = torch.cat([batch_shape, torch.as_tensor([1], device=device)], dim=0) - indices = indices.repeat(multiples.tolist()) - # equivalent (in Numpy:) - # indices = torch.as_tensor(np.tile(indices.numpy(), multiples.tolist())) + # create a rank 1 vector with num_segments elements, then broadcast it over the batch dimensions + indices = torch.arange(start=0, end=num_segments, device=device) + indices = indices.view([1] * len(batch_shape) + [num_segments]) + indices = indices.repeat(list(batch_shape) + [1]) - return IndexMap(indices=indices, num_segments=num_segments, batch_dims=list(batch_shape.size())[0]) + return IndexMap(indices=indices, num_segments=num_segments, batch_dims=len(batch_shape)) def _segment_reduce(values, index, segment_reduce_fn, name): @@ -1455,30 +1439,20 @@ def _segment_reduce(values, index, segment_reduce_fn, name): # unflattened. Segmented ops support vector-valued operations. flat_index = flatten(index) vector_shape = values.size()[len(index.indices.size()) :] # torch.Size object - flattened_shape = torch.cat( - [torch.as_tensor([-1], dtype=torch.long), torch.as_tensor(vector_shape, dtype=torch.long)], dim=0 - ) + flattened_shape = [-1] + list(vector_shape) # changed "view" by "reshape" in the following line - flat_values = values.reshape(flattened_shape.tolist()) + flat_values = values.reshape(flattened_shape) - out = torch.zeros(int(flat_index.num_segments), dtype=torch.float, device=flat_values.device) + out = torch.zeros(flat_index.num_segments, dtype=torch.float, device=flat_values.device) segment_means = out.scatter_reduce( dim=0, index=flat_index.indices.long(), src=flat_values.float(), reduce=segment_reduce_fn, include_self=False ) - device = index.num_segments.device # Unflatten the values. - new_shape = torch.cat( - [ - torch.as_tensor(index.batch_shape(), dtype=torch.long, device=device), - torch.as_tensor(index.num_segments, dtype=torch.long, device=device).unsqueeze(dim=0), - torch.as_tensor(vector_shape, dtype=torch.long, device=device), - ], - dim=0, - ) + new_shape = list(index.batch_shape()) + [index.num_segments] + list(vector_shape) - output_values = segment_means.clone().view(new_shape.tolist()).to(values.dtype) - output_index = range_index_map(index.batch_shape(), index.num_segments) + output_values = segment_means.clone().view(new_shape).to(values.dtype) + output_index = range_index_map(index.batch_shape(), index.num_segments, device=index.indices.device) return output_values, output_index @@ -1689,7 +1663,9 @@ def _single_column_cell_selection_loss(token_logits, column_logits, labels, cell no_cell_selected.view(column_label.size()), torch.zeros_like(column_label), column_label ) - column_dist = torch.distributions.Categorical(logits=column_logits) # shape (batch_size, max_num_cols) + column_dist = torch.distributions.Categorical( + logits=column_logits, validate_args=False + ) # shape (batch_size, max_num_cols) column_loss_per_example = -column_dist.log_prob(column_label) # Part 2: cell loss @@ -1713,7 +1689,7 @@ def _single_column_cell_selection_loss(token_logits, column_logits, labels, cell ) # Compute the log-likelihood for cells, but only for the selected column. - cell_dist = torch.distributions.Bernoulli(logits=logits_per_cell) # shape (batch_size, 64*32) + cell_dist = torch.distributions.Bernoulli(logits=logits_per_cell, validate_args=False) # shape (batch_size, 64*32) cell_log_prob = cell_dist.log_prob(labels_per_cell.type(torch.float32)) # shape(batch_size, 64*32) cell_loss = -torch.sum(cell_log_prob * column_mask * cell_mask, dim=1) @@ -1804,7 +1780,7 @@ def _calculate_aggregate_mask(answer, pooled_output, cell_selection_preference, # torch.FloatTensor(batch_size,) aggregate_mask_init = torch.logical_not(torch.isnan(answer)).type(torch.FloatTensor).to(answer.device) logits_aggregation = aggregation_classifier(pooled_output) - dist_aggregation = torch.distributions.categorical.Categorical(logits=logits_aggregation) + dist_aggregation = torch.distributions.categorical.Categorical(logits=logits_aggregation, validate_args=False) # Index 0 corresponds to "no aggregation". aggregation_ops_total_mass = torch.sum(dist_aggregation.probs[:, 1:], dim=1) @@ -1885,7 +1861,7 @@ def _calculate_aggregation_loss_unknown(logits_aggregation, aggregate_mask): aggregation_loss_unknown (`torch.FloatTensor` of shape `(batch_size,)`): Aggregation loss (in case of answer supervision) per example. """ - dist_aggregation = torch.distributions.categorical.Categorical(logits=logits_aggregation) + dist_aggregation = torch.distributions.categorical.Categorical(logits=logits_aggregation, validate_args=False) # Index 0 corresponds to "no aggregation". aggregation_ops_total_mass = torch.sum(dist_aggregation.probs[:, 1:], dim=1) # Predict some aggregation in case of an answer that needs aggregation. diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 137ce1d53e59..42acee1d3dfb 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -82,7 +82,13 @@ # Every backend, generate path only. "generate": {}, # ONNX, every variant. - "onnx": {}, + "onnx": { + "HunYuanVLModel": ( + "ONNX export trips an int32-overflow `GuardOnDataDependentSymNode` (`64*h*w`) in the vision " + "patch-merger conv on symbolic spatial dims. Plain `torch.export` (dynamo) exports fine." + ), + "HunYuanVLForConditionalGeneration": "Same ONNX vision-conv int32 guard as `HunYuanVLModel`.", + }, # ONNX, generate path only. "onnx.generate": {}, # ONNX, dynamic-shape only. @@ -106,11 +112,33 @@ "Deformable-attention pixel decoder exceeds the 1000s test timeout under dynamic shapes." ), "Mask2FormerForUniversalSegmentation": "Same `timeout` as `Mask2FormerModel`.", + "FunnelModel": ( + "onnxscript's constant-folding optimizer raises `Bitwidth not available for ONNX data type: " + "STRING` on funnel's dynamic-shape graph. `torch.export`/OpenVINO and static ONNX all export fine." + ), + "FunnelForMaskedLM": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelForPreTraining": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelForQuestionAnswering": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelForTokenClassification": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelBaseModel": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelForMultipleChoice": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", + "FunnelForSequenceClassification": "Same onnxscript optimizer `STRING` failure as `FunnelModel`.", }, # ExecuTorch, every variant. "executorch": { "GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."), "GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.", + # These export fine via `torch.export` (dynamo) but hit ExecuTorch edge-lowering limits. + "FunnelModel": "ExecuTorch edge lowering rejects a broadcast in the relative-position attention.", + "FunnelForMaskedLM": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelForPreTraining": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelForQuestionAnswering": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelForTokenClassification": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelBaseModel": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelForMultipleChoice": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "FunnelForSequenceClassification": "Same ExecuTorch broadcast limit as `FunnelModel`.", + "HunYuanVLModel": "ExecuTorch edge lowering fails on the vision stack (same family as the ONNX/OpenVINO gaps).", + "HunYuanVLForConditionalGeneration": "Same ExecuTorch limit as `HunYuanVLModel`.", "Kimi_K25Model": ( "ExecuTorch `to_edge` spec failure in the DeepseekV3 language model; the vision encoder exports fine." ), @@ -136,7 +164,14 @@ "DiffusionGemmaForBlockDiffusion": "Same `timeout` failure as `DiffusionGemmaModel`.", }, # OpenVINO, every variant. - "openvino": {}, + "openvino": { + "TapasModel": "OpenVINO has no conversion rule for `aten.scatter_reduce.two` (tapas segment reduction).", + "TapasForMaskedLM": "Same OpenVINO `scatter_reduce` gap as `TapasModel`.", + "TapasForQuestionAnswering": "Same OpenVINO `scatter_reduce` gap as `TapasModel`.", + "TapasForSequenceClassification": "Same OpenVINO `scatter_reduce` gap as `TapasModel`.", + "HunYuanVLModel": "OpenVINO conversion of the vision stack fails (same family as the ONNX/ExecuTorch gaps).", + "HunYuanVLForConditionalGeneration": "Same OpenVINO gap as `HunYuanVLModel`.", + }, # OpenVINO, generate path only. "openvino.generate": {}, # OpenVINO, dynamic-shape only. diff --git a/tests/models/funnel/test_modeling_funnel.py b/tests/models/funnel/test_modeling_funnel.py index 2cf064850bbb..3561ee5d07d0 100644 --- a/tests/models/funnel/test_modeling_funnel.py +++ b/tests/models/funnel/test_modeling_funnel.py @@ -375,7 +375,7 @@ class FunnelModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): else {} ) - test_torch_exportable = False # pending unbacked symbols from chunked pooling + test_torch_exportable = True # pending unbacked symbols from chunked pooling # special case for ForPreTraining model def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): @@ -434,7 +434,7 @@ class FunnelBaseModelTest(ModelTesterMixin, unittest.TestCase): (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else () ) - test_torch_exportable = False # pending unbacked symbols from chunked pooling + test_torch_exportable = True # pending unbacked symbols from chunked pooling def setUp(self): self.model_tester = FunnelModelTester(self, base=True) diff --git a/tests/models/hunyuan_vl/test_modeling_hunyuan_vl.py b/tests/models/hunyuan_vl/test_modeling_hunyuan_vl.py index e37ffa293e4a..54b970687784 100644 --- a/tests/models/hunyuan_vl/test_modeling_hunyuan_vl.py +++ b/tests/models/hunyuan_vl/test_modeling_hunyuan_vl.py @@ -146,7 +146,7 @@ def prepare_config_and_inputs(self): class HunYuanVLModelTest(VLMModelTest, unittest.TestCase): model_tester_class = HunYuanVLVisionText2TextModelTester test_all_params_have_gradient = False - test_torch_exportable = False + test_torch_exportable = True # HunYuanVL packs all images into one flat patch stream; pixel_values.shape[0] is total patches, not batch size. skip_test_image_features_output_shape = True diff --git a/tests/models/oneformer/test_modeling_oneformer.py b/tests/models/oneformer/test_modeling_oneformer.py index b3029c641cca..6e354f51fd03 100644 --- a/tests/models/oneformer/test_modeling_oneformer.py +++ b/tests/models/oneformer/test_modeling_oneformer.py @@ -236,7 +236,7 @@ class OneFormerModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCas is_encoder_decoder = False test_missing_keys = False - test_torch_exportable = False # detection head uses data-dependent filtering + test_torch_exportable = True # detection head uses data-dependent filtering # TODO: Fix the failed tests when this model gets more usage def is_pipeline_test_to_skip( diff --git a/tests/models/tapas/test_modeling_tapas.py b/tests/models/tapas/test_modeling_tapas.py index f0cb28453964..ec485e21d2b2 100644 --- a/tests/models/tapas/test_modeling_tapas.py +++ b/tests/models/tapas/test_modeling_tapas.py @@ -432,7 +432,7 @@ class TapasModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): ) test_resize_embeddings = True - test_torch_exportable = False # data-dependent aggregation logic + test_torch_exportable = True # data-dependent aggregation logic def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): inputs_dict = copy.deepcopy(inputs_dict)