`bridge.models.conversion.modelopt_utils`#

Module Contents#

Classes#

`QuantMeta`	ModelOpt quantization metadata for one Megatron parameter.
`_Nvfp4InputQuantizerView`	Minimal quantizer interface for ModelOpt activation-scale export.
`_Nvfp4WeightQuantizerView`	CPU-only quantizer state needed by ModelOpt’s canonical scale export.
`_LocalExpertMappingMixin`	Use expert TP while leaving the expert dimension local.
`_LocalExpertColumnMapping`
`_LocalExpertRowMapping`
`_LocalExpertReplicatedMapping`
`_LocalExpertGatedMLPMapping`
`_ModelOptFusedExpertMapping`	Build one E/EP-local expert batch without a BF16 EP gather.
`_ModelOptFusedGatedExpertMapping`	Fuse gate/up locally without a BF16 EP gather.

Functions#

`_iter_quant_ignore_name_candidates`
`matches_quant_ignore_pattern`	Return whether a parameter name matches any ModelOpt ignore pattern.
`is_modelopt_quantizable_weight_name`	Return whether an exported HF tensor name should be ModelOpt-quantized.
`_is_same_tensor`
`_iter_modelopt_weight_quantizers`
`_is_enabled_quantizer`
`find_modelopt_weight_quantizer_and_module`	Find the enabled weight quantizer and owning module for `param_weight`.
`_with_quant_meta_tensors`
`_clone_cpu`
`_clone_positive_cpu`
`_get_modelopt_weight_amax`	Read the raw global amax without deriving a scale on its CUDA device.
`_compute_modelopt_weight_scale_2_cpu`	Run ModelOpt’s quantizer-specific global-scale derivation on CPU.
`_get_modelopt_tp_process_group`	Return the raw TP process group attached to a ModelOpt quantized module.
`_max_reduce_modelopt_tp_scalar`	MAX-reduce one per-tensor quantization value over its owning TP group.
`_slice_optional_quant_tensor`
`_slice_gated_quant_meta`	Slice fused `[gate; up]` metadata to match a split HF tensor.
`_stack_optional_quant_tensors`
`_stack_grouped_quant_meta`
`_expert_param_template`
`_iter_grouped_quant_meta`	Yield all synced per-expert metadata entries for a grouped export task.
`build_hf_modelopt_quant_metadata`	Map Megatron ModelOpt metadata onto exported Hugging Face names.
`_build_hf_modelopt_pre_ep_quant_metadata`	Build metadata for selected local expert batches before EP gathering.
`collect_modelopt_quant_metadata`	Collect ModelOpt quantization metadata from conversion task modules.
`sync_modelopt_quant_metadata`	Synchronize ModelOpt quantization metadata across a distributed group.
`_reshape_nvfp4_weight_scale_2_for_compute`
`compute_nvfp4_weight_scale`	Compute the NVFP4 per-block weight scale tensor for ModelOpt export.
`compute_nvfp4_input_scale`	Compute a static NVFP4 activation scale from synchronized input amax.
`_nvfp4_export_names`
`_format_nvfp4_weight_scale_2_for_export`
`_format_nvfp4_input_scale_for_export`	Shape static activation scales for dense and fused-MoE vLLM loaders.
`quantize_nvfp4_weight`	Yield NVFP4 quantized weight tensors and associated scale tensors.
`get_modelopt_quant_exporter`	Return the ModelOpt quantization format and exporter for a quantization mode.
`_grouped_expert_projection_name`	Remove a resolved expert index and terminal weight suffix from an HF name.
`_fuse_grouped_projection_names`	Merge two projection leaves while preserving their common underscore suffix.
`_modelopt_pre_ep_mapping`	Build a fused local-expert mapping for ModelOpt expert projections.
`_stage_tensor_for_collective`	Move a CPU tensor to CUDA only when its collective backend requires it.
`_compose_export_hooks`
`build_modelopt_export_plan`	Prepare mapped conversion tasks and hooks for a ModelOpt export.

Data#

`HFExportHook`
`_NVFP4_AMAX_DENOMINATOR`
`_NVFP4_MAXBOUND`
`_FP8_E4M3_MIN`
`_FP8_E4M3_MAX`
`_QUANT_IGNORE_NAME_SUFFIXES`
`_EXPERT_NUMBER_PATTERNS`
`_FUSED_MOE_NVFP4_NAME_MAP`

API#

bridge.models.conversion.modelopt_utils.HFExportHook#: None

bridge.models.conversion.modelopt_utils._NVFP4_AMAX_DENOMINATOR#: None

bridge.models.conversion.modelopt_utils._NVFP4_MAXBOUND#: 6.0

bridge.models.conversion.modelopt_utils._FP8_E4M3_MIN#: None

bridge.models.conversion.modelopt_utils._FP8_E4M3_MAX#: 448.0

bridge.models.conversion.modelopt_utils._QUANT_IGNORE_NAME_SUFFIXES#: (‘.weight’, ‘.weight_scale’, ‘.weight_scale_2’)

bridge.models.conversion.modelopt_utils._EXPERT_NUMBER_PATTERNS#: ()

bridge.models.conversion.modelopt_utils._FUSED_MOE_NVFP4_NAME_MAP#: None

class bridge.models.conversion.modelopt_utils.QuantMeta#

ModelOpt quantization metadata for one Megatron parameter.

qformat: str#: None

block_size: int#: None

weight_amax: torch.Tensor | None#: None

weight_scale_2: torch.Tensor | None#: None

input_amax: torch.Tensor | None#: None

class bridge.models.conversion.modelopt_utils._Nvfp4InputQuantizerView#

Minimal quantizer interface for ModelOpt activation-scale export.

input_amax: torch.Tensor#: None

is_enabled: bool#: True

maxbound: float#: None

export_amax() → torch.Tensor#

class bridge.models.conversion.modelopt_utils._Nvfp4WeightQuantizerView#

CPU-only quantizer state needed by ModelOpt’s canonical scale export.

_amax: torch.Tensor#: None

block_sizes: object | None#: None

bridge.models.conversion.modelopt_utils._iter_quant_ignore_name_candidates(name: str) → Iterator[str]#

bridge.models.conversion.modelopt_utils.matches_quant_ignore_pattern(name: str, patterns: list[str]) → bool#: Return whether a parameter name matches any ModelOpt ignore pattern.

bridge.models.conversion.modelopt_utils.is_modelopt_quantizable_weight_name(name: str) → bool#: Return whether an exported HF tensor name should be ModelOpt-quantized.

bridge.models.conversion.modelopt_utils._is_same_tensor(param_weight: object, weight: object) → bool#

bridge.models.conversion.modelopt_utils._iter_modelopt_weight_quantizers( module: torch.nn.Module, ) → Iterator[tuple[object, object, bool]]#

bridge.models.conversion.modelopt_utils._is_enabled_quantizer(quantizer: object) → bool#

bridge.models.conversion.modelopt_utils.find_modelopt_weight_quantizer_and_module( module: torch.nn.Module, param_weight: object, ) → tuple[object | None, torch.nn.Module | None]#: Find the enabled weight quantizer and owning module for param_weight.

bridge.models.conversion.modelopt_utils._with_quant_meta_tensors( meta: bridge.models.conversion.modelopt_utils.QuantMeta, *, weight_amax: torch.Tensor | None, weight_scale_2: torch.Tensor | None, ) → bridge.models.conversion.modelopt_utils.QuantMeta#

bridge.models.conversion.modelopt_utils._clone_cpu(value: torch.Tensor | None) → torch.Tensor | None#

bridge.models.conversion.modelopt_utils._clone_positive_cpu(value: torch.Tensor | None) → torch.Tensor | None#

bridge.models.conversion.modelopt_utils._get_modelopt_weight_amax( weight_quantizer: object, ) → tuple[torch.Tensor | None, bool]#: Read the raw global amax without deriving a scale on its CUDA device.

bridge.models.conversion.modelopt_utils._compute_modelopt_weight_scale_2_cpu( weight_amax: torch.Tensor | None, weight_quantizer: object, ) → torch.Tensor | None#: Run ModelOpt’s quantizer-specific global-scale derivation on CPU.

bridge.models.conversion.modelopt_utils._get_modelopt_tp_process_group(module: object) → object | None#: Return the raw TP process group attached to a ModelOpt quantized module.

bridge.models.conversion.modelopt_utils._max_reduce_modelopt_tp_scalar( value: torch.Tensor | None, module: object, field_name: str, ) → torch.Tensor | None#: MAX-reduce one per-tensor quantization value over its owning TP group.

bridge.models.conversion.modelopt_utils._slice_optional_quant_tensor( value: torch.Tensor | None, split: slice, leading_dim: int, ) → torch.Tensor | None#

bridge.models.conversion.modelopt_utils._slice_gated_quant_meta( meta: bridge.models.conversion.modelopt_utils.QuantMeta, hf_key: str, ) → bridge.models.conversion.modelopt_utils.QuantMeta#: Slice fused [gate; up] metadata to match a split HF tensor.

bridge.models.conversion.modelopt_utils._stack_optional_quant_tensors( values: list[torch.Tensor | None], *, hf_name: str, field_name: str, ) → torch.Tensor | None#

bridge.models.conversion.modelopt_utils._stack_grouped_quant_meta( hf_name: str, expert_meta: dict[int, bridge.models.conversion.modelopt_utils.QuantMeta], ) → bridge.models.conversion.modelopt_utils.QuantMeta#

bridge.models.conversion.modelopt_utils._expert_param_template(param_name: str) → str | None#

bridge.models.conversion.modelopt_utils._iter_grouped_quant_meta( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, metadata: dict[str, bridge.models.conversion.modelopt_utils.QuantMeta], ) → Iterator[tuple[int, bridge.models.conversion.modelopt_utils.QuantMeta]]#: Yield all synced per-expert metadata entries for a grouped export task.

bridge.models.conversion.modelopt_utils.build_hf_modelopt_quant_metadata( conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask], metadata: dict[str, bridge.models.conversion.modelopt_utils.QuantMeta], ) → dict[str, bridge.models.conversion.modelopt_utils.QuantMeta]#: Map Megatron ModelOpt metadata onto exported Hugging Face names.

bridge.models.conversion.modelopt_utils._build_hf_modelopt_pre_ep_quant_metadata( conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask], metadata: dict[str, bridge.models.conversion.modelopt_utils.QuantMeta], ) → dict[str, bridge.models.conversion.modelopt_utils.QuantMeta]#: Build metadata for selected local expert batches before EP gathering.

bridge.models.conversion.modelopt_utils.collect_modelopt_quant_metadata( conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask], ) → dict[str, bridge.models.conversion.modelopt_utils.QuantMeta]#: Collect ModelOpt quantization metadata from conversion task modules.

bridge.models.conversion.modelopt_utils.sync_modelopt_quant_metadata( metadata: dict[str, bridge.models.conversion.modelopt_utils.QuantMeta], group=None, ) → None#: Synchronize ModelOpt quantization metadata across a distributed group.

bridge.models.conversion.modelopt_utils._reshape_nvfp4_weight_scale_2_for_compute( weight: torch.Tensor, weight_scale_2: torch.Tensor, ) → torch.Tensor#

bridge.models.conversion.modelopt_utils.compute_nvfp4_weight_scale( weight: torch.Tensor, block_size: int, weight_amax: torch.Tensor | None = None, weight_scale_2: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor]#: Compute the NVFP4 per-block weight scale tensor for ModelOpt export.

bridge.models.conversion.modelopt_utils.compute_nvfp4_input_scale( input_amax: torch.Tensor | None, ) → torch.Tensor#: Compute a static NVFP4 activation scale from synchronized input amax.

bridge.models.conversion.modelopt_utils._nvfp4_export_names(name: str) → tuple[str, str, str, str]#

bridge.models.conversion.modelopt_utils._format_nvfp4_weight_scale_2_for_export( source_name: str, weight_name: str, weight: torch.Tensor, weight_scale_2: torch.Tensor, ) → torch.Tensor#

bridge.models.conversion.modelopt_utils._format_nvfp4_input_scale_for_export( source_name: str, weight_name: str, weight: torch.Tensor, input_scale: torch.Tensor, ) → torch.Tensor#: Shape static activation scales for dense and fused-MoE vLLM loaders.

bridge.models.conversion.modelopt_utils.quantize_nvfp4_weight( name: str, weight: torch.Tensor, meta: bridge.models.conversion.modelopt_utils.QuantMeta, ) → Iterator[tuple[str, torch.Tensor]]#: Yield NVFP4 quantized weight tensors and associated scale tensors.

bridge.models.conversion.modelopt_utils.get_modelopt_quant_exporter(quant_mode: str)#: Return the ModelOpt quantization format and exporter for a quantization mode.

bridge.models.conversion.modelopt_utils._grouped_expert_projection_name( hf_name: str, ) → tuple[str, int] | None#: Remove a resolved expert index and terminal weight suffix from an HF name.

bridge.models.conversion.modelopt_utils._fuse_grouped_projection_names( gate_name: str, up_name: str, ) → str | None#: Merge two projection leaves while preserving their common underscore suffix.

class bridge.models.conversion.modelopt_utils._LocalExpertMappingMixin#

Use expert TP while leaving the expert dimension local.

property tp_group#

property is_expert: bool#

class bridge.models.conversion.modelopt_utils._LocalExpertColumnMapping#: Bases: bridge.models.conversion.modelopt_utils._LocalExpertMappingMixin, megatron.bridge.models.conversion.param_mapping.ColumnParallelMapping

class bridge.models.conversion.modelopt_utils._LocalExpertRowMapping#: Bases: bridge.models.conversion.modelopt_utils._LocalExpertMappingMixin, megatron.bridge.models.conversion.param_mapping.RowParallelMapping

class bridge.models.conversion.modelopt_utils._LocalExpertReplicatedMapping#: Bases: bridge.models.conversion.modelopt_utils._LocalExpertMappingMixin, megatron.bridge.models.conversion.param_mapping.ReplicatedMapping

class bridge.models.conversion.modelopt_utils._LocalExpertGatedMLPMapping#: Bases: bridge.models.conversion.modelopt_utils._LocalExpertMappingMixin, megatron.bridge.models.conversion.param_mapping.GatedMLPMapping

class bridge.models.conversion.modelopt_utils._ModelOptFusedExpertMapping#

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

Build one E/EP-local expert batch without a BF16 EP gather.

is_grouped_export#: False

is_modelopt_pre_ep_export#: True

_get_or_create_mapping(parallelism_type: str)#

class bridge.models.conversion.modelopt_utils._ModelOptFusedGatedExpertMapping( megatron_param: str, hf_param: str, permute_dims: tuple[int, ...] | None = None, transpose_on_export: bool = False, )#

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

Fuse gate/up locally without a BF16 EP gather.

Initialization

is_grouped_export#: False

is_modelopt_pre_ep_export#: True

bridge.models.conversion.modelopt_utils._modelopt_pre_ep_mapping( mapping: Any, pg_collection: Any = None, ) → tuple[Any, tuple[str, ...]] | None#: Build a fused local-expert mapping for ModelOpt expert projections.

bridge.models.conversion.modelopt_utils._stage_tensor_for_collective( tensor: torch.Tensor, group: Any, ) → torch.Tensor#: Move a CPU tensor to CUDA only when its collective backend requires it.

bridge.models.conversion.modelopt_utils._compose_export_hooks( exporter: bridge.models.conversion.modelopt_utils.HFExportHook, finalizer: bridge.models.conversion.modelopt_utils.HFExportHook | None, ) → bridge.models.conversion.modelopt_utils.HFExportHook#

bridge.models.conversion.modelopt_utils.build_modelopt_export_plan( conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask], *, model: list[torch.nn.Module], bridge: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge, quant_mode: str, ignore_patterns: list[str], ) → list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask]#: Prepare mapped conversion tasks and hooks for a ModelOpt export.

bridge.models.conversion.modelopt_utils#