wehub-resource-sync
e06fe8e8c6
Secret Leaks / trufflehog (push) Failing after 1s
Build documentation / build (push) Failing after 1s
Build documentation / build_other_lang (push) Failing after 0s
CodeQL Security Analysis / CodeQL Analysis (push) Failing after 0s
PR CI / pr-ci (push) Failing after 1s
Slow tests on important models (on Push - A10) / Get all modified files (push) Failing after 1s
Slow tests on important models (on Push - A10) / Model CI (push) Has been skipped
Self-hosted runner (benchmark) / Benchmark (aws-g5-4xlarge-cache) (push) Has been cancelled
New model PR merged notification / Notify new model (push) Has been cancelled
Update Transformers metadata / build_and_package (push) Has been cancelled
671 lines
35 KiB
Python
671 lines
35 KiB
Python
# 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.
|
||
|
||
import copy
|
||
import inspect
|
||
import re
|
||
|
||
import pytest
|
||
import torch
|
||
from parameterized import parameterized
|
||
|
||
from transformers import set_seed
|
||
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.utils import (
|
||
decompose_for_generation,
|
||
decompose_multimodal,
|
||
get_leaf_tensors,
|
||
is_multimodal,
|
||
)
|
||
from transformers.testing_utils import (
|
||
require_executorch,
|
||
require_onnxruntime,
|
||
require_onnxscript,
|
||
require_torch_greater_or_equal,
|
||
set_config_for_less_flaky_test,
|
||
set_model_for_less_flaky_test,
|
||
slow,
|
||
torch_device,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 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.
|
||
#
|
||
# 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.
|
||
|
||
|
||
EXPORT_SKIPS: dict[str, dict[str, str]] = {
|
||
# Every backend, every variant.
|
||
"all": {
|
||
"VideoMAEForPreTraining": (
|
||
"Computes loss even when `return_loss=False`, hitting a data-dependent guard in "
|
||
"`mse_loss`. TODO: skip loss when labels aren't provided."
|
||
),
|
||
"OpenAIPrivacyFilterModel": (
|
||
"`get_correct_experts_implementation` defaults to `eager` because the model is "
|
||
"sensitive to accumulation order. Eager experts forward iterates over "
|
||
"`expert_hit.nonzero()` (data-dependent shape). Users can opt into "
|
||
"`set_experts_implementation('batched_mm')` to export."
|
||
),
|
||
"OpenAIPrivacyFilterForTokenClassification": (
|
||
"Same root cause as `OpenAIPrivacyFilterModel` — eager experts implementation."
|
||
),
|
||
},
|
||
# Every backend, generate path only.
|
||
"generate": {
|
||
"Blip2ForConditionalGeneration": (
|
||
"`generate()` delegates to the inner language model without calling top-level "
|
||
"`forward()`, so `decompose_prefill_decode` can't capture inputs. "
|
||
"TODO: route generate through top-level `forward()`."
|
||
),
|
||
"InstructBlipForConditionalGeneration": "Same `generate()`-delegation as Blip2.",
|
||
"InstructBlipVideoForConditionalGeneration": "Same `generate()`-delegation as Blip2.",
|
||
"Kosmos2ForConditionalGeneration": "Same `generate()`-delegation as Blip2.",
|
||
"RecurrentGemmaForCausalLM": (
|
||
"Stores recurrent/conv state as module attributes (not a `Cache` object); "
|
||
"`torch.export` can't carry that state between calls. "
|
||
"TODO: refactor to a cache-based SSM pattern (like Mamba/Mamba2)."
|
||
),
|
||
"MoshiForConditionalGeneration": (
|
||
"`generate()` creates `blank_user_audio_codes` outside the traced forward and "
|
||
"passes it as a kwarg; the resulting ONNX input has mismatched rank (scalar vs 3D). "
|
||
"TODO: make `blank_user_audio_codes` part of the model state."
|
||
),
|
||
"UdopForConditionalGeneration": (
|
||
"Exported decoder output is missing `attention_mask` vs eager — encoder-decoder "
|
||
"cross-attention mask doesn't flow through the generate decomposition correctly."
|
||
),
|
||
"VoxtralRealtimeForConditionalGeneration": (
|
||
"Exported prefill drops `past_key_values.*.{keys,values,_sliding_window_tensor}` "
|
||
"tensors that eager returns. Plain forward exports work. "
|
||
"TODO: align generate-decomposition path with the realtime KV-cache shape."
|
||
),
|
||
"Gemma3nForConditionalGeneration": (
|
||
"KV-shared layers (`num_kv_shared_layers`) reuse cache entries from earlier layers; "
|
||
"exported prefill returns only `logits` while eager surfaces the populated KV cache. "
|
||
"Same shape as Voxtral. TODO: align the generate-decomposition path."
|
||
),
|
||
},
|
||
# ONNX, every variant.
|
||
"onnx": {
|
||
"CHMv2ForDepthEstimation": (
|
||
"`run_decompositions` retraces through aot_autograd which emits a `detach_(alias(...))` "
|
||
"pair the functional-graph assertion rejects (independent of any source `.detach()` — "
|
||
"verified). Torch export works. TODO: file upstream `torch.export` issue."
|
||
),
|
||
},
|
||
# ONNX, generate path only.
|
||
"onnx.generate": {
|
||
"ReformerModelWithLMHead": (
|
||
"Chunked local attention exports a Constant idx that exceeds the cached-keys axis "
|
||
"length under static decode (prefill+1 token, seq=17 vs chunked axis of 16). The same "
|
||
"computation stays symbolic under dynamic so ORT can't pre-validate it. The other "
|
||
"three Reformer-local-attn ONNX variants pass."
|
||
),
|
||
},
|
||
# ONNX, dynamic-shape only.
|
||
"onnx.dynamic": {
|
||
"GroundingDinoModel": (
|
||
"Same `detach_(alias(...))` retrace bug as CHMv2, but only triggered under dynamic "
|
||
"shapes — `aot_autograd`'s decomposition pipeline emits the detach itself (verified "
|
||
"by guarding all three modeling-side detaches with `if self.training`). Static works."
|
||
),
|
||
"GroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.",
|
||
"MMGroundingDinoModel": "Same as `GroundingDinoModel`.",
|
||
"MMGroundingDinoForObjectDetection": "Same as `GroundingDinoModel`.",
|
||
"Sam2VisionModel": (
|
||
"`torch.export` of the Hiera vision backbone under dynamic shapes takes ~7.5 min "
|
||
"even after simplifying `window_partition`/`window_unpartition` (12 attention blocks "
|
||
"× 3 Q-pool stage transitions on symbolic H/W). ONNX + ORT push past 1000s timeout."
|
||
),
|
||
"Sam2Model": "Same Hiera-backbone dynamic-shape budget overrun as `Sam2VisionModel`.",
|
||
},
|
||
# ExecuTorch — lowering failures grouped by root cause; see the first entry of each
|
||
# `Same ... as` chain for the full description.
|
||
"executorch": {
|
||
"BarkFineModel": (
|
||
"ExecuTorch memory planning miscomputes the tensor spec (`buffer of size N, expected nbytes of M`) — a dtype-size mismatch in the lowered program."
|
||
),
|
||
"ClvpModelForConditionalGeneration": (
|
||
"A pass-through output aliases an input (`Output node is already in the inputs`)."
|
||
),
|
||
"ColQwen2ForRetrieval": (
|
||
"ExecuTorch dim-order lowering requires a copying view (`Cannot view a tensor ... with shape/strides`)."
|
||
),
|
||
"DabDetrModel": ("XNNPACK partitioner: `Attempting to convert non-NHWC compatible node to NHWC`."),
|
||
"DabDetrForObjectDetection": "Same `nhwc` failure as `DabDetrModel`.",
|
||
"Ernie4_5_VLMoeModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Ernie4_5_VLMoeForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"FlavaForPreTraining": ("XNNPACK partitioner: `Invalid partition, found dependency cycles`."),
|
||
"GPT2Model": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GPT2LMHeadModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GPT2DoubleHeadsModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GPT2ForQuestionAnswering": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GPT2ForSequenceClassification": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GPT2ForTokenClassification": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Gemma3nModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"Gemma3nForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"Glm46VModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Glm46VForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Glm4vModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Glm4vForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Glm4vMoeModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Glm4vMoeForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GlmImageModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GlmImageForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GlmOcrModel": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GlmOcrForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"GroundingDinoModel": ("Lowering exceeds the test timeout under dynamic shapes."),
|
||
"GroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.",
|
||
"InstructBlipModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"InstructBlipForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"InstructBlipVideoForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"InstructBlipVideoModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"MMGroundingDinoModel": "Same `timeout` failure as `GroundingDinoModel`.",
|
||
"MMGroundingDinoForObjectDetection": "Same `timeout` failure as `GroundingDinoModel`.",
|
||
"MiniMaxM3VLModel": ("Serialization rejects an i64 constant (`bad number for type int32`)."),
|
||
"MiniMaxM3SparseForConditionalGeneration": "Same `int32` failure as `MiniMaxM3VLModel`.",
|
||
"PPDocLayoutV3ForObjectDetection": ("Delegation drops a referenced weight (`KeyError` on a state-dict key)."),
|
||
"PaddleOCRVLForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"PerceptionLMModel": "Same `passthrough` failure as `ClvpModelForConditionalGeneration`.",
|
||
"PerceptionLMForConditionalGeneration": "Same `passthrough` failure as `ClvpModelForConditionalGeneration`.",
|
||
"Qwen2VLModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen2VLForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen2_5OmniThinkerForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Qwen2_5_VLModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen2_5_VLForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3OmniMoeThinkerForConditionalGeneration": "Same `view` failure as `ColQwen2ForRetrieval`.",
|
||
"Qwen3_5Model": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3_5ForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3_5ForSequenceClassification": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3_5ForTokenClassification": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3_5MoeModel": "Same `spec` failure as `BarkFineModel`.",
|
||
"Qwen3_5MoeForConditionalGeneration": "Same `spec` failure as `BarkFineModel`.",
|
||
},
|
||
"executorch.generate": {
|
||
"PPFormulaNetForConditionalGeneration": (
|
||
"ExecuTorch memory planning miscomputes the tensor spec (`buffer of size N, expected nbytes of M`) — a dtype-size mismatch in the lowered program."
|
||
),
|
||
},
|
||
"executorch.dynamic": {
|
||
"BigBirdModel": ("Lowering exceeds the test timeout under dynamic shapes."),
|
||
"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`.",
|
||
"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`.",
|
||
},
|
||
}
|
||
|
||
|
||
# ──────────────────────────── 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.",
|
||
},
|
||
# 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`.",
|
||
},
|
||
}
|
||
|
||
|
||
# Parameterization for export tests: runs once with dynamic=True and once with dynamic=False.
|
||
DYNAMIC_EXPORT_PARAMS = parameterized.expand(
|
||
[(False,), (True,)],
|
||
name_func=lambda f, _, p: f"{f.__name__}_{'dynamic' if p.args[0] else 'static'}",
|
||
)
|
||
|
||
# Maximum time (in seconds) for a single export test before it is killed.
|
||
EXPORT_TEST_TIMEOUT = 1000
|
||
|
||
# Minimum torch version the exporters target — older releases lack `torch.export` features the
|
||
# exporters rely on, so the export sweep is skipped (not failed) below this. Sourced from the
|
||
# exporter itself so the test and the runtime check can't drift apart.
|
||
MIN_EXPORT_TORCH_VERSION = DynamoExporter.min_versions["torch"]
|
||
|
||
|
||
# ──────────────────────────── helpers ────────────────────────────
|
||
|
||
|
||
def _cast_inputs(obj, device, dtype):
|
||
"""Recursively move tensors to `device`, casting floating-point tensors to `dtype`."""
|
||
if isinstance(obj, torch.Tensor):
|
||
return obj.to(device=device, dtype=dtype) if obj.is_floating_point() else obj.to(device=device)
|
||
if isinstance(obj, dict):
|
||
return {k: _cast_inputs(v, device, dtype) for k, v in obj.items()}
|
||
if isinstance(obj, (list, tuple)):
|
||
return type(obj)(_cast_inputs(v, device, dtype) for v in obj)
|
||
return obj
|
||
|
||
|
||
def _clean_inputs_for_export(inputs_dict, config):
|
||
"""Strip None values and export-incompatible keys from an inputs dict. Mutates config in-place."""
|
||
inputs_dict = {k: v for k, v in inputs_dict.items() if v is not None}
|
||
for key in ("labels", "future_values", "return_loss"):
|
||
inputs_dict.pop(key, None)
|
||
config.return_loss = False
|
||
return inputs_dict
|
||
|
||
|
||
def _run_onnx_program(onnx_program, inputs) -> dict:
|
||
"""Run an ONNX program and return outputs as a `{name: tensor}` dict."""
|
||
set_seed(1234)
|
||
onnx_inputs = get_leaf_tensors(inputs)
|
||
onnx_outputs = onnx_program(**onnx_inputs)
|
||
onnx_names = (re.sub(r"^output\.", "", node.name) for node in onnx_program.model_proto.graph.output)
|
||
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.
|
||
|
||
Mirrors ``_should_skip``'s scope walk on ``ONNX_DISABLE_OPTIMIZE`` — ``"all"`` always
|
||
applies; ``"dynamic"`` adds the dynamic-only entries.
|
||
"""
|
||
name = model_class.__name__
|
||
scopes = ["all"] + (["dynamic"] if dynamic else [])
|
||
return not any(name in ONNX_DISABLE_OPTIMIZE.get(scope, {}) for scope in scopes)
|
||
|
||
|
||
# ──────────────────────────── mixins ────────────────────────────
|
||
|
||
|
||
class ExportTesterMixin:
|
||
"""Mixin providing non-generative export tests for Dynamo, ONNX, and ExecuTorch backends.
|
||
|
||
Mixed into [`ModelTesterMixin`] so every model test class that inherits from it
|
||
automatically runs these export tests against all entries in `all_model_classes`.
|
||
|
||
Expected attributes provided by [`ModelTesterMixin`]:
|
||
- `all_model_classes` — iterable of model class objects to test.
|
||
- `model_tester` — object with `prepare_config_and_inputs_for_common()` (and optionally
|
||
`prepare_config_and_inputs_for_model_class()`).
|
||
- `test_torch_exportable` — bool; set to `False` to skip all export tests for the model.
|
||
- `_prepare_for_class(inputs_dict, model_class)` — adjusts inputs per model class.
|
||
|
||
Tests are parameterised over `dynamic=True` / `dynamic=False` via `DYNAMIC_EXPORT_PARAMS`.
|
||
Multi-modal models (detected by `is_multimodal`) are automatically decomposed and each
|
||
submodule is tested independently.
|
||
"""
|
||
|
||
def _skip_if_not_exportable(self):
|
||
"""Skip the test if the model architecture is not exportable."""
|
||
if not self.test_torch_exportable:
|
||
self.skipTest(reason="Model architecture is not Dynamo exportable/traceable")
|
||
|
||
with open(inspect.getfile(self.all_model_classes[0]), "r") as f:
|
||
source_code = f.read()
|
||
# TODO: add use_experts_implementation support to remaining MoE models
|
||
if "for expert" in source_code and "use_experts_implementation" not in source_code:
|
||
self.skipTest(reason="Model architecture uses eager MoE implementation which is not torch exportable")
|
||
|
||
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, ``"<backend>"`` for that backend, and ``"<backend>.<variant>"`` for
|
||
the more-specific intersections.
|
||
"""
|
||
name = model_class.__name__
|
||
scopes = ["all"]
|
||
if generate:
|
||
scopes.append("generate")
|
||
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)
|
||
|
||
def _prepare_export_model_and_inputs(self, model_class):
|
||
"""Create model and forward inputs ready for export.
|
||
|
||
Returns:
|
||
Dict of `{name: (model, inputs)}` — one entry per component.
|
||
"""
|
||
if hasattr(self.model_tester, "prepare_config_and_inputs_for_model_class"):
|
||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_model_class(model_class)
|
||
else:
|
||
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
|
||
inputs_dict = self._prepare_for_class(inputs_dict, model_class)
|
||
inputs_dict = _clean_inputs_for_export(inputs_dict, config)
|
||
|
||
set_config_for_less_flaky_test(config)
|
||
model = model_class(config).eval().to(torch_device)
|
||
set_model_for_less_flaky_test(model)
|
||
|
||
# Cast inputs to model device/dtype
|
||
try:
|
||
model_param = next(iter(model.parameters()))
|
||
inputs_dict = _cast_inputs(inputs_dict, model_param.device, model_param.dtype)
|
||
except StopIteration:
|
||
pass
|
||
|
||
if is_multimodal(model):
|
||
return decompose_multimodal(model, inputs_dict)
|
||
return {"model": (model, inputs_dict)}
|
||
|
||
def _collect_eager_outputs(self, components):
|
||
"""Run eager forward for each component and return a ``{name: leaf_tensors}`` dict."""
|
||
eager_outputs = {}
|
||
for name, (model, inputs) in components.items():
|
||
with torch.no_grad():
|
||
set_seed(1234)
|
||
eager_outputs[name] = get_leaf_tensors(model(**copy.deepcopy(inputs)))
|
||
assert eager_outputs[name], f"Eager outputs are empty for {name}."
|
||
return eager_outputs
|
||
|
||
def _check_outputs_close(self, actual, expected, atol, rtol, check_device=True):
|
||
"""Assert outputs are close, allowing up to 5% element-level mismatch."""
|
||
try:
|
||
torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol, check_device=check_device)
|
||
except AssertionError as e:
|
||
mismatched_percentage = re.findall(r"Mismatched elements: (\d+) / (\d+)", str(e))
|
||
if mismatched_percentage:
|
||
mismatched, total = map(int, mismatched_percentage[0])
|
||
if mismatched / total < 0.05:
|
||
return # allow up to 5%
|
||
raise e
|
||
|
||
# ──────────────────── torch.export tests ─────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@pytest.mark.torch_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_torch_export(self, dynamic, atol=1e-4, rtol=1e-4):
|
||
"""Export each model class with ``torch.export`` and verify outputs match eager within tolerance."""
|
||
self._skip_if_not_exportable()
|
||
|
||
exporter = DynamoExporter()
|
||
config = DynamoConfig(dynamic=dynamic)
|
||
|
||
for model_class in self.all_model_classes:
|
||
if self._should_skip(model_class):
|
||
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}"):
|
||
exported_program = exporter.export(model, inputs, config=config)
|
||
|
||
with torch.no_grad():
|
||
set_seed(1234)
|
||
exported_outputs = get_leaf_tensors(exported_program.module()(**copy.deepcopy(inputs)))
|
||
self.assertTrue(exported_outputs, f"Exported outputs are empty for {name}.")
|
||
|
||
self._check_outputs_close(exported_outputs, eager_outputs[name], atol=atol, rtol=rtol)
|
||
|
||
# ──────────────────────── ONNX tests ─────────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@require_onnxscript
|
||
@require_onnxruntime
|
||
@pytest.mark.onnx_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_onnx_export(self, dynamic):
|
||
"""Export each model class to ONNX and verify output names match eager."""
|
||
self._skip_if_not_exportable()
|
||
|
||
for model_class in self.all_model_classes:
|
||
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)
|
||
|
||
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}"):
|
||
onnx_program = exporter.export(model, inputs, config=config)
|
||
onnx_outputs = _run_onnx_program(onnx_program, inputs)
|
||
self.assertTrue(onnx_outputs, f"ONNX outputs are empty for {name}.")
|
||
self.assertEqual(set(onnx_outputs.keys()), set(eager_outputs[name].keys()))
|
||
|
||
# ──────────────────── ExecuTorch tests ───────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@require_executorch
|
||
@pytest.mark.executorch_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_executorch_export(self, dynamic):
|
||
"""Export each model class to ExecuTorch (xnnpack on CPU, cuda on GPU) and verify no errors."""
|
||
|
||
self._skip_if_not_exportable()
|
||
exporter = ExecutorchExporter()
|
||
config = ExecutorchConfig(dynamic=dynamic)
|
||
|
||
for model_class in self.all_model_classes:
|
||
if self._should_skip(model_class, dynamic=dynamic, backend="executorch"):
|
||
continue
|
||
|
||
components = self._prepare_export_model_and_inputs(model_class)
|
||
|
||
for name, (model, inputs) in components.items():
|
||
with self.subTest(f"{model_class.__name__}/{name}"):
|
||
exporter.export(model, inputs, config=config)
|
||
|
||
|
||
class ExportGenerateTesterMixin(ExportTesterMixin):
|
||
"""Mixin providing generation-aware export tests for torch.export, ONNX, and ExecuTorch backends.
|
||
|
||
Inherits ``ExportTesterMixin`` for the shared exportability gate / skip logic / input prep, and
|
||
is mixed into a model test class alongside ``GenerationTesterMixin``.
|
||
|
||
Required attributes on the host class (in addition to those from ``ExportTesterMixin``):
|
||
- ``all_generative_model_classes`` — iterable of generative model class objects to test.
|
||
- ``prepare_config_and_inputs_for_generate()`` — returns ``(config, inputs_dict)`` suitable
|
||
for ``model.generate()``.
|
||
|
||
Each generative model is decomposed into prefill and decode components via
|
||
:func:`decompose_prefill_decode`. Multi-modal models additionally decompose the prefill
|
||
stage into individual submodules via :func:`decompose_multimodal`.
|
||
"""
|
||
|
||
def _prepare_export_generate_model_and_inputs(self, model_class):
|
||
"""Decompose a generative model into exportable components.
|
||
|
||
For multi-modal models: decomposes the prefill stage into individual submodules plus the decode stage.
|
||
For decoder-only models: returns prefill and decode components.
|
||
|
||
Returns:
|
||
Dict of `{name: (model, inputs)}` — one entry per component.
|
||
"""
|
||
config, inputs_dict = self.prepare_config_and_inputs_for_generate()
|
||
inputs_dict = _clean_inputs_for_export(inputs_dict, config)
|
||
|
||
set_config_for_less_flaky_test(config)
|
||
model = model_class(config).eval().to(torch_device)
|
||
set_model_for_less_flaky_test(model)
|
||
|
||
return decompose_for_generation(model, inputs_dict)
|
||
|
||
# ──────────────────── torch.export tests ─────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@pytest.mark.torch_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_torch_export_generate(self, dynamic, atol=1e-4, rtol=1e-4):
|
||
"""Export prefill and decode stages with ``torch.export`` and verify outputs match eager."""
|
||
self._skip_if_not_exportable()
|
||
|
||
exporter = DynamoExporter()
|
||
config = DynamoConfig(dynamic=dynamic)
|
||
|
||
for model_class in self.all_generative_model_classes:
|
||
if self._should_skip(model_class, generate=True):
|
||
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}"):
|
||
exported_program = exporter.export(model, inputs, config=config)
|
||
|
||
with torch.no_grad():
|
||
set_seed(1234)
|
||
exported_outputs = get_leaf_tensors(exported_program.module()(**copy.deepcopy(inputs)))
|
||
self.assertTrue(exported_outputs, "Exported outputs are empty.")
|
||
|
||
self._check_outputs_close(exported_outputs, eager_outputs[name], atol=atol, rtol=rtol)
|
||
|
||
# ──────────────────────── ONNX tests ─────────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@require_onnxscript
|
||
@require_onnxruntime
|
||
@pytest.mark.onnx_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_onnx_export_generate(self, dynamic):
|
||
"""Export prefill and decode stages to ONNX and verify output names match eager."""
|
||
self._skip_if_not_exportable()
|
||
|
||
for model_class in self.all_generative_model_classes:
|
||
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)
|
||
|
||
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}"):
|
||
onnx_program = exporter.export(model, inputs, config=config)
|
||
set_seed(1234)
|
||
onnx_outputs = _run_onnx_program(onnx_program, inputs)
|
||
self.assertTrue(onnx_outputs, "ONNX outputs are empty.")
|
||
self.assertEqual(set(onnx_outputs.keys()), set(eager_outputs[name].keys()))
|
||
|
||
# ──────────────────── ExecuTorch tests ───────────────────────
|
||
|
||
@DYNAMIC_EXPORT_PARAMS
|
||
@slow
|
||
@require_executorch
|
||
@pytest.mark.executorch_export_test
|
||
@pytest.mark.timeout(EXPORT_TEST_TIMEOUT)
|
||
@require_torch_greater_or_equal(MIN_EXPORT_TORCH_VERSION)
|
||
def test_executorch_export_generate(self, dynamic):
|
||
"""Export prefill and decode stages to ExecuTorch and verify no errors."""
|
||
|
||
self._skip_if_not_exportable()
|
||
exporter = ExecutorchExporter()
|
||
config = ExecutorchConfig(dynamic=dynamic)
|
||
|
||
for model_class in self.all_generative_model_classes:
|
||
if self._should_skip(model_class, generate=True, dynamic=dynamic, backend="executorch"):
|
||
continue
|
||
|
||
components = self._prepare_export_generate_model_and_inputs(model_class)
|
||
|
||
for name, (model, inputs) in components.items():
|
||
with self.subTest(f"{model_class.__name__}/{name}"):
|
||
exporter.export(model, inputs, config=config)
|