bridge.models.conversion.quant_bridge#

Module Contents#

Classes#

MegatronQuantizationBridge

Mixin providing quantization-aware utilities for Megatron model bridges.

Data#

API#

bridge.models.conversion.quant_bridge.MegatronModel#

‘TypeVar(…)’

bridge.models.conversion.quant_bridge.HFPreTrained#

‘TypeVar(…)’

class bridge.models.conversion.quant_bridge.MegatronQuantizationBridge#

Mixin providing quantization-aware utilities for Megatron model bridges.

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.