layerwise_export#

Write each decoder layer’s quantized checkpoint shard as soon as it is calibrated.

Classes

LayerwiseExporter

Writes one decoder layer's quantized shard per call, then the tail and index.

Functions

`assert_formats_supported`	Raise unless every format in `module` can be reproduced per layer.
`assert_layerwise_export_supported`	Raise unless per-layer export is valid for this model.
`layer_shard_name`	Shard filename for one decoder layer, keyed by index so a re-export overwrites.

LAYERWISE_EXPORTER_ATTR = '_layerwise_exporter'#: Set on the model handed to mtq.quantize, so calibration and the export that follows it reach the same exporter.

class LayerwiseExporter#

Bases: object

Writes one decoder layer’s quantized shard per call, then the tail and index.

Built before calibration, driven per layer, finalized after the last. finalize indexes the shards on disk, so an earlier run’s layers are picked up as they are.

__init__(model, export_dir, dtype=None)#

Name the model the checkpoint describes and where it goes.

Nothing is inspected: the caller builds this before mtq.quantize, when there is no quantizer yet to validate or read a config from. bind() does that.

Parameters:

model (Module)
export_dir (Path | str)
dtype (dtype | None)

Return type:

None

announce(module)#

Publish this exporter on module for a later pass to pick up.

Calibration and export are handed different models – a VLM calibrates its language model but exports the whole thing – so each end is told separately.

Parameters:: module (Module)
Return type:: None

assert_no_orphan_shards()#

Refuse to redo work when shards exist but no usable resume record does.

Return type:: None

assert_shards_present(upto)#

Require shards for layers [0, upto), which a resume intends to skip.

Otherwise a mismatched checkpoint/export pair only surfaces after a full run.

Parameters:: upto (int)
Return type:: None

bind(calibrated_layers)#

Validate the model and snapshot what the tail pass needs.

Called from calibration, after quantizer insertion and before any layer is converted – the only window where both hold.

Parameters:: calibrated_layers (list[Module])
Return type:: None

completed_layers()#

How many leading layers have a shard. Contiguous: a gap means the rest never ran.

Return type:: int

property export_dir: Path#: Where the shards go. The exporter owns this, not the caller’s config.

export_layer(layer_idx, layer_module, layer_inputs=None)#

Pack one calibrated layer into its shard, converting it in place.

layer_inputs are the layer’s cached calibration activations, replayed once so a fusing format can rediscover which modules share an input; omit them only when nothing fuses.

Parameters:

layer_idx (int)
layer_module (Module)
layer_inputs (list | None)

Return type:

None

finalize(extra_state_dict=None)#

Export the tail, write the config artifacts, and index all shards.

extra_state_dict carries tensors with no slot in model.state_dict() – MTP weights, whichever convention the checkpoint uses. They are already in export form, so only the hub-name reversal applies, and they win on a name clash exactly as they do in export_hf_checkpoint.

Parameters:: extra_state_dict (dict[str, Tensor] | None)
Return type:: dict

assert_formats_supported(module, scope)#

Raise unless every format in module can be reproduced per layer.

Called before calibration to fail early, and again per exported layer – AWQ and SVDQuant only become visible once the calibrator registers their discriminators.

Parameters:

module (Module)
scope (str)

Return type:

None

assert_layerwise_export_supported(model)#

Raise unless per-layer export is valid for this model.

Parameters:: model (Module)
Return type:: None

layer_shard_name(layer_idx)#

Shard filename for one decoder layer, keyed by index so a re-export overwrites.

Parameters:: layer_idx (int)
Return type:: str