bridge.models.llama.llama_bridge

Module Contents

Classes

LlamaBridge

Megatron Bridge for Llama Causal LM.

Data

logger

API

bridge.models.llama.llama_bridge.logger

‘getLogger(…)’

class bridge.models.llama.llama_bridge.LlamaBridge

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

Megatron Bridge for Llama Causal LM.

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(“meta-llama/Llama-3.1-8B-Instruct”) 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 Llama config to Megatron GPTModelProvider.

Uses base class implementation for common conversion, then sets Llama-specific config and preserves supported RoPE scaling.

Parameters:

hf_pretrained – HuggingFace PreTrainedCausalLM containing the Llama config

Returns:

GPTModelProvider configured for Llama architecture

classmethod megatron_to_hf_config(

provider: megatron.bridge.models.gpt_provider.GPTModelProvider,

) → dict

Convert Megatron GPTModelProvider config to HuggingFace Llama config dict.

Uses base class implementation, then adds supported Llama RoPE scaling.

Parameters:

provider – GPTModelProvider with Llama configuration

Returns:

Dictionary of HuggingFace LlamaConfig parameters

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

maybe_modify_converted_hf_weight(

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

converted_weights_dict: dict[str, torch.Tensor],

hf_state_dict: collections.abc.Mapping[str, torch.Tensor],

) → dict[str, torch.Tensor]

Preserve persisted Llama rotary inverse-frequency buffers on export.