hf_export_prep

Model-level preparation shared by every unified HF exporter.

Everything here runs on the whole model before any weight is packed: dtype resolution, MoE input-quantizer preparation, resmoothing and shared-input fusion, quant-config adjustments, and the transformers patches needed while writing artifacts.

It imports only leaf helpers (layer_utils, model_config, model_utils, quant_utils, registry, diffusers_utils) and never an exporter, which is what lets all three exporters import it without a cycle.

Functions

collect_shared_input_modules	Collect modules that share the same input using forward hooks.
requantize_resmooth_fused_llm_layers	Group modules that take the same input and register shared parameters in module.

collect_shared_input_modules(model, dummy_forward_fn, collect_layernorms=False)

Collect modules that share the same input using forward hooks.

This is a common helper for both LLM and diffusion model fusion.

Parameters:

• model (Module) – The model to analyze.

• dummy_forward_fn (Callable[[], None]) – A callable that runs a dummy forward pass on the model. Should be a function that takes no arguments.

• collect_layernorms (bool) – If True, also collect layernorm output mappings (for AWQ).

Returns:

A tuple of (input_to_linear, output_to_layernorm). input_to_linear: Dict mapping input tensor to list of modules sharing that input. output_to_layernorm: Dict mapping layernorm output to the layernorm module (or None).

Return type:

tuple[dict, dict | None]

requantize_resmooth_fused_llm_layers(model)

Group modules that take the same input and register shared parameters in module.

Parameters:

model (Module)