bridge.models.olmoe.olmoe_bridge#

Module Contents#

Classes#

OlMoEBridge

Megatron Bridge for OlMoE Models.

Data#

logger

API#

bridge.models.olmoe.olmoe_bridge.logger#

‘getLogger(…)’

class bridge.models.olmoe.olmoe_bridge.OlMoEBridge#

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

Megatron Bridge for OlMoE Models.

This bridge handles the conversion between HuggingFace OlMoEForCausalLM and Megatron-Core GPTModel formats. OlMoE models use mixture of experts architecture with QK layernorm.

.. rubric:: Example

from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained(“allenai/OLMoE-1B-7B-0125”) provider = bridge.to_megatron_provider()

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.gpt_provider.GPTModelProvider#

Convert HuggingFace OlMoE config to Megatron GPTModelProvider.

Uses base class implementation for common conversion, then sets OlMoE-specific config. OlMoE uses QK layernorm and mixture of experts.

Parameters:

hf_pretrained – HuggingFace PreTrainedCausalLM containing the OlMoE config

Returns:

GPTModelProvider configured for OlMoE architecture

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