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.

The tied-weight dedup caches must be scoped to one export invocation: a process-global cache would carry stale entries whose data_ptr keys can be recycled by PyTorch’s allocator across exports, causing silent false-positive aliasing. tied_cache (int keys) holds dense Linear / per-expert wrapper dedup; moe_tied_cache (tuple keys) holds MoE fused-experts module dedup.

__init__(model, dtype, is_modelopt_qlora=False, tied_cache=<factory>, moe_tied_cache=<factory>)
Parameters:

  • model (Module)

  • dtype (dtype)

  • is_modelopt_qlora (bool)

  • tied_cache (dict[int, Module])

  • moe_tied_cache (dict[tuple[int, int], Module])

Return type:

None

dtype: dtype
is_modelopt_qlora: bool = False
model: Module
moe_tied_cache: dict[tuple[int, int], Module]
tied_cache: dict[int, Module]