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_layerwise_export_supported`	Raise `NotImplementedError` unless per-layer export is valid for this model.
`layer_shard_name`	Shard filename for one decoder layer.
`transient_module_state`	Undo everything export does to `module`, so calibration can continue through it.

class LayerwiseExporter

Bases: object

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

Constructed before calibration begins, driven once per layer from inside the window calibration already opens, and finalized after the last one:

exporter = LayerwiseExporter(model, export_dir)
...
with persistent_materialization(layer, writeback=False):
    calib_func(layer, ...)
    exporter.export_layer(layer_idx, layer)
...
quant_config = exporter.finalize(extra_state_dict=mtp_state_dict)

finalize() rebuilds the index from the shards present on disk, so layers exported by an earlier run that this one skipped are picked up without being re-exported.

__init__(model, export_dir, dtype=None, is_modelopt_qlora=False)

Validate support and capture model-level state, before calibration runs.

Only quantizer configuration is read here, which mtq.quantize fixes when it swaps modules; anything amax-dependent belongs in finalize().

Parameters:

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

Return type:

None

assert_shards_present(upto)

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

Calibration resumes from its own checkpoint directory, which knows nothing about what was exported. If the two were produced by different runs, the skipped layers have no shards and the gap would only surface at finalize(), after the whole calibration had run. Fail before any of that work instead.

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

export_layer(layer_idx, layer_module)

Pack one calibrated layer into its shard, leaving the layer itself untouched.

Parameters:

layer_idx (int)
layer_module (Module)

Return type:

None

finalize(extra_state_dict=None)

Export the tail, write every config artifact, and index all shards.

Leaves export_dir a complete, loadable checkpoint, so no separate export_hf_checkpoint() call is needed. Returns the quant config.

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

assert_layerwise_export_supported(model)

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

Each case would otherwise produce a checkpoint differing from a whole-model export without failing, so all are rejected before the first shard is written.

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

layer_shard_name(layer_idx)

Shard filename for one decoder layer.

Derived from the index rather than a running counter so that re-exporting a layer overwrites its shard instead of leaving a stale copy behind for the index to pick up.

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

transient_module_state(module)

Undo everything export does to module, so calibration can continue through it.

Export is destructive – packed weights, new scale buffers, grafted per-expert submodules. An offloaded model discards that when its materialization window closes; a resident one has no window, and calibration still has every later layer to run.

Restoring the dicts suffices, and costs references rather than a deep copy, because export rebinds them instead of mutating tensors in place.

Parameters:: module (Module)