Enable Deepseek-v4 hybrid_model in dev branch Part (1/N)

hybrid_builders.py

@@ -2,8 +2,8 @@
 
 from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_inference_stack_spec
 from megatron.core.models.hybrid.hybrid_model import HybridModel
-from megatron.core.transformer import TransformerConfig
-from megatron.core.transformer.spec_utils import import_module
+from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.transformer.spec_utils import ModuleSpec, import_module
 from megatron.training import print_rank_0
 from megatron.training.arguments import core_transformer_config_from_args
 from model_provider import count_parameters_in_layer
@@ -13,6 +13,52 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None,
     print_rank_0('building Hybrid model ...')
     if config is None:
         config = core_transformer_config_from_args(args, TransformerConfig)
+    # MLA (and DSv4 hybrid) require MLATransformerConfig so that its __post_init__ runs
+    # the dsv4_hybrid derivation. The hybrid pretrain path can hand us a plain
+    # TransformerConfig, which silently skips that derivation; rebuild as MLA to match GPT.
+    if args.multi_latent_attention and not isinstance(config, MLATransformerConfig):
+        config = core_transformer_config_from_args(args)
+    # DSv4-hybrid head-dim contract: qk_head_dim and kv_lora_rank are derived from
+    # v_head_dim and qk_pos_emb_head_dim (MLATransformerConfig.__post_init__ does this for the
+    # GPT path). The hybrid config can reach here without that derivation applied, which breaks
+    # the MLA up-proj / fused-rope contract (q head dim must equal qk_head_dim + qk_pos_emb_head
+    # _dim == v_head_dim). Apply it for any DSv4 MLA attention: experimental_attention_variant
+    # == dsv4_hybrid, OR the layer pattern uses a DSv4 attention symbol (D/C/H/W). Idempotent.
+    _pattern = getattr(args, "hybrid_layer_pattern", None) or ""
+    _uses_dsv4_attn = (
+        getattr(args, "experimental_attention_variant", None) == "dsv4_hybrid"
+        or any(sym in _pattern for sym in ("C", "H", "W"))
+    )
+    if _uses_dsv4_attn:
+        derived = config.v_head_dim - config.qk_pos_emb_head_dim
+        if config.qk_head_dim != derived or config.kv_lora_rank != derived:
+            print_rank_0(
+                f"[hybrid dsv4] deriving qk_head_dim/kv_lora_rank = {config.v_head_dim} - "
+                f"{config.qk_pos_emb_head_dim} = {derived} (was qk_head_dim={config.qk_head_dim}, "
+                f"kv_lora_rank={config.kv_lora_rank})"
+            )
+            config.qk_head_dim = derived
+            config.kv_lora_rank = derived
+        # 'C'/'H'/'W' layers carry their compress ratio via the spec, but array-driven 'D' layers
+        # AND the indexer-loss logger (which counts ratio==4 layers) read
+        # config.csa_compress_ratios. When not given explicitly, derive it from the pattern
+        # symbols (C->4, H->128, W/D/other->0) so the array is consistent with the symbols and
+        # the indexer loss is normalized correctly; pad MTP depths with 0. An explicit
+        # --csa-compress-ratios is always respected.
+        if config.csa_compress_ratios is None:
+            ratio_map = {"C": 4, "H": 128}
+            # One entry per ACTUAL layer: main layers, then every MTP layer of every MTP depth
+            # (a depth can hold multiple hybrid layers, e.g. "/MD-E"), mirroring the arguments.py
+            # derivation. Padding by mtp_num_layers (depth count) would be too short and an MTP
+            # attention that isn't first would IndexError at num_layers + layer_number - 1.
+            sections = _pattern.split("/")
+            ratios = [ratio_map.get(c, 0) for c in sections[0].replace("|", "")]
+            for mtp_sec in sections[1:]:
+                ratios += [ratio_map.get(c, 0) for c in mtp_sec.replace("|", "")]
+            config.csa_compress_ratios = ratios
+            print_rank_0(
+                f"[hybrid dsv4] derived csa_compress_ratios from pattern symbols: {ratios}"
+            )
 
     if config.transformer_impl == "inference_optimized":
         hybrid_stack_spec = hybrid_inference_stack_spec
@@ -21,6 +67,11 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None,
         ), "inference_fuse_tp_communication is not supported for HybridModel"
     elif args.spec is not None:
         hybrid_stack_spec = import_module(args.spec)
+        # Allow config-aware specs: if --spec resolves to a callable (not a ModuleSpec),
+        # call it with config to build the stack spec (e.g. hybrid_dsv4_stack_spec, which
+        # wires the DSv4 CompressedSparseAttention into the 'D' layer per config).
+        if not isinstance(hybrid_stack_spec, ModuleSpec) and callable(hybrid_stack_spec):
+            hybrid_stack_spec = hybrid_stack_spec(config)
     else:
         raise ValueError("You must provide a valid hybrid layer spec via --spec")
 

megatron/core/extensions/transformer_engine.py

@@ -2011,12 +2011,28 @@ def __init__(
                 tp_group_for_te = None
 
             if is_te_min_version("2.14.0"):
-                extra_kwargs["single_grouped_weight"] = getattr(
-                    config, "moe_single_grouped_weight", False
-                )
-                extra_kwargs["single_grouped_bias"] = getattr(
-                    config, "moe_single_grouped_bias", False
-                )
+                # nemo_26.04 ships TE 2.14.0+71bbefbf whose GroupedLinear.__init__ does NOT
+                # yet accept single_grouped_{weight,bias}, even though the version string
+                # passes is_te_min_version("2.14.0"). Introspect the signature instead of
+                # version-gating, mirroring the patch in dsv4_fused_attn / main_megatron.
+                # The GroupedLinear.__init__ signature is constant for a given TE install, so
+                # introspect once and cache at module scope rather than on every TEGroupedLinear
+                # instantiation (matters for large MoE models with many expert groups).
+                global _TE_GROUPED_LINEAR_INIT_PARAMS
+                try:
+                    _gl_params = _TE_GROUPED_LINEAR_INIT_PARAMS
+                except NameError:
+                    _gl_params = _TE_GROUPED_LINEAR_INIT_PARAMS = set(
+                        inspect.signature(te.pytorch.GroupedLinear.__init__).parameters
+                    )
+                if "single_grouped_weight" in _gl_params:
+                    extra_kwargs["single_grouped_weight"] = getattr(
+                        config, "moe_single_grouped_weight", False
+                    )
+                if "single_grouped_bias" in _gl_params:
+                    extra_kwargs["single_grouped_bias"] = getattr(
+                        config, "moe_single_grouped_bias", False
+                    )
 
             self.te_quant_params: Optional[TEQuantizationParams] = None
             quant_config = get_quant_config_or_none(name, config.quant_recipe)