Bridge with 🤗Hugging Face

The bridge provides seamless bidirectional conversion between 🤗Hugging Face Transformers and megatron-core model formats, handling the complexities of distributed model parallelism transparently.

Quick Start

Loading a 🤗Hugging Face Model into Megatron

from megatron.bridge import AutoBridge

# Load Llama from HuggingFace Hub and convert to Megatron
bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B")
provider = bridge.to_megatron_provider()

# The provider is lazy - configure parallelism before creating models
provider.tensor_model_parallel_size = 8
provider.pipeline_model_parallel_size = 2

model = provider.provide_distributed_model(wrap_with_ddp=False)

Converting Megatron Models back to 🤗Hugging Face

# Export a trained Megatron model to HuggingFace format
bridge.save_hf_pretrained(model, "./my-fine-tuned-llama")

# Or stream weights for memory efficiency
for name, weight in bridge.export_to_hf(model):
    print(f"Exporting {name}: {weight.shape}")

Auto-detection with AutoBridge

from megatron.bridge import AutoBridge

# Automatically detect model architecture and create appropriate bridge
bridge = AutoBridge.from_hf_pretrained("any-supported-huggingface/model")

Architecture Overview

The bridge framework uses a layered architecture with clear separation of concerns:

┌─────────────────────────────────────────────────────────┐
│                   User API Layer                        │
│         (AutoBridge)                                    │
├─────────────────────────────────────────────────────────┤
│                Orchestration Layer                      │
│            (MegatronModelBridge)                        │
├─────────────────────────────────────────────────────────┤
│    Mapping Layer         │    Transformation Layer      │
│ (MegatronMappingRegistry)│   (MegatronParamMapping)     │
└─────────────────────────────────────────────────────────┘

Core Components

MegatronModelBridge: High-level orchestrator that coordinates the conversion process
MegatronMappingRegistry: Registry of parameter name mappings between formats
MegatronParamMapping: Handles weight transformations and distributed communication
AutoBridge: User-friendly API for causal language models

Design Patterns

Multi-Dispatch Registration

The bridge uses decorators to register bridge implementations, enabling automatic routing:

@MegatronModelBridge.register_bridge(source=LlamaForCausalLM, target=GPTModel)
class MegatronCausalLlamaBridge(MegatronModelBridge):
    def provider_bridge(self, hf_pretrained):
        # Convert HF config to Megatron provider
        return LlamaModelProvider(
            num_layers=hf_pretrained.config.num_hidden_layers,
            hidden_size=hf_pretrained.config.hidden_size,
            # ... more config mapping
        )
    
    def mapping_registry(self):
        # Define weight mappings
        return MegatronMappingRegistry(
            AutoMapping(
                megatron_param="embedding.word_embeddings.weight",
                hf_param="model.embed_tokens.weight"
            ),
            # ... more mappings
        )

Param Mapping Strategies

Different weight transformation strategies handle various parallelism patterns:

DirectMapping: Simple 1:1 mapping
ColumnParallelMapping: Splits along output dimension
RowParallelMapping: Splits along input dimension
QKVMapping: Handles QKV matrix interleaving
GatedMLPMapping: Manages gated activation concatenation
AutoMapping: Auto-detects and applies correct strategy

Conversion Process

🤗Hugging Face → Megatron

Planning Phase
- Iterate through Megatron model parameters
- Query StateBridge for HF source mappings
- Build execution task with resolved names
Execution Phase
- Load HF weights according to mappings
- Apply format transformations (e.g., QKV fusion)
- Distribute via tensor parallelism if needed
- Copy to target Megatron parameters
Distribution
- Column-parallel: scatter along dim 0
- Row-parallel: scatter along dim 1
- Replicated: broadcast to all ranks

Megatron → 🤗Hugging Face

Collection Phase
- Gather parameter locations across pipeline ranks
- Build export task respecting ordering preferences
Export Phase
- Broadcast from owning pipeline rank
- Gather from tensor parallel ranks
- Apply inverse transformations (e.g., QKV splitting)
- Yield HF-format weight tensors

Parallelism Handling

Tensor Parallelism (TP)

The bridge automatically handles three TP patterns:

# Column-parallel (output split)
embeddings: [vocab, hidden] → [vocab, hidden/TP]
qkv_proj: [hidden, 3*hidden] → [hidden, 3*hidden/TP]

# Row-parallel (input split)  
out_proj: [hidden, hidden] → [hidden/TP, hidden]
down_proj: [intermediate, hidden] → [intermediate/TP, hidden]

# Replicated (no split)
layer_norm: [hidden] → [hidden] (same on all ranks)

Pipeline Parallelism (PP)

Parameters exist on specific pipeline ranks
Cross-rank broadcasting ensures collective participation
Metadata broadcast precedes tensor data for efficiency

Virtual Pipeline Parallelism (VPP)

Multiple model chunks per device
Automatic layer offset calculation
Transparent to weight transformation logic

Advanced Usage

Custom Weight Transformations

Create custom param mappings for special formats:

class MyCustomMapping(MegatronParamMapping):
   def hf_to_megatron(self, hf_weights, megatron_module):
        # Custom transformation logic
        transformed = my_transform(hf_weights)
        # Use provided helpers for distribution
        return self.scatter_to_tp_ranks(transformed, dim=0)

   def megatron_to_hf(self, megatron_weights, module):
        # Gather and inverse transform
        gathered = self.gather_from_tp_ranks(megatron_weights, dim=0)
        return {"custom_weight": my_inverse_transform(gathered)}

Exporting Weights

Export weights to 🤗Hugging Face format (all ranks receive full weights):

# Export weights
weights = bridge.export_to_hf(model)

Streaming Large Models

Handle massive models without loading all weights:

# Stream weights during conversion
with bridge.stream_conversion(src_model) as stream:
    for weight_batch in stream:
        # Process weights in batches
        process_batch(weight_batch)

Adding New Model Support

To add support for a new model architecture:

Create a Bridge Class

@MegatronModelBridge.register_bridge(source=YourHFModel, target=YourMegatronModel)
class YourModelBridge(MegatronModelBridge):
    pass

Implement Configuration Mapping

def provider_bridge(self, hf_pretrained):
    return YourModelProvider(
        # Map HF config to Megatron config
    )

Define Weight Mappings

def mapping_registry(self):
    return MegatronMappingRegistry(
        # Define all weight mappings
    )

Register Custom Modules (if needed)

AutoMapping.register_module_type(
    "YourColumnParallelLinear", "column"
)

Implementation Examples

Llama Bridge

The Llama implementation demonstrates key patterns:

# Special handling for model variations
if is_llama_3_1_config(config):
    # Handle RoPE scaling
    rotary_base = config.rope_scaling["rope_type"] == "llama3"
    
# QKV fusion with proper head ordering
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"
)

# Gated MLP handling
GatedMLPMapping(
    megatron_param="decoder.layers.*.mlp.linear_fc1.weight",
    gate="model.layers.*.mlp.gate_proj.weight",
    up="model.layers.*.mlp.up_proj.weight"
)

Best Practices

Always Use High-Level APIs: Prefer AutoBridge over direct bridge usage
Configure Before Creating: Set parallelism parameters on providers before model creation
Handle Missing Weights: Check for None returns in custom bridges
Test Bidirectionality: Ensure HF→Megatron→HF preserves weights exactly
Profile Large Models: Use streaming APIs to avoid memory issues

Troubleshooting

Common Issues

Shape Mismatches: Usually indicates incorrect parallelism configuration
Missing Weights: Check StateBridge mappings cover all parameters
Memory Errors: Use streaming APIs or increase distribution

Debugging Tips

# Enable verbose logging
import logging
logging.getLogger("megatron.bridge.models").setLevel(logging.DEBUG)

# Inspect mappings
bridge = AutoBridge.from_hf_pretrained("model")
mapping_registry = bridge.mapping_registry()
print(mapping_registry.get_all_mappings())

# Verify weight shapes
for task in bridge._build_conversion_tasks(hf_model, meg_models):
    print(f"{task.param_name}: {task.megatron_param.shape}")

Performance Considerations

Memory: Streaming prevents full model loads during conversion
Communication: Optimized collective operations for distributed settings
Caching: Compiled regex patterns for efficient name matching
Progress: Shown only on rank 0 to avoid output clutter

Future Enhancements

PEFT model support (LoRA, QLoRA)
Automatic optimal parallelism detection
Checkpoint format conversions
Multi-GPU to single-GPU consolidation utilities

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!