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.
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)#

Validate support and capture model-level state.

Runs before calibration, so nothing amax-dependent exists yet.

Parameters:

  • model (Module)

  • export_dir (Path | str)

  • dtype (dtype | None)

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

completed_layers()#

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

Return type:

int

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()#

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

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

Return type:

dict

MTP_EXTRA_STATE_ATTR = '_mtp_extra_state_dict'#

Set by the caller on the model, holding MTP tensors that have no slot in state_dict(). finalize() runs inside calibration, so the caller cannot pass them as an argument; this follows the _mtp_layer_prefixes convention already used to hand MTP information across the same boundary.

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