quantized_weight

Non-mutating export helpers for quantized checkpoint weights.

Classes

QuantizedWeightExport

Packed checkpoint tensors for one quantized weight.

QuantizedWeightExportState

Scalar ModelOpt state for one logical floating-point weight.

Functions

build_hf_quantization_config

Build the canonical ModelOpt HF config from canonical layer formats.

capture_quantized_weight_export_state

Capture scalar quantized-weight export state without changing the module.

export_quantized_weight

Pack one canonical weight without mutating it or its quantizers.
class QuantizedWeightExport

Bases: object

Packed checkpoint tensors for one quantized weight.

__init__(weight, weight_scale, weight_scale_2, input_scale=None)
Parameters:

  • weight (Tensor)

  • weight_scale (Tensor)

  • weight_scale_2 (Tensor)

  • input_scale (Tensor | None)

Return type:

None

input_scale: Tensor | None = None
named_tensors(weight_name='weight')

Return checkpoint tensors using ModelOpt’s canonical relative names.

Parameters:

weight_name (str)

Return type:

OrderedDict[str, Tensor]

weight: Tensor
weight_scale: Tensor
weight_scale_2: Tensor
class QuantizedWeightExportState

Bases: object

Scalar ModelOpt state for one logical floating-point weight.

__init__(quantization_format, block_size, weight_amax, input_amax=None)
Parameters:

  • quantization_format (str)

  • block_size (int)

  • weight_amax (Tensor)

  • input_amax (Tensor | None)

Return type:

None

block_size: int
input_amax: Tensor | None = None
quantization_format: str
weight_amax: Tensor
build_hf_quantization_config(layer_states)

Build the canonical ModelOpt HF config from canonical layer formats.

Parameters:

layer_states (Mapping[str, QuantizedWeightExportState | None] | Iterable[tuple[str, QuantizedWeightExportState | None]])

Return type:

dict

capture_quantized_weight_export_state(module, weight_name='weight', *, weight_quantizer=None, input_quantizer=None)

Capture scalar quantized-weight export state without changing the module.

Parameters:

  • module (Module)

  • weight_name (str)

  • weight_quantizer (object | None)

  • input_quantizer (object | None)

Return type:

QuantizedWeightExportState

export_quantized_weight(weight, state, *, dtype)

Pack one canonical weight without mutating it or its quantizers.

Parameters:
Return type:

QuantizedWeightExport