refactor(exaone): configure provider in bridge

Drop the EXAONE-specific GPTModelProvider subclasses and configure the plain GPTModelProvider returned by the base bridge instead. This follows the Qwen3 provider_bridge pattern and keeps only the custom EXAONE layer spec in the provider module.

Signed-off-by: Bias92 <pewpewplay315@gmail.com>

Author: Bias92

src/megatron/bridge/models/__init__.py

@@ -38,8 +38,6 @@
 )
 from megatron.bridge.models.exaone import (
     Exaone4Bridge,
-    Exaone4ModelProvider,
-    Exaone4ModelProvider1P2B,
 )
 from megatron.bridge.models.gemma import (
     CodeGemmaModelProvider2B,
@@ -231,8 +229,6 @@
     "OlMoEBridge",
     "OlMoEModelProvider",
     "Exaone4Bridge",
-    "Exaone4ModelProvider",
-    "Exaone4ModelProvider1P2B",
     "NemotronHBridge",
     "MambaModelProvider",
     "MimoBridge",

src/megatron/bridge/models/exaone/__init__.py

@@ -1,4 +1,3 @@
 from megatron.bridge.models.exaone.exaone4_bridge import Exaone4Bridge
-from megatron.bridge.models.exaone.exaone4_provider import Exaone4ModelProvider, Exaone4ModelProvider1P2B
 
-__all__ = ["Exaone4Bridge", "Exaone4ModelProvider", "Exaone4ModelProvider1P2B"]
+__all__ = ["Exaone4Bridge"]

src/megatron/bridge/models/exaone/exaone4_bridge.py

@@ -35,6 +35,7 @@
 """
 
 import torch
+import torch.nn.functional as F
 from megatron.core.models.gpt.gpt_model import GPTModel
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
@@ -44,7 +45,8 @@
     GatedMLPMapping,
     QKVMapping,
 )
-from megatron.bridge.models.exaone.exaone4_provider import Exaone4ModelProvider
+from megatron.bridge.models.exaone.exaone4_provider import exaone4_layer_spec
+from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 
 
@@ -57,7 +59,6 @@
 @MegatronModelBridge.register_bridge(
     source="Exaone4ForCausalLM",  # HF architecture string (auto_map / trust_remote_code)
     target=GPTModel,
-    provider=Exaone4ModelProvider,
     model_type="exaone4",
 )
 class Exaone4Bridge(MegatronModelBridge):
@@ -84,8 +85,8 @@ class Exaone4Bridge(MegatronModelBridge):
         >>> provider = bridge.to_megatron_provider()
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Exaone4ModelProvider:
-        """Convert HuggingFace EXAONE 4.0 config to Megatron Exaone4ModelProvider.
+    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
+        """Convert HuggingFace EXAONE 4.0 config to Megatron GPTModelProvider.
 
         Maps HF config fields to Megatron TransformerConfig parameters and sets
         EXAONE-specific options including Post-LN, QK norm, and RoPE scaling.
@@ -94,28 +95,25 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Exaone4ModelProv
             hf_pretrained: HuggingFace PreTrainedCausalLM containing the EXAONE config
 
         Returns:
-            Exaone4ModelProvider configured for EXAONE 4.0 architecture
+            GPTModelProvider configured for EXAONE 4.0 architecture
         """
         hf_config = hf_pretrained.config
 
-        provider = Exaone4ModelProvider(
-            num_layers=hf_config.num_hidden_layers,
-            hidden_size=hf_config.hidden_size,
-            ffn_hidden_size=hf_config.intermediate_size,
-            num_attention_heads=hf_config.num_attention_heads,
-            num_query_groups=hf_config.num_key_value_heads,
-            seq_length=hf_config.max_position_embeddings,
-            init_method_std=hf_config.initializer_range,
-            layernorm_epsilon=hf_config.rms_norm_eps,
-            rotary_base=getattr(hf_config, "rope_theta", getattr(hf_config, "rotary_base", 1000000.0)),
-            make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(hf_config.vocab_size),
-            share_embeddings_and_output_weights=getattr(hf_config, "tie_word_embeddings", True),
-            kv_channels=getattr(hf_config, "head_dim", None),
-            fp16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.float16),
-            bf16=(self.dtype_from_hf(hf_config, default=torch.float32) == torch.bfloat16),
-            params_dtype=self.dtype_from_hf(hf_config, default=torch.float32),
-            vocab_size=hf_config.vocab_size,
-        )
+        provider = super().provider_bridge(hf_pretrained)
+
+        # EXAONE-specific architecture settings. Size-dependent fields are
+        # populated by the shared HF config mapping in the base bridge.
+        provider.normalization = "RMSNorm"
+        provider.activation_func = F.silu
+        provider.gated_linear_unit = True
+        provider.position_embedding_type = "rope"
+        provider.add_bias_linear = False
+        provider.add_qkv_bias = False
+        provider.qk_layernorm = True
+        provider.hidden_dropout = 0.0
+        provider.attention_dropout = 0.0
+        provider.transformer_layer_spec = exaone4_layer_spec
+        provider.autocast_dtype = torch.bfloat16
 
         # RoPE scaling for EXAONE 4.0 (llama3-style)
         hf_rope_scaling = getattr(hf_config, "rope_scaling", None)
@@ -131,11 +129,11 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Exaone4ModelProv
         return provider
 
     @classmethod
-    def megatron_to_hf_config(cls, provider: Exaone4ModelProvider) -> dict:
-        """Convert Megatron Exaone4ModelProvider config to HuggingFace config dict.
+    def megatron_to_hf_config(cls, provider: GPTModelProvider) -> dict:
+        """Convert Megatron GPTModelProvider config to HuggingFace config dict.
 
         Args:
-            provider: Exaone4ModelProvider with EXAONE configuration
+            provider: GPTModelProvider with EXAONE configuration
 
         Returns:
             Dictionary of HuggingFace Exaone4Config parameters
@@ -202,22 +200,20 @@ def mapping_registry(self) -> MegatronMappingRegistry:
         # =====================================================================
         # Composite Mappings (require concatenation/splitting)
         # =====================================================================
-        mapping_list.extend(
-            [
-                # QKV: Merge separate Q, K, V projections into single QKV matrix
-                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: Merge gate and up projections into single FC1 matrix
-                GatedMLPMapping(
-                    megatron_param="decoder.layers.*.mlp.linear_fc1.weight",
-                    gate="model.layers.*.mlp.gate_proj.weight",
-                    up="model.layers.*.mlp.up_proj.weight",
-                ),
-            ]
-        )
+        mapping_list.extend([
+            # QKV: Merge separate Q, K, V projections into single QKV matrix
+            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: Merge gate and up projections into single FC1 matrix
+            GatedMLPMapping(
+                megatron_param="decoder.layers.*.mlp.linear_fc1.weight",
+                gate="model.layers.*.mlp.gate_proj.weight",
+                up="model.layers.*.mlp.up_proj.weight",
+            ),
+        ])
 
         return MegatronMappingRegistry(*mapping_list)

src/megatron/bridge/models/exaone/exaone4_provider.py

@@ -30,16 +30,10 @@
 established by Gemma2 bridge.
 """
 
-from dataclasses import dataclass
-from typing import Callable
-
-import torch
-import torch.nn.functional as F
 from megatron.core.extensions.transformer_engine import TEColumnParallelLinear
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.transformer import (
     ModuleSpec,
-    TransformerConfig,
     TransformerLayer,
     TransformerLayerSubmodules,
 )
@@ -97,81 +91,3 @@ def exaone4_layer_spec(config: "GPTModelProvider") -> ModuleSpec:  # noqa: ARG00
             mlp_bda=get_bias_dropout_add,
         ),
     )
-
-
-# =============================================================================
-# EXAONE 4.0 Model Providers
-# =============================================================================
-
-
-@dataclass
-class Exaone4ModelProvider(GPTModelProvider):
-    """Base configuration for EXAONE 4.0 models (LG AI Research).
-
-    Architecture features:
-    - Pure Post-LayerNorm (no input_layernorm / pre_feedforward_layernorm)
-    - QK RMSNorm after Q/K projection
-    - GQA (Grouped Query Attention)
-    - SwiGLU activation
-    - RoPE with llama3-style scaling
-    - Tied word embeddings (inherits share_embeddings_and_output_weights=True from GPTModelProvider)
-    """
-
-    # Architecture defaults common across EXAONE 4.0 model sizes
-    normalization: str = "RMSNorm"
-    activation_func: Callable = F.silu
-    gated_linear_unit: bool = True
-    position_embedding_type: str = "rope"
-    add_bias_linear: bool = False
-    add_qkv_bias: bool = False
-    qk_layernorm: bool = True  # EXAONE 4.0 uses QK RMSNorm
-    hidden_dropout: float = 0.0
-    attention_dropout: float = 0.0
-
-    # Custom layer spec for Post-LN architecture
-    transformer_layer_spec: ModuleSpec | Callable[["GPTModelProvider"], ModuleSpec] = exaone4_layer_spec
-
-    # Dtype defaults
-    autocast_dtype: torch.dtype = torch.bfloat16
-    params_dtype: torch.dtype = torch.bfloat16
-    bf16: bool = True
-
-
-@dataclass
-class Exaone4ModelProvider1P2B(Exaone4ModelProvider):
-    """Configuration for EXAONE 4.0 1.2B.
-
-    Model: LGAI-EXAONE/EXAONE-4.0-1.2B
-    - 30 layers, 2048 hidden, 32 attention heads, 8 KV heads
-    - Full attention only (no sliding window / hybrid attention)
-    - RoPE: theta=1M, llama3 scaling factor=16, original_max_pos=8192
-    - Vocab: 102,400 tokens, Context: 65,536 tokens
-    """
-
-    num_layers: int = 30
-    hidden_size: int = 2048
-    ffn_hidden_size: int = 4096
-    num_attention_heads: int = 32
-    num_query_groups: int = 8
-    kv_channels: int = 64
-    seq_length: int = 65536
-    vocab_size: int = 102400
-    rotary_base: float = 1000000.0
-    layernorm_epsilon: float = 1e-5
-    init_method_std: float = 0.02
-
-    # RoPE scaling (llama3-style)
-    rope_scaling: bool = True
-    rope_scaling_factor: float = 16.0
-    rope_scaling_low_freq_factor: float = 1.0
-    rope_scaling_high_freq_factor: float = 4.0
-    rope_scaling_original_max_position_embeddings: int = 8192
-
-
-# TODO: Add Exaone4ModelProvider32B when 32B model details are confirmed
-# The 32B model introduces hybrid attention (LLLG pattern: 3 local + 1 global)
-# with sliding_window_pattern and layer_types configuration.
-# This will require:
-# - Layer-wise attention type branching (local vs global)
-# - RoPE disable for global attention layers
-# - Sliding window size configuration per layer

tests/functional_tests/models/exaone/test_exaone4_provider.py

@@ -13,13 +13,51 @@
 # limitations under the License.
 
 import torch
+import torch.nn.functional as F
 from transformers import PretrainedConfig
 
 from megatron.bridge.models.conversion.auto_bridge import AutoBridge
-from megatron.bridge.models.exaone import Exaone4ModelProvider1P2B
+from megatron.bridge.models.exaone.exaone4_provider import exaone4_layer_spec
+from megatron.bridge.models.gpt_provider import GPTModelProvider
 from tests.functional_tests.utils import compare_provider_configs
 
 
+def make_expected_exaone4_1p2b_provider():
+    provider = GPTModelProvider(
+        num_layers=30,
+        hidden_size=2048,
+        ffn_hidden_size=4096,
+        num_attention_heads=32,
+        num_query_groups=8,
+        kv_channels=64,
+        seq_length=65536,
+        vocab_size=102400,
+        rotary_base=1000000.0,
+        layernorm_epsilon=1e-5,
+        init_method_std=0.02,
+        normalization="RMSNorm",
+        activation_func=F.silu,
+        gated_linear_unit=True,
+        position_embedding_type="rope",
+        add_bias_linear=False,
+        add_qkv_bias=False,
+        qk_layernorm=True,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        transformer_layer_spec=exaone4_layer_spec,
+        autocast_dtype=torch.bfloat16,
+        params_dtype=torch.bfloat16,
+        bf16=True,
+        fp16=False,
+        rope_scaling=True,
+        rope_scaling_factor=16.0,
+    )
+    provider.rope_scaling_low_freq_factor = 1.0
+    provider.rope_scaling_high_freq_factor = 4.0
+    provider.rope_scaling_original_max_position_embeddings = 8192
+    return provider
+
+
 class TestExaone4ProviderMapping:
     """Test EXAONE 4.0 config-only AutoBridge provider mapping."""
 
@@ -53,7 +91,7 @@ def test_bridge_vs_predefined_provider_config_from_config_only(self):
         converted_provider = bridge.to_megatron_provider(load_weights=False)
         converted_provider.finalize()
 
-        predefined_provider = Exaone4ModelProvider1P2B()
+        predefined_provider = make_expected_exaone4_1p2b_provider()
         predefined_provider.finalize()
 
         compare_provider_configs(converted_provider, predefined_provider, "exaone4-config-only")