quant_aware_conversion

Quantization-aware reverse weight conversion for unified HF export.

Background

transformers may apply a conversion_mapping when loading a model, so the in-memory parameter names differ from the original model-hub checkpoint (e.g. fused mlp.gate_up_proj, renamed MoE leaves, reordered model/language_model prefix). On save, transformers reverses this via revert_weight_conversion so the on-disk names match the hub checkpoint again.

ModelOpt’s unified export disables that reverse (it raises RuntimeError on 0-d scalar scale tensors such as weight_scale_2/input_scale), so a quantized export emits the in-memory (post-conversion) names — violating the unified checkpoint contract that names stay aligned with the original hub checkpoint.

This module performs the reverse in a quantization-aware way: it carries each weight’s companion scale tensors (weight_scale, weight_scale_2, input_scale, weight_scale_inv, bias) through the rename and un-fuse operations.

Scope

Two reverse primitives cover the conversion_mapping cases:

• Rename — a key-level string substitution. Because a quantized linear stores every tensor under <module>.<leaf>, renaming the module substring rewrites the weight and all its scale siblings together with no tensor manipulation.

• Split — un-fuse an output-dim concatenation (e.g. dense gate_up_proj -> gate_proj + up_proj). weight/weight_scale/weight_scale_inv/ bias are chunked along the fused (output) dim; 0-d scalar weight_scale_2/ input_scale are duplicated to each part (they are per-tensor and shared).

MoE experts need only Rename: ModelOpt’s export already expands the fused, stacked in-memory experts (experts.gate_up_proj of shape [E, 2F, H]) into per-expert 2-D linears (experts.<i>.gate_proj / up_proj / down_proj) before save, so the reverse just maps those per-expert leaf names back to the hub leaves (e.g. gate_proj -> w1, up_proj -> w3, down_proj -> w2).

Reverse rules are derived from the model’s conversion mapping via transformers’ reverse_transform(). Any op shape not covered raises QuantConversionUnsupportedError so the caller falls back to the legacy (in-memory-name) behavior rather than emit a silently-wrong checkpoint.

Classes

RenameRule	Reverse of a WeightRenaming: re.sub(pattern, repl, key) on every key.
SplitRule	Reverse of an output-dim Concatenate: un-fuse one module into parts.

Functions

apply_reverse_rules	Apply quant-aware reverse conversion: splits first, then renames.
build_reverse_name_mapper	Build a str -> str mapper that applies the quant-aware reverse rename rules.
revert_quant_config_names	Revert exclude_modules / quantized_layers keys to hub names, in place.
revert_weight_conversion_quant_aware	Reverse a transformers conversion_mapping on a quantized state dict.

exception QuantConversionUnsupportedError

Bases: Exception

Raised when a conversion op cannot be reversed quant-aware (caller falls back).

class RenameRule

Bases: object

Reverse of a WeightRenaming: re.sub(pattern, repl, key) on every key.

__init__(pattern, repl)

Parameters:

• pattern (str)

• repl (str)

Return type:

None

pattern: str

repl: str

class SplitRule

Bases: object

Reverse of an output-dim Concatenate: un-fuse one module into parts.

Parameters:

• fused_suffix – module suffix of the fused tensor, e.g. ".gate_up_proj".

• part_suffixes – ordered replacements, e.g. (".gate_proj", ".up_proj").

• dim – the fused (output) dim along which weight/weight_scale/bias are chunked. NVFP4 weight is [out, in//2] and weight_scale is [out, in//block] so the output dim is 0 for both.

__init__(fused_suffix, part_suffixes, dim=0)

Parameters:

• fused_suffix (str)

• part_suffixes (tuple[str, ...])

• dim (int)

Return type:

None

dim: int = 0

fused_suffix: str

part_suffixes: tuple[str, ...]

apply_reverse_rules(state_dict, split_rules, rename_rules)

Apply quant-aware reverse conversion: splits first, then renames.

Splits run on the in-memory (post-conversion) names; renames then map the resulting keys back to the original hub names. Renames are applied in order.

Parameters:

• state_dict (dict[str, Tensor])

• split_rules (list[SplitRule])

• rename_rules (list[RenameRule])

Return type:

dict[str, Tensor]

build_reverse_name_mapper(model)

Build a str -> str mapper that applies the quant-aware reverse rename rules.

The exported weight tensors are reverted to the original hub names by revert_weight_conversion_quant_aware(), but the quantization config’s module references (exclude_modules and, for mixed precision, quantized_layers keys) are built from the in-memory module names and would otherwise stay in the post-conversion namespace – so a deployment loader matching those patterns against the (reverted) hub-named modules finds no match, silently loads an excluded BF16 layer as quantized, and fails. Applying the same rename rules to those name strings keeps them aligned with the weights. Only the rename rules apply (splits act on tensors, not names).

Returns None when no renaming applies. Raises QuantConversionUnsupportedError when the mapping can’t be reversed, so the caller can keep the in-memory names for BOTH weights and config (mutually consistent).

revert_quant_config_names(quantization, mapper)

Revert exclude_modules / quantized_layers keys to hub names, in place.

mapper is the callable from build_reverse_name_mapper() (a no-op when None). Applies to the ModelOpt {"quantization": {...}} sub-dict before it is written / format-converted, so both hf_quant_config.json and the embedded config.json quantization_config inherit the reverted names.

Parameters:

quantization (dict)

Return type:

None

revert_weight_conversion_quant_aware(model, state_dict)

Reverse a transformers conversion_mapping on a quantized state dict.

Builds reverse rules from the model’s conversion mapping and applies them carrying companion scale tensors. Raises QuantConversionUnsupportedError when the mapping uses an op that cannot be reversed quant-aware yet, so the caller can fall back to the legacy behavior.

Parameters:

state_dict (dict[str, Tensor])