bridge.models.conversion.mapping_registry#
Module Contents#
Classes#
Registry for weight mappings between model formats with pattern matching support. |
API#
- class bridge.models.conversion.mapping_registry.MegatronMappingRegistry(
- *mappings: megatron.bridge.models.conversion.param_mapping.MegatronParamMapping,
Registry for weight mappings between model formats with pattern matching support.
This class serves as a registry of weight mappings between Megatron and external (typically HuggingFace) model formats. It provides efficient pattern matching for parameter names using glob-like wildcards (*) and supports both forward (Megatron → HF) and reverse (HF → Megatron) lookups.
The registry pre-compiles regex patterns for efficient repeated lookups and handles the resolution of wildcards in parameter names.
- Parameters:
*mappings – Variable number of MegatronParamMapping objects defining the individual weight mappings
.. rubric:: Example
Create a mapping registry with various mappings
mapping_registry = MegatronMappingRegistry( … AutoMapping( … megatron_param=”embedding.word_embeddings.weight”, … hf_param=”model.embed_tokens.weight”, … ), … QKVMapping( … megatron_param=”decoder.layers..self_attention.linear_qkv.weight”, … q=”model.layers..self_attn.q_proj.weight”, … k=”model.layers..self_attn.k_proj.weight”, … v=”model.layers..self_attn.v_proj.weight”, … ), … )
Query for a specific layer (wildcards are resolved)
mapping = mapping_registry.megatron_to_hf_lookup(“decoder.layers.0.self_attention.linear_qkv.weight”) print(mapping.hf_param) # Will show resolved HF names for layer 0
Reverse lookup from HF name
mapping = mapping_registry.hf_to_megatron_lookup(“model.layers.5.self_attn.q_proj.weight”) print(mapping.megatron_param) # Shows corresponding Megatron name
Build from a list
mappings = [bridge1, bridge2, bridge3] mapping_registry = MegatronMappingRegistry(*mappings)
.. note::
Wildcard patterns use ‘*’ which matches any sequence of digits (0-9). This is specifically designed for layer indices in transformer models.
Initialization
Initialize MegatronMappingRegistry with weight mappings.
- Parameters:
*mappings – MegatronParamMapping objects
- megatron_to_hf_lookup(
- megatron_param_name: str,
Get mapping for a Megatron parameter name.
This method performs efficient lookups by first checking for exact matches, then falling back to pattern matching using pre-compiled regex patterns. When a pattern match is found, wildcards are automatically resolved.
- Parameters:
megatron_param_name – Megatron parameter name to look up Example: “decoder.layers.0.self_attention.linear_qkv.weight”
- Returns:
Bridge instance with resolved wildcards, or None if no matching mapping is found. The returned bridge will have all wildcards replaced with actual values.
- Return type:
.. rubric:: Example
Query with exact layer number
bridge = state_map.megatron_to_hf_lookup(“decoder.layers.5.mlp.linear_fc1.weight”) if bridge: … print(f”Maps to: {bridge.hf_param}”) # Shows HF name for layer 5
- hf_to_megatron_lookup(
- hf_param_name: str,
Get mapping for a destination parameter name (reverse lookup).
This is useful when you have a destination name and want to find the corresponding megatron name.
- Parameters:
hf_param_name – Destination parameter name to look up
- Returns:
MegatronParamMapping with resolved wildcards, or None if no match found
- get_all_mappings() List[megatron.bridge.models.conversion.param_mapping.MegatronParamMapping]#
Get all mappings in this MegatronMappingRegistry.
- get_mappings_by_pattern(
- pattern: str,
Get all mappings that match a given pattern.
- Parameters:
pattern – Pattern to match (supports * wildcards)
- Returns:
List of matching MegatronParamMapping objects
- __len__() int#
Return number of mappings.
- __iter__()#
Iterate over mappings.
- __repr__() str#
String representation of MegatronMappingRegistry.
- describe() str#
Get a human-readable description of all mappings.
- Returns:
Formatted string describing all weight mappings