bridge.models.gemma.gemma2_bridge

Module Contents

Classes

Gemma2Bridge

Megatron Bridge for Gemma2 Causal LM. This bridge handles the conversion between HuggingFace Gemma2ForCausalLM and Megatron-Core GPTModel formats, including weight mappings and configuration translation. Gemma2 includes specific features like attention logit softcapping, sliding window attention, and additional layer normalization compared to the original Gemma model. As a user you would not use this bridge directly, but through AutoBridge. .. rubric:: Example

API

class bridge.models.gemma.gemma2_bridge.Gemma2Bridge

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

Megatron Bridge for Gemma2 Causal LM. This bridge handles the conversion between HuggingFace Gemma2ForCausalLM and Megatron-Core GPTModel formats, including weight mappings and configuration translation. Gemma2 includes specific features like attention logit softcapping, sliding window attention, and additional layer normalization compared to the original Gemma model. As a user you would not use this bridge directly, but through AutoBridge. .. rubric:: Example

from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained(“google/gemma-2-2b”) provider = bridge.to_megatron_provider()

provider_bridge(

hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,

) → megatron.bridge.models.gemma.gemma2_provider.Gemma2ModelProvider

Convert HuggingFace config to Gemma2ModelProvider.

Parameters:

hf_pretrained – HuggingFace pretrained model wrapper

Returns:

Configured provider for Megatron model

Return type:

Gemma2ModelProvider

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

Return MegatronMappingRegistry containing parameter mappings from HF to Megatron format.

Returns:

Registry of parameter mappings

Return type:

MegatronMappingRegistry