Merge branch 'main' into chcui/nvbug-4/fix-lora-plus-warmup

Author: cuichenx

.github/workflows/cicd-main.yml

@@ -1388,7 +1388,7 @@ jobs:
   notify-nightly-failure:
     name: Notify scheduled CI failure
     runs-on: ubuntu-latest
-    environment: main
+    environment: public
     needs:
       - ephemeral-runner-routing
       - pre-flight

src/megatron/bridge/models/bailing/bailing_moe2_bridge.py

@@ -33,10 +33,11 @@
 
 import logging
 from functools import partial
+from typing import Any
 
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
-from transformers import AutoConfig, AutoModelForCausalLM
+from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
 
 from megatron.bridge.models.bailing.configuration_bailing_moe_v2 import BailingMoeV2Config
 from megatron.bridge.models.bailing.modeling_bailing_moe_v2 import BailingMoeV2ForCausalLM
@@ -47,8 +48,7 @@
     ConcatenatedQKVMapping,
     GatedMLPMapping,
 )
-from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 
 
 # Register the Bailing MoE V2 config and model classes with transformers Auto classes.
@@ -76,38 +76,52 @@ class BailingMoeV2Bridge(MegatronModelBridge):
     Example:
         >>> from megatron.bridge import AutoBridge
         >>> bridge = AutoBridge.from_hf_pretrained("inclusionAI/Ling-mini-2.0")
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
-        provider = super().provider_bridge(hf_pretrained)
-        hf_config = hf_pretrained.config
-
-        provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
-        provider.normalization = "RMSNorm"
-        provider.gated_linear_unit = True
-        provider.add_bias_linear = False
-        provider.share_embeddings_and_output_weights = False
-        provider.qk_layernorm = True
-        provider.add_qkv_bias = getattr(hf_config, "use_qkv_bias", False)
-
-        provider.moe_grouped_gemm = True
-        provider.moe_router_pre_softmax = True
-        provider.moe_router_load_balancing_type = "none"
-        provider.moe_router_score_function = "sigmoid"
-        provider.moe_router_enable_expert_bias = True
-        provider.moe_router_dtype = "fp32"
-        provider.moe_token_dispatcher_type = "alltoall"
-        provider.moe_permute_fusion = True
-
-        provider.hidden_dropout = 0.0
-
-        provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * (
-            hf_config.num_hidden_layers - hf_config.first_k_dense_replace
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
+
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face Ling MoE V2 config to builder config kwargs.
+
+        The builder's default MoE layer spec already uses the GPT decoder block
+        spec, which honors the dense-first ``moe_layer_freq``.
+        """
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
+        config_kwargs.update(
+            normalization="RMSNorm",
+            gated_linear_unit=True,
+            add_bias_linear=False,
+            share_embeddings_and_output_weights=False,
+            qk_layernorm=True,
+            add_qkv_bias=getattr(hf_config, "use_qkv_bias", False),
+            masked_softmax_fusion=True,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+            moe_grouped_gemm=True,
+            moe_router_pre_softmax=True,
+            moe_router_load_balancing_type="none",
+            moe_router_score_function="sigmoid",
+            moe_router_enable_expert_bias=True,
+            moe_router_dtype="fp32",
+            moe_token_dispatcher_type="alltoall",
+            moe_permute_fusion=True,
+            hidden_dropout=0.0,
+            moe_layer_freq=[0] * hf_config.first_k_dense_replace
+            + [1] * (hf_config.num_hidden_layers - hf_config.first_k_dense_replace),
+            moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size,
         )
-        provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size
-
-        return provider
+        config_kwargs.setdefault("position_embedding_type", "rope")
+        return config_kwargs
+
+    def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Adapt the canonical builder mapping to the deprecated provider path."""
+        provider_kwargs = self.hf_config_to_model_config_kwargs(hf_config)
+        # The provider's default layer spec builds homogeneous layers; the decoder
+        # block spec is needed for the dense-first MoE layer pattern.
+        provider_kwargs["transformer_layer_spec"] = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
+        return provider_kwargs
 
     def mapping_registry(self) -> MegatronMappingRegistry:
         mapping_list = []

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

@@ -34,9 +34,13 @@
 - EXAONE bridge: QK layernorm mapping pattern
 """
 
+import logging
+from typing import Any
+
 import torch
 import torch.nn.functional as F
 from megatron.core.models.gpt.gpt_model import GPTModel
+from transformers import PretrainedConfig
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
@@ -46,10 +50,12 @@
     QKVMapping,
 )
 from megatron.bridge.models.exaone.exaone4.exaone4_provider import exaone4_layer_spec
+from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
 from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 
 
+logger = logging.getLogger(__name__)
+
 # Register custom EXAONE modules for AutoMapping weight distribution
 # TERowParallelLinearLayerNorm is a row-parallel linear with post-layernorm
 # (same pattern as other Post-LN architectures)
@@ -82,51 +88,64 @@ class Exaone4Bridge(MegatronModelBridge):
         ...     "LGAI-EXAONE/EXAONE-4.0-1.2B",
         ...     trust_remote_code=True,
         ... )
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
-        """Convert HuggingFace EXAONE 4.0 config to Megatron GPTModelProvider.
+    MODEL_CONFIG_CLASS = BridgeGPTModelConfig
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
 
-        Maps HF config fields to Megatron TransformerConfig parameters and sets
-        EXAONE-specific options including Post-LN, QK norm, and RoPE scaling.
+    # Megatron Core applies llama3 RoPE scaling with these fixed values; only the factor is configurable.
+    _MCORE_LLAMA3_ROPE_SCALING = {
+        "low_freq_factor": 1.0,
+        "high_freq_factor": 4.0,
+        "original_max_position_embeddings": 8192,
+    }
 
-        Args:
-            hf_pretrained: HuggingFace PreTrainedCausalLM containing the EXAONE config
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face EXAONE 4.0 config to builder config kwargs.
 
-        Returns:
-            GPTModelProvider configured for EXAONE 4.0 architecture
+        Sets EXAONE-specific options including Post-LN, QK norm, and RoPE scaling.
+        Size-dependent fields are populated by the shared HF config mapping.
         """
-        hf_config = hf_pretrained.config
-
-        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
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
+        config_kwargs.update(
+            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,
+            masked_softmax_fusion=True,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+        )
 
         # RoPE scaling for EXAONE 4.0 (llama3-style)
         hf_rope_scaling = getattr(hf_config, "rope_scaling", None)
         if hf_rope_scaling is not None and hf_rope_scaling.get("rope_type") == "llama3":
-            provider.rope_scaling = True
-            provider.rope_scaling_factor = hf_rope_scaling.get("factor", 16.0)
-            provider.rope_scaling_low_freq_factor = hf_rope_scaling.get("low_freq_factor", 1.0)
-            provider.rope_scaling_high_freq_factor = hf_rope_scaling.get("high_freq_factor", 4.0)
-            provider.rope_scaling_original_max_position_embeddings = hf_rope_scaling.get(
-                "original_max_position_embeddings", 8192
-            )
+            config_kwargs.update(rope_scaling=True, rope_scaling_factor=hf_rope_scaling.get("factor", 16.0))
+            unsupported = {
+                key: hf_rope_scaling[key]
+                for key, value in self._MCORE_LLAMA3_ROPE_SCALING.items()
+                if key in hf_rope_scaling and hf_rope_scaling[key] != value
+            }
+            if unsupported:
+                logger.warning(
+                    "EXAONE 4.0 llama3 RoPE scaling values %s differ from the fixed values Megatron Core applies (%s).",
+                    unsupported,
+                    self._MCORE_LLAMA3_ROPE_SCALING,
+                )
+
+        return config_kwargs
 
-        return provider
+    def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Adapt the canonical builder mapping to the deprecated provider path."""
+        return self.hf_config_to_model_config_kwargs(hf_config)
 
     @classmethod
     def megatron_to_hf_config(cls, provider: GPTModelProvider) -> dict:
@@ -149,9 +168,11 @@ def megatron_to_hf_config(cls, provider: GPTModelProvider) -> dict:
             hf_config["rope_scaling"] = {
                 "rope_type": "llama3",
                 "factor": provider.rope_scaling_factor,
-                "low_freq_factor": provider.rope_scaling_low_freq_factor,
-                "high_freq_factor": provider.rope_scaling_high_freq_factor,
-                "original_max_position_embeddings": provider.rope_scaling_original_max_position_embeddings,
+                # Megatron Core's fixed llama3 values, unless a caller set them explicitly.
+                **{
+                    key: getattr(provider, f"rope_scaling_{key}", value)
+                    for key, value in cls._MCORE_LLAMA3_ROPE_SCALING.items()
+                },
             }
 
         return hf_config

src/megatron/bridge/models/hf_pretrained/state.py

@@ -196,13 +196,26 @@ def __init__(self, source: Dict[str, torch.Tensor] | "StateSource"):
             raise TypeError(f"StateDict source must be a dict or a StateSource, got {type(source)}")
 
         self.source = source
+        self._key_set: set[str] | None = None
+        self._key_set_source: List[str] | None = None
 
     def _get_all_keys(self) -> List[str]:
         """
         Get all available tensor keys from the underlying source.
         """
         return self.source.get_all_keys()
 
+    def _get_key_set(self) -> set[str]:
+        # Membership checks run once per tensor during conversion; a list scan made
+        # them quadratic in the number of checkpoint keys. Sources return a cached
+        # list, so the set is rebuilt only if they hand back a different one (e.g.
+        # while a bridge temporarily wraps ``get_all_keys`` with virtual keys).
+        keys = self._get_all_keys()
+        if self._key_set_source is not keys:
+            self._key_set = set(keys)
+            self._key_set_source = keys
+        return self._key_set
+
     def _load_tensors(self, keys_to_load: List[str]) -> Dict[str, torch.Tensor]:
         """
         Load specified tensors from the underlying source.
@@ -221,7 +234,7 @@ def _match_keys(self, pattern: Union[str, Pattern]) -> List[str]:
             return [k for k in all_keys if fnmatch.fnmatch(k, pattern)]
         else:
             # Exact match
-            return [pattern] if pattern in all_keys else []
+            return [pattern] if pattern in self._get_key_set() else []
 
     @overload
     def __getitem__(self, key: str) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: ...
@@ -299,11 +312,11 @@ def __getitem__(self, key: Union[str, List[str], Pattern]) -> Union[torch.Tensor
                     raise KeyError(f"No keys match pattern: {key}")
                 return self._load_tensors(matched_keys)
             else:
-                if key not in self._get_all_keys():
+                if key not in self._get_key_set():
                     raise KeyError(f"Key not found: {key}")
                 return self._load_tensors([key])[key]
         elif isinstance(key, list):
-            all_keys_set = set(self._get_all_keys())
+            all_keys_set = self._get_key_set()
             missing_keys = [k for k in key if k not in all_keys_set]
             if missing_keys:
                 raise KeyError(f"Keys not found: {missing_keys}")
@@ -334,7 +347,7 @@ def keys(self) -> List[str]:
 
     def __contains__(self, key: str) -> bool:
         """Check if a key exists in the state dict."""
-        return key in self._get_all_keys()
+        return key in self._get_key_set()
 
     def __repr__(self) -> str:
         """String representation."""
@@ -350,7 +363,7 @@ def get(self, key: str, default=None) -> Optional[torch.Tensor]:
         Returns `default` if the key is not found.
         Note: This method is for single key lookup and does not support patterns.
         """
-        if key in self._get_all_keys():
+        if key in self._get_key_set():
             return self._load_tensors([key])[key]
         return default
 
@@ -500,6 +513,7 @@ def __init__(
         # DeepSeek-V4-Pro). Both can be disabled through environment variables.
         self._small_tensor_cache: Dict[str, torch.Tensor] = {}
         self._open_handles: Dict[str, Any] = {}
+        self._open_handle_keys: dict[str, set[str]] = {}
         self._cache_small = os.environ.get("MBRIDGE_STATE_CACHE_SMALL_TENSORS", "1") == "1"
         self._reuse_handles = os.environ.get("MBRIDGE_STATE_REUSE_HANDLES", "1") == "1"
         self._small_numel = int(os.environ.get("MBRIDGE_STATE_SMALL_TENSOR_NUMEL", str(16 * 1024 * 1024)))
@@ -687,6 +701,17 @@ def _get_handle(self, file_path: Path):
             self._open_handles[key] = handle
         return handle
 
+    def _get_handle_keys(self, file_path: Path, handle: Any) -> set[str]:
+        """Return the key set of ``handle``, cached alongside reused handles."""
+        if not self._reuse_handles:
+            return set(handle.keys())
+        key = str(file_path)
+        keys = self._open_handle_keys.get(key)
+        if keys is None:
+            keys = set(handle.keys())
+            self._open_handle_keys[key] = keys
+        return keys
+
     def _maybe_cache(self, key: str, tensor: torch.Tensor) -> None:
         # Keep a private copy: callers may modify the tensors they get back in place.
         if self._cache_small and tensor.numel() <= self._small_numel:
@@ -695,6 +720,7 @@ def _maybe_cache(self, key: str, tensor: torch.Tensor) -> None:
     def _invalidate_handle(self, file_path: Path) -> None:
         """Drop a reused handle so the next access re-reads the shard header (e.g. after a retry)."""
         self._open_handles.pop(str(file_path), None)
+        self._open_handle_keys.pop(str(file_path), None)
 
     def load_tensors(self, keys_to_load: List[str]) -> Dict[str, torch.Tensor]:
         if not keys_to_load:
@@ -744,7 +770,7 @@ def load_tensors(self, keys_to_load: List[str]) -> Dict[str, torch.Tensor]:
                     for attempt in range(max_retries):
                         f = self._get_handle(file_path)
                         if True:  # the handle may be shared across calls; do not close it here
-                            file_keys = set(f.keys())
+                            file_keys = self._get_handle_keys(file_path, f)
                             for key in keys_in_file:
                                 source_key = (self.key_prefix or "") + key
                                 if source_key in file_keys and key not in loaded_tensors:

src/megatron/bridge/recipes/exaone/h100/exaone4.py

@@ -35,7 +35,7 @@ def _set_optimizer_precision(cfg: ConfigContainer) -> None:
 
 
 def _apply_exaone4_common(cfg: ConfigContainer) -> None:
-    cfg.model = AutoBridge.from_hf_pretrained(_HF_PATH).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(_HF_PATH).get_model_config()
     cfg.tokenizer.tokenizer_model = _HF_PATH
 
     cfg.model.tensor_model_parallel_size = 1

src/megatron/bridge/recipes/qwen/__init__.py

@@ -23,6 +23,7 @@
     qwen35_text_9b_pretrain_8gpu_gb200_bf16_deterministic_config,
     qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config,
     qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config,
+    qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_bf16_config,
     qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config,
 )
 
@@ -182,5 +183,6 @@
     "qwen35_text_35b_a3b_pretrain_config",
     "qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config",
     "qwen35_text_35b_a3b_sft_8gpu_gb200_bf16_dynamic_cp_config",
+    "qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_bf16_config",
     "qwen35_text_35b_a3b_sft_long_context_16gpu_gb200_fp8mx_config",
 ]