HfExporters: Native, Unified export for PyTorch / ONNX / ExecuTorch

Edit: pieces of this PR were merged separately (#42697, #42317), so this one is now focused on the HfExporters machinery itself.

What this adds

A native, in-Transformers export pipeline — one base class (HfExporter), three subclasses for the runtimes we care about, one unified API:

Same call shape across all three. Dynamic shapes by default. Generation-style models split automatically into prefill + decode (+ vision/audio sub-encoders for VLMs).

from transformers import AutoModelForMaskedLM, AutoTokenizer
from transformers.exporters import OnnxExporter, OnnxConfig

model_id = "hf-internal-testing/tiny-random-BertForMaskedLM"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(model_id).eval()
inputs = tokenizer(["Hello, my dog is cute"] * 2, return_tensors="pt")
onnx_program = OnnxExporter().export(model, inputs, config=OnnxConfig(dynamic=True))

new_input = tokenizer("Hello, my cat is so adorable!", return_tensors="pt")
torch.testing.assert_close(
    onnx_program.call_reference(**new_input)[0],   # numpy reference
    onnx_program(**new_input)[0],                  # onnxruntime
    rtol=1e-4, atol=1e-4,
)

Swap one line for another runtime — DynamoExporter() / DynamoConfig or ExecutorchExporter() / ExecutorchConfig(backend=...).

For generative models the prefill/decode split is captured automatically:

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.exporters import OnnxExporter, OnnxConfig

model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).eval()
inputs = tokenizer(["Hello, my dog is cute"] * 2, return_tensors="pt")

artifacts = OnnxExporter().export_for_generation(model, inputs, config=OnnxConfig(dynamic=True))
# {"prefill": ONNXProgram, "decode": ONNXProgram}
# For VLMs: also vision_encoder, audio_encoder, multi_modal_projector, language_model, lm_head

Why native, not via Optimum

torch.export is fast and clean but strict compared to TorchScript — data-dependent ifs that TorchScript would just trace through (with a warning) become hard guard failures, hence the if not torch.compiler.is_exporting() opt-outs landing throughout modeling code. Living downstream in Optimum means every new architecture (a new attention pattern, a new cache type, a new modality) waits for someone to write export glue. By colocating the exporters with the models, the export evolves with the model and breakages surface in the model's own test file, not weeks later in a separate repo.

Optimum will keep building on this — generating inputs, dynamic axes, splitting components, wiring inference. This PR is the API and the registry it builds on.

How it works

export_for_generation runs the model's real generate(..., max_new_tokens=2) once, hooks forward to capture the actual prefill and decode kwargs, then exports each independently. Multimodal models decompose further into per-component sub-graphs. Decoder-only, SSM, encoder-decoder — same machinery.

Multi-stage patching

Every backend follows the same pipeline (5 stages for ONNX, 4 for the others). The principle: fix exports without touching modeling code wherever possible. Each fix lives at the lowest stage that can express it cleanly.

@register_export_input_preparer(*markers) covers a related case: precompute data-dependent kwargs (cu_seqlens, position ids, window indices, bilinear interpolation grids, audio chunking) outside the traced graph.

Adding support for a new model is a function and a decorator. No new files, no plumbing, no per-backend boilerplate.

Other moving parts

• Dynamic shapes by default. dynamic=True walks tensors / caches / dicts / lists / ModelOutputs and assigns Dim.AUTO everywhere — no per-model dynamic_shapes dict.
• Cache pytree registration is reflection-driven. Cache.__subclasses__() is walked at export time and each registered as a pytree node; model-specific cache types (anything ending in Cache) get picked up too. past_key_values becomes just another tensor structure to the exporter.
• State cleanup after export. Stateful tensor attrs (sliding-window buffers, etc.) that torch.export may leave as FakeTensors are reset so a follow-up eager forward on the same module stays safe.
• Forward-signature flattening. **kwargs: Unpack[TransformersKwargs] is rewritten to a flat signature derived from the user's inputs — otherwise torch.export bundles **kwargs into one blob that doesn't line up with dynamic_shapes.

What's tested

The export tests are wired into ModelTesterMixin directly, so every model that uses the standard test mixin gets test_torch_export_*, test_onnx_export_*, and test_executorch_export_* for free (forward + generation variants where relevant). That's the bulk of the model zoo — LLMs, VLMs, audio encoders, vision backbones, detection / segmentation.

Skip lists in tests/exporters/test_utils.py are short, scoped, and each entry points to a specific upstream torch.export / onnxscript issue or modeling bug we want to fix at the source — not paper over here.

What's not perfect yet

• export_for_generation returns independent graphs, not a turnkey inference pipeline. Wiring them together (encoder → projector → language model → lm_head, decode loop) is downstream's job. We're actively working to reduce that glue.
• Some workarounds in the registries will disappear when PyTorch or onnxscript lands the equivalent fix. Each is marked with the issue it works around — keeping these visible is the point.

Reproducing

pip install -e . torch onnxscript executorch onnxruntime

Then any of the snippets above; full per-model tests:

RUN_SLOW=1 pytest tests/models/<model>/test_modeling_<model>.py \
    -k "test_onnx_export or test_torch_export or test_executorch_export"