diff --git a/docs/source/en/exporters.md b/docs/source/en/exporters.md
index 9b2871a56e3a..901b3132e3d7 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}
+```
+
@@ -380,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)
@@ -466,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/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/cache_utils.py b/src/transformers/cache_utils.py
index a0ea5b2b8c6c..faff0602790d 100644
--- a/src/transformers/cache_utils.py
+++ b/src/transformers/cache_utils.py
@@ -916,8 +916,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/__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..c2d44dcd6278 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,
}
@@ -69,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 713448debacd..e05d79cbae64 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"
@@ -99,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
@@ -141,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
@@ -168,3 +168,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 60a9ce62808b..516762acfd01 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
@@ -351,8 +352,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)
@@ -382,6 +383,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",
@@ -449,6 +452,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}
@@ -471,12 +490,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")
@@ -523,9 +551,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")
@@ -566,6 +597,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
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..793d323256b7 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__)
@@ -336,7 +354,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.
@@ -347,14 +365,43 @@ 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.
+
+ 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.
"""
+ 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]:
@@ -367,13 +414,16 @@ 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):
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
@@ -410,10 +460,34 @@ 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])
- return original(self, *sizes).clone(memory_format=torch.contiguous_format)
+ 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():
+ return result.clone(memory_format=torch.contiguous_format)
+ return result
+
+ 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
@@ -433,22 +507,32 @@ 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
@@ -465,7 +549,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)]
@@ -526,7 +609,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:
@@ -560,7 +642,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
@@ -584,6 +665,125 @@ 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.
+ """
+
+ 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``."""
+
+ 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",
@@ -609,12 +809,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)
@@ -637,6 +831,343 @@ 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.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
@@ -734,6 +1265,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"]``.
@@ -875,10 +1430,53 @@ 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
+
+
+@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/exporters/exporter_onnx.py b/src/transformers/exporters/exporter_onnx.py
index ef65aaa27044..7af8349aa783 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,
@@ -188,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:
@@ -245,6 +270,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 +365,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)."""
@@ -417,7 +497,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)
@@ -668,8 +749,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)
@@ -746,7 +829,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]))
@@ -780,6 +865,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`.
@@ -808,10 +898,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
new file mode 100644
index 000000000000..3a1f142d06c2
--- /dev/null
+++ b/src/transformers/exporters/exporter_openvino.py
@@ -0,0 +1,2096 @@
+# 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_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.
+ 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_runtime_asserts(graph_module) -> None:
+ """Drop ``_assert_tensor_metadata`` / ``_assert_scalar`` runtime asserts before the replay.
+
+ ``_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 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()
+
+
+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 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). 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.
+ 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 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))
+ 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 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"))
+ 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:
+ # ``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:
+ 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):
+ # ``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=(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))
+ # 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.self, args=(any_match, maxes, self_arg))
+ 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
+
+
+@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
+
+
+@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
+
+
+@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
+# 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)
+ # ``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")
+ # 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_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.
+
+ 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":
+ # 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))
+ # ``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):
+ """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.
+
+ 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)]
+
+
+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(
+ [
+ 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.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_truediv),
+ 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..6a80d55c2722 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):
@@ -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,11 +444,24 @@ 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)
- # 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)
+ # 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)
+
+ # 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")
@@ -458,9 +472,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")
@@ -482,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`,
@@ -613,12 +662,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 +688,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]),
}
@@ -684,9 +745,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]]:
@@ -741,7 +806,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/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/fsmt/modeling_fsmt.py b/src/transformers/models/fsmt/modeling_fsmt.py
index 9388d4962dfe..b5c14227a4d2 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/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/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/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/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/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/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/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/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 ee9a2e739cfa..37725a09ffc2 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 054347ef5b82..ede80422a952 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/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/perceiver/modeling_perceiver.py b/src/transformers/models/perceiver/modeling_perceiver.py
index bc5fdff0a308..51cef1d63423 100755
--- a/src/transformers/models/perceiver/modeling_perceiver.py
+++ b/src/transformers/models/perceiver/modeling_perceiver.py
@@ -2484,7 +2484,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.
@@ -2493,13 +2493,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")
@@ -2578,7 +2582,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, ...).
@@ -2589,12 +2593,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)
@@ -2641,14 +2647,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
@@ -3262,8 +3268,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/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/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py b/src/transformers/models/pp_doclayout_v3/modeling_pp_doclayout_v3.py
index 33d445632566..2dc3de0b4807 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
@@ -1533,23 +1533,14 @@ 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)
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 4fc9088a652d..0d08b7706c60 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
@@ -578,23 +578,14 @@ 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)
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/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 4db2c92c7251..035ff132e8b4 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,
)
@@ -619,20 +603,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):
@@ -676,36 +646,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/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py
index e9049a761042..06c88a20b271 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 358dba3308d6..1653bb22939f 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 dcd41b7e552c..7e053ea63b56 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/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/src/transformers/models/tvp/modeling_tvp.py b/src/transformers/models/tvp/modeling_tvp.py
index 688f41a67721..2d58f4cd08bc 100644
--- a/src/transformers/models/tvp/modeling_tvp.py
+++ b/src/transformers/models/tvp/modeling_tvp.py
@@ -745,9 +745,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 ed0505d2f1b7..b8c9c87c057f 100755
--- a/src/transformers/models/unispeech/modeling_unispeech.py
+++ b/src/transformers/models/unispeech/modeling_unispeech.py
@@ -1131,11 +1131,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 edeae0a3c9cb..c2c5a0f09b29 100644
--- a/src/transformers/models/unispeech/modular_unispeech.py
+++ b/src/transformers/models/unispeech/modular_unispeech.py
@@ -376,11 +376,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..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/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 9f32ed760900..b6361095253b 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 155d023243cb..75631ae10636 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/testing_utils.py b/src/transformers/testing_utils.py
index 287386b1d5e7..4f08162d0cb3 100644
--- a/src/transformers/testing_utils.py
+++ b/src/transformers/testing_utils.py
@@ -128,6 +128,7 @@
is_onnxruntime_available,
is_onnxscript_available,
is_openai_available,
+ is_openvino_available,
is_optimum_available,
is_optimum_quanto_available,
is_pandas_available,
@@ -626,6 +627,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 2f697261e0d4..86cd08cb78b8 100644
--- a/src/transformers/utils/__init__.py
+++ b/src/transformers/utils/__init__.py
@@ -182,6 +182,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 1770eeb7b2f5..a31c8f6aed13 100644
--- a/src/transformers/utils/import_utils.py
+++ b/src/transformers/utils/import_utils.py
@@ -916,6 +916,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/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 fd1d22c53c71..42acee1d3dfb 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,
require_torch_greater_or_equal,
set_config_for_less_flaky_test,
set_model_for_less_flaky_test,
@@ -45,169 +47,108 @@
# ──────────────────────────── 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]] = {
# 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."
+ "all": {},
+ # 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."
),
- "OpenAIPrivacyFilterForTokenClassification": (
- "Same root cause as `OpenAIPrivacyFilterModel` — eager experts implementation."
+ "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": {
- "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."
- ),
- },
+ "generate": {},
# 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."
+ "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": {
- "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.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."
+ "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."
),
- "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."
+ "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."
),
- "Sam2Model": "Same Hiera-backbone dynamic-shape budget overrun as `Sam2VisionModel`.",
+ "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 — lowering failures grouped by root cause; see the first entry of each
- # `Same ... as` chain for the full description.
+ # ExecuTorch, every variant.
"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`.",
+ # 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."
+ ),
+ "Kimi_K25ForConditionalGeneration": "Same ExecuTorch `spec` failure as `Kimi_K25Model` (DeepseekV3 LM).",
"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, 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`.",
@@ -217,82 +158,40 @@
"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)."),
- "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`.",
- "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`.",
+ "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`.",
},
-}
-
-
-# ──────────────────────────── 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 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.",
+ # OpenVINO, every variant.
+ "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`.",
},
- # 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`.",
+ # OpenVINO, generate path only.
+ "openvino.generate": {},
+ # OpenVINO, dynamic-shape only.
+ "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`.",
},
}
@@ -344,15 +243,76 @@ def _run_onnx_program(onnx_program, inputs) -> dict:
return dict(zip(onnx_names, onnx_outputs))
-def _onnx_optimize_enabled(model_class, dynamic: bool) -> bool:
- """Return whether onnxscript optimisation should run for this model under this shape mode.
+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()
- Mirrors ``_should_skip``'s scope walk on ``ONNX_DISABLE_OPTIMIZE`` — ``"all"`` always
- applies; ``"dynamic"`` adds the dynamic-only entries.
+ return outputs
+
+
+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``.
+
+ 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 ────────────────────────────
@@ -390,22 +350,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.
@@ -472,7 +427,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)
@@ -506,9 +461,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)
@@ -520,6 +474,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 ───────────────────────
@DYNAMIC_EXPORT_PARAMS
@@ -630,9 +611,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)
@@ -645,6 +625,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 ───────────────────────
@DYNAMIC_EXPORT_PARAMS
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/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/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/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/kimi_k25/test_modeling_kimi_k25.py b/tests/models/kimi_k25/test_modeling_kimi_k25.py
index 10a11acaa02d..34c92edceb14 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):
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/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/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,
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)
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)