bridge.models.conversion.quant_mapping#

Module Contents#

Classes#

AmaxMapping

Amax mapping for quantization.

AmaxFanoutMapping

Replicated amax mapping that fans out one Megatron amax to multiple HF targets.

Functions#

convert_to_amax_map

Convert weight mappings to amax mappings for quantization.

API#

class bridge.models.conversion.quant_mapping.AmaxMapping(megatron_param: str, hf_param: str)#

Bases: megatron.bridge.models.conversion.param_mapping.ReplicatedMapping

Amax mapping for quantization.

Initialization

Initialize the Amax mapping.

class bridge.models.conversion.quant_mapping.AmaxFanoutMapping(megatron_param: str, hf_params: list[str])#

Bases: bridge.models.conversion.quant_mapping.AmaxMapping

Replicated amax mapping that fans out one Megatron amax to multiple HF targets.

Used for QKV and gate/up where the amax values are shared but need to be written/read under multiple HF parameter names.

Initialization

Initialize the Amax mapping.

megatron_to_hf(megatron_weights, megatron_module)#
resolve(captures: tuple[str, ...])#

Resolve wildcards for both megatron_param and all HF targets.

bridge.models.conversion.quant_mapping.convert_to_amax_map(
mappings: list[megatron.bridge.models.conversion.param_mapping.MegatronParamMapping],
mapped_name='.weight_quantizer._amax',
) → list[megatron.bridge.models.conversion.param_mapping.MegatronParamMapping]#

Convert weight mappings to amax mappings for quantization.

This function converts parameter mappings for weights to their corresponding amax (absolute maximum) parameter mappings used in quantization. For example:

  • “layer.weight” -> “layer.weight_quantizer._amax”

Parameters:

mappings – List of MegatronParamMapping objects for weight parameters

Returns:

List of new MegatronParamMapping objects for amax parameters

.. note::

Only mappings with parameter names ending in ‘.weight’ are converted. Other mappings are ignored.