export_utils#

Utilities for exporting LLM models to ONNX.

Classes

ModelLoader

A class to handle HuggingFace model loading and configuration.

RopeType

Rope type enum.

WrapperModelForCausalLM

Wrapper Model to ensure all models have the same I/O.

Functions

llm_to_onnx

Export the WrapperModelForCausalLM to ONNX with fixed I/O names and shape definitions and save to output_dir.

torch_to_onnx

Export the model to ONNX.
class ModelLoader#

Bases: object

A class to handle HuggingFace model loading and configuration.

__init__(hf_model_path, config_path)#

Initialize the ModelLoader.

Parameters:

  • hf_model_path (str)

  • config_path (str)

get_model_type()#

Get model type from config file.

get_rope_type()#

Get rope type.

load_model(trust_remote_code=False)#

Load HuggingFace model based on model type.

Parameters:

trust_remote_code (bool)

Return type:

AutoModelForCausalLM

class RopeType#

Bases: Enum

Rope type enum.

K_MROPE = 3#
K_NONE = 0#
K_ROPE_ROTATE_GPTJ = 1#
K_ROPE_ROTATE_NEOX = 2#
class WrapperModelForCausalLM#

Bases: Module

Wrapper Model to ensure all models have the same I/O.

__init__(model)#

Initialize the WrapperModelForCausalLM.

forward(input_ids, past_key_values)#

Forward pass.

Parameters:

  • input_ids (Tensor | None)

  • past_key_values (tuple)

llm_to_onnx(model, output_dir, extra_inputs={}, extra_dyn_axes={})#

Export the WrapperModelForCausalLM to ONNX with fixed I/O names and shape definitions and save to output_dir.

Parameters:

  • model – torch.Module

  • output_dir – str, the output_dir of the original ONNX.

  • extra_inputs – dict, append additional inputs after kv_cache. Usually for VL models

  • extra_dyn_axes – dict. Usually for VL models

torch_to_onnx(model, inputs, onnx_dir, onnx_name, input_names, output_names, dynamic_axes)#

Export the model to ONNX.