`bridge.models.conversion.quant_bridge`#

Module Contents#

Classes#

`_NativeMXFP8Storage`
`_NativeMXFP8StorageSpec`
`_LocalMXFP8ParamSpec`
`_NativeMXFP8TaskPlan`
`MegatronQuantizationBridge`	Mixin providing quantization-aware utilities for Megatron model bridges.

Functions#

`_validate_native_mxfp8_storage`	Validate native MXFP8 backing metadata without projecting its payload.
`_inspect_native_mxfp8_storage`	Return an allocation-free description of native MXFP8 storage.
`_extract_native_mxfp8_storage`	Return validated, unconverted native MXFP8 storage views.
`_make_meta_native_mxfp8_storage`	Create storage-free tensors for structural mapping validation.
`_lookup_grouped_expert_mapping`	Resolve a grouped-member mapping without relying on a parameter-name suffix.
`_supports_native_grouped_mxfp8`
`_supports_native_mxfp8_mapping`	Return whether a mapping class explicitly implements the native contract.
`_uses_dtensor_or_fsdp`	Return whether local materialization would inspect distributed storage.
`_validate_local_native_mxfp8_param`	Validate one final mapping result before it can leave the Bridge.
`_validate_local_native_mxfp8_params`	Materialize and validate all mapping results for one source parameter.
`_local_native_mxfp8_param_specs`	Describe validated mapping results without retaining their tensors.
`_describe_local_native_mxfp8_params`	Validate mapping results and retain only their structural metadata.
`_validate_local_native_mxfp8_params_against_specs`	Validate real mapping results against their preflight structure.

Data#

`MegatronModel`
`HFPreTrained`

API#

bridge.models.conversion.quant_bridge.MegatronModel#: ‘TypeVar(…)’

bridge.models.conversion.quant_bridge.HFPreTrained#: ‘TypeVar(…)’

class bridge.models.conversion.quant_bridge._NativeMXFP8Storage#

weight: torch.Tensor#: None

weight_scale: torch.Tensor#: None

class bridge.models.conversion.quant_bridge._NativeMXFP8StorageSpec#

weight_shape: torch.Size#: None

weight_scale_shape: torch.Size#: None

class bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec#

name: str#: None

weight_shape: torch.Size#: None

weight_dtype: torch.dtype#: None

weight_scale_shape: torch.Size#: None

weight_scale_dtype: torch.dtype#: None

global_weight_shape: torch.Size#: None

shard_group: str#: None

shard_dim: Optional[int]#: None

class bridge.models.conversion.quant_bridge._NativeMXFP8TaskPlan#

task: Any#: None

grouped: bool#: None

output_specs: tuple[bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec, ...]#: None

bridge.models.conversion.quant_bridge._validate_native_mxfp8_storage( param: torch.Tensor, global_param_name: str, ) → tuple[bridge.models.conversion.quant_bridge._NativeMXFP8StorageSpec, torch.Tensor, torch.Tensor]#: Validate native MXFP8 backing metadata without projecting its payload.

bridge.models.conversion.quant_bridge._inspect_native_mxfp8_storage( param: torch.Tensor, global_param_name: str, ) → bridge.models.conversion.quant_bridge._NativeMXFP8StorageSpec#: Return an allocation-free description of native MXFP8 storage.

bridge.models.conversion.quant_bridge._extract_native_mxfp8_storage( param: torch.Tensor, global_param_name: str, ) → bridge.models.conversion.quant_bridge._NativeMXFP8Storage#: Return validated, unconverted native MXFP8 storage views.

bridge.models.conversion.quant_bridge._make_meta_native_mxfp8_storage( storage_spec: bridge.models.conversion.quant_bridge._NativeMXFP8StorageSpec, ) → bridge.models.conversion.quant_bridge._NativeMXFP8Storage#: Create storage-free tensors for structural mapping validation.

bridge.models.conversion.quant_bridge._lookup_grouped_expert_mapping( mapping_registry: Any, global_param_name: str, ) → Optional[megatron.bridge.models.conversion.param_mapping.MegatronParamMapping]#: Resolve a grouped-member mapping without relying on a parameter-name suffix.

bridge.models.conversion.quant_bridge._supports_native_grouped_mxfp8( mapping: megatron.bridge.models.conversion.param_mapping.MegatronParamMapping, ) → bool#

bridge.models.conversion.quant_bridge._supports_native_mxfp8_mapping( mapping: megatron.bridge.models.conversion.param_mapping.MegatronParamMapping, ) → bool#: Return whether a mapping class explicitly implements the native contract.

bridge.models.conversion.quant_bridge._uses_dtensor_or_fsdp( param_weight: torch.Tensor, megatron_module: Any, ) → bool#: Return whether local materialization would inspect distributed storage.

bridge.models.conversion.quant_bridge._validate_local_native_mxfp8_param( param: object, global_param_name: str, ) → megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param#: Validate one final mapping result before it can leave the Bridge.

bridge.models.conversion.quant_bridge._validate_local_native_mxfp8_params( params: Iterable[object], global_param_name: str, ) → tuple[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param, ...]#: Materialize and validate all mapping results for one source parameter.

bridge.models.conversion.quant_bridge._local_native_mxfp8_param_specs( params: tuple[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param, ...], ) → tuple[bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec, ...]#: Describe validated mapping results without retaining their tensors.

bridge.models.conversion.quant_bridge._describe_local_native_mxfp8_params( params: Iterable[object], global_param_name: str, ) → tuple[bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec, ...]#: Validate mapping results and retain only their structural metadata.

bridge.models.conversion.quant_bridge._validate_local_native_mxfp8_params_against_specs( params: Iterable[object], expected_specs: tuple[bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec, ...], global_param_name: str, ) → tuple[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param, ...]#: Validate real mapping results against their preflight structure.

class bridge.models.conversion.quant_bridge.MegatronQuantizationBridge#

Mixin providing quantization-aware utilities for Megatron model bridges.

build_export_mxfp8_tasks( hf_pretrained: bridge.models.conversion.quant_bridge.HFPreTrained, megatron_model: List[bridge.models.conversion.quant_bridge.MegatronModel], ) → List[megatron.bridge.models.conversion.model_bridge.WeightConversionTask]#

Build deterministic export tasks for native MXFP8 parameters.

Singular grouped-expert parameters remain one task when their storage is native MXFP8. Grouped parameters kept in BF16 are expanded into ordinary per-expert tasks so existing conversion mappings continue to apply.

Parameters:

hf_pretrained – Hugging Face model metadata used for mapping validation.
megatron_model – Virtual-pipeline model chunks on the current rank.

Returns:

Conversion tasks in global Megatron parameter order.

iter_local_mxfp8_params( tasks: Iterable[megatron.bridge.models.conversion.model_bridge.WeightConversionTask], ) → Iterator[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param]#

Yield validated local views for native MXFP8 conversion tasks.

Parameters:: tasks – Conversion tasks in the order they should be materialized.
Yields:: Canonical local MXFP8 value and scale pairs in task, local-expert, and mapping order after every task passes preflight validation.
Raises:: ValueError – If a native task cannot be represented without conversion.

_preflight_local_native_mxfp8_task( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, ) → Optional[bridge.models.conversion.quant_bridge._NativeMXFP8TaskPlan]#: Validate one task using storage-free mapping inputs.

_preflight_native_mxfp8_mapping( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, param_weight: torch.Tensor, global_param_name: str, ) → tuple[bridge.models.conversion.quant_bridge._LocalMXFP8ParamSpec, ...]#: Validate one mapping projection with meta tensors.

_materialize_local_native_mxfp8_task( plan: bridge.models.conversion.quant_bridge._NativeMXFP8TaskPlan, ) → tuple[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param, ...]#: Materialize and validate one native MXFP8 conversion task.

_get_cached_grouped_native_mxfp8_members( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, ) → Optional[Iterable[torch.Tensor]]#: Return cached grouped members with task-qualified failures.

_grouped_native_mxfp8_member_contexts( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, members: Optional[Iterable[torch.Tensor]], ) → tuple[tuple[str, torch.Tensor], ...]#: Validate and name cached grouped members in local expert order.

_iter_grouped_native_mxfp8_params( task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask, members: Optional[Iterable[torch.Tensor]], ) → Iterator[megatron.bridge.models.conversion.param_mapping.LocalMXFP8Param]#: Project cached grouped MXFP8 members in local expert order.

static _is_mtp_param(param_name: str) → bool#: Return whether a Megatron parameter belongs to a co-trained MTP module.

stream_weights_megatron_to_hf_quant( megatron_model: Union[bridge.models.conversion.quant_bridge.MegatronModel, List[bridge.models.conversion.quant_bridge.MegatronModel]], hf_pretrained: bridge.models.conversion.quant_bridge.HFPreTrained, quantization_checker: Callable[[str], bool], quant_fn: Callable[..., Tuple[torch.Tensor, torch.Tensor]], quant_block_size: Optional[Tuple[int, int]] = None, cpu: bool = True, show_progress: bool = True, conversion_tasks: Optional[List[megatron.bridge.models.conversion.model_bridge.WeightConversionTask]] = None, merge_adapter_weights: bool = False, ) → Iterable[megatron.bridge.models.conversion.model_bridge.HFWeightTuple]#: Export Megatron weights to HuggingFace format with quantization.

bridge.models.conversion.quant_bridge#