registry#

Registries dispatching per-module logic for the unified HF export path.

This mirrors the registration-and-dispatch idiom of QuantModuleRegistry, but not its mechanism: quantization registers replacement classes and converts modules in place, whereas export registers functions that emit compressed weights and scale buffers for a module without changing its class.

Preparation and export use separate registries because they have independent matching precedence. Registering a handler for a new module type replaces what previously required editing if/elif chains inside unified_export_hf.py.

Classes

ExportContext

Shared state for a single export invocation, passed to every handler call.
class ExportContext#

Bases: object

Shared state for a single export invocation, passed to every handler call.

Tied-weight dedup is not a handler concern: the driver builds one name-based TiedWeightMap and feeds it to sync_tied_input_amax and postprocess_state_dict directly. Both dense and fused-MoE tied weights are packed independently and their duplicate keys are dropped by name there, so the context carries no tied-weight map (handlers never consulted it).

__init__(model, dtype, is_modelopt_qlora=False, model_type=None, keep_fused_experts=False, module_names=None, hold_fsdp2_windows=False, fsdp2_window=None)#
Parameters:

  • model (Module)

  • dtype (dtype)

  • is_modelopt_qlora (bool)

  • model_type (str | None)

  • keep_fused_experts (bool)

  • module_names (ModuleNames | None)

  • hold_fsdp2_windows (bool)

  • fsdp2_window (_HeldWindow | None)

Return type:

None

dtype: dtype#
fsdp2_window: _HeldWindow | None = None#

The FSDP2 unit’s weight-update window the handlers hold open across its modules.

hold_fsdp2_windows: bool = False#

set only inside held_fsdp2_weight_windows(), which closes the last window. Elsewhere each module is packed in a window of its own.

Type:

Whether the handlers hold an FSDP2 unit’s weight-update window open across its modules

is_modelopt_qlora: bool = False#
keep_fused_experts: bool = False#

Marks the distributed no-gather export, whose handlers pack while the model is still sharded: an FSDP2-owned weight is packed inside its unit’s weight-update window, and fused MoE experts are folded in place and kept fused and sharded rather than split into whole per-expert tensors, which split_fused_experts_state_dict() does later.

model: Module#
model_type: str | None = None#

The model’s HF model type (model.config.model_type), used to resolve the model’s spec in modelopt.torch.models. None means unknown: spec lookups then fail loudly instead of guessing.

module_names: ModuleNames | None = None#

model’s name maps, built once per export by the first handler that looks a module up (module_name_maps()).