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.
`build_legacy_name_mapper`	Hub-name mapper for transformers < 5, or None if the model declares no mapping.
`layer_shard_name`	Shard filename for one decoder layer.
`resolve_export_parent`	Return `(model the checkpoint describes, key prefix for its tensors)`.
`transient_module_state`	Undo the parameter, buffer and child-module changes export makes to `module`.

EXPORT_PARENT_ATTR = '_layerwise_export_parent': Set by the caller on the calibrated submodel to name the model the checkpoint should describe. Multimodal pipelines calibrate the extracted language model, but the exported checkpoint has to describe the whole VLM or its config and tensor names disagree.

class LayerwiseExporter

Bases: object

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

Built before calibration, driven per layer from inside the window calibration opens, finalized after the last. finalize indexes the shards on disk, so layers an earlier run exported 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; anything amax-dependent belongs in finalize().

Parameters:

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

Return type:

None

assert_no_orphan_shards(manifest_present)

Refuse to silently redo work when shards exist but the resume record does not.

The resume point comes from the manifest – the shards cannot supply it. Without one start_layer is 0, assert_shards_present() checks an empty range, and calibration overwrites every finished layer without a word.

Parameters:: manifest_present (bool)
Return type:: None

assert_shards_present(upto)

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

The checkpoint directory knows nothing about what was exported, so a mismatched pair would only surface at finalize() – after a full calibration.

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

completed_layers()

How many leading layers already have a shard on disk.

A contiguous run from layer 0: a gap means the layers after it were never finished, and resuming past one would leave the checkpoint missing them.

Return type:: int

export_layer(layer_idx, layer_module, probe_forward=None)

Pack one calibrated layer into its shard, leaving that layer untouched.

probe_forward runs the layer on real activations so a fusing format can rediscover which modules share an input; only omit it when nothing fuses.

Parameters:

layer_idx (int)
layer_module (Module)
probe_forward (Callable[[Module], None] | None)

Return type:

None

finalize(extra_state_dict=None)

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

Leaves export_dir a complete checkpoint; no export_hf_checkpoint() needed.

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 twice, because AWQ and SVDQuant only become visible once the calibrator has registered _pre_quant_scale / svdquant_lora_a: before calibration to fail early, and on each exported layer, which is the authority.

Parameters:

module (Module)
scope (str)

Return type:

None

assert_layerwise_export_supported(model)

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

Structural cases only; each would otherwise differ from a whole-model export without failing. Formats are settled by assert_formats_supported().

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

build_legacy_name_mapper(model)

Hub-name mapper for transformers < 5, or None if the model declares no mapping.

save_pretrained is what reverses _checkpoint_conversion_mapping on 4.x. Per-layer export writes shards with save_file and never passes through it, so without this it would publish in-memory names where the whole-model path publishes hub names.

Parameters:: model (Module)

layer_shard_name(layer_idx)

Shard filename for one decoder layer.

Keyed by index, not a counter, so re-exporting overwrites rather than leaving a stale shard for the index to find.

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

resolve_export_parent(model)

Return (model the checkpoint describes, key prefix for its tensors).

Without EXPORT_PARENT_ATTR this is (model, ""). With it, the prefix is the dotted path from parent down to the calibrated submodel, found by identity so a module that merely looks like the language model cannot be mistaken for it.

Parameters:: model (Module)
Return type:: tuple[Module, str]

transient_module_state(module)

Undo the parameter, buffer and child-module changes export makes to module.

A resident model has no materialization window to discard what export leaves behind. Restoring the dicts covers rebinding; buffers are additionally clone-restored because scale fusion writes amax in place. Parameters are not – that would copy the layer. Attributes outside those three dicts are not restored (e.g. the layer-relative name that collect_shared_input_modules sets), so handlers must keep their state there.

Parameters:: module (Module)