Implementation:Huggingface Optimum Model Decomposition Utils

Field	Value
Page Type	Implementation
Source Repository	https://github.com/huggingface/optimum
Source File	`optimum/exporters/utils.py`
Domains	NLP, Computer_Vision, Export
Last Updated	2026-02-15 00:00 GMT

Overview

This implementation provides the utility functions that decompose complex multi-component models into individually exportable sub-models. Each function returns a dictionary mapping component names to (model, ExporterConfig) tuples, ready for independent export.

API Reference

`get_encoder_decoder_models_for_export`

Source location: optimum/exporters/utils.py, lines 190-221

Purpose: Splits an encoder-decoder model (e.g., T5, BART, Marian) into its encoder and decoder components, each paired with an appropriate export configuration.

Signature:

def get_encoder_decoder_models_for_export(
    model: "PreTrainedModel",
    config: "ExporterConfig",
) -> Dict[str, Tuple["PreTrainedModel", "ExporterConfig"]]:

Parameters:

Parameter	Type	Description
`model`	`PreTrainedModel`	The encoder-decoder model to decompose.
`config`	`ExporterConfig`	The export configuration for the full model.

Returns: Dict[str, Tuple[PreTrainedModel, ExporterConfig]] -- A dictionary with keys:

"encoder_model" -- Tuple of (encoder sub-model, encoder export config)
"decoder_model" -- Tuple of (full model, decoder export config without past inputs)
"decoder_with_past_model" (if config.use_past is True) -- Tuple of (full model, decoder export config with past inputs)

Internal logic:

Calls _get_submodels_for_export_encoder_decoder to extract sub-models using model.get_encoder()
Creates encoder config via config.with_behavior("encoder")
Creates decoder config via config.with_behavior("decoder", use_past=..., use_past_in_inputs=False)
Optionally creates decoder-with-past config via config.with_behavior("decoder", use_past=True, use_past_in_inputs=True)

`get_decoder_models_for_export`

Source location: optimum/exporters/utils.py, lines 224-258

Purpose: Prepares a decoder-only model (e.g., GPT-2, LLaMA) for export with appropriate KV-cache handling.

Signature:

def get_decoder_models_for_export(
    model: "PreTrainedModel",
    config: "ExporterConfig",
) -> Dict[str, Tuple["PreTrainedModel", "ExporterConfig"]]:

Parameters:

Parameter	Type	Description
`model`	`PreTrainedModel`	The decoder-only model to prepare.
`config`	`ExporterConfig`	The export configuration for the model.

Returns: Dict[str, Tuple[PreTrainedModel, ExporterConfig]] -- A dictionary with key:

"model" -- Tuple of (model, export config with use_past and use_past_in_inputs set appropriately)

Internal logic:

Calls _get_submodels_for_export_decoder to extract the model
Creates a new ExporterConfig instance from the same config class, with use_past_in_inputs=config.use_past

`get_diffusion_models_for_export`

Source location: optimum/exporters/utils.py, lines 261-362

Purpose: Decomposes a diffusion pipeline (e.g., Stable Diffusion, SDXL, SD3) into its constituent sub-models, each paired with its own export configuration.

Signature:

def get_diffusion_models_for_export(
    pipeline: "DiffusionPipeline",
    int_dtype: str = "int64",
    float_dtype: str = "fp32",
    exporter: str = "onnx",
) -> Dict[str, Tuple[Union["PreTrainedModel", "ModelMixin"], "ExporterConfig"]]:

Parameters:

Parameter	Type	Default	Description
`pipeline`	`DiffusionPipeline`	required	The diffusion pipeline to decompose.
`int_dtype`	`str`	`"int64"`	Data type for integer tensors (`"int64"`, `"int32"`, `"int8"`).
`float_dtype`	`str`	`"fp32"`	Data type for float tensors (`"fp32"`, `"fp16"`, `"bf16"`).
`exporter`	`str`	`"onnx"`	The export backend to target.

Returns: Dict[str, Tuple[Union[PreTrainedModel, ModelMixin], ExporterConfig]] -- A dictionary with keys for each pipeline component (varies by pipeline type):

Component Key	Task Used for Config Lookup	Description
`"text_encoder"`	`"feature-extraction"`	CLIP or other text encoder
`"text_encoder_2"`	`"feature-extraction"`	Second text encoder (SDXL)
`"text_encoder_3"`	`"feature-extraction"`	Third text encoder (SD3)
`"unet"`	`"semantic-segmentation"`	Denoising UNet
`"transformer"`	`"semantic-segmentation"`	Denoising transformer (DiT)
`"vae_encoder"`	`"semantic-segmentation"`	VAE encoder (model_type: `"vae-encoder"`)
`"vae_decoder"`	`"semantic-segmentation"`	VAE decoder (model_type: `"vae-decoder"`)

Internal logic:

Calls _get_submodels_for_export_diffusion to extract raw sub-models from the pipeline
For each component, looks up the appropriate ExporterConfig via TasksManager.get_exporter_config_constructor
Constructs the config with the specified int_dtype and float_dtype
Pairs each sub-model with its config in the output dictionary

Import

from optimum.exporters.utils import (
    get_encoder_decoder_models_for_export,
    get_decoder_models_for_export,
    get_diffusion_models_for_export,
)

Input/Output Summary

Function	Input	Output
`get_encoder_decoder_models_for_export`	Model + ExporterConfig	Dict of `(sub-model, sub-config)` tuples keyed by component name
`get_decoder_models_for_export`	Model + ExporterConfig	Dict with `"model"` key mapping to `(model, config)` tuple
`get_diffusion_models_for_export`	DiffusionPipeline + dtype/exporter options	Dict of `(sub-model, sub-config)` tuples for each pipeline component

Usage Example

from transformers import AutoModelForSeq2SeqLM
from optimum.exporters import TasksManager
from optimum.exporters.utils import get_encoder_decoder_models_for_export

# Load an encoder-decoder model
model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")

# Get the export config
config_constructor = TasksManager.get_exporter_config_constructor(
    exporter="onnx",
    model=model,
    task="text2text-generation",
    library_name="transformers",
)
export_config = config_constructor(model.config)

# Decompose into exportable components
models_for_export = get_encoder_decoder_models_for_export(model, export_config)
# Returns: {"encoder_model": (encoder, enc_config),
#           "decoder_model": (model, dec_config),
#           "decoder_with_past_model": (model, dec_past_config)}

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment