Support dpsk-4.1 sft and rl

src/mcore_bridge/__init__.py

@@ -9,15 +9,15 @@
 
 if TYPE_CHECKING:
     from .bridge import GPTBridge
-    from .config import ModelConfig, hf_to_mcore_config
+    from .config import MLAModelConfig, ModelConfig, hf_to_mcore_config
     from .model import get_mcore_model
     from .tuners import LoraParallelLinear
     from .utils import get_logger, set_random_seed, split_cp_inputs, unwrap_model
     from .version import __release_datetime__, __version__
 else:
     _import_structure = {
         'bridge': ['GPTBridge'],
-        'config': ['ModelConfig', 'hf_to_mcore_config'],
+        'config': ['MLAModelConfig', 'ModelConfig', 'hf_to_mcore_config'],
         'model': ['get_mcore_model'],
         'tuners': ['LoraParallelLinear'],
         'utils': ['get_logger', 'set_random_seed', 'split_cp_inputs', 'unwrap_model'],

src/mcore_bridge/bridge/gpt_bridge.py

@@ -1756,7 +1756,11 @@ def _set_hyper_connection(self, mg_layer, hf_state_dict, layer_idx, to_mcore):
             self._set_state_dict(hyper_connection, 'bias', hf_state_dict, f'hc_{hf_key}_base', to_mcore)
             has_hyper_connection = hyper_connection is not None
             has_hyper_connection = self._reduce_tensor_pp_group(has_hyper_connection, to_mcore)
-            if has_hyper_connection:
+            # ``alpha_*`` are frozen base parameters written outside ``_set_state_dict``, so they
+            # need the peft guard the mapping_proj/bias calls above get for free -- otherwise a
+            # LoRA export writes base weights into ``adapter_model.safetensors`` and a LoRA load
+            # demands a key the adapter does not carry. Same shape as the Engram export guard.
+            if has_hyper_connection and not self._peft_format:
                 if to_mcore:
                     alpha = hf_state_dict[f'hc_{hf_key}_scale'].load()
                     for i, alpha_suffix in enumerate(['pre', 'post', 'res']):
@@ -1844,6 +1848,10 @@ def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore):
         self._set_state_dict(lm_model, 'decoder.final_layernorm.weight', hf_state_dict, self.hf_final_layernorm_key,
                              to_mcore)
 
+    def _convert_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage):
+        """Extension point for model-specific auxiliary stacks outside standard MTP."""
+        return ()
+
     def _convert_hf_state_dict(self, hf_state_dict, to_mcore):
         res = {}
         for k, v in hf_state_dict.items():
@@ -1929,6 +1937,8 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd
                     res = self._convert_hf_state_dict(res, to_mcore)
                     yield from list(self._add_prefix(res, hf_prefix).items())
                     hf_state_dict = {}
+        yield from self._convert_additional_layers(
+            mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage)
         if not to_mcore or is_pp_last_stage:
             hf_state_dict.update(self._convert_post_process(mg_model, hf_state_dict, '', to_mcore))
         if to_mcore:
@@ -2049,6 +2059,7 @@ def export_weights(
         tqdm_desc: str = 'Exporting: ',
         disable_tqdm: bool = True,
         _is_saving: bool = False,
+        skip_unsupported_export: bool = False,
     ):
         """Export Megatron model weights to safetensors (HuggingFace) format as a generator.
 
@@ -2066,6 +2077,11 @@ def export_weights(
             converter: Used to perform key-value conversion on the newly exported state_dict.
             tqdm_desc: Description text for the progress bar. Defaults to 'Exporting: '.
             disable_tqdm: Whether to disable the tqdm progress bar. Defaults to True.
+            skip_unsupported_export: When True, weights whose Megatron->HF export is not implemented
+                (e.g. DeepSeek-V4.1 Engram tables, which are frozen during on-policy RL and already
+                loaded in the rollout engine) are silently skipped instead of raising. Used by the RL
+                weight-sync path; the checkpoint-save path keeps the default (False) so a saved HF
+                checkpoint stays complete.
 
         Yields:
             Tuple[str, torch.Tensor]: Key-value pairs of parameter names and tensors.
@@ -2076,6 +2092,7 @@ def export_weights(
         self._adapter_name = adapter_name
         self._disable_tqdm = disable_tqdm
         self._is_saving = _is_saving
+        self._skip_unsupported_export = skip_unsupported_export
         self._peft_target_modules = set()
         self._peft_modules_to_save = set()
         self._fp8_skip_modules = set()
@@ -2146,6 +2163,10 @@ def save_weights(
         saver.finalize()
         dist.barrier()  # Ensure all weights are saved completely
 
+    def _normalize_missing_weight_key(self, key: str) -> str:
+        """Return the identity used to detect aliases while restoring source-only weights."""
+        return key
+
     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.
 
@@ -2162,7 +2183,9 @@ def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> Non
             return
         with SafetensorLazyLoader(source_model_dir) as loader:
             state_dict = loader.get_state_dict()
-            missing_keys = sorted(set(state_dict.keys()) - saved_keys)
+            saved_identities = {self._normalize_missing_weight_key(key) for key in saved_keys}
+            missing_keys = sorted(
+                key for key in state_dict if self._normalize_missing_weight_key(key) not in saved_identities)
             if not missing_keys:
                 return
             logger.info(f'Restoring {len(missing_keys)} weights from the source checkpoint '

src/mcore_bridge/config/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .model_config import ModelConfig
+from .model_config import MLAModelConfig, ModelConfig
 from .parser import hf_to_mcore_config

src/mcore_bridge/config/model_config.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from megatron.core import mpu
 from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.transformer_config import MLATransformerConfig
 from transformers import PretrainedConfig
 from transformers.utils import is_torch_npu_available
 from transformers.utils.versions import require_version
@@ -173,6 +174,10 @@ class ModelConfig(TransformerConfig):
     moe_router_score_function: Literal['sigmoid', 'softmax'] = 'softmax'
     moe_router_bias_update_rate: float = 1e-3
     moe_router_enable_expert_bias: bool = False
+    # Model-specific VL routing belongs to mcore-bridge rather than Megatron-Core.
+    # DeepSeek-V4.1 selects a separately checkpointed correction bias for image tokens.
+    moe_router_enable_vl_bias: bool = False
+    image_token_id: Optional[int] = None
     moe_router_topk_scaling_factor: Optional[float] = None
     # 'aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'
     moe_router_load_balancing_type: Union[str, List[str]] = 'aux_loss'
@@ -251,6 +256,31 @@ class ModelConfig(TransformerConfig):
     mhc_init_gating_factor: float = 0.01
     moe_n_hash_layers: int = 0
 
+    # deepseek-v4.1 engram (HF layer IDs are 0-based)
+    # Declared here as well so the bridge remains importable on the PR #7224 baseline,
+    # where NVIDIA's optional Engram extension is not installed.
+    engram_enabled: bool = False
+    engram_layer_ids: Optional[List[int]] = None
+    engram_num_embeddings: Optional[List[int]] = None
+    engram_max_ngram_size: Optional[int] = None
+    engram_vocab_size: Optional[int] = None
+    engram_n_heads: Optional[int] = None
+    engram_head_dim: Optional[int] = None
+    engram_pad_token_id: Optional[int] = None
+    engram_compressed_vocab_size: Optional[int] = None
+    engram_tokenizer_map: Optional[str] = None
+
+    # DeepSeek-V4.1 DSpark. This is intentionally separate from mtp_num_layers:
+    # DSpark drafts a block in parallel and adds Markov/confidence heads, whereas
+    # Megatron MTP predicts successive tokens autoregressively.
+    dspark_num_layers: Optional[int] = None
+    dspark_block_size: int = 0
+    dspark_noise_token_id: Optional[int] = None
+    dspark_target_layer_ids: Optional[List[int]] = None
+    dspark_markov_rank: Optional[int] = None
+    dspark_num_experts: Optional[int] = None
+    dspark_router_topk: Optional[int] = None
+
     # mtp
     mtp_decoder_input_detach: bool = False
     mtp_shared_weights: bool = False
@@ -288,7 +318,6 @@ def _augment_mindspeed_defaults(self):
             defaults = {}
             try:
                 import mindspeed.features_manager as mfm
-                import sys
                 from argparse import ArgumentParser
                 from mindspeed.arguments import process_args
 
@@ -333,6 +362,11 @@ def __post_init__(self):
         if self.num_moe_experts is not None:
             if self.moe_ffn_hidden_size is None:
                 self.moe_ffn_hidden_size = self.ffn_hidden_size
+        if self.moe_router_enable_vl_bias:
+            if not self.moe_router_enable_expert_bias:
+                raise ValueError('VL expert bias requires moe_router_enable_expert_bias.')
+            if self.image_token_id is None:
+                raise ValueError('VL expert bias requires image_token_id.')
         if self.rope_scaling is not None:
             self.rope_scaling = json_parse_to_dict(self.rope_scaling)
             if 'type' in self.rope_scaling and 'rope_type' not in self.rope_scaling:
@@ -367,6 +401,39 @@ def __post_init__(self):
             self.mtp_num_layers = 1
         else:
             self.mtp_unroll_steps = self.mtp_num_layers
+        # ``num_nextn_predict_layers`` counts the draft layers for both DeepSeek's standard MTP and
+        # V4.1's DSpark, so it alone cannot tell them apart: on a plain V3/V4 checkpoint it means MTP
+        # and there is no DSpark at all. ``dspark_block_size`` is what actually marks a DSpark
+        # checkpoint (V4.1-Flash and the V4 Vision experiment carry it; plain V4-Flash does not).
+        if not self.dspark_block_size:
+            self.dspark_num_layers = None
+        if self.dspark_block_size:
+            required_dspark = {
+                'dspark_num_layers': self.dspark_num_layers,
+                'dspark_block_size': self.dspark_block_size,
+                'dspark_noise_token_id': self.dspark_noise_token_id,
+                'dspark_target_layer_ids': self.dspark_target_layer_ids,
+                'dspark_markov_rank': self.dspark_markov_rank,
+            }
+            missing_dspark = [name for name, value in required_dspark.items() if value is None]
+            if missing_dspark:
+                raise ValueError(f'DSpark config is missing required fields: {missing_dspark}.')
+            if self.dspark_num_layers <= 0 or self.dspark_block_size <= 0 or self.dspark_markov_rank <= 0:
+                raise ValueError('DSpark layer count, block size and Markov rank must all be positive.')
+            if not self.dspark_target_layer_ids:
+                raise ValueError('DSpark requires at least one target layer ID.')
+            if len(set(self.dspark_target_layer_ids)) != len(self.dspark_target_layer_ids):
+                raise ValueError('DSpark target layer IDs must be unique.')
+            if min(self.dspark_target_layer_ids) < 0 or max(self.dspark_target_layer_ids) >= self.num_layers:
+                raise ValueError('DSpark target layer IDs must refer to decoder layers.')
+            if self.dspark_noise_token_id < 0 or self.dspark_noise_token_id >= self.padded_vocab_size:
+                raise ValueError('DSpark noise token ID must be inside the padded vocabulary.')
+            # The draft stack's own expert counts are optional -- a checkpoint that omits them (as the
+            # V4 Vision experiment does) means its draft layers reuse the backbone's MoE shape, which
+            # is where the builder falls back to.
+            if self.dspark_num_experts is not None and self.dspark_router_topk is not None:
+                if self.dspark_num_experts <= 0 or not 0 < self.dspark_router_topk <= self.dspark_num_experts:
+                    raise ValueError('DSpark router top-k must be positive and no larger than its expert count.')
         if self.csa_compress_ratios is not None and self.mtp_num_layers is not None:
             self.csa_compress_ratios += [0] * self.mtp_num_layers
         if self.multi_latent_attention:
@@ -432,3 +499,8 @@ def __deepcopy__(self, memo):
             else:
                 setattr(new_obj, k, copy.deepcopy(v, memo))
         return new_obj
+
+
+@dataclass
+class MLAModelConfig(ModelConfig, MLATransformerConfig):
+    """ModelConfig variant for models requiring native Megatron MLA semantics."""

src/mcore_bridge/config/parser.py

@@ -25,6 +25,7 @@
     'add_bias_linear': ['mlp_bias'],
     'kv_channels': ['head_dim'],
     'hf_model_type': ['model_type'],
+    'image_token_id': ['image_token_id'],
     # moe
     'moe_ffn_hidden_size': ['moe_intermediate_size'],
     'moe_shared_expert_intermediate_size': ['shared_expert_intermediate_size', 'moe_shared_expert_intermediate_size'],
@@ -85,11 +86,34 @@
     # deepseek_v4
     'csa_compress_ratios': ['compress_rates'],
     'csa_compress_rotary_base': ['compress_rope_theta'],
+    # deepseek_v41 / CSA2 source-layer routing
+    'csa2_kv_source_layers': ['kv_source_layer_ids'],
+    'csa2_index_source_layers': ['index_source_layer_ids'],
+    'csa2_candidate_source_layer': ['candidate_source_layer_id'],
+    'csa2_candidate_topk_blocks': ['candidate_topk_blocks'],
+    'csa2_candidate_block_size': ['candidate_block_size'],
     'o_groups': ['o_groups'],
     'o_lora_rank': ['o_lora_rank'],
     'num_residual_streams': ['hc_mult'],
     'mhc_sinkhorn_iterations': ['hc_sinkhorn_iters'],
     'moe_n_hash_layers': ['mlp_layer_types'],
+    'engram_layer_ids': ['engram_layer_ids'],
+    'engram_num_embeddings': ['engram_num_embeddings'],
+    'engram_max_ngram_size': ['engram_max_ngram_size'],
+    'engram_vocab_size': ['engram_vocab_size'],
+    'engram_n_heads': ['engram_n_heads'],
+    'engram_head_dim': ['engram_head_dim'],
+    'engram_pad_token_id': ['engram_pad_token_id'],
+    'engram_compressed_vocab_size': ['engram_compressed_vocab_size'],
+    'engram_tokenizer_map': ['engram_tokenizer_map'],
+    # DeepSeek-V4.1 DSpark is a parallel draft stack, not Megatron's autoregressive MTP.
+    'dspark_num_layers': ['num_nextn_predict_layers'],
+    'dspark_block_size': ['dspark_block_size'],
+    'dspark_noise_token_id': ['dspark_noise_token_id'],
+    'dspark_target_layer_ids': ['dspark_target_layer_ids'],
+    'dspark_markov_rank': ['dspark_markov_rank'],
+    'dspark_num_experts': ['dspark_n_routed_experts'],
+    'dspark_router_topk': ['dspark_num_experts_per_tok', 'dspark_n_activated_experts'],
     'activation_func_clamp_value': ['swiglu_limit'],
     # nemotron_h / mamba2
     'mamba_num_heads': ['mamba_num_heads'],
@@ -191,8 +215,8 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         res.pop('ffn_hidden_size', None)
         if llm_model_type in {'qwen2_moe', 'qwen3_next'} or hf_model_type == 'qwen3_5_moe':
             res['moe_shared_expert_gate'] = True
-    if llm_model_type in {'deepseek', 'deepseek_v2', 'deepseek_v3', 'kimi_k2', 'deepseek_v32', 'dots1', 'deepseek_v4'
-                          } or hf_model_type == 'kimi_vl':
+    if llm_model_type in {'deepseek', 'deepseek_v2', 'deepseek_v3', 'kimi_k2', 'deepseek_v32', 'dots1', 'deepseek_v4',
+                           'deepseek_v41_text'} or hf_model_type == 'kimi_vl':
         if llm_model_type != 'deepseek':
             res['qk_layernorm'] = True
         res['moe_router_load_balancing_type'] = 'seq_aux_loss'
@@ -210,6 +234,28 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
             csa_compress_ratios = res.pop('csa_compress_ratios', None)
             res['csa_compress_ratios'] = [csa_compress_ratios.get(layer_type, 0) for layer_type in layer_types]
             res['moe_n_hash_layers'] = len([layer for layer in moe_n_hash_layers if layer == 'hash_moe'])
+        elif llm_model_type == 'deepseek_v41_text':
+            if 'v_head_dim' not in res:
+                res['v_head_dim'] = res['kv_channels']
+            res['experimental_attention_variant'] = 'dsv4_hybrid'
+            res['dsv4_version'] = 'v4.1'
+            # Native V4.1 uses unrotated indexer activations and no YaRN
+            # amplitude scaling (generic MLA defaults differ).
+            res['dsa_indexer_rotate_activation'] = False
+            res['mscale'] = 0.0
+            res['mscale_all_dim'] = 0.0
+            res['moe_router_enable_expert_bias'] = True
+            res['moe_router_enable_vl_bias'] = getattr(hf_config, 'vision_config', None) is not None
+            res['csa_window_size'] = window_size
+            res['enable_hyper_connections'] = True
+            res['mhc_single_pass'] = True
+            # CSA2 consumes raw 0/1/2 ratios; drop any trailing MTP entries.
+            res.pop('csa_compress_ratios', None)
+            text_config = getattr(hf_config, 'text_config', hf_config)
+            res['csa_compress_ratios'] = list(text_config.compress_ratios)[:res['num_layers']]
+            # V4.1 has no Hash-MoE bootstrap layers.
+            res['moe_n_hash_layers'] = 0
+            res['engram_enabled'] = bool(res.get('engram_layer_ids'))
     elif llm_model_type == 'hunyuan':
         # Since HunYuan’s attention applies RoPE before using q/k_layernorm,
         # which is incompatible with megatron-core, support is not provided here.

src/mcore_bridge/model/constant.py

@@ -11,6 +11,7 @@ class LLMModelType:
     bailing_moe = 'bailing_moe'
     bailing_hybrid = 'bailing_hybrid'
     deepseek_v4 = 'deepseek_v4'
+    deepseek_v41 = 'deepseek_v41'
     glm_moe_dsa = 'glm_moe_dsa'
     nemotron_h = 'nemotron_h'
 

src/mcore_bridge/model/gpt_model.py

@@ -309,7 +309,8 @@ def forward(
                 padding_mask = torch.chunk(padding_mask, tp_size, dim=1)[mpu.get_tensor_model_parallel_rank()]
             extra_block_kwargs['padding_mask'] = padding_mask.contiguous()
 
-        if self.config.moe_n_hash_layers > 0 or getattr(self.config, 'ple_layer_ids', None):
+        if self.config.moe_n_hash_layers > 0 or getattr(self.config, 'ple_layer_ids', None) \
+                or getattr(self.config, 'moe_router_enable_vl_bias', False):
             extra_block_kwargs['input_ids'] = input_ids
         if getattr(self.config, 'indexer_n_heads', None) is not None:
             extra_block_kwargs['position_ids'] = position_ids

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, nemotron_h, olmoe,
+from . import (bailing_hybrid, bailing_moe, deepseek_v4, deepseek_v41, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe,
                qwen3_emb, qwen3_next, qwen4_exp)