bridge.models.mimo.mimo_bridge

Module Contents

Classes

MimoBridge

Megatron Bridge for MiMo Causal LM.

API

class bridge.models.mimo.mimo_bridge.MimoBridge

Bases: megatron.bridge.models.qwen.qwen2_bridge.Qwen2Bridge

Megatron Bridge for MiMo Causal LM.

hf_config_to_model_config_kwargs(

hf_config: transformers.PretrainedConfig,

) → dict[str, Any]

Convert a Hugging Face MiMo config to builder config kwargs.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry

static _swap_input_proj_halves(weight: torch.Tensor) → torch.Tensor

maybe_modify_loaded_hf_weight(

hf_param: str | dict[str, str],

hf_state_dict: Mapping[str, torch.Tensor],

) → torch.Tensor

maybe_modify_converted_hf_weight(

task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask,

converted_weights_dict: dict[str, torch.Tensor],

hf_state_dict: Mapping[str, torch.Tensor],

) → dict[str, torch.Tensor]