Merge remote-tracking branch 'origin/main' into codex/pr5479-hybrid-model

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/models/conversion/model_bridge.py

@@ -1897,17 +1897,17 @@ def build_conversion_tasks(
         hf_pretrained: HFPreTrained,
         megatron_model: List[MegatronModel],
         weight_dtype: Optional[torch.dtype] = None,
-    ) -> List[None | WeightConversionTask]:
+    ) -> List[WeightConversionTask]:
         """Construct the conversion tasks between HF and megatron.
 
         Args:
             weight_dtype: Export dtype recorded on each task. Overrides must forward it.
 
         The algorithm walks over every parameter of every destination model,
         asks the :class:`MegatronMappingRegistry` whether it has a mapping for that
-        parameter, and – if the corresponding HF weights actually exist – yields
-        an :class:`_HFLoadTask` describing exactly how that parameter will be
-        populated.
+        parameter and returns a concrete task describing exactly how that
+        parameter will be populated. Missing mappings or source weights are
+        conversion errors, not empty task slots.
         """
 
         has_hf_state = hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")
@@ -1934,9 +1934,15 @@ def build_conversion_tasks(
                 name for name in sorted_global_param_names_all_pp_ranks if "output_layer" not in name
             ]
 
+        mappings_by_global_name = self._validate_conversion_mappings(
+            mapping_registry,
+            sorted_global_param_names_all_pp_ranks,
+            hf_keys,
+        )
+
         global_names_index_dict = {name: idx for idx, name in enumerate(sorted_global_param_names_all_pp_ranks)}
 
-        tasks = [None] * len(sorted_global_param_names_all_pp_ranks)
+        pending_tasks: list[WeightConversionTask | None] = [None] * len(sorted_global_param_names_all_pp_ranks)
         for vp_stage, model in enumerate(megatron_model):
             # persistent buffers are part of the model's state_dict, but not the named_parameters, so we must include them here separately
             for local_name, _ in itertools.chain(model.named_parameters(), persistent_buffers(model)):
@@ -1950,34 +1956,15 @@ def build_conversion_tasks(
                     print_rank_0(f"WARNING: {global_name} not in global_names_index_dict")
                     continue
                 global_name_idx = global_names_index_dict[global_name]
-                mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
-
-                if not mapping:
-                    logger.warning(f"WARNING: No mapping found for megatron_param: {global_name}")
-                    continue
-                # Ensure hf weights exist (skip for config-only export where hf_keys is None)
-                if hf_keys is not None and not mapping.allow_hf_name_mismatch:
-                    if isinstance(mapping.hf_param, str):
-                        if mapping.hf_param not in hf_keys:
-                            logger.warning(f"WARNING: Can't find {mapping.hf_param} in hf_keys")
-                            continue
-                    else:
-                        missing_params = [
-                            hf_param for hf_param in mapping.hf_param.values() if hf_param not in hf_keys
-                        ]
-                        if missing_params:
-                            logger.warning(
-                                f"WARNING: Can't find the following HF parameters in hf_keys: {missing_params}"
-                            )
-                            continue
+                mapping = mappings_by_global_name[global_name]
 
                 local_module, local_weights = get_module_and_param_from_name(megatron_model, local_name, vp_stage)
                 if local_module is not None and not hasattr(local_module, "config"):
                     # If module is not a MegatronModule (e.g. torch.nn.Conv1d or a module list) we need
                     # to get the config from the model
                     setattr(local_module, "config", model_config)
 
-                tasks[global_name_idx] = WeightConversionTask(
+                pending_tasks[global_name_idx] = WeightConversionTask(
                     pp_rank=pp_rank,
                     vp_stage=vp_stage,
                     param_name=local_name,
@@ -1990,15 +1977,12 @@ def build_conversion_tasks(
 
         # Fill the remaining ones for pp communications
         for idx, global_name in enumerate(sorted_global_param_names_all_pp_ranks):
-            if tasks[idx] is None:
-                mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
-                # Skip tasks with no mapping found
-                if mapping is None:
-                    continue
+            if pending_tasks[idx] is None:
+                mapping = mappings_by_global_name[global_name]
                 # This is an exception here we pass in global name
                 # we are not using global_name to extract module and weights
                 # only use it for param mapping auto dispatch checks
-                tasks[idx] = WeightConversionTask(
+                pending_tasks[idx] = WeightConversionTask(
                     pp_rank=pp_rank,
                     vp_stage=None,
                     param_name=global_name,
@@ -2009,6 +1993,68 @@ def build_conversion_tasks(
                     weight_dtype=weight_dtype,
                 )
 
+        return self._require_concrete_tasks(pending_tasks)
+
+    def _validate_conversion_mappings(
+        self,
+        mapping_registry: MegatronMappingRegistry,
+        global_param_names: Iterable[str],
+        hf_keys: Iterable[str] | None = None,
+    ) -> dict[str, MegatronParamMapping]:
+        """Resolve and validate mappings for the full cross-PP parameter list."""
+        mappings_by_global_name: dict[str, MegatronParamMapping] = {}
+        missing_mappings: list[str] = []
+        missing_hf_weights: list[tuple[str, str]] = []
+        hf_key_set = set(hf_keys) if hf_keys is not None else None
+
+        for global_name in global_param_names:
+            mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
+            if mapping is None:
+                missing_mappings.append(global_name)
+                continue
+
+            mappings_by_global_name[global_name] = mapping
+            if hf_key_set is None or mapping.allow_hf_name_mismatch:
+                continue
+
+            expected_hf_names = (
+                [mapping.hf_param] if isinstance(mapping.hf_param, str) else list(mapping.hf_param.values())
+            )
+            missing_hf_weights.extend(
+                (global_name, hf_name) for hf_name in expected_hf_names if hf_name not in hf_key_set
+            )
+
+        if missing_mappings:
+            missing_names = "\n  ".join(missing_mappings)
+            raise ValueError(
+                "No mapping found for the following Megatron parameter(s):\n"
+                f"  {missing_names}\n"
+                "Every global Megatron parameter must have a concrete mapping so import and export remain strict."
+            )
+
+        if missing_hf_weights:
+            missing_names = "\n  ".join(f"{global_name} -> {hf_name}" for global_name, hf_name in missing_hf_weights)
+            raise ValueError(
+                "Hugging Face checkpoint is missing mapped parameter(s):\n"
+                f"  {missing_names}\n"
+                "If the HF config determines whether the weight exists, register the mapping "
+                "conditionally on that config instead. If it does not, and the name is synthesized "
+                "or the weight is absent on only some layers, set allow_hf_name_mismatch on the "
+                "mapping."
+            )
+
+        return mappings_by_global_name
+
+    @staticmethod
+    def _require_concrete_tasks(
+        pending_tasks: Iterable[WeightConversionTask | None],
+    ) -> list[WeightConversionTask]:
+        """Return tasks after enforcing the internal no-empty-slot invariant."""
+        tasks: list[WeightConversionTask] = []
+        for task in pending_tasks:
+            if task is None:
+                raise RuntimeError("Internal error: conversion task construction left an empty slot")
+            tasks.append(task)
         return tasks
 
     def _detect_fp8_params(
@@ -2103,7 +2149,7 @@ def build_export_fp8_tasks(
         *,
         scale_inv_suffix: str = "_scale_inv",
         fp8_scale_inv_attr: str = "_rowwise_scale_inv",
-    ) -> List[None | WeightConversionTask]:
+    ) -> List[WeightConversionTask]:
         """
         Build Megatron→(export) conversion tasks, inserting extra *scale_inv* tasks for blockwise FP8 params.
         """
@@ -2128,6 +2174,11 @@ def build_export_fp8_tasks(
                 name for name in sorted_global_param_names_all_pp_ranks if "output_layer" not in name
             ]
 
+        mappings_by_global_name = self._validate_conversion_mappings(
+            mapping_registry,
+            sorted_global_param_names_all_pp_ranks,
+        )
+
         # 1) Determine which global params are blockwise FP8 and gather flags across PP ranks
         global_fp8_flags = self._detect_fp8_params(
             megatron_model,
@@ -2160,10 +2211,7 @@ def build_export_fp8_tasks(
                 if global_name not in global_names_index_dict:
                     continue
 
-                mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
-                if not mapping:
-                    logger.warning(f"WARNING: No mapping found for megatron_param: {global_name}")
-                    continue
+                mapping = mappings_by_global_name[global_name]
                 local_module, local_weights = get_module_and_param_from_name(megatron_model, local_name, vp_stage)
                 if local_module is not None and not hasattr(local_module, "config"):
                     setattr(local_module, "config", model_config)
@@ -2221,22 +2269,16 @@ def build_export_fp8_tasks(
             # For scale_inv entries, reuse the base param's mapping type.
             if global_name.endswith(scale_inv_suffix):
                 base_global_name = global_name[: -len(scale_inv_suffix)]
-                base_mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(base_global_name))
-                if base_mapping is not None:
-                    # clone mapping instance to avoid sharing state across tasks.
-                    base_mapping_for_scale = mapping_registry.resolve_mapping(base_mapping, ())
-                    mapping = _HFNameSuffixMapping(
-                        base_mapping_for_scale,
-                        scale_inv_suffix,
-                        self._fp8_scale_block_size(global_fp8_flags.get(base_global_name)),
-                    )
-                else:
-                    mapping = None
+                base_mapping = mappings_by_global_name[base_global_name]
+                # clone mapping instance to avoid sharing state across tasks.
+                base_mapping_for_scale = mapping_registry.resolve_mapping(base_mapping, ())
+                mapping = _HFNameSuffixMapping(
+                    base_mapping_for_scale,
+                    scale_inv_suffix,
+                    self._fp8_scale_block_size(global_fp8_flags.get(base_global_name)),
+                )
             else:
-                mapping = mapping_registry.megatron_to_hf_lookup(self._get_lora_unwrapped_name(global_name))
-            if mapping is None:
-                logger.warning(f"No mapping found for global_name: {global_name}")
-                continue
+                mapping = mappings_by_global_name[global_name]
 
             tasks[idx] = WeightConversionTask(
                 pp_rank=pp_rank,
@@ -2248,7 +2290,7 @@ def build_export_fp8_tasks(
                 mapping=mapping,
             )
 
-        return tasks
+        return self._require_concrete_tasks(tasks)
 
     @staticmethod
     def _fp8_scale_block_size(fp8_flag: bool | int | None) -> int | None:

src/megatron/bridge/models/conversion/param_mapping.py

@@ -125,8 +125,11 @@ def __init__(self, megatron_param: str, hf_param: Union[str, Dict[str, str]]):
             self._tp_group = None
             self._etp_group = None
 
-        # if a param mapping class takes in modified HF weight name from maybe_modify_loaded_hf_weight,
-        # allow_hf_name_mismatch should be set to True to bypass a check in `build_conversion_tasks`
+        # Set allow_hf_name_mismatch to True when the declared HF name will not be found verbatim
+        # in the checkpoint's key set. That covers two cases: a name that is rewritten or
+        # synthesized (see maybe_modify_loaded_hf_weight), and a weight that is legitimately
+        # absent for some layers or configurations. Both bypass the hf_keys check in
+        # `build_conversion_tasks`, which raises otherwise.
         self.allow_hf_name_mismatch = False
 
     def set_process_groups_from_pg_collection(self, pg_collection: Any) -> None:

src/megatron/bridge/models/deepseek/attention.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MLA attention spec helpers for the DeepSeek family."""
+
+from dataclasses import replace
+from typing import Optional
+
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mla_qk_norm_config import get_backend
+from megatron.core.transformer.multi_latent_attention import MLASelfAttention
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class MLASelfAttentionWithoutQueryNorm(MLASelfAttention):
+    """MLA self-attention that does not add a query norm when there is no query LoRA.
+
+    MCore derives Q and KV normalization from a single ``qk_layernorm`` flag. DeepSeek
+    needs it enabled for ``kv_a_layernorm``, which every checkpoint ships. When
+    ``q_lora_rank`` is None, that same flag also makes MCore fuse a query normalization
+    into ``linear_q_proj`` (``QKNormConfigResolver._resolve_mla_qk_layernorm``), but the
+    HF architecture defines no query-side norm in that case: ``DeepseekV3Attention``
+    builds a bare ``q_proj``.
+
+    The result is a trainable parameter with no HF counterpart, which cannot be loaded
+    and is silently dropped on export. This subclass keeps the KV norm and drops the
+    query norm so the converted model matches the source architecture.
+
+    Transformer Engine is required for the no-query-LoRA case. MCore builds
+    ``linear_q_proj`` from the backend's fused norm+linear implementation, which only
+    Transformer Engine provides, so the local backend is rejected with an explicit
+    message rather than an internal one.
+    """
+
+    def _resolve_qk_norm_config(self, submodules):
+        """Replace the fused query projection with a plain one when there is no query LoRA.
+
+        The standalone-``q_layernorm`` case is neutralised *before* delegating: an MLA spec
+        may set ``q_layernorm`` to a real norm whenever ``qk_layernorm`` is on, and the
+        parent resolver rejects that outright when there is no query LoRA to consume it
+        (``_raise_unused_q_norm``). Dropping the query norm is exactly what this class
+        exists to do, so the rejection would fire on a configuration this class already
+        knows how to satisfy.
+        """
+        if self.config.q_lora_rank is not None:
+            return super()._resolve_qk_norm_config(submodules)
+
+        backend = get_backend(self.config.transformer_impl)
+        if backend.column_parallel_layer_norm_linear() is None:
+            raise ValueError(
+                "DeepSeek without a query LoRA (`q_lora_rank=None`) requires "
+                f"`transformer_impl='transformer_engine'`; `{self.config.transformer_impl}` "
+                "provides no fused norm+linear projection. MCore's MLA resolver builds "
+                "`linear_q_proj` from that fused implementation whenever `qk_layernorm` is "
+                "on, and DeepSeek needs `qk_layernorm` on for `kv_a_layernorm`, so this "
+                "backend cannot express the architecture."
+            )
+
+        if submodules.q_layernorm not in (None, IdentityOp):
+            submodules = replace(submodules, q_layernorm=IdentityOp)
+
+        layer_classes = super()._resolve_qk_norm_config(submodules)
+        layer_classes["linear_q_proj"] = backend.column_parallel_linear()
+        return layer_classes
+
+
+def get_deepseek_decoder_block_spec(
+    config: TransformerConfig,
+    use_transformer_engine: bool,
+    normalization: Optional[str] = None,
+    qk_l2_norm: Optional[bool] = False,
+    vp_stage: Optional[int] = None,
+    pp_rank: Optional[int] = None,
+) -> ModuleSpec:
+    """Build the decoder block spec, omitting the query norm when ``q_lora_rank`` is None.
+
+    The signature mirrors ``get_gpt_decoder_block_spec`` exactly, including ``vp_stage``
+    and ``pp_rank``. ``GPTModelProvider.provide()`` inspects the callable's parameters and
+    only forwards ``vp_stage`` when it is declared, so dropping it here would leave
+    interleaved pipeline parallelism calling MCore's layer-offset helper without a virtual
+    stage, which asserts.
+
+    Args:
+        config: The model provider / transformer config.
+        use_transformer_engine: Whether to build Transformer Engine submodules.
+        normalization: Optional normalization override, forwarded unchanged.
+        qk_l2_norm: Optional QK L2 norm flag, forwarded unchanged.
+        vp_stage: Virtual pipeline stage, forwarded unchanged.
+        pp_rank: Pipeline rank, forwarded unchanged.
+
+    Returns:
+        The decoder block spec, with MLA self-attention replaced by
+        :class:`MLASelfAttentionWithoutQueryNorm` when there is no query LoRA.
+    """
+    spec = get_gpt_decoder_block_spec(
+        config,
+        use_transformer_engine=use_transformer_engine,
+        normalization=normalization,
+        qk_l2_norm=qk_l2_norm,
+        vp_stage=vp_stage,
+        pp_rank=pp_rank,
+    )
+    return replace_mla_self_attention(config, spec)
+
+
+def replace_mla_self_attention(config: TransformerConfig, spec: ModuleSpec) -> ModuleSpec:
+    """Swap MLA self-attention for the query-norm-free variant, in place, on every layer.
+
+    Shared with the MTP path: a standalone MTP pipeline stage owns no decoder layers, so
+    the provider re-derives a layer spec straight from MCore and never passes through
+    :func:`get_deepseek_decoder_block_spec`. Without this the MTP layer regains the query
+    norm that the decoder layers just dropped.
+
+    Accepts either a block spec (``.layer_specs``) or a single layer spec.
+    """
+    if getattr(config, "q_lora_rank", None) is not None:
+        return spec
+
+    layer_specs = getattr(spec, "layer_specs", None)
+    for layer_spec in layer_specs if layer_specs is not None else [spec]:
+        self_attention = getattr(layer_spec.submodules, "self_attention", None)
+        if self_attention is not None and self_attention.module is MLASelfAttention:
+            self_attention.module = MLASelfAttentionWithoutQueryNorm
+    return spec