Merge branch 'main' into release/1.6

Author: tastelikefeet

requirements.txt

@@ -1,6 +1,6 @@
 megatron-core>=0.16,<0.20
 modelscope
-peft>=0.11,<0.20
+peft>=0.11,<0.21
 safetensors
 tqdm
-transformers>=4.33,<5.13.0
+transformers>=4.33,<5.15.0

src/mcore_bridge/bridge/gpt_bridge.py

@@ -1,5 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import math
+import os
 import re
 import torch
 import torch.distributed as dist
@@ -61,6 +62,7 @@ def __init__(self, config: ModelConfig):
         self._peft_format = False
         self._adapter_name = 'default'
         self._is_saving = False
+        self._source_model_dir = None
         self.model_type = config.hf_model_type
         self.llm_model_type = config.llm_model_type
         self.is_multimodal = config.is_multimodal
@@ -1915,6 +1917,8 @@ def load_weights(
         """
         self._peft_format = peft_format
         self._adapter_name = adapter_name
+        if not peft_format:
+            self._source_model_dir = hf_model_dir
         mg_models = unwrap_model(mg_models)
         self._disable_tqdm = False
         self._is_saving = False
@@ -1997,6 +2001,7 @@ def save_weights(
         adapter_name: str = 'default',
         converter: Optional[Callable] = None,
         max_shard_size: str = '5GB',
+        save_missing_weights: Union[bool, str] = False,
     ) -> None:
         """Save Megatron model checkpoint in safetensors (HuggingFace) format.
 
@@ -2013,10 +2018,15 @@ def save_weights(
             adapter_name: Name of the adapter for PEFT models. Defaults to 'default'.
             converter: Used to perform key-value conversion on the newly exported state_dict.
             max_shard_size: Maximum size of a single storage file, default is '5GB'.
+            save_missing_weights: Whether to copy tensors that exist in the source checkpoint but are
+                absent from the exported weights, such as submodules Megatron does not support. Pass a
+                path to specify the source checkpoint, otherwise the one recorded by `load_weights` is
+                used. Ignored when `peft_format` is True.
         """
         gc_collect()
         saver = StreamingSafetensorSaver(save_dir=output_dir, max_shard_size=max_shard_size, peft_format=peft_format)
         mg_models = unwrap_model(mg_models)
+        saved_keys = set()
         for k, v in self.export_weights(
                 mg_models,
                 target_device='cpu',
@@ -2028,9 +2038,37 @@ def save_weights(
                 disable_tqdm=False,
                 _is_saving=True):
             saver.add_tensor(k, v)
+            saved_keys.add(k)
+        if save_missing_weights and not peft_format:
+            source_model_dir = save_missing_weights if isinstance(save_missing_weights, str) else None
+            self._save_missing_weights(saver, saved_keys, source_model_dir)
         saver.finalize()
         dist.barrier()  # Ensure all weights are saved completely
 
+    def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None:
+        """Copy tensors present in the source checkpoint but absent from the exported ones.
+
+        Megatron only materializes the modules it knows about, so weights of unsupported
+        submodules (for instance the DSpark stages under `mtp.*`) would silently vanish
+        from the exported checkpoint. Restoring them verbatim keeps the saved model
+        functionally complete.
+        """
+        source_model_dir = source_model_dir or self._source_model_dir
+        if source_model_dir is None or not is_master():
+            return
+        if not os.path.isdir(source_model_dir):
+            logger.warning(f'Source model dir does not exist, skip restoring missing weights: {source_model_dir}')
+            return
+        with SafetensorLazyLoader(source_model_dir) as loader:
+            state_dict = loader.get_state_dict()
+            missing_keys = sorted(set(state_dict.keys()) - saved_keys)
+            if not missing_keys:
+                return
+            logger.info(f'Restoring {len(missing_keys)} weights from the source checkpoint '
+                        f'that were not exported by Megatron, e.g. {missing_keys[:3]}.')
+            for key in missing_keys:
+                saver.add_tensor(key, state_dict[key].load())
+
     @contextmanager
     def _patch_hf_initialize_weight(self):
 

src/mcore_bridge/config/model_config.py

@@ -196,6 +196,9 @@ class ModelConfig(TransformerConfig):
     attention_output_gate: bool = False
     linear_decoupled_in_proj: bool = False
 
+    # nemotron_h (hybrid mamba2 + attention + moe)
+    hybrid_layer_pattern: Optional[str] = None
+
     # dsa
     experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = None
     dsa_indexer_n_heads: Optional[int] = None

src/mcore_bridge/config/parser.py

@@ -1,6 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import torch.nn.functional as F
 from functools import partial
+from megatron.core.activations import squared_relu
 from transformers import PretrainedConfig
 from typing import Any, Dict
 
@@ -14,7 +15,7 @@
     'num_attention_heads': ['num_attention_heads'],
     'num_query_groups': ['num_key_value_heads'],
     'max_position_embeddings': ['max_position_embeddings'],
-    'layernorm_epsilon': ['rms_norm_eps'],
+    'layernorm_epsilon': ['rms_norm_eps', 'layer_norm_epsilon'],
     'rotary_base': ['rope_theta'],
     'padded_vocab_size': ['vocab_size'],
     'attention_dropout': ['attention_dropout'],
@@ -26,7 +27,7 @@
     'hf_model_type': ['model_type'],
     # moe
     'moe_ffn_hidden_size': ['moe_intermediate_size'],
-    'moe_shared_expert_intermediate_size': ['shared_expert_intermediate_size'],
+    'moe_shared_expert_intermediate_size': ['shared_expert_intermediate_size', 'moe_shared_expert_intermediate_size'],
     'moe_router_topk': ['num_experts_per_tok', 'moe_topk', 'moe_k', 'top_k_experts'],
     'moe_router_num_groups': ['n_group'],
     'moe_router_group_topk': ['topk_group'],
@@ -67,6 +68,14 @@
     'mhc_sinkhorn_iterations': ['hc_sinkhorn_iters'],
     'moe_n_hash_layers': ['mlp_layer_types'],
     'activation_func_clamp_value': ['swiglu_limit'],
+    # nemotron_h / mamba2
+    'mamba_num_heads': ['mamba_num_heads'],
+    'mamba_head_dim': ['mamba_head_dim'],
+    'mamba_state_dim': ['ssm_state_size', 'mamba_state_dim'],
+    'mamba_num_groups': ['n_groups', 'mamba_num_groups'],
+    'hybrid_layer_pattern': ['hybrid_override_pattern'],
+    'fp32_residual_connection': ['residual_in_fp32'],
+    'mtp_hybrid_override_pattern': ['mtp_hybrid_override_pattern'],
     # other
     'original_max_position_embeddings': ['original_max_position_embeddings'],
     'partial_rotary_factor': ['partial_rotary_factor'],
@@ -198,6 +207,15 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         res['swiglu'] = False
         res['gated_linear_unit'] = True
         res['activation_func'] = partial(F.gelu, approximate='tanh')
+    elif hf_model_type == 'muse_glimmer':
+        # 39 sliding layers (window 2048) interleaved with 13 full-attention layers; the latter are
+        # exactly the NoPE layers (`layer_rope_theta == 0`).
+        res['window_size'] = f'{window_size - 1},0'
+        window_attn_skip_freq = ','.join(['1' if lt == 'sliding_attention' else '0' for lt in layer_types])
+        res['window_attn_skip_freq'] = f'[{window_attn_skip_freq}]'
+        # The four per-layer norms are `CenteredRMSNorm` (`x * (1.0 + w)`). The final `norm` is a plain
+        # RMSNorm, so `MuseGlimmerLoader.build_model` opts that single module back out.
+        res['layernorm_zero_centered_gamma'] = True
     elif llm_model_type == 'gpt_oss':
         res['add_bias_linear'] = True
         res['bias_dropout_fusion'] = False
@@ -255,6 +273,19 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         res['add_qkv_bias'] = False
         res['moe_router_score_function'] = 'sigmoid'
         res['moe_router_load_balancing_type'] = 'seq_aux_loss'
+    elif llm_model_type == 'nemotron_h':
+        res['is_hybrid_model'] = True
+        res['position_embedding_type'] = 'none'
+        # relu^2 ("relu2") activation: non-gated, so fc1 is a single up_proj (no gate_proj).
+        res['swiglu'] = False
+        res['gated_linear_unit'] = False
+        res['activation_func'] = squared_relu
+        res['add_bias_linear'] = False
+        res['add_qkv_bias'] = False
+        res['qk_layernorm'] = False
+        res['moe_router_score_function'] = 'sigmoid'
+        res['moe_router_enable_expert_bias'] = True
+        res['moe_router_load_balancing_type'] = 'seq_aux_loss'
 
     if 'partial_rotary_factor' not in res and 'partial_rotary_factor' in rope_scaling:
         res['partial_rotary_factor'] = rope_scaling['partial_rotary_factor']

src/mcore_bridge/model/constant.py

@@ -12,6 +12,7 @@ class LLMModelType:
     bailing_hybrid = 'bailing_hybrid'
     deepseek_v4 = 'deepseek_v4'
     glm_moe_dsa = 'glm_moe_dsa'
+    nemotron_h = 'nemotron_h'
 
     qwen3_emb = 'qwen3_emb'
 
@@ -41,6 +42,8 @@ class MLLMModelType:
 
     minicpmv4_6 = 'minicpmv4_6'
 
+    muse_glimmer = 'muse_glimmer'
+
 
 class ModelType(LLMModelType, MLLMModelType):
     pass

src/mcore_bridge/model/gpt_model.py

@@ -203,6 +203,8 @@ def _preprocess(
         return decoder_input, mtp_decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
 
     def _set_inv_freq(self):
+        if getattr(self, 'rotary_pos_emb', None) is None:
+            return
         new_inv_freq, self.config.attention_scaling = get_rope_inv_freq(self.config)
         self.rotary_pos_emb.inv_freq = new_inv_freq.to(self.rotary_pos_emb.inv_freq.device)
 

src/mcore_bridge/model/gpts/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, olmoe, qwen3_emb,
-               qwen3_next)
+from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe,
+               qwen3_emb, qwen3_next)

src/mcore_bridge/model/gpts/deepseek_v4.py

@@ -60,6 +60,38 @@ def _patch_YarnRotaryEmbedding(config):
             delattr(config, attr)
 
 
+def _apply_mla_rope(t, freqs, *, config, cu_seqlens, cp_group, inverse=False):
+    """Apply DSv4's MLA RoPE to a tensor whose frequencies are already expanded per token.
+
+    `GPTModel` pre-indexes the rotary table by `position_ids`, so `freqs` is row-aligned with
+    `t`: row i holds the frequency of token i. That holds for every layout DSv4 supports --
+    unpacked, packed (thd), and packed CP: swift CP-splits `position_ids` with the same
+    partition mode as the hidden states, so the pre-indexed frequencies come out rank-local
+    while still carrying absolute positions. The multiply is therefore purely elementwise and
+    needs no `cu_seqlens`-based segment alignment, under any `cp_partition_mode`.
+
+    Enforcing that invariant here matters: when it does not hold, the generic
+    `apply_rotary_pos_emb` thd path re-derives positions from `cu_seqlens` assuming a *zigzag*
+    CP split, which is wrong for DSv4's contiguous split and would corrupt positions silently.
+    Asserting row alignment turns any future layout change into an immediate, explicit failure.
+    """
+    assert freqs.shape[0] == t.shape[0], (
+        f'DSv4 MLA RoPE expects per-token frequencies row-aligned with the input, got '
+        f'freqs.shape[0]={freqs.shape[0]} vs tokens={t.shape[0]}. `GPTModel` must pre-index the '
+        'rotary table by `position_ids` (requires `apply_rope_fusion=False`), and under CP the '
+        '`position_ids` must be split with the same partition mode as the hidden states.')
+    return apply_rotary_pos_emb(
+        t,
+        freqs,
+        config=config,
+        cu_seqlens=cu_seqlens,
+        cp_group=cp_group,
+        mla_rotary_interleaved=True,
+        mla_output_remove_interleaving=True,
+        inverse=inverse,
+    )
+
+
 class DSv4HybridSelfAttention(McoreDSv4HybridSelfAttention):
 
     def __init__(self, config, *args, **kwargs):
@@ -153,6 +185,15 @@ def qkv_up_proj_and_rope_apply(q_compressed,
             When sequence packing enabled, the input tensors adopt a packed shape of [t, ...];
             otherwise, they maintain the unpacked shape [s, b, ...]. In subsequent code comments,
             we uniformly use [num_tokens, ...] to denote [s, b, ...] or [t, ...] for two cases.
+
+            RoPE frequency layout: `GPTModel` pre-indexes the rotary table by `position_ids`
+            (see gpt_model.py, the `not apply_rope_fusion` branch), so `rotary_pos_emb` here is
+            already expanded per token -- row i belongs to token i -- rather than being a
+            position->frequency lookup table. Every RoPE call below is therefore an elementwise
+            multiply; see `_apply_mla_rope` for why that invariant is asserted. Under CP the
+            frequencies arrive rank-local because `position_ids` is split alongside the hidden
+            states, which is also why the boundary rows carry their own frequencies
+            (`boundary_rotary_pos_emb`) instead of being re-derived from positions.
             """
             # q_compressed: [num_tokens, q_lora_rank]
             # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)]
@@ -166,6 +207,8 @@ def qkv_up_proj_and_rope_apply(q_compressed,
             if boundary_kv_compressed is not None:
                 boundary_rows = boundary_kv_compressed.shape[0]
                 kv_projection_input = torch.cat([boundary_kv_compressed, kv_compressed], dim=0)
+                # The boundary rows precede this rank's block, so their frequencies must precede
+                # too -- keeping kv_rotary_pos_emb row-aligned with kv_projection_input.
                 kv_rotary_pos_emb = torch.cat([boundary_rotary_pos_emb, rotary_pos_emb], dim=0)
             else:
                 kv_projection_input = kv_compressed
@@ -182,29 +225,25 @@ def qkv_up_proj_and_rope_apply(q_compressed,
 
             # RoPE and query (shared for wkv and latent)
             # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim]
-            q_pos_emb = apply_rotary_pos_emb(
+            q_pos_emb = _apply_mla_rope(
                 q_pos_emb,
                 rotary_pos_emb,
                 config=self.config,
                 cu_seqlens=cu_seqlens_q,
                 cp_group=self.pg_collection.cp,
-                mla_rotary_interleaved=True,
-                mla_output_remove_interleaving=True,
             )
             # query: [num_tokens, n, (qk_head_dim + v_head_dim)]
             query = torch.cat([q_no_pe, q_pos_emb], dim=-1)
 
             kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1)
 
             # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim]
-            k_pos_emb = apply_rotary_pos_emb(
+            k_pos_emb = _apply_mla_rope(
                 k_pos_emb,
                 kv_rotary_pos_emb,
                 config=self.config,
                 cu_seqlens=cu_seqlens_kv,
                 cp_group=self.pg_collection.cp,
-                mla_rotary_interleaved=True,
-                mla_output_remove_interleaving=True,
             )
 
             # Single head: key = value = [num_tokens, 1, v_head_dim]
@@ -384,15 +423,13 @@ def forward(
             rot_part_in = rot_part.squeeze(1)
         else:
             rot_part_in = rot_part
-        rot_part_out = apply_rotary_pos_emb(
+        rot_part_out = _apply_mla_rope(
             rot_part_in,
             rotary_pos_emb,
-            self.config,
+            config=self.config,
             cu_seqlens=cu_seqlens_kv,
             cp_group=self.pg_collection.cp,
-            mla_rotary_interleaved=True,
             inverse=True,
-            mla_output_remove_interleaving=True,
         )
         if packed_seq:
             rot_part = rot_part_out.unsqueeze(1)