diff --git a/src/mcore_bridge/__init__.py b/src/mcore_bridge/__init__.py index 50100df..e0fbc16 100644 --- a/src/mcore_bridge/__init__.py +++ b/src/mcore_bridge/__init__.py @@ -9,7 +9,7 @@ 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 @@ -17,7 +17,7 @@ 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'], diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 68e75ba..72d9a22 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/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 ' diff --git a/src/mcore_bridge/config/__init__.py b/src/mcore_bridge/config/__init__.py index 204ce71..09fa6b7 100644 --- a/src/mcore_bridge/config/__init__.py +++ b/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 diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index 8e10c96..942f2cf 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/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.""" diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index 91fae30..6cc204e 100644 --- a/src/mcore_bridge/config/parser.py +++ b/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. diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py index a7c99ee..9c1c386 100644 --- a/src/mcore_bridge/model/constant.py +++ b/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' diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py index 5294610..cf96a05 100644 --- a/src/mcore_bridge/model/gpt_model.py +++ b/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 diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py index 655415b..40a1bfb 100644 --- a/src/mcore_bridge/model/gpts/__init__.py +++ b/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) diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index de282bd..c58f0aa 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -100,7 +100,15 @@ def __init__(self, config, *args, **kwargs): '`pip install git+https://github.com/NVIDIA/Megatron-LM@dev`') with _patch_YarnRotaryEmbedding(config): super().__init__(config, *args, **kwargs) - self.layer_type = self.config.hf_config.layer_types[self.layer_number - 1] + # ``layer_types`` is an HF-space (length ``num_layers``) list. On the HybridStack the layer + # index space is doubled -- HF layer ``i`` becomes attention layer ``2*i`` (1-based + # ``layer_number`` ``2*i + 1``) and MLP layer ``2*i + 1`` -- so map the hybrid layer_number + # back to the HF index. On the GPT stack ``layer_number - 1`` is already the HF index. + if getattr(self.config, 'is_hybrid_model', False): + hf_layer_idx = (self.layer_number - 1) // 2 + else: + hf_layer_idx = self.layer_number - 1 + self.layer_type = self.config.hf_config.layer_types[hf_layer_idx] self.rope_layer_type = 'main' if self.layer_type == 'sliding_attention' else 'compress' if config.fp8_param: group_proj_in_size = self.query_projection_size // config.o_groups @@ -162,6 +170,12 @@ def get_query_key_value_tensors( # In Megatron-Core, the qkv shape is [t, 1, h, d]. # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. q_compressed = q_compressed.squeeze(1) + # The KV latent (and any CP boundary rows) must drop the dummy batch axis too; + # otherwise linear_kv_proj emits [t, 1, 1, d] and the V4.1 CSA2 core attention + # rejects the layout (it requires key/value shaped [t, 1, d]). + kv_compressed = kv_compressed.squeeze(1) + if boundary_hidden is not None: + boundary_hidden = boundary_hidden.squeeze(1) # ========================================= # Apply norm @@ -201,7 +215,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, # q: [num_tokens, n, q_head_dim] q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) - q = _q_rms_norm(q, self.config.layernorm_epsilon) + # Per-head query RMS norm is a V4-only step: V4.1 normalizes the query latent + # (``q_layernorm``) and feeds ``wq_b``'s output straight into RoPE, so applying it + # here would rescale every head to unit RMS and change the attention scores. + if self.config.dsv4_version == 'v4': + q = _q_rms_norm(q, self.config.layernorm_epsilon) boundary_rows = 0 if boundary_kv_compressed is not None: @@ -304,6 +322,7 @@ def forward( sequence_len_offset=None, *, inference_params=None, + csa2_state=None, ): """Forward pass for DeepSeek-v4 Hybrid Attention""" rotary_pos_emb = rotary_pos_emb[self.rope_layer_type] @@ -382,6 +401,8 @@ def forward( if boundary_hidden is not None: core_attn_kwargs['boundary_hidden'] = boundary_hidden core_attn_kwargs['boundary_kv'] = boundary_kv + if csa2_state is not None: + core_attn_kwargs['csa2_state'] = csa2_state core_attn_out = self.core_attention( query, key, @@ -531,6 +552,12 @@ class DeepseekV4Bridge(GPTBridge): hf_post_attention_layernorm_key = 'ffn_norm.weight' hf_expert_bias_key = 'gate.bias' + def _normalize_missing_weight_key(self, key: str) -> str: + # Native checkpoints use `mtp.*`, while Megatron exports the same stack as + # `model.mtp.*`. Treat them as one identity so save_missing_weights never + # writes both namespaces into the same checkpoint. + return key[len('model.'):] if key.startswith('model.mtp.') else key + def _set_o_group_proj_grouped(self, mg_attn, hf_state_dict, to_mcore): """Handle GroupedLinear state dict for linear_o_group_proj in fp8 mode. diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py new file mode 100644 index 0000000..f82002d --- /dev/null +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -0,0 +1,1804 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1-Flash bridge for megatron-core. + +Wires the V4.1 model (composite HF ``model_type='deepseek_v41'`` with text sub-config +``deepseek_v41_text``) into mcore-bridge. V4.1 reuses DeepSeek-V4's DSv4 hybrid MLA + +hyper-connection stack but swaps the sparse-attention core to CSA2 (selected by +``config.dsv4_version == 'v4.1'``). Differences vs V4/CSA: + + * Compressor (``CSA2Compressor``): no absolute-position embedding (``ape``); the + gate projection (``linear_wgate``) exists only on ratio-2 layers. + * Indexer (``CSA2Indexer``): flattened -- kv-source (``owns_k``) layers own + ``linear_wk`` + ``k_norm`` instead of a nested compressor. + * mHC is single-pass (``mhc_single_pass=True``): there are no learned final + ``hc_head_*`` params; per-layer ``hc_attn_*``/``hc_ffn_*`` are mapped by the + base ``GPTBridge``. + +The backbone is megatron-core's native ``HybridModel``, which splits every HF layer into an +attention-only and an MLP-only hybrid layer. That split is what makes pipeline parallelism +work: a plain ``TransformerBlock`` cannot carry the hyper-connection payload across PP +stages. :func:`derive_hybrid_layer_config` re-expands the HF-space layer config into the +doubled hybrid space, and :class:`DeepseekV41Bridge` fans each HF layer out onto its two +hybrid layers. + +The text integration also attaches the native trainable Engram modules -- loading their +EP-local table rows directly from the official flat FP8 tensors -- and the DSpark (``mtp.*``) +draft stack. Vision is wired in :class:`DeepseekV41MultimodalModel`. + +Requires a megatron-core that ships ``megatron.core.models.hybrid``; without it V4.1 is not +registered at all (see ``_HYBRID_MODEL_AVAILABLE``). +""" +import copy +import os +from dataclasses import dataclass +from types import SimpleNamespace +from typing import List, Optional, Sequence, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +import transformer_engine +from megatron.core import mpu, parallel_state +from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding +from megatron.core.tensor_parallel.layers import VocabParallelEmbedding +from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 +from megatron.core.transformer.spec_utils import build_module +from megatron.core.transformer.transformer_block import TransformerBlock as McoreTransformerBlock +from torch import nn +from tqdm import tqdm + +from mcore_bridge.config import MLAModelConfig +from mcore_bridge.model.modules.dspark import DeepseekV41DSparkStack +from mcore_bridge.model.modules.engram import ( + build_deepseek_v41_engram_config, + DeepseekV41Engram, + DeepseekV41TransformerLayer, + has_native_engram, +) +from mcore_bridge.utils import is_master + +from ..constant import ModelType +from ..mm_gpt_model import MultimodalGPTModel +from ..register import ModelMeta, register_model +from ..rope import get_rope_inv_freq +from .deepseek_v4 import _apply_mla_rope, DeepseekV4Bridge, DeepseekV4Loader, DSv4HybridSelfAttention + +try: + from megatron.core.transformer.experimental_attention_variant.csa2 import CSA2Compressor as McoreCSA2Compressor + from megatron.core.transformer.experimental_attention_variant.csa2 import CSA2Indexer as McoreCSA2Indexer +except ImportError: + McoreCSA2Compressor = object + McoreCSA2Indexer = object + +try: + from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer + from megatron.core.models.hybrid.hybrid_layer_specs import hybrid_dsv4_stack_spec + + from ..hybrid_model import HybridModel + + _HYBRID_MODEL_AVAILABLE = True +except ImportError as error: + if not (error.name or '').startswith('megatron.core.models.hybrid'): + raise + HybridModel = hybrid_dsv4_stack_spec = HyperConnectionHybridLayer = None + _HYBRID_MODEL_AVAILABLE = False + + +def _duplicated_linear_kwargs(config): + return dict( + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode='duplicated', + ) + + +class CSA2Compressor(McoreCSA2Compressor): + """CSA2 compressor keeping its bf16 projections out of fp8 under fp8_param. + + The V4.1 checkpoint stores ``compressor.wkv``/``compressor.wgate`` in bf16, so + rebuild them with fp8 disabled to match (mirrors the V4 ``Compressor`` wrapper). + """ + + def __init__(self, config, submodules, *args, **kwargs): + super().__init__(config, submodules, *args, **kwargs) + if getattr(config, 'fp8_param', False): + linear_kwargs = _duplicated_linear_kwargs(config) + with transformer_engine.pytorch.fp8_model_init(enabled=False): + self.linear_wkv = build_module(submodules.linear_wkv, config.hidden_size, config.v_head_dim, + **linear_kwargs) + if self.compress_ratio == 2: + self.linear_wgate = build_module(submodules.linear_wgate, config.hidden_size, config.v_head_dim, + **linear_kwargs) + + +class DeepseekV41DSparkCoreAttention(MegatronModule): + """Parameter holder for DSpark's latent attention sink. + + The actual attention computation belongs to ``DeepseekV41DSparkAttention``; + this module intentionally has no CSA compressor or indexer parameters. + """ + + def __init__(self, config, *args, **kwargs): + super().__init__(config=config) + world_size = parallel_state.get_tensor_model_parallel_world_size() + if config.num_attention_heads % world_size: + raise ValueError('DSpark attention heads must be divisible by tensor parallel size.') + device = 'cpu' if config.use_cpu_initialization else torch.cuda.current_device() + self.attn_sink = mark_keep_in_fp32( + nn.Parameter(torch.zeros(config.num_attention_heads // world_size, dtype=torch.float32, device=device))) + + +class CSA2Indexer(McoreCSA2Indexer): + """CSA2 indexer keeping its bf16 projections out of fp8 under fp8_param. + + ``linear_weights_proj`` and (on ``owns_k`` layers) ``linear_wk`` are bf16 in the + V4.1 checkpoint; ``linear_wq_b`` stays fp8. Mirrors the V4 ``CSAIndexer`` wrapper. + """ + + def __init__(self, config, submodules, *args, **kwargs): + super().__init__(config, submodules, *args, **kwargs) + if getattr(config, 'fp8_param', False): + linear_kwargs = _duplicated_linear_kwargs(config) + with transformer_engine.pytorch.fp8_model_init(enabled=False): + self.linear_weights_proj = build_module(submodules.linear_weights_proj, config.hidden_size, + self.n_heads, **linear_kwargs) + if self.owns_k: + self.linear_wk = build_module(submodules.linear_wk, config.v_head_dim, self.head_dim, + **linear_kwargs) + + +class DeepseekV41DSparkAttention(DSv4HybridSelfAttention): + """DSpark latent attention over a main-token ring window and draft block.""" + + def __init__(self, config, *args, **kwargs): + super().__init__(config, *args, **kwargs) + self.window_size = config.csa_window_size + self._dspark_window_kv_cache = None + + @staticmethod + def _select_rotary(rotary_pos_emb): + if isinstance(rotary_pos_emb, dict): + return rotary_pos_emb['main'] + return rotary_pos_emb + + def _project_kv(self, hidden_states, rotary_pos_emb): + kv, _ = self.linear_kv_proj(hidden_states) + kv = self.kv_layernorm(kv) + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, kv_pos_emb = torch.split(kv, [kv.shape[-1] - pos_dim, pos_dim], dim=-1) + kv_pos_emb = _apply_mla_rope( + kv_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + cp_group=self.pg_collection.cp, + ) + return torch.cat((kv_no_pe, kv_pos_emb), dim=-1).unsqueeze(-2).contiguous() + + def _project_query(self, hidden_states, rotary_pos_emb): + query_compressed, _ = self.linear_q_down_proj(hidden_states) + query_compressed = self.q_layernorm(query_compressed) + query, _ = self.linear_q_up_proj(query_compressed) + query = query.view( + *query.shape[:-1], + self.num_attention_heads_per_partition, + self.q_head_dim, + ) + pos_dim = self.config.qk_pos_emb_head_dim + query_no_pe, query_pos_emb = torch.split(query, [query.shape[-1] - pos_dim, pos_dim], dim=-1) + query_pos_emb = _apply_mla_rope( + query_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + cp_group=self.pg_collection.cp, + ) + return torch.cat((query_no_pe, query_pos_emb), dim=-1).contiguous() + + def _ensure_cache(self, slot_count, hidden_size, dtype, device): + expected = (slot_count, self.window_size, hidden_size) + cache = self._dspark_window_kv_cache + if cache is None or cache.device != device or cache.dtype != dtype or cache.shape[-1] != hidden_size: + cache = torch.zeros(expected, dtype=dtype, device=device) + elif cache.shape[0] < slot_count: + expanded = torch.zeros(expected, dtype=dtype, device=device) + expanded[:cache.shape[0]].copy_(cache) + cache = expanded + self._dspark_window_kv_cache = cache + return cache + + @staticmethod + def _normalize_cache_inputs(main_kv, start_pos, cache_slots): + batch_size = main_kv.shape[1] + device = main_kv.device + if cache_slots is None: + cache_slots = torch.arange(batch_size, dtype=torch.long, device=device) + else: + cache_slots = cache_slots.to(device=device, dtype=torch.long) + if cache_slots.shape != (batch_size, ): + raise ValueError(f'DSpark cache slots must be [b={batch_size}], got {tuple(cache_slots.shape)}.') + start_positions = torch.as_tensor(start_pos, dtype=torch.long, device=device) + if start_positions.ndim == 0: + start_positions = start_positions.expand(batch_size) + if start_positions.shape != (batch_size, ): + raise ValueError(f'DSpark start positions must be scalar or [b={batch_size}], got ' + f'{tuple(start_positions.shape)}.') + return start_positions, cache_slots + + def _write_main_cache(self, main_kv, start_pos, cache_slots=None): + main_kv = main_kv.squeeze(-2) + start_positions, cache_slots = self._normalize_cache_inputs(main_kv, start_pos, cache_slots) + cache = self._ensure_cache( + int(cache_slots.max().item()) + 1, + main_kv.shape[-1], + main_kv.dtype, + main_kv.device, + ) + if main_kv.shape[0] > self.window_size: + offset = main_kv.shape[0] - self.window_size + main_kv = main_kv[offset:] + start_positions = start_positions + offset + sequence_offsets = torch.arange(main_kv.shape[0], device=main_kv.device) + positions = (start_positions.unsqueeze(0) + sequence_offsets.unsqueeze(1)) % self.window_size + slots = cache_slots.unsqueeze(0).expand_as(positions) + cache[slots, positions] = main_kv.detach() + active_cache = cache.index_select(0, cache_slots).transpose(0, 1).contiguous() + valid_lengths = (start_positions + main_kv.shape[0]).clamp(max=self.window_size) + return active_cache, valid_lengths + + def prefill_dspark( + self, + main_hidden, + rotary_pos_emb, + inference_context=None, + start_pos=0, + cache_slots=None, + ): + del inference_context + rotary_pos_emb = self._select_rotary(rotary_pos_emb) + if rotary_pos_emb is None: + raise ValueError('DSpark prefill requires main-token rotary embeddings.') + main_kv = self._project_kv(main_hidden, rotary_pos_emb) + self._write_main_cache(main_kv, start_pos, cache_slots) + + def reset_dspark_cache(self): + self._dspark_window_kv_cache = None + + def _latent_attention(self, query, key_value, valid_main_lengths=None): + key_value = key_value.expand(-1, -1, query.shape[-2], -1) + scores = torch.einsum('sbhd,tbhd->bhst', query.float(), key_value.float()) + scores.mul_(self.config.v_head_dim**-0.5) + if valid_main_lengths is not None: + main_width = self.window_size + main_indices = torch.arange(main_width, device=scores.device) + invalid_main = main_indices.unsqueeze(0) >= valid_main_lengths.unsqueeze(1) + invalid = F.pad(invalid_main, (0, key_value.shape[0] - main_width), value=False) + scores = scores.masked_fill(invalid[:, None, None, :], float('-inf')) + sink = self.core_attention.attn_sink.view(1, -1, 1, 1) + probabilities = torch.softmax( + torch.cat((scores, sink.expand(scores.shape[:-1] + (1, ))), dim=-1), + dim=-1, + dtype=torch.float32, + )[..., :-1] + return torch.einsum('bhst,tbhd->sbhd', probabilities.to(key_value.dtype), key_value) + + def _project_output(self, output): + seq_len, batch_size = output.shape[:2] + output = output.view(seq_len, batch_size, self.o_local_groups, -1) + if self._o_group_proj_is_grouped_linear: + output = output.permute(2, 0, 1, 3).contiguous().reshape(-1, output.shape[-1]) + output = self.linear_o_group_proj(output, [seq_len * batch_size] * self.o_local_groups) + output = output.view(self.o_local_groups, seq_len, batch_size, -1) + output = output.permute(1, 2, 0, 3).contiguous().reshape(seq_len, batch_size, -1) + else: + weight = self.linear_o_group_proj.view(self.o_local_groups, self.config.o_lora_rank, -1) + output = torch.einsum('...gd,grd->...gr', output, weight).flatten(-2) + return self.linear_proj(output) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + dspark_main_hidden=None, + dspark_main_rotary_pos_emb=None, + dspark_cache_slots=None, + ): + del attention_mask, key_value_states, rotary_pos_cos, rotary_pos_sin + del rotary_pos_cos_sin, attention_bias, packed_seq_params, position_ids, inference_params + if dspark_main_hidden is None: + raise ValueError('DSpark attention requires target-layer main hidden states.') + if sequence_len_offset is None: + raise ValueError('DSpark attention requires the main-token sequence offset.') + start_pos = sequence_len_offset + draft_rotary = self._select_rotary(rotary_pos_emb) + main_rotary = self._select_rotary(dspark_main_rotary_pos_emb) + if draft_rotary is None or main_rotary is None: + raise ValueError('DSpark attention requires main and draft rotary embeddings.') + + main_kv = self._project_kv(dspark_main_hidden, main_rotary) + main_window, valid_main = self._write_main_cache(main_kv, start_pos, dspark_cache_slots) + main_window = main_window.unsqueeze(-2) + query = self._project_query(hidden_states, draft_rotary) + draft_kv = self._project_kv(hidden_states, draft_rotary) + key_value = torch.cat((main_window, draft_kv), dim=0) + output = self._latent_attention(query, key_value, valid_main) + + pos_dim = self.config.qk_pos_emb_head_dim + output_no_pe, output_pos_emb = torch.split(output, [output.shape[-1] - pos_dim, pos_dim], dim=-1) + output_pos_emb = _apply_mla_rope( + output_pos_emb, + draft_rotary, + config=self.config, + cu_seqlens=None, + cp_group=self.pg_collection.cp, + inverse=True, + ) + return self._project_output(torch.cat((output_no_pe, output_pos_emb), dim=-1)) + + +def _vision_cos_sin(n_h: int, n_w: int, dim: int, theta: float, device: torch.device): + """Build the official row-major 2D RoPE table for one image.""" + inv_freq = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + hpos = torch.arange(n_h, device=device).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w, device=device).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack((hpos, wpos), dim=-1).reshape(-1, 2, 1).float() * inv_freq + freqs = freqs.flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def _apply_vision_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1).to(dtype) + + +class DeepseekV41VisionRMSNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class DeepseekV41PatchEmbed(nn.Module): + + def __init__(self, patch_size: int, hidden_size: int): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Linear(3 * patch_size**2, hidden_size) + + def forward(self, patches: torch.Tensor): + if patches.ndim not in (2, 4): + raise ValueError('DeepSeek-V4.1 pixel_values must be [num_patches, 3, patch, patch] ' + f'or flattened [num_patches, 3 * patch ** 2], got {tuple(patches.shape)}.') + return self.proj(patches.flatten(1)) + + +class DeepseekV41VisionAttention(nn.Module): + + def __init__(self, hidden_size: int, num_heads: int): + super().__init__() + if hidden_size % num_heads: + raise ValueError(f'vision hidden_size {hidden_size} must be divisible by num_heads {num_heads}.') + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.wqkv = nn.Linear(hidden_size, 3 * hidden_size) + self.wo = nn.Linear(hidden_size, hidden_size) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + num_tokens = x.shape[0] + q, k, v = (tensor.view(num_tokens, self.num_heads, self.head_dim) for tensor in self.wqkv(x).chunk(3, dim=-1)) + q = _apply_vision_rotary(q, cos, sin) + k = _apply_vision_rotary(k, cos, sin) + output = F.scaled_dot_product_attention(q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1)) + return self.wo(output.transpose(0, 1).reshape(num_tokens, -1)) + + +class DeepseekV41VisionMLP(nn.Module): + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.w1 = nn.Linear(hidden_size, 2 * intermediate_size, bias=False) + self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x: torch.Tensor): + gate, up = self.w1(x).chunk(2, dim=-1) + return self.w2(F.silu(gate) * up) + + +class DeepseekV41VisionBlock(nn.Module): + + def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int): + super().__init__() + self.norm1 = DeepseekV41VisionRMSNorm(hidden_size) + self.attn = DeepseekV41VisionAttention(hidden_size, num_heads) + self.norm2 = DeepseekV41VisionRMSNorm(hidden_size) + self.mlp = DeepseekV41VisionMLP(hidden_size, intermediate_size) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class DeepseekV41VisionTransformer(nn.Module): + + def __init__(self, vision_config): + super().__init__() + hidden_size = vision_config.hidden_size + num_heads = vision_config.num_attention_heads + self.rope_dim = hidden_size // num_heads // 2 + self.rope_theta = vision_config.rope_theta + self.patch_embed = DeepseekV41PatchEmbed(vision_config.patch_size, hidden_size) + self.blocks = nn.ModuleList([ + DeepseekV41VisionBlock(hidden_size, num_heads, vision_config.intermediate_size) + for _ in range(vision_config.num_hidden_layers) + ]) + self.norm = DeepseekV41VisionRMSNorm(hidden_size) + + def forward(self, patches: torch.Tensor, n_h: int, n_w: int): + if patches.shape[0] != n_h * n_w: + raise ValueError(f'Image grid {n_h}x{n_w} requires {n_h * n_w} patches, got {patches.shape[0]}.') + x = self.patch_embed(patches) + cos, sin = _vision_cos_sin(n_h, n_w, self.rope_dim, self.rope_theta, x.device) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class DeepseekV41Aligner(nn.Module): + + def __init__(self, vision_config, text_hidden_size: int): + super().__init__() + self.downsample_ratio = vision_config.downsample_ratio + in_dim = vision_config.hidden_size * self.downsample_ratio**2 + self.w1 = nn.Linear(in_dim, text_hidden_size) + self.w2 = nn.Linear(text_hidden_size, text_hidden_size) + + def forward(self, x: torch.Tensor, n_h: int, n_w: int): + ratio = self.downsample_ratio + x = x.view(n_h, n_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_w % ratio, 0, -n_h % ratio)) + x = F.unfold(x.unsqueeze(0), ratio, stride=ratio).squeeze(0).transpose(0, 1) + return self.w2(F.gelu(self.w1(x))) + + +class DeepseekV41Vision(nn.Module): + """Trainable V4.1 vision tower, aligner and image-span embedding merger.""" + + # DeepseekV4Bridge normalizes root-level official keys under an internal + # ``model.`` prefix before conversion and strips it again on export. + module_mapping = {'model.vision': 'vision', 'model.aligner': 'aligner'} + _vision_tower = ['vision'] + _aligner = ['aligner'] + test_mm_type = 'image' + + IMAGE_START = 0 + IMAGE = 1 + IMAGE_NEW_LINE = 2 + IMAGE_END = 3 + + def __init__(self, config): + super().__init__() + self.config = config + self.image_token_id = config.hf_config.image_token_id + if config.language_model_only: + self.vision = None + self.aligner = None + self.register_parameter('image_start', None) + self.register_parameter('image_end', None) + self.register_parameter('image_newline', None) + return + vision_config = config.hf_config.vision_config + self.vision = DeepseekV41VisionTransformer(vision_config) + self.aligner = DeepseekV41Aligner(vision_config, config.hidden_size) + self.image_start = nn.Parameter(torch.empty(config.hidden_size)) + self.image_end = nn.Parameter(torch.empty(config.hidden_size)) + self.image_newline = nn.Parameter(torch.empty(config.hidden_size)) + target_device = torch.cuda.current_device() if torch.cuda.is_available() else None + self.to(device=target_device, dtype=config.params_dtype) + # Official RMSNorm scales remain fp32 even when the remaining ViT is bf16. + for module in self.modules(): + if isinstance(module, DeepseekV41VisionRMSNorm): + module.weight.data = module.weight.data.float() + + def get_inputs_embeds_language_model(self, inputs_embeds, **kwargs): + return inputs_embeds + + @staticmethod + def _grid_hw(image_grid_thw: torch.Tensor): + if image_grid_thw.ndim != 2 or image_grid_thw.shape[1] not in (2, 3): + raise ValueError(f'image_grid_thw must have shape [num_images, 2 or 3], got {tuple(image_grid_thw.shape)}.') + if image_grid_thw.shape[1] == 3: + if not torch.all(image_grid_thw[:, 0] == 1): + raise ValueError('DeepSeek-V4.1 supports still images only; every temporal grid size must be 1.') + image_grid_thw = image_grid_thw[:, 1:] + return image_grid_thw.to(dtype=torch.long, device='cpu') + + def encode_images(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor): + grids = self._grid_hw(image_grid_thw) + outputs = [] + patch_offset = 0 + for n_h, n_w in grids.tolist(): + patch_count = n_h * n_w + patches = pixel_values[patch_offset:patch_offset + patch_count] + outputs.append(self.aligner(self.vision(patches, n_h, n_w), n_h, n_w)) + patch_offset += patch_count + if patch_offset != pixel_values.shape[0]: + raise ValueError( + f'Image grids describe {patch_offset} patches, but pixel_values has {pixel_values.shape[0]}.') + if not outputs: + return pixel_values.new_empty((0, self.image_start.numel())) + return torch.cat(outputs, dim=0) + + def _zero_parameter_dependency(self, inputs_embeds: torch.Tensor): + zero = inputs_embeds.new_zeros(()) + for parameter in self.parameters(): + zero = zero + parameter.reshape(-1)[0].to(inputs_embeds.dtype) * 0 + return inputs_embeds + zero + + def get_inputs_embeds(self, inputs_embeds, **kwargs): + pixel_values = kwargs.get('pixel_values') + image_grid_thw = kwargs.get('image_grid_thw') + token_types = kwargs.get('image_token_types', kwargs.get('token_types')) + if pixel_values is None: + return self._zero_parameter_dependency(inputs_embeds) + if image_grid_thw is None or token_types is None: + raise ValueError('DeepSeek-V4.1 vision requires image_grid_thw and image_token_types/token_types.') + if token_types.shape != kwargs['input_ids'].shape: + raise ValueError(f'image token types shape {tuple(token_types.shape)} must match input_ids ' + f'{tuple(kwargs["input_ids"].shape)}.') + + image_mask = token_types >= 0 + input_image_mask = kwargs['input_ids'] == self.image_token_id + if not torch.equal(image_mask.to(input_image_mask.device), input_image_mask): + raise ValueError( + 'Every DeepSeek-V4.1 image-span position must carry image_token_id, and no text position may use it.') + image_features = self.encode_images(pixel_values.to(self.vision.patch_embed.proj.weight), image_grid_thw) + flat_types = token_types[image_mask].to(device=inputs_embeds.device) + if int((flat_types == self.IMAGE).sum()) != image_features.shape[0]: + raise ValueError(f'Image spans contain {int((flat_types == self.IMAGE).sum())} patch slots, ' + f'but the aligner produced {image_features.shape[0]} rows.') + replacements = inputs_embeds.new_empty((flat_types.numel(), inputs_embeds.shape[-1])) + replacements[flat_types == self.IMAGE_START] = self.image_start.to(inputs_embeds.dtype) + replacements[flat_types == self.IMAGE_END] = self.image_end.to(inputs_embeds.dtype) + replacements[flat_types == self.IMAGE_NEW_LINE] = self.image_newline.to(inputs_embeds.dtype) + replacements[flat_types == self.IMAGE] = image_features.to(inputs_embeds.dtype) + if not torch.all((flat_types >= self.IMAGE_START) & (flat_types <= self.IMAGE_END)): + raise ValueError('image_token_types values must be TEXT=-1 or one of START=0, IMAGE=1, NEW_LINE=2, END=3.') + expanded_mask = image_mask.to(inputs_embeds.device).unsqueeze(-1).expand_as(inputs_embeds) + return inputs_embeds.masked_scatter(expanded_mask, replacements) + + +@dataclass +class HybridLayerConfig: + """Per-layer config re-expanded from HF layer space into hybrid (2x) layer space. + + All source-layer / candidate fields are 0-based indices in the doubled hybrid space + (i.e. ``layer_number - 1`` as CSA2 reads them), where GPT layer ``i`` maps to hybrid + attention layer ``2 * i``. + """ + + hybrid_layer_pattern: str + num_layers: int + csa_compress_ratios: List[int] + csa2_kv_source_layers: List[int] + csa2_index_source_layers: List[int] + csa2_candidate_source_layer: Optional[int] + + +def _normalize_moe_layer_freq(moe_layer_freq: Union[int, Sequence[int], None], num_layers: int) -> List[int]: + """Return a length-``num_layers`` 0/1 list marking MoE layers. + + Mirrors megatron-core's own interpretation (moe_logging.py:660-664): an ``int`` N means + layer ``i`` is MoE iff ``i % N == 0``; a list is used verbatim. ``None`` (no experts) + means every layer is dense. + """ + if moe_layer_freq is None: + return [0] * num_layers + if isinstance(moe_layer_freq, int): + return [1 if i % moe_layer_freq == 0 else 0 for i in range(num_layers)] + freq = list(moe_layer_freq) + if len(freq) != num_layers: + raise ValueError(f'moe_layer_freq length {len(freq)} does not match num_layers {num_layers}.') + return [1 if x else 0 for x in freq] + + +def derive_hybrid_layer_config( + num_layers: int, + csa_compress_ratios: Sequence[int], + moe_layer_freq: Union[int, Sequence[int], None], + csa2_kv_source_layers: Sequence[int] = (), + csa2_index_source_layers: Sequence[int] = (), + csa2_candidate_source_layer: Optional[int] = None, +) -> HybridLayerConfig: + """Translate an HF-space V4.1 config into the doubled hybrid layer space. + + Each HF transformer layer ``i`` (0-based) becomes two hybrid layers: + + * hybrid index ``2*i`` -- attention-only, always the array-driven ``D`` symbol so it + reads its ratio from ``csa_compress_ratios[2*i]`` and stays numerically identical to + the GPT ``dsv4_hybrid`` attention layer (baking a fixed ratio via ``C``/``H``/``W`` + would instead trip the ``compress_ratio != ratio`` guard in csa2.py:1280). + * hybrid index ``2*i + 1`` -- MLP-only, ``E`` if the layer is MoE else ``-``. + + Because CSA2 indexes every per-layer array by ``layer_number - 1``, the compress ratios + and the kv/index/candidate source layers are re-expanded so a GPT source layer ``j`` + lands on hybrid attention index ``2*j`` (MLP slots get ratio 0 and are never read). + """ + if len(csa_compress_ratios) != num_layers: + raise ValueError( + f'csa_compress_ratios length {len(csa_compress_ratios)} does not match num_layers {num_layers}.') + moe_mask = _normalize_moe_layer_freq(moe_layer_freq, num_layers) + + pattern_chars: List[str] = [] + hybrid_ratios: List[int] = [] + for i in range(num_layers): + # attention-only layer (array-driven DSv4 attention) + pattern_chars.append('D') + hybrid_ratios.append(int(csa_compress_ratios[i])) + # MLP-only layer + pattern_chars.append('E' if moe_mask[i] else '-') + hybrid_ratios.append(0) + + def _remap(layers: Sequence[int]) -> List[int]: + return [2 * int(j) for j in layers] + + return HybridLayerConfig( + hybrid_layer_pattern=''.join(pattern_chars), + num_layers=2 * num_layers, + csa_compress_ratios=hybrid_ratios, + csa2_kv_source_layers=_remap(csa2_kv_source_layers), + csa2_index_source_layers=_remap(csa2_index_source_layers), + csa2_candidate_source_layer=(None if csa2_candidate_source_layer is None else 2 + * int(csa2_candidate_source_layer)), + ) + + +if _HYBRID_MODEL_AVAILABLE: + + class DeepseekV41HyperConnectionHybridLayer(HyperConnectionHybridLayer): + """Hyper-connection wrapper that applies Engram on the n-stream residual, matching GPT. + + The GPT single-pass path (``HyperConnectionTransformerLayer._forward_attention``, upstream + transformer_layer.py) applies Engram to the *n-stream* residual (width + ``num_residual_streams * hidden_size``) **before** the self-attention hyper-connection + aggregates it to a single stream:: + + hidden_states = self._maybe_apply_engram(hidden_states, input_ids) # n-stream, 20480 + hidden_states, ... = self.self_attention_hyper_connection(hidden_states, ...) # -> 1 stream + + HybridStack inverts that order: :meth:`HyperConnectionHybridLayer.forward` aggregates first + (``self.hyper_connection(hidden_states)``) and runs the inner layer on the single aggregated + stream, and its eager fast path (``_call_inner_transformer_layer_without_local_bda``, taken + for the attention-only 'D' layer) calls ``_forward_self_attention_output_with_bias`` + directly, which skips ``_maybe_apply_engram`` entirely. So the base wrapper either drops + Engram (fast path) or -- if the fast path is declined -- applies it on the aggregated + *single*-stream tensor, which is both the wrong width (``hidden_size`` vs. + ``num_streams * hidden_size``) and the wrong point in the residual. + + We therefore apply Engram here, on the incoming n-stream ``hidden_states``, before + delegating to the base wrapper forward (aggregation + fast-path attention), so the Engram + contribution lands on the pre-aggregation streams. The inner ``DeepseekV41TransformerLayer`` keeps its + ``engram`` module only so the bridge can load/export its weights; the base fast path never + calls it, so there is no double add. Non-Engram layers keep the base wrapper untouched (this + subclass is only swapped onto Engram-carrying wrappers, see + :meth:`DeepseekV41Loader._rewrap_engram_hyper_connection_layers`). + + The fast path is also invoked by the CUDA-graph capture body, which is out of scope for this + change (plan: no CUDA Graph). + """ + + def forward(self, + hidden_states, + attention_mask=None, + inference_context=None, + rotary_pos_emb=None, + sequence_len_offset=None, + packed_seq_params=None, + padding_mask=None, + input_ids=None, + mhc_recompute_manager=None, + mhc_state=None, + **layer_kwargs): + engram = getattr(self.inner_layer, 'engram', None) + if engram is not None: + if input_ids is None: + raise ValueError('DeepSeek-V4.1 hybrid Engram requires input token IDs on the layer forward.') + # ``Engram.forward`` reads the THD / inference context off the module itself + # (mirrors ``_DeepseekV41EngramLayerMixin._forward_attention``), so stash it for + # the duration of this call and add the n-stream Engram delta like + # ``TransformerLayer._maybe_apply_engram``. + previous_ctx = getattr(engram, '_bridge_inference_context', None) + previous_pack = getattr(engram, '_bridge_packed_seq_params', None) + engram._bridge_inference_context = inference_context + engram._bridge_packed_seq_params = packed_seq_params + try: + hidden_states = hidden_states + engram(hidden_states, input_ids, inference_context) + finally: + engram._bridge_inference_context = previous_ctx + engram._bridge_packed_seq_params = previous_pack + return super().forward( + hidden_states, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + sequence_len_offset=sequence_len_offset, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + input_ids=input_ids, + mhc_recompute_manager=mhc_recompute_manager, + **({ + 'mhc_state': mhc_state + } if mhc_state is not None else {}), + **layer_kwargs, + ) + + class DeepseekV41HybridStackModel(HybridModel): + """``HybridModel`` that splits PP / VPP stages on complete attention+FFN blocks. + + Upstream ``select_pipeline_segment`` (called inside ``HybridModel.__init__``, + hybrid_model.py:265) handles the split, but for a pattern *without* ``|`` separators it + (a) refuses VPP outright and (b) slices the ``2 * num_layers`` sublayers evenly, which + cuts a ``D``/``E`` block across a stage boundary whenever ``2N // stages`` is odd. Either + breaks :class:`DeepseekV41Bridge`, whose 1-HF-layer -> 2-hybrid-layer fan-out + assumes the attention half (``2 * i``) and its MLP half (``2 * i + 1``) are co-resident. + + Mirroring GLM-5.3 (``Glm5NextHybridModel``), we pre-segment the *main* pattern on block + boundaries into ``|``-delimited, PP*VPP-ordered stages before it reaches upstream, and + assert after build that this rank holds whole blocks -- failing loudly instead of + mis-mapping weights. The MTP suffix is still appended by the base resolver, so a + segmented main becomes ``seg0|seg1|.../mtp``. + """ + + # This backbone is text-only, but ``deepseek_v41`` is a multimodal model_type, so the + # trainer's ``is_multimodal`` path reads ``model.visual`` (expecting ``None`` for text). + # Expose it so that guard short-circuits; the real vision tower arrives with + # :class:`DeepseekV41MultimodalModel`. + visual = None + + # ``MultimodalGPTModel.forward`` (the multimodal wrapper) reads ``language_model.extra_forward_keys`` + # to forward a whitelist of extra kwargs into the decoder. ``McoreHybridModel`` has no such + # attribute (it lives on the mcore-bridge ``GPTModel``, default ``[]``); expose the same + # empty default so the wrapper can treat this backbone like any mcore-bridge ``GPTModel``. + extra_forward_keys: List[str] = [] + + @staticmethod + def _segment_main_pattern(config) -> Optional[str]: + pattern = config.hybrid_layer_pattern + if getattr(config, 'pipeline_model_parallel_layout', None) is not None: + raise ValueError('DeepSeek-V4.1 hybrid splits pipeline stages by hybrid_layer_pattern, so ' + 'pipeline_model_parallel_layout does not apply; use ' + 'num_layers_in_first_pipeline_stage / num_layers_in_last_pipeline_stage for an ' + 'uneven split.') + # An explicit layout is respected as-is; upstream + the post-build guard validate it. + if (not pattern or '|' in pattern or config.num_layers_in_first_pipeline_stage is not None + or config.num_layers_in_last_pipeline_stage is not None): + return pattern + stages = config.pipeline_model_parallel_size + if config.virtual_pipeline_model_parallel_size: + stages *= config.virtual_pipeline_model_parallel_size + if stages <= 1: + return pattern + blocks, extra = divmod(len(pattern) // 2, stages) + if blocks == 0: + raise ValueError('DeepSeek-V4.1 hybrid needs at least one attention+FFN block per pipeline stage, ' + f'but {len(pattern) // 2} blocks cannot cover {stages} stages; lower ' + 'pipeline_model_parallel_size / virtual_pipeline_model_parallel_size.') + # Consecutive segments map to (vp0,pp0),(vp0,pp1),... matching upstream's + # segment_index = vp_stage * pp_size + pp_rank (hybrid_layer_allocation.py:478). + segments, offset = [], 0 + for stage in range(stages): + count = 2 * (blocks + int(stage < extra)) + segments.append(pattern[offset:offset + count]) + offset += count + return '|'.join(segments) + + @staticmethod + def _resolve_hybrid_layer_pattern(config) -> Optional[str]: + segmented = DeepseekV41HybridStackModel._segment_main_pattern(config) + if segmented == config.hybrid_layer_pattern: + return HybridModel._resolve_hybrid_layer_pattern(config) + seg_config = copy.copy(config) + seg_config.hybrid_layer_pattern = segmented + return HybridModel._resolve_hybrid_layer_pattern(seg_config) + + def __init__(self, config, transformer_layer_spec, pre_process=True, post_process=True, vp_stage=None): + super().__init__(config, transformer_layer_spec, pre_process, post_process, vp_stage) + # A stage holding a partial block would break the HF-layer fan-out in the bridge. + layers = getattr(self.decoder, 'layers', None) or [] + if layers: + offset = layers[0].layer_number - 1 + count = len(layers) + if offset % 2 or count % 2: + raise ValueError('DeepSeek-V4.1 hybrid pipeline stage boundaries must fall on complete ' + f'attention+FFN blocks, but this stage starts at sublayer {offset} and holds ' + f'{count} sublayers (both must be even, since one block is two sublayers). ' + 'Leave num_layers_in_first_pipeline_stage / num_layers_in_last_pipeline_stage ' + 'unset for an even block-aligned split, or pass even values.') + # ``HybridModel.forward`` builds no model-level RoPE for ``multi_latent_attention`` and + # hard-sets ``rotary_pos_emb=None`` when calling the decoder. The reused DSv4 attention + # (shared with V4) instead expects the decoupled ``{'main', 'compress'}`` dict + # that ``DeepseekV4GPTModel`` builds. Build the same two RoPE tables here and inject the + # dict into the decoder via a forward pre-hook, keeping the attention numerically + # identical to the GPTModel baseline. ``get_rotary_seq_len`` reads ``decoder.input_tensor`` + # when the local ``hidden_states`` is ``None``, so this also covers PP intermediate/last + # stages. + self._dsv4_position_ids = None + self._build_dsv4_rotary_tables() + self.decoder.register_forward_pre_hook(self._inject_dsv4_rotary_pos_emb, with_kwargs=True) + + def _build_dsv4_rotary_tables(self): + """Build the MLA decoupled-RoPE ``main``/``compress`` tables (mirrors + ``mcore_bridge.model.gpt_model.GPTModel`` MLA setup + ``DeepseekV4GPTModel._set_inv_freq``).""" + self.rotary_pos_emb = RotaryEmbedding( + kv_channels=self.config.qk_pos_emb_head_dim, + rotary_percent=1, + rotary_interleaved=self.config.rotary_interleaved, + rotary_base=self.config.rotary_base, + use_cpu_initialization=self.config.use_cpu_initialization, + ) + rope_scaling = self.config.rope_scaling + self.config.rope_scaling = rope_scaling['main'] + new_inv_freq, 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) + self.config.attention_scaling = attention_scaling + # compress + self.compress_rotary_pos_emb = copy.copy(self.rotary_pos_emb) + self.config.rope_scaling = rope_scaling['compress'] + new_inv_freq, attention_scaling = get_rope_inv_freq(self.config) + self.compress_rotary_pos_emb.inv_freq = new_inv_freq + self.config.compress_attention_scaling = attention_scaling + self.config.rope_scaling = rope_scaling + + def _dsv4_rotary_pos_emb(self, transformer_input, packed_seq_params, inference_context=None): + """Return the ``{'main', 'compress'}`` RoPE dict the DSv4 attention indexes by + ``rope_layer_type`` (mirrors ``DeepseekV4GPTModel._get_rotary_pos_emb`` plus the + packed pre-indexing ``GPTModel.forward`` does). + + The DSv4 attention consumes *per-token* frequencies row-aligned with the hidden states + (see ``_apply_mla_rope``), not a position->frequency table. For one sequence per row the + table is already row-aligned, but under ``thd`` packing a row holds several sequences + whose positions restart, so the table (sized by the longest sequence) must be indexed by + ``position_ids`` here -- exactly what the GPT path does in ``GPTModel.forward``. Under CP + ``position_ids`` arrives already split with the hidden states' partition mode, so the + indexed frequencies come out rank-local while keeping absolute positions. + """ + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(inference_context, self.decoder, transformer_input, + self.config, packed_seq_params) + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + rotary_pos_emb = { + 'main': self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq), + 'compress': self.compress_rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq), + } + if packed_seq and not self.config.apply_rope_fusion: + position_ids = self._dsv4_position_ids + if position_ids is None: + raise ValueError('DeepSeek-V4.1 hybrid needs position_ids on every pipeline ' + 'stage to pre-index the MLA rotary table under sequence ' + 'packing.') + assert position_ids.shape[0] == 1, f'position_ids.shape: {position_ids.shape}' + rotary_pos_emb = {k: v[position_ids[0]] for k, v in rotary_pos_emb.items()} + return rotary_pos_emb + + def _inject_dsv4_rotary_pos_emb(self, module, args, kwargs): + if kwargs.get('rotary_pos_emb') is not None: + return None + transformer_input = kwargs.get('hidden_states') + if transformer_input is None and args: + transformer_input = args[0] + kwargs['rotary_pos_emb'] = self._dsv4_rotary_pos_emb(transformer_input, kwargs.get('packed_seq_params'), + kwargs.get('inference_context')) + return args, kwargs + + # Visual kwargs are injected into the embeddings by the multimodal wrapper and then + # cleared before the language model runs; the base HybridModel.forward never accepts them. + # This backbone is text-only, so strip them here. For a text batch + # ``DeepseekV41Vision.get_inputs_embeds`` is a numeric no-op (``_zero_parameter_dependency`` + # adds ``0 * vision_params``), so dropping them changes nothing numerically. + _visual_forward_keys = ('pixel_values', 'image_grid_thw', 'image_token_types', 'token_types') + + def forward(self, *args, **kwargs): + # The multimodal wrapper (``MultimodalGPTModel.forward``) always funnels the + # decoder's extra kwargs through ``extra_block_kwargs`` -- the mcore-bridge ``GPTModel`` + # calling convention. Upstream ``HybridModel.forward`` has no such parameter (it threads + # ``input_ids`` into the decoder itself, hybrid_model.py), so unpack the container here + # and let the visual-key strip below drop anything the text backbone does not consume. + extra_block_kwargs = kwargs.pop('extra_block_kwargs', None) + if extra_block_kwargs: + kwargs.update(extra_block_kwargs) + if kwargs.get('pixel_values') is not None: + raise NotImplementedError('The DeepSeek-V4.1 hybrid backbone is text-only; multimodal inputs must go ' + 'through DeepseekV41MultimodalModel.') + for key in self._visual_forward_keys: + kwargs.pop(key, None) + # Upstream ``HybridModel.forward`` never threads position_ids into the decoder, so stash + # it for the rotary pre-hook (see :meth:`_dsv4_rotary_pos_emb`). + position_ids = kwargs.get('position_ids') + if position_ids is None and len(args) > 1: + position_ids = args[1] + self._dsv4_position_ids = position_ids + try: + return super().forward(*args, **kwargs) + finally: + self._dsv4_position_ids = None + + class DeepseekV41MultimodalModel(MultimodalGPTModel): + """Multimodal wrapper hosting the ``HybridModel`` backbone. + + ``MultimodalGPTModel`` consumes its ``language_model`` through a backbone-agnostic + interface -- ``embedding(input_ids, position_ids)`` / ``vp_stage`` / + ``share_embeddings_and_output_weights`` / ``extra_forward_keys`` / ``set_input_tensor`` / + ``get_input_tensor`` / ``shared_embedding_or_output_weight`` plus the standard forward + signature -- all of which :class:`DeepseekV41HybridStackModel` provides + (``extra_forward_keys`` is added on it for exactly this). The vision tower, image-embed + injection (``_patch_word_embeddings``) and the vision/aligner weight bridging + (``MultimodalGPTBridge._convert_pre_process``) are inherited unchanged. + + The wrapper injects image embeddings into the embedding output and clears the visual + kwargs before the language model runs, so the hybrid backbone only ever sees a text batch + (its ``forward`` strips ``_visual_forward_keys`` as a defensive backstop). + """ + + language_model_cls = DeepseekV41HybridStackModel + + @property + def vocab_size(self): + return self.language_model.vocab_size +else: + DeepseekV41HyperConnectionHybridLayer = None + DeepseekV41HybridStackModel = None + DeepseekV41MultimodalModel = None + + +class DeepseekV41Loader(DeepseekV4Loader): + """Build DeepSeek-V4.1 on megatron-core's native ``HybridModel``. + + Extends the V4 loader's MLA/CSA2 knowledge with the V4.1 Engram config resolution and + rewrites the layer config into the doubled hybrid layer space (see + :func:`derive_hybrid_layer_config`); the derivation works on a config copy so the caller's + config is never mutated. + + On top of the text backbone it attaches the DSpark (``mtp.*``) draft stack (in + :meth:`build_model`, mapped in :meth:`DeepseekV41Bridge._convert_additional_layers`). + Autoregressive MTP (``mtp_num_layers`` / ``MultiTokenPredictionBlock``) does not apply to + V4.1 -- its ``mtp.*`` checkpoint keys *are* DSpark -- so it stays disabled here. The backbone + is wrapped in :class:`DeepseekV41MultimodalModel` for the vision tower + image-embed + injection. + """ + + # ``HybridModel`` builds its own stack, so leave megatron-core's TransformerBlock + # unpatched (``register.py`` would otherwise swap in mcore-bridge's variant). + transformer_block = McoreTransformerBlock + + model_cls = DeepseekV41MultimodalModel + + def _engram_placement_layer_ids(self, hf_layer_ids): + """On HybridStack, HF layer ``e`` becomes the attention-only 'D' layer at hybrid index + ``2 * e`` (0-based) -- 1-based ``layer_number`` ``2 * e + 1``. Engram is placed there + (never on the MLP-only 'E'/'-' layer), so ``TransformerLayer``'s + ``layer_number in engram_config.layer_ids`` gate builds it on the right hybrid layers. + ``hash_layer_ids`` stays 0-based HF so the tokenizer artifact / hash multipliers are + looked up unchanged.""" + return tuple(2 * layer_id + 1 for layer_id in hf_layer_ids) + + def _build_hybrid_config(self): + # Shallow copy + per-field reassignment (mirrors ``get_dspark_layer_spec``); every field + # written below is replaced by a fresh object, so the original config is never mutated. + cfg = copy.copy(self.config) + derived = derive_hybrid_layer_config( + self.config.num_layers, + list(self.config.csa_compress_ratios), + self.config.moe_layer_freq, + csa2_kv_source_layers=self.config.csa2_kv_source_layers or [], + csa2_index_source_layers=self.config.csa2_index_source_layers or [], + csa2_candidate_source_layer=self.config.csa2_candidate_source_layer, + ) + cfg.num_layers = derived.num_layers + cfg.hybrid_layer_pattern = derived.hybrid_layer_pattern + cfg.csa_compress_ratios = derived.csa_compress_ratios + cfg.csa2_kv_source_layers = derived.csa2_kv_source_layers + cfg.csa2_index_source_layers = derived.csa2_index_source_layers + cfg.csa2_candidate_source_layer = derived.csa2_candidate_source_layer + cfg.is_hybrid_model = True + # HybridStack picks E/- from the pattern; keep moe_layer_freq consistent with the doubled + # space so any layer-count validation that reads it still agrees with num_layers. + cfg.moe_layer_freq = [1 if symbol == 'E' else 0 for symbol in derived.hybrid_layer_pattern] + # Autoregressive MTP does not apply to V4.1: the parser never sets ``mtp_num_layers`` + # (it maps ``num_nextn_predict_layers`` to ``dspark_num_layers`` instead), and the + # ``mtp.*`` checkpoint keys are the DSpark draft stack (attached in ``build_model``). + # Keep it disabled so no ``MultiTokenPredictionBlock`` is built. + cfg.mtp_num_layers = None + return cfg + + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + # Build the spec from the *hybrid* config so the CSA2 attention sees the doubled-space + # csa arrays. ``build_model`` caches it on ``self._hybrid_config`` first. + spec = hybrid_dsv4_stack_spec(self._hybrid_config) + # Apply the fp8-parity module swaps on the array-driven 'D' + # attention layer (the only attention symbol V4.1 emits). + attn = spec.submodules.dsa_layer.submodules.self_attention + attn.module = DSv4HybridSelfAttention + core = attn.submodules.core_attention.submodules + if getattr(core, 'compressor', None) is not None: + core.compressor.module = CSA2Compressor + if getattr(core, 'indexer', None) is not None: + # CSA2 indexer is flat (no nested compressor). + core.indexer.module = CSA2Indexer + # Attach Engram to the 'D' (attention-only) layer spec. Because HybridStack shares one + # ``dsa_layer`` spec across every 'D' layer, per-layer placement is handled by + # ``TransformerLayer.__init__`` (only ``layer_number in engram_config.layer_ids`` builds + # it) rather than by editing per-layer specs. + engram_config = self._get_engram_config() + if engram_config is not None: + from megatron.core.transformer.spec_utils import ModuleSpec + dsa = spec.submodules.dsa_layer + # The inference-aware subclass adds the ``_forward_attention`` Engram hook. + dsa.module = DeepseekV41TransformerLayer + dsa.submodules.engram = ModuleSpec(module=DeepseekV41Engram, params={'engram_config': engram_config}) + # HybridStack exposes MoE via ``moe_layer`` (symbol 'E') instead of GPT's ``layer_specs``, + # so ``ModelLoader._replace_router`` never sees it. Swap the stock ``McoreTopKRouter`` for + # the project ``TopKRouter`` here too, otherwise the MoE ``router`` has no ``expert_bias_vl`` + # buffer and the V4.1 bridge fails to load ``gate.bias_vl`` (mirrors ``_replace_router``). + self._replace_hybrid_router(spec) + return spec + + @staticmethod + def _replace_hybrid_router(spec): + from functools import partial + + from megatron.core.transformer.moe.router import TopKRouter as McoreTopKRouter + + from ..modules import TopKRouter + moe_layer = getattr(spec.submodules, 'moe_layer', None) + mlp_spec = getattr(getattr(moe_layer, 'submodules', None), 'mlp', None) + # ``get_moe_module_spec_for_backend`` hands back a ``functools.partial(MoELayer, ...)`` + # here (not a plain ``ModuleSpec``), so read its ``submodules`` from ``keywords`` -- same + # dual handling as ``ModelLoader._replace_router``. + if isinstance(mlp_spec, partial): + mlp_submodules = mlp_spec.keywords.get('submodules') + else: + mlp_submodules = getattr(mlp_spec, 'submodules', None) + if getattr(mlp_submodules, 'router', None) is McoreTopKRouter: + mlp_submodules.router = TopKRouter + + def _rewrap_engram_hyper_connection_layers(self, model): + """Retrofit Engram-carrying ``HyperConnectionHybridLayer`` wrappers with the V4.1 + subclass that declines the fast path (see + :class:`DeepseekV41HyperConnectionHybridLayer`). + + HybridStack hard-codes ``HyperConnectionHybridLayer`` (hybrid_block.py:1120-1121) with no + spec hook, so the swap is done in place after build. ``DeepseekV41HyperConnectionHybridLayer`` + only overrides one method and adds no state, making the ``__class__`` reassignment safe. + Only wrappers whose inner layer actually built an Engram module (``layer_number in + engram_config.layer_ids``) are touched; every other layer keeps the base fast path and + stays numerically identical to a plain hybrid stack. + """ + if not self.config.enable_hyper_connections or DeepseekV41HyperConnectionHybridLayer is None: + return + decoder = getattr(model, 'decoder', None) + for layer in getattr(decoder, 'layers', []) or []: + inner = getattr(layer, 'inner_layer', None) + if (isinstance(layer, HyperConnectionHybridLayer) and inner is not None + and getattr(inner, 'engram', None) is not None): + layer.__class__ = DeepseekV41HyperConnectionHybridLayer + + def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): + """Build the multimodal wrapper around ``HybridModel``, skipping ``ModelLoader.build_model``'s + GPT layer-spec post-processing (MLA / router / TransformerLayer substitution): a + ``HybridStack`` spec exposes per-symbol submodules instead, and the DSv4 attention swap is + done in ``get_transformer_layer_spec`` above. + + ``model`` is :class:`DeepseekV41MultimodalModel` (vision tower + wrapper); the hybrid + text backbone -- which owns the decoder / MoE / Engram layers the fix-ups below touch -- + is nested under ``model.language_model``, so they target that.""" + self._hybrid_config = self._build_hybrid_config() + model = self.model_cls( + config=self._hybrid_config, + transformer_layer_spec=self.get_transformer_layer_spec(vp_stage=vp_stage), + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + language_model = getattr(model, 'language_model', model) + self._rewrap_engram_hyper_connection_layers(language_model) + self._set_linear_is_expert(language_model) + # DSpark: the ``mtp.*`` draft stack is backbone-agnostic (plain experimental-attention + # layers), so :meth:`_attach_dspark` builds it unchanged and attaches it to the hybrid text + # backbone (``language_model.dspark``). Inference-time target-layer capture on HybridStack is + # not implemented (it is not exercised by training / weight round-trip); the stack only needs + # to exist so its parameters are loaded / saved via ``mtp.*``. + self._attach_dspark(language_model, post_process, vp_stage=vp_stage) + return model + + def _get_engram_config(self): + hf_layer_ids = tuple(self.config.engram_layer_ids or ()) + if not hf_layer_ids: + return None + if not has_native_engram(): + raise RuntimeError('DeepSeek-V4.1 Engram requires NVIDIA Megatron-LM Engram support. ' + 'The PR #7224 text-backbone baseline intentionally does not provide it; ' + 'install the official Engram extension or disable Engram explicitly.') + required = ( + 'engram_num_embeddings', + 'engram_max_ngram_size', + 'engram_vocab_size', + 'engram_n_heads', + 'engram_head_dim', + 'engram_pad_token_id', + ) + missing = [name for name in required if getattr(self.config, name, None) is None] + if missing: + raise ValueError(f'DeepSeek-V4.1 Engram config is missing required fields: {missing}.') + + tokenizer_map = self.config.engram_tokenizer_map + if tokenizer_map is None: + model_dir = getattr(self.config.hf_config, 'name_or_path', '') + candidate = os.path.join(model_dir, 'engram_tokenizer_map.json') if model_dir else '' + if candidate and os.path.isfile(candidate): + tokenizer_map = candidate + if not tokenizer_map: + raise ValueError('DeepSeek-V4.1 Engram requires engram_tokenizer_map. Generate it with ' + 'Megatron-LM/tools/engram/generate_tokenizer_map.py using the HF 0-based ' + f'layer IDs {list(hf_layer_ids)}.') + + max_ngram_order = self.config.engram_max_ngram_size + image_token_id = getattr(self.config.hf_config, 'image_token_id', None) + engram_config = build_deepseek_v41_engram_config( + global_vocab_sizes=(self.config.engram_vocab_size, ) * (max_ngram_order - 1), + # TransformerLayer numbers are 1-based, while the official checkpoint and + # tokenizer artifact use the original 0-based HF layer IDs. + placement_layer_ids=self._engram_placement_layer_ids(hf_layer_ids), + hash_layer_ids=hf_layer_ids, + max_ngram_order=max_ngram_order, + num_hash_heads=self.config.engram_n_heads, + memory_dim=self.config.engram_n_heads * self.config.engram_head_dim, + kernel_size=1, + hash_seed=0, + boundary_token_id=self.config.engram_pad_token_id, + tokenizer_map_path=tokenizer_map, + excluded_token_ids=(() if image_token_id is None else (image_token_id, )), + ) + actual_rows = tuple(sum(engram_config.table_sizes(layer_id)) for layer_id in engram_config.layer_ids) + expected_rows = tuple(self.config.engram_num_embeddings) + if actual_rows != expected_rows: + raise ValueError('DeepSeek-V4.1 Engram table layout does not match engram_num_embeddings: ' + f'computed {actual_rows}, checkpoint declares {expected_rows}.') + if (self.config.engram_compressed_vocab_size is not None + and engram_config.tokenizer_remap.max().item() + 1 != self.config.engram_compressed_vocab_size): + raise ValueError('DeepSeek-V4.1 compressed tokenizer vocabulary mismatch: artifact has ' + f'{engram_config.tokenizer_remap.max().item() + 1}, config declares ' + f'{self.config.engram_compressed_vocab_size}.') + engram_config.validate_startup(self.config, expected_tokenizer_vocab_size=self.config.padded_vocab_size) + return engram_config + + def get_dspark_layer_spec(self): + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + _get_backend_spec_provider, + get_transformer_layer_with_experimental_attention_variant_spec, + ) + + dspark_config = copy.copy(self.config) + dspark_config.hf_config = getattr(self.config.hf_config, 'text_config', self.config.hf_config) + dspark_config.num_layers = self.config.dspark_num_layers + # A checkpoint may leave the draft stack's expert counts out, which means its draft layers are + # shaped like the backbone's MoE rather than carrying their own shape. + if self.config.dspark_num_experts is not None: + dspark_config.num_moe_experts = self.config.dspark_num_experts + if self.config.dspark_router_topk is not None: + dspark_config.moe_router_topk = self.config.dspark_router_topk + dspark_config.moe_layer_freq = [1] * self.config.dspark_num_layers + dspark_config.first_pipeline_num_layers = None + dspark_config.last_pipeline_num_layers = None + dspark_config.num_layers_in_first_pipeline_stage = None + dspark_config.num_layers_in_last_pipeline_stage = None + dspark_config.sequence_parallel = False + dspark_config.csa_compress_ratios = [0] * self.config.dspark_num_layers + dspark_config.csa2_kv_source_layers = [] + dspark_config.csa2_index_source_layers = [] + dspark_config.csa2_candidate_source_layer = None + backend = _get_backend_spec_provider(config=dspark_config) + layer_specs = get_transformer_layer_with_experimental_attention_variant_spec( + config=dspark_config, backend=backend) + for layer_spec in layer_specs: + attention_spec = layer_spec.submodules.self_attention + attention_spec.module = DeepseekV41DSparkAttention + attention_spec.submodules.core_attention.module = DeepseekV41DSparkCoreAttention + # DSpark specs bypass ModelLoader.build_model, so apply the same router + # override the main layers get: swap the vendored McoreTopKRouter for the + # custom TopKRouter, otherwise the draft MoE has no ``expert_bias_vl`` and + # loading the checkpoint's ``mtp.*.ffn.gate.bias_vl`` asserts. + self._replace_router(SimpleNamespace(layer_specs=layer_specs)) + return dspark_config, layer_specs + + def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = None): + """Build the DSpark (``mtp.*``) draft stack and attach it to ``language_model`` on the + final pipeline stage. + + Backbone-agnostic: the draft layers are plain experimental-attention-variant + ``TransformerLayer`` instances (see :meth:`get_dspark_layer_spec`), independent of whether + the main model is a ``GPTModel`` or ``HybridModel``; the caller passes whichever object + owns the stack -- here the ``HybridModel`` backbone, which exposes ``pg_collection`` / + ``vocab_size`` / ``config`` all the same. + The stack is never part of the training forward (capture is inference-only), so it only + needs to exist here so its parameters are loaded / saved through the ``mtp.*`` bridge -- which + is also why it is frozen at the end of this method. + + ``vp_stage`` must be threaded into ``build_module`` because the draft layers reuse the + experimental-attention ``TransformerLayer``, whose ``__init__`` calls + ``get_transformer_layer_offset`` -- and that helper asserts ``vp_stage is not None`` under + VPP. The draft stack keeps its own local 1-based numbering; the pipeline offset the helper + adds is the same value the last stage already applied under plain PP (0 for the tiny draft + stack), so this only satisfies the VPP assertion without changing placement. + """ + if not self.config.dspark_num_layers or not post_process: + return + dspark_config, dspark_layer_specs = self.get_dspark_layer_spec() + layers = [ + build_module( + layer_spec, + config=dspark_config, + layer_number=index + 1, + pg_collection=language_model.pg_collection, + vp_stage=vp_stage, + ) for index, layer_spec in enumerate(dspark_layer_specs) + ] + language_model.dspark = DeepseekV41DSparkStack(dspark_config, layers) + self._set_linear_is_expert(language_model.dspark) + # ``Float16Module`` casts every unmarked float buffer, and mcore's ``TopKRouter`` only + # restores the aux-loss-free bias to fp32 lazily -- from ``forward`` and from + # ``_save_to_state_dict``. The draft stack never runs in the training forward, and the + # bridge copies ``param.data`` directly instead of going through the state-dict hooks, so + # without this marker the checkpoint's fp32 ``mtp.*.ffn.gate.bias`` would round-trip + # through bf16. The main layers escape it only because their routers do run. + for layer in layers: + expert_bias = getattr(getattr(layer.mlp, 'router', None), 'expert_bias', None) + if expert_bias is not None: + mark_keep_in_fp32(expert_bias) + # DSpark embeds its draft seed with the base input embedding. On a PP>1 last + # stage with untied embeddings the base model has no ``embedding`` here, so + # build a dedicated replicated DSpark embedding; the bridge loads it from + # ``model.embed_tokens.weight`` (same source as the first-stage embedding). + if not hasattr(language_model, 'embedding'): + language_model.dspark_word_embeddings = VocabParallelEmbedding( + language_model.vocab_size, + language_model.config.hidden_size, + init_method=language_model.config.init_method, + config=language_model.config, + tp_group=language_model.pg_collection.tp, + ) + # Nothing here runs in the training forward, so none of these parameters can ever receive a + # gradient -- and leaving them trainable does more than waste the optimizer state and the + # gradient buffers. Adam's weight decay is decoupled from the gradient, so a parameter whose + # gradient stays zero is still multiplied by ``1 - lr * wd`` on every step with nothing to + # balance it, and the ``mtp.*`` weights loaded from the checkpoint would decay away over a + # long run (bf16 rounding only hides this until the drift crosses one ulp of the parameter; + # the fp32 main weights shrink from the first step). Freezing keeps them out of the optimizer + # altogether. The bridge reads and writes ``param.data`` directly, so the checkpoint still + # round-trips the draft stack unchanged -- which is the whole reason it is built here. + for module in (language_model.dspark, getattr(language_model, 'dspark_word_embeddings', None)): + if module is not None: + for param in module.parameters(): + param.requires_grad = False + + +class DeepseekV41Bridge(DeepseekV4Bridge): + """Weight bridge for the ``HybridModel`` backbone. + + An HF layer owns both attention and MLP; on ``HybridModel`` it is split in two (see + :func:`derive_hybrid_layer_config`), so this bridge fans a single HF layer ``i`` out onto + two hybrid layers: + + * hybrid layer ``2*i`` -- attention half: MLA / CSA2 state + ``attn_norm`` (+ Engram when + ``i in engram_layer_ids``) + the ``hc_attn_*`` hyper-connection channel. + * hybrid layer ``2*i + 1`` -- MLP half: MoE / dense state + ``ffn_norm`` + the ``hc_ffn_*`` + hyper-connection channel. + + When ``enable_hyper_connections`` is set each hybrid layer is wrapped in a + ``HyperConnectionHybridLayer`` whose real payload lives under ``inner_layer`` and which owns + a *single* ``hyper_connection`` module, so the HF ``hc_{attn,ffn}_*`` keys split across the + two wrappers. + + ``self.config`` is seen in two layer spaces depending on direction: on load it is the + original HF-space config (``num_layers == N``, no ``hybrid_layer_pattern``); on export it is + the doubled hybrid megatron config used to build the model (``num_layers == 2 * N``, + ``hybrid_layer_pattern`` populated). :meth:`_convert` normalizes this so it always iterates + the decoder's ``2 * N`` hybrid layers. + """ + + _ENGRAM_LOAD_CHUNK_ROWS = 65536 + additional_dim0_keys = DeepseekV4Bridge.additional_dim0_keys | {'embed', 'head'} + additional_dim1_keys = DeepseekV4Bridge.additional_dim1_keys | {'main_proj'} + + @staticmethod + def _lm(mg_model): + """Resolve the language model. A bare HybridModel is text-only (no ``language_model`` + wrapper); :class:`DeepseekV41MultimodalModel` nests it under a multimodal container.""" + language_model = getattr(mg_model, 'language_model', None) + return mg_model if language_model is None else language_model + + @staticmethod + def _num_hybrid_layers(config) -> int: + """Decoder layer count in the doubled hybrid space, regardless of which layer space + ``config`` is currently in. + + The two conversion entrypoints hand :meth:`_convert` a config in *different* spaces: + load (``to_mcore=True``) passes the original HF-space config (``num_layers == N``, no + ``hybrid_layer_pattern``) whose built decoder holds ``2 * N`` layers; export + (``to_mcore=False``) passes the doubled hybrid megatron config used to build the model + (``num_layers == 2 * N`` with ``hybrid_layer_pattern`` populated). Discriminating by the + pattern makes both directions iterate exactly the decoder's layer count -- using the raw + ``2 * num_layers`` on export would over-count and dereference ``None`` layers past the + decoder end (see PP-availability window in :meth:`_convert`).""" + if getattr(config, 'hybrid_layer_pattern', None): + return config.num_layers + return 2 * config.num_layers + + def _engram_hf_layer_id(self, engram): + # Engram lives on the doubled-space attention layer ``2 * hf_id + 1`` (see + # ``DeepseekV41Loader._engram_placement_layer_ids``), so map it back to HF space. + return (engram.layer_number - 1) // 2 + + def _set_word_embeddings(self, mg_model, hf_state_dict, to_mcore): + # The base ``MultimodalGPTBridge`` resolves the language model with a raw + # ``getattr(mg_model, 'language_model')``; route it through :meth:`_lm` so both the + # multimodal wrapper and a bare backbone resolve correctly. + self._set_state_dict( + self._lm(mg_model), 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key, to_mcore) + + def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): + # Runs the vision/aligner + image_* block unconditionally instead of branching on *this* + # rank's ``mg_model.visual``. On export every pipeline stage runs ``_convert`` -> + # ``_convert_pre_process``, and this path issues the *same* pp-group collective sequence on + # all ranks: word-embeddings (routed through ``_lm`` by :meth:`_set_word_embeddings`), then + # the config-guarded vision/aligner block and the image_* markers, all driven via + # ``_set_module``/``_set_state_dict`` which stay in lockstep even where the submodule is + # ``None`` (see ``_set_module``'s ``src_rank`` all-reduce / ``_set_state_dict``'s ``state`` + # all-reduce). A per-rank ``visual is not None`` guard would skip that whole block on + # non-first stages (where the wrapper built ``visual=None``), desynchronizing the + # collectives so the last stage's later per-layer ``has_model`` all-reduce reads a stale + # value -> ``next(mg_models)`` -> ``StopIteration``. On load only the first stage reaches + # this method (see ``_convert``'s ``is_pp_first_stage`` guard). + result = super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) + target = hf_state_dict if to_mcore else result + for name in ('image_start', 'image_end', 'image_newline'): + self._set_state_dict(mg_model, f'visual.{name}', target, f'model.{name}', to_mcore) + return result + + def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + lm_model = self._lm(mg_model) + if self.config.task_type != 'embedding': + if self.config.untie_embeddings_and_output_weights: + hf_lm_head_key = self.hf_lm_head_key + if self.config.task_type == 'seq_cls': + hf_lm_head_key = self.hf_score_key + if not to_mcore or hf_lm_head_key in hf_state_dict: + self._set_state_dict(lm_model, 'output_layer.weight', hf_state_dict, hf_lm_head_key, to_mcore) + elif to_mcore and lm_model.output_layer.weight is not None: + self._set_state_dict(lm_model, 'output_layer.weight', hf_state_dict, self.hf_embed_key, to_mcore) + self._set_final_layernorm(lm_model, hf_state_dict, to_mcore) + if to_mcore: + return {} + return self._add_prefix(hf_state_dict, hf_prefix) + + def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): + # HybridStack names its trailing norm ``final_norm`` (vs the GPT block's + # ``final_layernorm``). Like the GPT V4.1 bridge, single-pass mHC has no learned + # ``hc_head_*`` output head (only built when ``not mhc_single_pass``), so nothing else + # is mapped here. + self._set_state_dict(lm_model, 'decoder.final_norm.weight', hf_state_dict, self.hf_final_layernorm_key, + to_mcore) + + def _set_one_hyper_connection(self, hyper_connection, hf_state_dict, hf_key, to_mcore): + """Bridge a single ``HyperConnectionModule`` (one wrapper == one channel). + + Same parameter layout as the GPT ``_set_hyper_connection`` per-channel body, but keyed + by an explicit ``hf_key`` ('attn' or 'ffn') because each hybrid wrapper owns exactly one + connection instead of the GPT layer's attention + FFN pair. + """ + self._set_state_dict(hyper_connection, 'mapping_proj.weight', hf_state_dict, f'hc_{hf_key}_fn', 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) + # ``alpha_*`` bypass ``_set_state_dict``, so mirror the peft guard the GPT + # ``_set_hyper_connection`` applies -- these are frozen base weights and must stay out of + # ``adapter_model.safetensors``. + 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']): + getattr(hyper_connection, f'alpha_{alpha_suffix}').data[:] = alpha[i] + else: + alpha = None + if hyper_connection is not None: + alpha = torch.concat( + [getattr(hyper_connection, f'alpha_{suffix}') for suffix in ['pre', 'post', 'res']], dim=0) + hf_state_dict[f'hc_{hf_key}_scale'] = self._get_weight(alpha, 'alpha')[0] + + def _set_hybrid_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, hybrid_idx: int, to_mcore: bool): + """Map HF layer ``hybrid_idx // 2`` onto one half of the hybrid pair. + + Even ``hybrid_idx`` is the attention half, odd is the MLP half; both read/write the + same ``model.layers.{hf_idx}.`` prefix so the HF checkpoint stays single-layer-per-index. + """ + hf_idx = hybrid_idx // 2 + is_attn = (hybrid_idx % 2 == 0) + layer_prefix = f'{hf_prefix}{hf_idx}.' + local_state = self._remove_prefix(hf_state_dict, layer_prefix) if to_mcore else {} + # The wrapper carries the payload under ``inner_layer`` and the single hyper-connection + # under ``hyper_connection``; without mHC the layer is the payload itself. + inner = None if mg_layer is None else getattr(mg_layer, 'inner_layer', mg_layer) + hyper_connection = None if mg_layer is None else getattr(mg_layer, 'hyper_connection', None) + if is_attn: + local_state.update(self._set_layer_attn(inner, local_state, hf_idx, to_mcore)) + if hf_idx in (self.config.engram_layer_ids or []): + # ``_get_layer_engram`` already unwraps ``inner_layer.engram``. + self._set_layer_engram(mg_layer, local_state, to_mcore) + if self.config.enable_hyper_connections: + self._set_one_hyper_connection(hyper_connection, local_state, 'attn', to_mcore) + else: + local_state.update(self._set_layer_mlp(inner, local_state, hf_idx, to_mcore)) + if self.config.enable_hyper_connections: + self._set_one_hyper_connection(hyper_connection, local_state, 'ffn', to_mcore) + if to_mcore: + return {} + return self._add_prefix(local_state, layer_prefix) + + def _convert_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage): + """Map the DSpark (``mtp.*``) draft stack. + + The draft layers are plain experimental-attention ``TransformerLayer`` instances -- + backbone-agnostic -- so :meth:`_convert_dspark_stack` owns the mapping and this method + only locates the stack. It is attached to the text backbone (no ``language_model`` wrapper on + a bare model, so use :meth:`_lm`) and only + on the final pipeline stage. During export non-last stages use an empty structural proxy so + every PP rank executes the same collective sequence. MTP (``mtp_num_layers``) is skipped in + :meth:`_convert` and does not apply to V4.1.""" + if not self.config.dspark_num_layers or (to_mcore and not is_pp_last_stage): + return + language_model = self._lm(mg_model) + dspark = getattr(language_model, 'dspark', None) + if dspark is None: + if to_mcore or is_pp_last_stage: + raise RuntimeError('DSpark weights require the draft stack on the final pipeline stage.') + dspark = SimpleNamespace(layers=[None] * self.config.dspark_num_layers) + yield from self._convert_dspark_stack(language_model, dspark, hf_state_dict, hf_prefix, to_mcore) + + def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqdm_desc: str = 'Converting: '): + """Backbone conversion with a 1->2 layer fan-out. + + Mirrors :meth:`GPTBridge._convert` but iterates the doubled hybrid layer space + (``2 * num_layers``) and dispatches each hybrid layer to :meth:`_set_hybrid_layer_state`. + MTP is intentionally skipped: V4.1's ``mtp.*`` keys are DSpark, mapped separately. + """ + self._pending_export_iter = None + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore) + else: + hf_state_dict = {} + mg_models = iter(mg_models) + mg_model = next(mg_models) + is_pp_first_stage = mpu.is_pipeline_first_stage(ignore_virtual=False, vp_stage=mg_model.vp_stage) + is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=mg_model.vp_stage) + if not to_mcore or is_pp_first_stage: + hf_state_dict.update(self._convert_pre_process(mg_model, hf_state_dict, '', to_mcore)) + if to_mcore: + yield + else: + hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore) + yield from list(self._add_prefix(hf_state_dict, hf_prefix).items()) + hf_state_dict = {} + # Total hybrid (attention + MLP) layer count in the doubled space; ``_num_hybrid_layers`` + # normalizes the two layer spaces ``self.config`` may be in (see its docstring) so both + # load and export iterate exactly the decoder's layer count. HybridStack layer_number + # spans this same space (i + 1 + pp_offset), matching this loop's hybrid index so the + # PP-availability window below stays correct. + num_hybrid_layers = self._num_hybrid_layers(self.config) + layer_idx = 0 + disable_tqdm = self._disable_tqdm or not is_master() + prog_bar = tqdm(range(num_hybrid_layers), dynamic_ncols=True, desc=tqdm_desc, disable=disable_tqdm) + while layer_idx < num_hybrid_layers: + lm_model = self._lm(mg_model) + if len(lm_model.decoder.layers) > 0: + start_idx = lm_model.decoder.layers[0].layer_number - 1 + mg_layer_available = (start_idx <= layer_idx < lm_model.decoder.layers[-1].layer_number) + else: + mg_layer_available = False + if mg_layer_available: + mg_layer = lm_model.decoder.layers[layer_idx - start_idx] + else: + if to_mcore: + layer_idx += 1 + prog_bar.update() + continue + else: + mg_layer = None + if not to_mcore and self.pp_size > 1: + has_model = torch.tensor([mg_layer is not None], dtype=torch.bool, device='cuda') + dist.all_reduce(has_model, group=self.pp_group) + if not has_model: + mg_model = next(mg_models) # compat vpp + continue + res = self._set_hybrid_layer_state(mg_layer, hf_state_dict, f'{self.hf_layers_prefix}.', layer_idx, + to_mcore) + layer_idx += 1 + prog_bar.update() + if to_mcore: + yield + else: + res = self._convert_hf_state_dict(res, to_mcore) + yield from self._drain_pending_export(hf_prefix) + yield from self._add_prefix(res, hf_prefix).items() + hf_state_dict = {} + prog_bar.close() + 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: + yield + else: + hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore) + yield from list(self._add_prefix(hf_state_dict, hf_prefix).items()) + + def _set_router(self, mg_mlp, hf_state_dict, to_mcore, **kwargs): + super()._set_router(mg_mlp, hf_state_dict, to_mcore, **kwargs) + if self.config.moe_router_enable_vl_bias: + self._set_state_dict(mg_mlp, 'router.expert_bias_vl', hf_state_dict, 'gate.bias_vl', to_mcore) + + @staticmethod + def _get_layer_engram(mg_layer): + if mg_layer is None: + return None + engram = getattr(mg_layer, 'engram', None) + if engram is None: + engram = getattr(getattr(mg_layer, 'inner_layer', None), 'engram', None) + return engram + + @staticmethod + def _load_lazy_slice(lazy_tensor, row_start, row_end): + if hasattr(lazy_tensor, 'load_slice'): + return lazy_tensor.load_slice(slice(row_start, row_end)) + return lazy_tensor.load()[row_start:row_end] + + @staticmethod + def _dequantize_engram_rows(weight, scale): + if scale is None: + return weight + if weight.ndim != 2 or scale.ndim != 2 or weight.shape[0] != scale.shape[0]: + raise ValueError(f'Invalid Engram FP8 weight/scale shapes: {tuple(weight.shape)} and {tuple(scale.shape)}.') + if weight.shape[1] % scale.shape[1] != 0: + raise ValueError(f'Engram weight width {weight.shape[1]} is not divisible by scale width {scale.shape[1]}.') + block_size = weight.shape[1] // scale.shape[1] + return (weight.float().unflatten(-1, (-1, block_size)) * scale.float().unsqueeze(-1)).flatten(-2) + + def _load_engram_embedding(self, engram, hf_state_dict): + weight = hf_state_dict['engram.embed.weight'] + scale = hf_state_dict.get('engram.embed.weight_scale_inv') + flat_offset = 0 + for table in engram.embedding.tables: + table_offset = flat_offset + flat_offset += table.global_num_embeddings + for local_start in range(0, table.local_num_embeddings, self._ENGRAM_LOAD_CHUNK_ROWS): + local_end = min(local_start + self._ENGRAM_LOAD_CHUNK_ROWS, table.local_num_embeddings) + source_start = table_offset + table.row_start + local_start + source_end = table_offset + table.row_start + local_end + rows = self._load_lazy_slice(weight, source_start, source_end) + row_scales = None if scale is None else self._load_lazy_slice(scale, source_start, source_end) + rows = self._dequantize_engram_rows(rows, row_scales) + table.weight.data[local_start:local_end].copy_( + rows.to(device=table.weight.device, dtype=table.weight.dtype)) + hf_layer_id = self._engram_hf_layer_id(engram) + expected_rows = self.config.engram_num_embeddings[self.config.engram_layer_ids.index(hf_layer_id)] + if flat_offset != expected_rows: + raise ValueError(f'Engram layer {hf_layer_id} expected {expected_rows} flat rows, ' + f'but its prime tables contain {flat_offset}.') + + def _set_layer_engram(self, mg_layer, hf_state_dict, to_mcore): + engram = self._get_layer_engram(mg_layer) + if to_mcore: + if engram is None: + return + self._load_engram_embedding(engram, hf_state_dict) + wkv = hf_state_dict['engram.wkv.weight'].load() + wkv_scale = hf_state_dict.get('engram.wkv.weight_scale_inv') + if wkv_scale is not None: + wkv_scale = wkv_scale.load() + wkv = self._dequantize_engram_rows(wkv, wkv_scale) + key_rows = engram.num_streams * engram.hidden_size + if tuple(wkv.shape) != (key_rows + engram.hidden_size, engram.engram_config.total_memory_dim): + raise ValueError(f'Unexpected DeepSeek-V4.1 Engram wkv shape: {tuple(wkv.shape)}.') + engram.key_projection.weight.data.copy_(wkv[:key_rows].to(engram.key_projection.weight)) + engram.value_projection.weight.data.copy_(wkv[key_rows:].to(engram.value_projection.weight)) + engram.query_norm.weight.data.copy_(hf_state_dict['engram.q_weight'].load().reshape(-1).to( + engram.query_norm.weight)) + engram.key_norm.weight.data.copy_(hf_state_dict['engram.k_weight'].load().reshape(-1).to( + engram.key_norm.weight)) + elif not self._peft_format: + if getattr(self, '_skip_unsupported_export', False): + # On-policy RL weight sync: Engram tables are frozen and already resident in the + # rollout engine from the base checkpoint, so skip re-exporting them (a full 183 GiB + # resync per step is infeasible) instead of raising. + return + if engram is None: + return + # --- export embedding tables to a single flat tensor (bf16, no FP8) --- + all_rows = [] + for table in engram.embedding.tables: + all_rows.append(table.weight.data.cpu()) + hf_state_dict['engram.embed.weight'] = torch.cat(all_rows, dim=0) + # --- export key+value projections as combined wkv --- + key_w = engram.key_projection.weight.data.cpu() # [stream_width, total_memory_dim] + val_w = engram.value_projection.weight.data.cpu() # [hidden, total_memory_dim] + hf_state_dict['engram.wkv.weight'] = torch.cat([key_w, val_w], dim=0) + # --- export norm weights as q_weight / k_weight --- + num_streams = engram.num_streams + hf_state_dict['engram.q_weight'] = engram.query_norm.weight.data.cpu().reshape(num_streams, -1) + hf_state_dict['engram.k_weight'] = engram.key_norm.weight.data.cpu().reshape(num_streams, -1) + + def _set_dspark_layer_state(self, mg_layer, hf_state_dict, layer_idx, to_mcore): + stage_prefix = f'{self.hf_mtp_prefix}.{layer_idx}.' + local_state = self._remove_prefix(hf_state_dict, stage_prefix) if to_mcore else {} + local_state.update(self._set_layer_attn(mg_layer, local_state, layer_idx, to_mcore)) + local_state.update(self._set_layer_mlp(mg_layer, local_state, layer_idx, to_mcore, is_mtp=True)) + self._set_hyper_connection(mg_layer, local_state, layer_idx, to_mcore) + if to_mcore: + return {} + return self._add_prefix(local_state, stage_prefix) + + def _set_dspark_endpoints(self, dspark, hf_state_dict, to_mcore): + first_prefix = f'{self.hf_mtp_prefix}.0.' + last_prefix = f'{self.hf_mtp_prefix}.{self.config.dspark_num_layers - 1}.' + first_state = self._remove_prefix(hf_state_dict, first_prefix) if to_mcore else {} + last_state = self._remove_prefix(hf_state_dict, last_prefix) if to_mcore else {} + self._set_state_dict(dspark, 'input.main_proj.weight', first_state, 'main_proj.weight', to_mcore) + self._set_state_dict(dspark, 'input.main_norm.weight', first_state, 'main_norm.weight', to_mcore) + self._set_state_dict(dspark, 'output.norm.weight', last_state, 'norm.weight', to_mcore) + self._set_state_dict(dspark, 'output.markov_head.embed.weight', last_state, 'markov_head.embed.weight', + to_mcore) + self._set_state_dict(dspark, 'output.markov_head.head.weight', last_state, 'markov_head.head.weight', to_mcore) + self._set_state_dict(dspark, 'output.confidence_head.proj.weight', last_state, 'confidence_head.proj.weight', + to_mcore) + if to_mcore: + return {} + result = self._add_prefix(first_state, first_prefix) + result.update(self._add_prefix(last_state, last_prefix)) + return result + + def _load_dspark_word_embeddings(self, embedding, hf_state_dict): + """Load the dedicated DSpark input embedding from ``model.embed_tokens.weight``. + + Mirrors the base embedding load: pad the HF rows up to ``padded_vocab_size`` + (already a multiple of TP) and take this rank's vocab-parallel shard. + """ + weight = hf_state_dict[self.hf_embed_key].load() + padded = self.config.padded_vocab_size + if weight.shape[0] < padded: + weight = F.pad(weight, (0, 0, 0, padded - weight.shape[0])) + if self.tp_size > 1: + weight = weight.chunk(self.tp_size, dim=0)[self.tp_rank] + embedding.weight.data.copy_(weight.to(device=embedding.weight.device, dtype=embedding.weight.dtype)) + + def _convert_dspark_stack(self, language_model, dspark, hf_state_dict, hf_prefix, to_mcore): + """Map the DSpark draft layers + endpoints between HF ``mtp.*`` keys and the megatron + stack. Backbone-agnostic (the draft layers are plain ``TransformerLayer`` instances); the + caller locates ``dspark`` and guards the pipeline stage.""" + # On a PP>1 last stage with untied embeddings, DSpark owns a dedicated input + # embedding (see build_model). Load it from the same HF source as the base + # first-stage embedding. On export the first stage already emits this tensor, + # so the redundant DSpark copy is not written back. + if to_mcore and getattr(language_model, 'dspark_word_embeddings', None) is not None: + self._load_dspark_word_embeddings(language_model.dspark_word_embeddings, hf_state_dict) + yield + + original_num_experts = self.config.num_moe_experts + # Mirrors get_dspark_layer_spec: without its own expert count the draft stack was built with the + # backbone's, so the weight mapping has to agree with what was built. + if self.config.dspark_num_experts is not None: + self.config.num_moe_experts = self.config.dspark_num_experts + try: + for layer_idx, layer in enumerate(dspark.layers): + result = self._set_dspark_layer_state(layer, hf_state_dict, layer_idx, to_mcore) + if to_mcore: + yield + else: + result = self._convert_hf_state_dict(result, to_mcore) + yield from self._add_prefix(result, hf_prefix).items() + result = self._set_dspark_endpoints(dspark, hf_state_dict, to_mcore) + if to_mcore: + yield + else: + result = self._convert_hf_state_dict(result, to_mcore) + yield from self._add_prefix(result, hf_prefix).items() + finally: + self.config.num_moe_experts = original_num_experts + + def _set_mla_attn_state(self, mg_attn, hf_state_dict, hf_prefix, layer_idx, to_mcore): + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + # --- shared MLA projections (identical to V4) --- + self._set_state_dict(mg_attn, 'linear_proj.weight', hf_state_dict, 'wo_b.weight', to_mcore) + if self.config.fp8_param: + self._set_o_group_proj_grouped(mg_attn, hf_state_dict, to_mcore) + else: + self._set_state_dict(mg_attn, 'linear_o_group_proj', hf_state_dict, 'wo_a.weight', to_mcore) + self._set_state_dict(mg_attn, 'linear_q_down_proj.weight', hf_state_dict, 'wq_a.weight', to_mcore) + self._set_state_dict(mg_attn, 'linear_q_up_proj.weight', hf_state_dict, 'wq_b.weight', to_mcore) + self._set_state_dict(mg_attn, 'linear_kv_proj.weight', hf_state_dict, 'wkv.weight', to_mcore) + self._set_state_dict(mg_attn, 'core_attention.attn_sink', hf_state_dict, 'attn_sink', to_mcore) + if self.config.qk_layernorm: + self._set_state_dict(mg_attn, 'q_layernorm.weight', hf_state_dict, 'q_norm.weight', to_mcore) + self._set_state_dict(mg_attn, 'kv_layernorm.weight', hf_state_dict, 'kv_norm.weight', to_mcore) + # --- CSA2 compressor / indexer (no `ape`; indexer owns wk/k_norm on owns_k) --- + core_attn = None if mg_attn is None else mg_attn.core_attention + compressor = None if core_attn is None else getattr(core_attn, 'compressor', None) + indexer = None if core_attn is None else getattr(core_attn, 'indexer', None) + has_compressor = self._reduce_tensor_pp_group(compressor is not None, to_mcore) + has_indexer = self._reduce_tensor_pp_group(indexer is not None, to_mcore) + # ratio-2 compressor layers additionally own a gate projection. + has_wgate = self._reduce_tensor_pp_group( + compressor is not None and getattr(compressor, 'linear_wgate', None) is not None, to_mcore) + # kv-source (owns_k) indexer layers additionally own wk + k_norm. + owns_k = self._reduce_tensor_pp_group( + indexer is not None and getattr(indexer, 'linear_wk', None) is not None, to_mcore) + if has_compressor: + self._set_state_dict(mg_attn, 'core_attention.compressor.linear_wkv.weight', hf_state_dict, + 'compressor.wkv.weight', to_mcore) + self._set_state_dict(mg_attn, 'core_attention.compressor.norm.weight', hf_state_dict, + 'compressor.norm.weight', to_mcore) + if has_wgate: + self._set_state_dict(mg_attn, 'core_attention.compressor.linear_wgate.weight', hf_state_dict, + 'compressor.wgate.weight', to_mcore) + if has_indexer: + self._set_state_dict(mg_attn, 'core_attention.indexer.linear_wq_b.weight', hf_state_dict, + 'indexer.wq_b.weight', to_mcore) + self._set_state_dict(mg_attn, 'core_attention.indexer.linear_weights_proj.weight', hf_state_dict, + 'indexer.weights_proj.weight', to_mcore) + if owns_k: + self._set_state_dict(mg_attn, 'core_attention.indexer.linear_wk.weight', hf_state_dict, + 'indexer.wk.weight', to_mcore) + self._set_state_dict(mg_attn, 'core_attention.indexer.k_norm.weight', hf_state_dict, + 'indexer.k_norm.weight', to_mcore) + if to_mcore: + hf_state_dict = {} + else: + hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) + return hf_state_dict + + +if _HYBRID_MODEL_AVAILABLE: + register_model( + ModelMeta( + ModelType.deepseek_v41, + ['deepseek_v41'], + bridge_cls=DeepseekV41Bridge, + visual_cls=DeepseekV41Vision, + loader=DeepseekV41Loader, + config_cls=MLAModelConfig, + )) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index e96f082..356c163 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -31,6 +31,16 @@ def __init__(self, **kwargs) self.vp_stage = self.language_model.vp_stage self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights + # Surface the language model's typed-pipeline payload interface on the wrapper. The PP + # schedulers locate it with ``get_attr_wrapped_model(chunk, 'pipeline_payload_factory')``, + # which only descends through ``.module`` wrappers and never reaches ``self.language_model``. + # Without this a HybridModel backbone that configures a custom cross-stage payload (e.g. + # DeepSeek-V4.1 CSA2 / single-pass mHC) is not recognised as typed, so both the 1F1B and the + # interleaved (VPP) schedules fall back to the shape-based ``P2PCommunicator`` -- which calls + # ``.size()`` on the payload object and crashes. Backbones without a payload (plain + # ``GPTModel``, GLM's HybridModel adapter) expose ``None`` here and keep the legacy path. + self.pipeline_payload_factory = getattr(self.language_model, 'pipeline_payload_factory', None) + self.pipeline_payload_spec = getattr(self.language_model, 'pipeline_payload_spec', None) self.model_meta = config.model_meta self.visual = None if pre_process and self.model_meta.visual_cls is not None: @@ -59,7 +69,8 @@ def forward(_self, input_): kwargs.update(res) res = inputs_embeds if self.config.context_parallel_size > 1: - res = split_cp_inputs(res, getattr(packed_seq_params, 'cu_seqlens_q', None), 1) + res = split_cp_inputs(res, getattr(packed_seq_params, 'cu_seqlens_q', None), 1, + cp_partition_mode=self.config.cp_partition_mode) if reduce_scatter_embeddings: res = res.transpose(0, 1).contiguous() res = scatter_to_sequence_parallel_region(res, group=_self.tp_group) @@ -87,14 +98,17 @@ def forward( runtime_gather_output: Optional[bool] = None, **kwargs, ) -> torch.Tensor: + inference_context = kwargs.pop('inference_context', None) extra_kwargs = {k: kwargs[k] for k in self.language_model.extra_forward_keys} # Compatible with legacy mcore-bridge behavior. cp_size = self.config.context_parallel_size + cp_partition_mode = self.config.cp_partition_mode needs_split = cp_size > 1 and input_ids is not None and position_ids.shape[-1] * cp_size == input_ids.shape[-1] if decoder_input is not None: pass elif self.pre_process: - input_ids_ = input_ids if needs_split else reconstruct_tensor_cp(input_ids, packed_seq_params, dim=1) + input_ids_ = input_ids if needs_split else reconstruct_tensor_cp( + input_ids, packed_seq_params, dim=1, cp_partition_mode=cp_partition_mode) kwargs.update({'input_ids': input_ids_, 'packed_seq_params': packed_seq_params}) with self._patch_word_embeddings(kwargs): decoder_input = self.language_model.embedding(input_ids=input_ids_, position_ids=position_ids) @@ -105,13 +119,15 @@ def forward( kwargs = {} kwargs.update(extra_kwargs) if needs_split: - input_ids = split_cp_inputs(input_ids, getattr(packed_seq_params, 'cu_seqlens_q', None), dim=1) + input_ids = split_cp_inputs(input_ids, getattr(packed_seq_params, 'cu_seqlens_q', None), dim=1, + cp_partition_mode=cp_partition_mode) return self.language_model( input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask, decoder_input=decoder_input, labels=labels, + inference_context=inference_context, inference_params=inference_params, packed_seq_params=packed_seq_params, runtime_gather_output=runtime_gather_output, diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py index 996a3bb..9eb3331 100644 --- a/src/mcore_bridge/model/modules/__init__.py +++ b/src/mcore_bridge/model/modules/__init__.py @@ -2,6 +2,18 @@ from .absorbed_mla import AbsorbedMLASelfAttention from .compressor import Compressor, CSAIndexer from .dsa_indexer import DSAIndexer +from .dspark import ( + DeepseekV41DSparkConfidenceHead, + DeepseekV41DSparkInput, + DeepseekV41DSparkMarkovHead, + DeepseekV41DSparkOutput, + DeepseekV41DSparkRMSNorm, + DeepseekV41DSparkStack, + DeepseekV41DSparkState, + DeepseekV41DSparkVerification, + dspark_sample, + verify_dspark_draft, +) from .gated_delta_net import GatedDeltaNet from .gated_self_attention import GatedSelfAttention from .hyper_connection_gated import Qwen4ExpTextGatedResidual, Qwen4ExpTextGroupedRMSNorm diff --git a/src/mcore_bridge/model/modules/dspark.py b/src/mcore_bridge/model/modules/dspark.py new file mode 100644 index 0000000..d70dec8 --- /dev/null +++ b/src/mcore_bridge/model/modules/dspark.py @@ -0,0 +1,456 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tensor-parallel building blocks for DeepSeek-V4.1 DSpark.""" +import copy +from dataclasses import dataclass +from typing import Callable, Optional, Sequence + +import torch +import torch.nn.functional as F +from megatron.core.tensor_parallel.layers import ( + ColumnParallelLinear, + RowParallelLinear, + VocabParallelEmbedding, +) +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +from megatron.core.transformer.hyper_connection import SinglePassMHCState +from torch import nn + + +class DeepseekV41DSparkRMSNorm(nn.Module): + """RMSNorm matching the fp32 accumulation used by the reference model.""" + + def __init__(self, hidden_size: int, eps: float, dtype: torch.dtype, device=None): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + + def forward(self, hidden_states: torch.Tensor): + dtype = hidden_states.dtype + hidden_states = hidden_states.float() + variance = hidden_states.square().mean(dim=-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return (hidden_states * self.weight.float()).to(dtype) + + +class DeepseekV41DSparkInput(nn.Module): + """Project target-layer states and construct the parallel draft-token block. + + DSpark receives target states in Megatron layout ``[s, b, targets * h]`` and + runs its draft stack on ``[block, b, streams * h]``. The feature projection + is row-parallel, while sequence-parallel sharding is deliberately disabled: + the target states are already local to the caller's sequence partition. + """ + + def __init__(self, config): + super().__init__() + target_count = len(config.dspark_target_layer_ids or ()) + if target_count == 0: + raise ValueError('DSpark input projection requires at least one target layer.') + if config.dspark_block_size <= 0: + raise ValueError('DSpark block size must be positive.') + self.hidden_size = config.hidden_size + self.num_streams = config.num_residual_streams + self.sequence_parallel = config.sequence_parallel + self.block_size = config.dspark_block_size + self.noise_token_id = config.dspark_noise_token_id + + projection_config = copy.copy(config) + projection_config.sequence_parallel = False + self.main_proj = RowParallelLinear( + config.hidden_size * target_count, + config.hidden_size, + config=projection_config, + init_method=config.init_method, + bias=False, + input_is_parallel=False, + skip_bias_add=False, + ) + device = None if config.use_cpu_initialization else torch.cuda.current_device() + self.main_norm = DeepseekV41DSparkRMSNorm( + config.hidden_size, + config.layernorm_epsilon, + config.params_dtype, + device=device, + ) + + def project_main_hidden(self, main_hidden: torch.Tensor): + if main_hidden.ndim != 3: + raise ValueError( + f'DSpark main hidden states must be [s, b, targets*h], got {tuple(main_hidden.shape)}.') + if self.sequence_parallel: + main_hidden = gather_from_sequence_parallel_region(main_hidden) + main_x, _ = self.main_proj(main_hidden) + return self.main_norm(main_x) + + def build_draft_hidden( + self, + input_ids: torch.Tensor, + embedding: Callable[[torch.Tensor], torch.Tensor], + ): + if input_ids.ndim != 1: + raise ValueError(f'DSpark input_ids must be [b], got {tuple(input_ids.shape)}.') + draft_input_ids = input_ids.new_full( + (self.block_size, input_ids.shape[0]), + self.noise_token_id, + ) + draft_input_ids[0] = input_ids + hidden_states = embedding(draft_input_ids) + expected = (self.block_size, input_ids.shape[0], self.hidden_size) + if tuple(hidden_states.shape) != expected: + raise ValueError( + f'DSpark embedding must return {expected}, got {tuple(hidden_states.shape)}.') + hidden_states = hidden_states.unsqueeze(-2).expand( + *hidden_states.shape[:-1], self.num_streams, self.hidden_size) + hidden_states = hidden_states.reshape( + self.block_size, + input_ids.shape[0], + self.num_streams * self.hidden_size, + ) + return hidden_states, draft_input_ids + + def forward( + self, + main_hidden: torch.Tensor, + input_ids: torch.Tensor, + embedding: Callable[[torch.Tensor], torch.Tensor], + ): + if input_ids.ndim != 1 or input_ids.shape[0] != main_hidden.shape[1]: + raise ValueError( + f'DSpark input_ids must be [b] matching main hidden batch {main_hidden.shape[1]}, ' + f'got {tuple(input_ids.shape)}.') + main_x = self.project_main_hidden(main_hidden) + hidden_states, draft_input_ids = self.build_draft_hidden(input_ids, embedding) + return hidden_states, main_x, draft_input_ids + + +class DeepseekV41DSparkMarkovHead(nn.Module): + """Low-rank Markov logit bias with vocab-row tensor parallelism.""" + + def __init__(self, config): + super().__init__() + rank = config.dspark_markov_rank + if rank is None or rank <= 0: + raise ValueError('DSpark Markov rank must be positive.') + self.vocab_size = config.padded_vocab_size + self.embed = VocabParallelEmbedding( + self.vocab_size, + rank, + init_method=config.init_method, + config=config, + ) + head_config = copy.copy(config) + head_config.params_dtype = torch.float32 + self.head = ColumnParallelLinear( + rank, + self.vocab_size, + config=head_config, + init_method=config.init_method, + bias=False, + gather_output=False, + skip_bias_add=False, + ) + + def forward(self, token_ids: torch.Tensor, gather_output: bool = True): + markov_embed = self.embed(token_ids) + logits, _ = self.head(markov_embed.float(), runtime_gather_output=gather_output) + return logits, markov_embed + + +class DeepseekV41DSparkConfidenceHead(nn.Module): + """FP32 acceptance-confidence projection from draft and Markov states.""" + + def __init__(self, config): + super().__init__() + input_size = config.hidden_size + config.dspark_markov_rank + device = None if config.use_cpu_initialization else torch.cuda.current_device() + self.proj = nn.Linear( + input_size, + 1, + bias=False, + dtype=torch.float32, + device=device, + ) + + def forward(self, hidden_states: torch.Tensor, markov_embed: torch.Tensor): + if hidden_states.shape[:-1] != markov_embed.shape[:-1]: + raise ValueError( + 'DSpark confidence inputs must have matching leading dimensions, got ' + f'{tuple(hidden_states.shape)} and {tuple(markov_embed.shape)}.') + hidden_states = torch.cat((hidden_states, markov_embed), dim=-1) + return F.linear(hidden_states.float(), self.proj.weight).squeeze(-1) + + +def dspark_sample(logits: torch.Tensor, temperature: float = 0.0): + """Sample one DSpark token, matching the reference Gumbel-max path.""" + if temperature == 0: + return logits.argmax(dim=-1) + logits = logits / max(temperature, 1e-5) + probabilities = torch.softmax(logits, dim=-1, dtype=torch.float32) + return probabilities.div(torch.empty_like(probabilities).exponential_()).argmax(dim=-1) + + +@dataclass +class DeepseekV41DSparkVerification: + """Verified prefix length and fallback/bonus token for each request.""" + + accepted_lengths: torch.Tensor + accepted_mask: torch.Tensor + next_tokens: torch.Tensor + + +def verify_dspark_draft( + draft_ids: torch.Tensor, + target_ids: torch.Tensor, + confidence_logits: Optional[torch.Tensor] = None, + confidence_threshold: Optional[float] = None, +): + """Verify proposals in strict prefix order against target-model tokens. + + ``draft_ids`` contains the committed seed followed by ``K`` proposals. + ``target_ids`` contains ``K`` verification tokens and one bonus token. + Confidence may shorten a proposal but cannot accept a target mismatch. + """ + if draft_ids.ndim != 2 or target_ids.ndim != 2: + raise ValueError('DSpark verification expects rank-2 token tensors.') + block_size = draft_ids.shape[1] - 1 + if block_size <= 0 or target_ids.shape != (draft_ids.shape[0], block_size + 1): + raise ValueError( + f'Expected draft [b, K+1] and target [b, K+1], got ' + f'{tuple(draft_ids.shape)} and {tuple(target_ids.shape)}.') + accepted_mask = draft_ids[:, 1:].eq(target_ids[:, :block_size]) + if confidence_threshold is not None: + if confidence_logits is None or confidence_logits.shape != accepted_mask.shape: + raise ValueError('Confidence logits must have shape [b, K] when a threshold is set.') + accepted_mask &= confidence_logits.sigmoid().ge(confidence_threshold) + accepted_mask = accepted_mask.cumprod(dim=-1).bool() + accepted_lengths = accepted_mask.sum(dim=-1) + next_tokens = target_ids.gather(1, accepted_lengths.unsqueeze(-1)).squeeze(-1) + return DeepseekV41DSparkVerification(accepted_lengths, accepted_mask, next_tokens) + + +@dataclass +class DeepseekV41DSparkState: + """Per-forward target state passed explicitly to every DSpark attention layer.""" + + main_hidden: torch.Tensor + main_rotary_pos_emb: Optional[torch.Tensor] = None + cache_slots: Optional[torch.Tensor] = None + + def attention_kwargs(self): + return { + 'dspark_main_hidden': self.main_hidden, + 'dspark_main_rotary_pos_emb': self.main_rotary_pos_emb, + 'dspark_cache_slots': self.cache_slots, + } + + def recompute_boundary_tensors(self): + if self.main_rotary_pos_emb is None: + return (self.main_hidden,) + return self.main_hidden, self.main_rotary_pos_emb + + def save_for_recompute(self): + has_rotary = self.main_rotary_pos_emb is not None + cache_slots = self.cache_slots + tensors = self.recompute_boundary_tensors() + + def restore(saved_tensors): + return type(self)( + main_hidden=saved_tensors[0], + main_rotary_pos_emb=saved_tensors[1] if has_rotary else None, + cache_slots=cache_slots, + ) + + return tensors, restore + + +class DeepseekV41DSparkOutput(nn.Module): + """Final DSpark norm, main logits, Markov recurrence and confidence.""" + + def __init__(self, config): + super().__init__() + device = None if config.use_cpu_initialization else torch.cuda.current_device() + self.block_size = config.dspark_block_size + self.norm = DeepseekV41DSparkRMSNorm( + config.hidden_size, + config.layernorm_epsilon, + config.params_dtype, + device=device, + ) + self.markov_head = DeepseekV41DSparkMarkovHead(config) + self.confidence_head = DeepseekV41DSparkConfidenceHead(config) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + output_layer: Callable, + temperature: float = 0.0, + sample_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + ): + if hidden_states.ndim != 3 or hidden_states.shape[0] != self.block_size: + raise ValueError( + f'DSpark output hidden states must be [block={self.block_size}, b, h], ' + f'got {tuple(hidden_states.shape)}.') + if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[1]: + raise ValueError( + f'DSpark output input_ids must be [b] matching hidden batch {hidden_states.shape[1]}, ' + f'got {tuple(input_ids.shape)}.') + + base_logits, _ = output_layer(self.norm(hidden_states), runtime_gather_output=True) + output_ids = input_ids.new_empty((input_ids.shape[0], self.block_size + 1)) + output_ids[:, 0] = input_ids + markov_embeds = [] + logits = [] + for index in range(self.block_size): + logits_bias, markov_embed = self.markov_head(output_ids[:, index]) + step_logits = base_logits[index] + logits_bias + logits.append(step_logits) + markov_embeds.append(markov_embed) + output_ids[:, index + 1] = ( + sample_fn(step_logits) if sample_fn is not None else dspark_sample(step_logits, temperature)) + logits = torch.stack(logits, dim=0) + markov_embed = torch.stack(markov_embeds, dim=0) + confidence = self.confidence_head(hidden_states, markov_embed) + return output_ids, logits.transpose(0, 1).contiguous(), confidence.transpose(0, 1).contiguous() + + +class DeepseekV41DSparkStack(nn.Module): + """Orchestrate the dedicated DSpark layers around the TP input/output modules. + + The supplied layers must be mHC-enabled TransformerLayer instances whose + attention accepts ``dspark_main_hidden`` and ``dspark_main_rotary_pos_emb``. + Keeping this stack separate prevents V4.1 checkpoint layers under ``mtp.*`` + from entering Megatron's serial MultiTokenPredictionBlock path. + """ + + def __init__(self, config, layers: Sequence[nn.Module]): + super().__init__() + if len(layers) != config.dspark_num_layers: + raise ValueError( + f'DSpark requires {config.dspark_num_layers} layers, got {len(layers)}.') + if not config.mhc_single_pass: + raise ValueError('DeepSeek-V4.1 DSpark requires single-pass mHC.') + self.config = config + self.input = DeepseekV41DSparkInput(config) + self.layers = nn.ModuleList(layers) + self.output = DeepseekV41DSparkOutput(config) + self._request_cache_slots = {} + + def reset_cache(self): + self._request_cache_slots.clear() + for layer in self.layers: + reset = getattr(layer.self_attention, 'reset_dspark_cache', None) + if reset is not None: + reset() + + def resolve_cache_slots(self, request_ids: torch.Tensor, live_request_ids: Optional[torch.Tensor] = None): + """Map scheduler request IDs to stable rows in every DSpark ring cache.""" + request_ids_cpu = request_ids.detach().to(device='cpu', dtype=torch.long).tolist() + if live_request_ids is not None: + live_ids = set(live_request_ids.detach().to(device='cpu', dtype=torch.long).tolist()) + self._request_cache_slots = { + request_id: slot for request_id, slot in self._request_cache_slots.items() + if request_id in live_ids + } + occupied = set(self._request_cache_slots.values()) + for request_id in request_ids_cpu: + if request_id in self._request_cache_slots: + continue + slot = 0 + while slot in occupied: + slot += 1 + self._request_cache_slots[request_id] = slot + occupied.add(slot) + return torch.tensor( + [self._request_cache_slots[request_id] for request_id in request_ids_cpu], + dtype=torch.long, + device=request_ids.device, + ) + + def _seed_main_kv( + self, + main_hidden: torch.Tensor, + main_rotary_pos_emb: Optional[torch.Tensor], + inference_context=None, + cache_slots: Optional[torch.Tensor] = None, + ): + for layer in self.layers: + attention = layer.self_attention + if not hasattr(attention, 'prefill_dspark'): + raise TypeError( + f'{type(attention).__name__} does not implement prefill_dspark().') + attention.prefill_dspark( + main_hidden, + rotary_pos_emb=main_rotary_pos_emb, + inference_context=inference_context, + cache_slots=cache_slots, + ) + + def update_main_cache( + self, + main_hidden: torch.Tensor, + main_rotary_pos_emb: torch.Tensor, + *, + start_pos, + cache_slots: torch.Tensor, + inference_context=None, + ): + main_x = self.input.project_main_hidden(main_hidden) + for layer in self.layers: + layer.self_attention.prefill_dspark( + main_x, + rotary_pos_emb=main_rotary_pos_emb, + inference_context=inference_context, + start_pos=start_pos, + cache_slots=cache_slots, + ) + return main_x + + def forward( + self, + main_hidden: torch.Tensor, + input_ids: torch.Tensor, + embedding: Callable[[torch.Tensor], torch.Tensor], + output_layer: Callable, + *, + start_pos: int, + rotary_pos_emb: Optional[torch.Tensor] = None, + main_rotary_pos_emb: Optional[torch.Tensor] = None, + inference_context=None, + temperature: float = 0.0, + sample_fn: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + cache_slots: Optional[torch.Tensor] = None, + prefill_only: Optional[bool] = None, + ): + hidden_states, main_x, _ = self.input(main_hidden, input_ids, embedding) + if prefill_only is None: + prefill_only = bool(torch.as_tensor(start_pos).eq(0).all().item()) + if prefill_only: + self._seed_main_kv( + main_x, main_rotary_pos_emb, inference_context, cache_slots=cache_slots) + return None + + mhc_state = SinglePassMHCState() + dspark_state = DeepseekV41DSparkState(main_x, main_rotary_pos_emb, cache_slots) + for layer in self.layers: + hidden_states, _ = layer( + hidden_states, + attention_mask=None, + rotary_pos_emb=rotary_pos_emb, + inference_context=inference_context, + sequence_len_offset=start_pos, + cross_layer_state=dspark_state, + mhc_state=mhc_state, + ) + hidden_states = mhc_state.contract( + hidden_states, + self.config.num_residual_streams, + use_fused=self.config.use_fused_mhc, + ) + return self.output( + hidden_states, + input_ids, + output_layer, + temperature=temperature, + sample_fn=sample_fn, + ) diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py new file mode 100644 index 0000000..8fc9b85 --- /dev/null +++ b/src/mcore_bridge/model/modules/engram.py @@ -0,0 +1,402 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1 adapters for NVIDIA Megatron-Core's optional Engram modules.""" + +import dataclasses + +import torch +from torch import Tensor + +from ...utils.megatron_utils import split_cp_inputs + +try: + from megatron.core import mpu + from megatron.core.models.engram.config import EngramConfig + from megatron.core.models.engram.engram import Engram + from megatron.core.models.engram.hashing import ( + compress_token_ids, + shift_right_reset_at_eos, + slice_hashes_for_sequence_parallel, + ) + from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayer, + ) + from megatron.core.utils import get_pg_size, nvtx_range_pop, nvtx_range_push +except ImportError: + EngramConfig = None + Engram = torch.nn.Module + HyperConnectionTransformerLayer = None + TransformerLayer = None + + +def has_native_engram() -> bool: + """Return whether the installed Megatron-Core includes the official Engram extension.""" + return EngramConfig is not None + + +class _RelaxedParallelismView: + """Read-only view of a transformer config that hides the CP and VPP guards. + + Lets the Engram validators reuse every upstream parallelism check except the two + DeepSeek-V4.1 has shown safe to drop, without mutating the shared transformer config: + + * ``context_parallel_size``: V4.1 hashes the full sequence locally and then slices its own + CP interval, so no n-gram window ever crosses a CP rank boundary. + * ``virtual_pipeline_model_parallel_size``: ``Engram.forward`` is self-contained (it never + exchanges state across pipeline/VP stages), and its layer placement keys off the + vp_stage-aware global ``layer_number`` that ``select_pipeline_segment`` assigns to each + chunk. The hybrid stack's ``__init__`` block-alignment guard rejects any partial-block + stage at build time, so an Engram-carrying ``D`` layer can never be split across stages. + """ + + context_parallel_size = 1 + virtual_pipeline_model_parallel_size = None + + def __init__(self, transformer_config): + self._transformer_config = transformer_config + + def __getattr__(self, name): + return getattr(self._transformer_config, name) + + +if EngramConfig is not None: + + class DeepseekV41EngramConfig(EngramConfig): + """Translate 0-based checkpoint hash IDs to 1-based Megatron placement IDs.""" + + def __init__(self, *, placement_layer_ids, hash_layer_ids, excluded_token_ids=(), **kwargs): + self.hash_layer_ids = tuple(hash_layer_ids) + placement_layer_ids = tuple(placement_layer_ids) + if len(placement_layer_ids) != len(self.hash_layer_ids): + raise ValueError('Engram placement and hash layer IDs must have the same length.') + self.excluded_token_ids = tuple(excluded_token_ids) + super().__init__(layer_ids=placement_layer_ids, **kwargs) + hash_multipliers = self.layer_multipliers + self.layer_multipliers = { + placement: hash_multipliers[source] + for placement, source in zip(placement_layer_ids, self.hash_layer_ids) + } + + def _load_tokenizer_map(self): + # The artifact and hash multipliers use checkpoint-native 0-based IDs, while the + # TransformerLayer composition point validates and selects 1-based layer numbers. + placement_layer_ids = self.layer_ids + self.layer_ids = self.hash_layer_ids + try: + return super()._load_tokenizer_map() + finally: + self.layer_ids = placement_layer_ids + + def _validate_parallelism(self, transformer_config, sequence_length): + # Drop only the upstream CP and VPP rejections (see _RelaxedParallelismView for why + # both are safe on the V4.1 hybrid path) and keep every other check -- etp==tp, the + # SP rank-local history length, etc. CP shortens the rank-local slice the SP check + # compares against, so apply it here before delegating. + context_parallel_size = transformer_config.context_parallel_size + if sequence_length is not None: + # The SP checks compare against a rank-local slice, which CP shortens first. + sequence_length = sequence_length // context_parallel_size + super()._validate_parallelism( + _RelaxedParallelismView(transformer_config), sequence_length) + + def _validate_packed_sequences(self, transformer_config, packed_sequences): + # DeepseekV41Engram restarts its n-gram windows at every cu_seqlens document + # boundary, so packed (THD) rows are safe even though the upstream DeepSeek + # variant advertises resets_windows_at_boundary_token=False. Keep the remaining + # packed checks (pipeline stages, padding alignment) from super(). + variant_spec = self.variant_spec + self.variant_spec = dataclasses.replace( + variant_spec, resets_windows_at_boundary_token=True) + try: + super()._validate_packed_sequences(transformer_config, packed_sequences) + finally: + self.variant_spec = variant_spec +else: + DeepseekV41EngramConfig = None + + +def build_deepseek_v41_engram_config(**kwargs): + """Build the DeepSeek adapter over NVIDIA's official Engram configuration.""" + if DeepseekV41EngramConfig is None: + raise RuntimeError( + 'DeepSeek-V4.1 Engram requires NVIDIA Megatron-LM Engram support. ' + 'Install the official Engram extension or disable Engram.') + return DeepseekV41EngramConfig(**kwargs) + + +def _hash_token_windows( + token_windows: Tensor, + tokenizer_remap: Tensor | None, + multipliers: Tensor, + table_sizes: Tensor, + max_ngram_order: int, + num_hash_heads: int, + boundary_token_id: int, + invalid_token_id: int | None = None, +) -> Tensor: + if token_windows.ndim != 3 or token_windows.shape[-1] != max_ngram_order: + raise ValueError( + 'Engram token windows must have shape [batch, sequence, max_ngram_order], ' + f'got {token_windows.shape}.') + tokens = token_windows.to(torch.int64) + compressed = tokens if tokenizer_remap is None else compress_token_ids(tokens, tokenizer_remap) + suffixes = [] + blocked = torch.zeros_like(compressed[..., 0], dtype=torch.bool) + for shift in range(max_ngram_order): + source = compressed[..., shift] + if invalid_token_id is not None: + blocked = blocked | (source == invalid_token_id) + source = torch.where(blocked, source.new_full((), boundary_token_id), source) + suffixes.append(source) + + hashes = [] + table_index = 0 + for order in range(2, max_ngram_order + 1): + mixed = suffixes[0] * multipliers[0] + for suffix_index in range(1, order): + mixed = torch.bitwise_xor(mixed, suffixes[suffix_index] * multipliers[suffix_index]) + for _ in range(num_hash_heads): + hashes.append(torch.remainder(mixed, table_sizes[table_index])) + table_index += 1 + return torch.stack(hashes, dim=-1) + + +def _positions_in_segment(cu_seqlens: Tensor | None, sequence_length: int, device) -> Tensor: + """Distance from each position to the start of the document that contains it. + + Without ``cu_seqlens`` the whole row is one document, so this is just the position + index and the window is only padded at the row start. + """ + positions = torch.arange(sequence_length, device=device, dtype=torch.int64) + if cu_seqlens is None: + return positions + boundaries = cu_seqlens.reshape(-1).to(device=device, dtype=torch.int64) + if int(boundaries[-1]) != sequence_length: + raise ValueError( + f'Engram cu_seqlens ends at {int(boundaries[-1])} but the hashed sequence has ' + f'{sequence_length} tokens; cu_seqlens must describe the full packed row.') + segment_index = torch.searchsorted(boundaries, positions, right=True) - 1 + return positions - boundaries[segment_index] + + +def _build_ngram_hashes( + input_ids: Tensor, + tokenizer_remap: Tensor | None, + multipliers: Tensor, + table_sizes: Tensor, + max_ngram_order: int, + num_hash_heads: int, + boundary_token_id: int, + reset_at_boundary: bool, + cu_seqlens: Tensor | None = None, +) -> Tensor: + tokens = input_ids.to(torch.int64) + compressed = tokens if tokenizer_remap is None else compress_token_ids(tokens, tokenizer_remap) + sequence_length = compressed.shape[1] + if reset_at_boundary: + suffixes = [ + shift_right_reset_at_eos(compressed, shift, boundary_token_id) + for shift in range(max_ngram_order) + ] + else: + # The DeepSeek variant carries no boundary token in the stream, so packed (THD) rows + # need the document starts from cu_seqlens to keep n-grams inside one document. With + # cu_seqlens=None this reduces exactly to padding the window at the row start. + if cu_seqlens is not None and compressed.shape[0] != 1: + raise ValueError( + 'Engram cu_seqlens-based window reset expects a single packed row, got ' + f'batch size {compressed.shape[0]}.') + position_in_segment = _positions_in_segment( + cu_seqlens, sequence_length, compressed.device).unsqueeze(0) + suffixes = [compressed] + for shift in range(1, max_ngram_order): + shifted = torch.nn.functional.pad( + compressed, (shift, 0), value=boundary_token_id)[:, :sequence_length] + suffixes.append( + torch.where(position_in_segment >= shift, shifted, + shifted.new_full((), boundary_token_id))) + return _hash_token_windows( + torch.stack(suffixes, dim=-1), + tokenizer_remap=None, + multipliers=multipliers, + table_sizes=table_sizes, + max_ngram_order=max_ngram_order, + num_hash_heads=num_hash_heads, + boundary_token_id=boundary_token_id, + invalid_token_id=-1, + ) + + +class DeepseekV41Engram(Engram): + """DeepSeek V4.1 Engram without projection bias or the Qwen short-conv branch.""" + + def __init__(self, *args, **kwargs): + if EngramConfig is None: + raise RuntimeError('The installed Megatron-Core does not provide Engram.') + super().__init__(*args, **kwargs) + # PR #7231's DeepSeek variant predates the released V4.1 checkpoint layout. + # Keep its initialized weights but remove parameters absent from that checkpoint. + self.value_projection.register_parameter('bias', None) + self.key_projection.register_parameter('bias', None) + self.conv_norm = None + self.short_conv = None + + def _mask_excluded_tokens(self, input_ids: Tensor) -> tuple[Tensor, Tensor]: + live = torch.ones_like(input_ids, dtype=torch.bool) + for token_id in self.engram_config.excluded_token_ids: + live = live & (input_ids != token_id) + return torch.where(live, input_ids, input_ids.new_full((), -1)), live + + def _build_hash_ids(self, input_ids: Tensor, cu_seqlens: Tensor | None = None) -> tuple[Tensor, Tensor]: + masked_ids, live = self._mask_excluded_tokens(input_ids) + hashes = _build_ngram_hashes( + masked_ids, + self.tokenizer_remap, + self.hash_multipliers, + self.table_sizes, + self.engram_config.max_ngram_order, + self.engram_config.num_hash_heads, + self.engram_config.hash_boundary_token_id, + self.engram_config.variant_spec.resets_windows_at_boundary_token, + cu_seqlens=cu_seqlens, + ) + return hashes, live + + def _cp_local_sequence_length(self, hidden_states: Tensor) -> int: + """Length of this rank's CP slice, undoing the innermost SP split first.""" + length = hidden_states.shape[0] + if self.config.sequence_parallel: + length *= get_pg_size(self.tp_group) + return length + + def _gather_input_ids_for_context_parallel(self, input_ids: Tensor, + local_sequence_length: int) -> Tensor: + """Restore the full token sequence so every rank hashes identical n-gram windows. + + The data pipeline hands us either a CP-sharded copy of ``input_ids`` (swift + ``get_batch_on_this_cp_rank`` splits it for text models) or the full sequence + (multimodal models keep it whole and split the embeddings instead), so re-align + only when the lengths disagree. Gathering int64 token IDs is far cheaper than the + hidden states the surrounding attention already exchanges. + """ + cp_size = self.config.context_parallel_size + present = input_ids.shape[1] + if present == local_sequence_length * cp_size: + return input_ids + if present != local_sequence_length: + raise ValueError( + f'Engram input_ids length {present} matches neither this CP rank slice ' + f'({local_sequence_length}) nor the full sequence ' + f'({local_sequence_length * cp_size}).') + shards = [torch.empty_like(input_ids) for _ in range(cp_size)] + torch.distributed.all_gather( + shards, input_ids.contiguous(), group=mpu.get_context_parallel_group()) + # Contiguous partitioning gives rank r the block [r * local, (r + 1) * local), so + # concatenating the gathered shards in rank order rebuilds the original token order. + return torch.cat(shards, dim=1) + + def _slice_for_context_parallel(self, hashes: Tensor) -> Tensor: + """Select this rank's CP interval after the hashes were computed globally.""" + if self.config.context_parallel_size == 1: + return hashes + return split_cp_inputs(hashes, None, 1, cp_partition_mode='contiguous') + + def forward(self, hidden_states: Tensor, input_ids: Tensor, inference_context=None) -> Tensor: + if inference_context is None: + inference_context = getattr(self, '_bridge_inference_context', None) + if inference_context is not None: + raise RuntimeError( + 'DeepSeek-V4.1 Engram inference is not wired into this integration; rollout runs ' + 'through vLLM, so the Engram module supports the training forward only.') + packed_seq_params = getattr(self, '_bridge_packed_seq_params', None) + if hidden_states.ndim != 3: + raise ValueError(f'Engram hidden_states must be [S,B,H], got {hidden_states.shape}.') + expected_hidden = self.num_streams * self.hidden_size + if hidden_states.shape[-1] != expected_hidden: + raise ValueError(f'Engram expected hidden width {expected_hidden}, got {hidden_states.shape[-1]}.') + context_parallel = self.config.context_parallel_size > 1 + + cu_seqlens = None + if packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', None) == 'thd': + # cu_seqlens_q stays global: the data pipeline builds it before the CP split. + cu_seqlens = getattr(packed_seq_params, 'cu_seqlens_q', None) + cp_partition_mode = getattr(packed_seq_params, 'cp_partition_mode', 'zigzag') + else: + # Non-packed CP carries the partition layout on the transformer config instead of on + # packed_seq_params (see mm_gpt_model's CP data path). + cp_partition_mode = getattr(self.config, 'cp_partition_mode', 'zigzag') + if context_parallel and cp_partition_mode != 'contiguous': + # Both _gather_input_ids_for_context_parallel (rank-order concat) and + # _slice_for_context_parallel (contiguous slice) assume contiguous CP blocks, so a + # zigzag layout would silently mis-align the hashes with the local hidden states. + # Fail loud on every CP path -- packed (THD) and non-packed alike -- matching the + # DSv4 THD CP forward. + raise ValueError( + "Engram with context parallelism requires cp_partition_mode='contiguous', " + 'matching the DSv4 THD CP forward.') + + nvtx_range_push('engram.hash') + try: + if context_parallel: + input_ids = self._gather_input_ids_for_context_parallel( + input_ids, self._cp_local_sequence_length(hidden_states)) + hash_ids, live_tokens = self._build_hash_ids(input_ids, cu_seqlens) + live_tokens = live_tokens.unsqueeze(-1) + # CP is the outer split and SP the inner one, so undo them in that order. + hash_ids = self._slice_for_context_parallel(hash_ids) + live_tokens = self._slice_for_context_parallel(live_tokens) + hash_ids = slice_hashes_for_sequence_parallel(hash_ids, hidden_states.shape[0], self.tp_group) + live_tokens = slice_hashes_for_sequence_parallel( + live_tokens, hidden_states.shape[0], self.tp_group).squeeze(-1) + finally: + nvtx_range_pop('engram.hash') + + nvtx_range_push('engram.lookup') + try: + memory = self.embedding(hash_ids).flatten(start_dim=-2).transpose(0, 1).contiguous() + finally: + nvtx_range_pop('engram.lookup') + streams = hidden_states.view( + hidden_states.shape[0], hidden_states.shape[1], self.num_streams, self.hidden_size) + shared_value = self.value_projection(memory) + key = self.key_norm(self.key_projection(memory)).view_as(streams) + query = self.query_norm(hidden_states).view_as(streams) + score = (key * query).sum(dim=-1) / self.hidden_size**0.5 + score = score.abs().clamp_min(1e-6).sqrt() * score.sign() + output = score.sigmoid().unsqueeze(-1) * shared_value.unsqueeze(2) + output = output * live_tokens.transpose(0, 1).unsqueeze(-1).unsqueeze(-1) + return output.reshape(hidden_states.shape) + + +class _DeepseekV41EngramLayerMixin: + + def _forward_attention(self, *args, **kwargs): + engram = getattr(self, 'engram', None) + if engram is None: + return super()._forward_attention(*args, **kwargs) + # `_maybe_apply_engram` only forwards hidden_states and input_ids, so stash the + # per-microbatch context the Engram needs (inference context, THD cu_seqlens) on the + # module itself for the duration of this attention call. + previous = getattr(engram, '_bridge_inference_context', None) + previous_packed = getattr(engram, '_bridge_packed_seq_params', None) + engram._bridge_inference_context = kwargs.get('inference_context') + engram._bridge_packed_seq_params = kwargs.get('packed_seq_params') + try: + return super()._forward_attention(*args, **kwargs) + finally: + engram._bridge_inference_context = previous + engram._bridge_packed_seq_params = previous_packed + + +if TransformerLayer is not None: + + class DeepseekV41TransformerLayer(_DeepseekV41EngramLayerMixin, TransformerLayer): + pass + + + class DeepseekV41HyperConnectionTransformerLayer( + _DeepseekV41EngramLayerMixin, HyperConnectionTransformerLayer): + pass +else: + DeepseekV41TransformerLayer = None + DeepseekV41HyperConnectionTransformerLayer = None diff --git a/src/mcore_bridge/model/modules/topk_router.py b/src/mcore_bridge/model/modules/topk_router.py index eae7a9f..3c26514 100644 --- a/src/mcore_bridge/model/modules/topk_router.py +++ b/src/mcore_bridge/model/modules/topk_router.py @@ -1,10 +1,62 @@ +import copy +from typing import Optional + import torch +from megatron.core import tensor_parallel from megatron.core.jit import jit_fuser from megatron.core.transformer.moe.router import TopKRouter as McoreTopKRouter -from typing import Optional class TopKRouter(McoreTopKRouter): + """mcore-bridge router extensions kept outside the vendored Megatron-LM tree.""" + + def __init__(self, config, *args, **kwargs): + enable_vl_bias = getattr(config, 'moe_router_enable_vl_bias', False) + if enable_vl_bias: + # The fused TE route accepts one shared correction vector, while VL routing + # selects a correction vector per token. Keep this override local to the router. + config = copy.copy(config) + config.moe_router_fusion = False + super().__init__(config, *args, **kwargs) + + # Stay compatible with a future Megatron version that grows native VL routing. + self._mcore_has_native_vl_bias = hasattr(self, 'expert_bias_vl') + if not self._mcore_has_native_vl_bias: + if enable_vl_bias: + if self.expert_bias is None: + raise ValueError('VL expert bias requires the standard expert bias.') + self.register_buffer('expert_bias_vl', torch.zeros_like(self.expert_bias, dtype=torch.float32)) + else: + self.expert_bias_vl = None + + def routing(self, logits, padding_mask=None, input_ids=None, packed_seq_params=None): + if self.expert_bias_vl is None or self._mcore_has_native_vl_bias: + return super().routing(logits, padding_mask, input_ids, packed_seq_params) + if input_ids is None: + raise ValueError('input_ids is required when VL expert bias is enabled.') + + seq_length, batch_size = logits.shape[:2] + image_mask = (input_ids == self.config.image_token_id).transpose(0, 1).contiguous() + if image_mask.shape != (seq_length, batch_size): + if (self.config.sequence_parallel and image_mask.shape[1] == batch_size + and image_mask.shape[0] % self.config.tensor_model_parallel_size == 0 + and image_mask.shape[0] // self.config.tensor_model_parallel_size == seq_length): + image_mask = tensor_parallel.scatter_to_sequence_parallel_region(image_mask) + else: + raise ValueError( + 'image token mask cannot be aligned with router logits: ' + f'input_ids={tuple(input_ids.shape)}, logits={tuple(logits.shape)}.') + + original_expert_bias = self.expert_bias + self.expert_bias = torch.where( + image_mask.reshape(-1, 1), + self.expert_bias_vl.unsqueeze(0), + original_expert_bias.unsqueeze(0), + ) + try: + return super().routing(logits, padding_mask, input_ids, packed_seq_params) + finally: + self.expert_bias = original_expert_bias @jit_fuser def _apply_expert_bias(self, routing_map: torch.Tensor, padding_mask: Optional[torch.Tensor] = None): diff --git a/src/mcore_bridge/model/modules/transformer_layer.py b/src/mcore_bridge/model/modules/transformer_layer.py index 68d35c1..7dd0796 100644 --- a/src/mcore_bridge/model/modules/transformer_layer.py +++ b/src/mcore_bridge/model/modules/transformer_layer.py @@ -79,6 +79,29 @@ def __init__( self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout self.is_mtp_layer = is_mtp_layer + # This mirrors the upstream composition point. It has to be built here, and unconditionally + # set to None otherwise, because the inherited `_forward_attention` reads `self.engram` on + # every layer -- this `__init__` replaces upstream's rather than extending it, so anything + # the inherited forward relies on has to be set up here too. `getattr` keeps this working + # against a Megatron-Core whose submodules predate the Engram composition point. + self.engram = None + engram_submodule = getattr(submodules, 'engram', IdentityOp) + if engram_submodule is not IdentityOp and not is_mtp_layer: + # MTP layers use their own local layer numbering, which would collide with the + # decoder's global Engram layer IDs; Engram never attaches to MTP layers. + if not isinstance(engram_submodule, ModuleSpec): + raise TypeError('The Engram composition point must be a ModuleSpec or IdentityOp.') + engram_config = engram_submodule.params.get('engram_config') + if engram_config is None: + raise ValueError('The Engram ModuleSpec must provide engram_config.') + if self.layer_number in engram_config.layer_ids: + self.engram = build_module( + engram_submodule, + config=self.config, + layer_number=self.layer_number, + pg_collection=pg_collection, + ) + # [Module 1: Input Layernorm] Optional Layernorm on the input data # TODO: add pytorch only layernorm self.input_layernorm = submodules.input_layernorm( diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index 847c9b0..6548c07 100644 --- a/src/mcore_bridge/model/register.py +++ b/src/mcore_bridge/model/register.py @@ -40,6 +40,7 @@ class ModelMeta: visual_cls: Optional[Type[nn.Module]] = None is_multimodal: bool = False loader: Optional[Type['ModelLoader']] = None + config_cls: Type[ModelConfig] = ModelConfig def __post_init__(self): if self.visual_cls is not None: diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index 419ed53..5108100 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -107,16 +107,18 @@ def get_num_samples(packed_seq_params) -> int: return int(packed_seq_params.cu_seqlens_q.numel()) - 1 -def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int) -> torch.Tensor: - """In CP mode, all-gather and undo the load-balanced (zigzag) chunking - produced by ``split_cp_inputs``, restoring the full sequence in original - token order along ``dim``. +def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int, cp_partition_mode: str = 'zigzag') -> torch.Tensor: + """In CP mode, all-gather and undo the chunking produced by + ``split_cp_inputs``, restoring the full sequence in original token order + along ``dim``. Args: tensor: CP-sharded local tensor whose sequence dim is at ``dim``. packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for regular ``[B, S, ...]`` inputs. dim: Sequence dimension index of ``tensor`` (default: 1). + cp_partition_mode: CP partition layout, either ``zigzag`` or ``contiguous``. + It must match the mode used by ``split_cp_inputs``. Returns: torch.Tensor: Full-sequence tensor with the same shape as ``tensor`` @@ -136,6 +138,11 @@ def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int) -> torch.Tensor: output_list[cp_rank] = tensor gathered = torch.cat(output_list, dim=dim) + if cp_partition_mode == 'contiguous': + # Rank r owns block r, so concatenating the shards in rank order already + # restores the original token order. + return gathered + # `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed. if dim != 0: gathered = gathered.transpose(0, dim).contiguous() diff --git a/src/mcore_bridge/utils/safetensors.py b/src/mcore_bridge/utils/safetensors.py index 2e22f0e..0ad3ea3 100644 --- a/src/mcore_bridge/utils/safetensors.py +++ b/src/mcore_bridge/utils/safetensors.py @@ -10,16 +10,25 @@ class LazyTensor: - def __init__(self, tensor=None, loader=None): + def __init__(self, tensor=None, loader=None, slice_loader=None): """You need to provide a tensor or loader""" self.tensor = tensor self.loader = loader + self.slice_loader = slice_loader def load(self): if self.tensor is None: return self.loader() return self.tensor + def load_slice(self, slices): + """Load only ``slices`` when the backing format supports partial reads.""" + if self.tensor is not None: + return self.tensor[slices] + if self.slice_loader is not None: + return self.slice_loader(slices=slices) + return self.loader()[slices] + class SafetensorLazyLoader: @@ -60,7 +69,10 @@ def _load_index(self): def get_state_dict(self): res = {} for k in self._weight_map.keys(): - res[k] = LazyTensor(loader=partial(self._load_tensor, key=k)) + res[k] = LazyTensor( + loader=partial(self._load_tensor, key=k), + slice_loader=partial(self._load_tensor_slice, key=k), + ) return res def _load_tensor(self, key): @@ -68,6 +80,11 @@ def _load_tensor(self, key): file_handle = self._open_file(filename) return file_handle.get_tensor(key) + def _load_tensor_slice(self, key, slices): + filename = self._weight_map[key] + file_handle = self._open_file(filename) + return file_handle.get_slice(key)[slices] + def close(self): self._file_handles.clear() diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py new file mode 100644 index 0000000..84356a0 --- /dev/null +++ b/tests/test_deepseek_v41_engram.py @@ -0,0 +1,993 @@ +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from megatron.core import mpu +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, VocabParallelEmbedding +from megatron.core.transformer import TransformerConfig +from safetensors.torch import save_file + +from mcore_bridge.config.parser import _convert_config +from mcore_bridge.model.gpts import deepseek_v41 as deepseek_v41_module +from mcore_bridge.model.gpts.deepseek_v41 import ( + DeepseekV41Aligner, + DeepseekV41Bridge, + DeepseekV41DSparkAttention, + DeepseekV41Loader, + DeepseekV41Vision, + DeepseekV41VisionTransformer, +) +from mcore_bridge.model.modules import engram as engram_adapter +from mcore_bridge.model.modules.dspark import ( + DeepseekV41DSparkConfidenceHead, + DeepseekV41DSparkInput, + DeepseekV41DSparkMarkovHead, + DeepseekV41DSparkOutput, + DeepseekV41DSparkStack, + DeepseekV41DSparkState, + dspark_sample, + verify_dspark_draft, +) +from mcore_bridge.utils.safetensors import SafetensorLazyLoader + + +def test_dspark_config_is_kept_separate_from_standard_mtp(): + text_config = SimpleNamespace( + model_type='deepseek_v41_text', + num_nextn_predict_layers=3, + dspark_block_size=5, + dspark_noise_token_id=128799, + dspark_target_layer_ids=[37, 38, 39], + dspark_markov_rank=256, + dspark_n_routed_experts=128, + dspark_num_experts_per_tok=3, + ) + + converted = _convert_config(SimpleNamespace(model_type='deepseek_v41', text_config=text_config)) + + assert converted['dspark_num_layers'] == 3 + assert converted['dspark_block_size'] == 5 + assert converted['dspark_noise_token_id'] == 128799 + assert converted['dspark_target_layer_ids'] == [37, 38, 39] + assert converted['dspark_markov_rank'] == 256 + assert converted['dspark_num_experts'] == 128 + assert converted['dspark_router_topk'] == 3 + assert 'mtp_num_layers' not in converted + + +def test_dspark_input_builds_parallel_noise_block(): + + class _Projection(torch.nn.Module): + + def forward(self, hidden_states): + return hidden_states[..., :2], None + + module = DeepseekV41DSparkInput.__new__(DeepseekV41DSparkInput) + torch.nn.Module.__init__(module) + module.hidden_size = 2 + module.num_streams = 3 + module.sequence_parallel = False + module.block_size = 4 + module.noise_token_id = 7 + module.main_proj = _Projection() + module.main_norm = torch.nn.Identity() + + main_hidden = torch.arange(2 * 2 * 4, dtype=torch.float32).view(2, 2, 4) + input_ids = torch.tensor([1, 2]) + + def embedding(token_ids): + return torch.nn.functional.one_hot(token_ids % 2, 2).float() + + hidden_states, main_x, draft_ids = module(main_hidden, input_ids, embedding) + + assert hidden_states.shape == (4, 2, 6) + assert torch.equal(draft_ids[0], input_ids) + assert torch.all(draft_ids[1:] == 7) + torch.testing.assert_close(main_x, main_hidden[..., :2]) + streams = hidden_states.view(4, 2, 3, 2) + torch.testing.assert_close(streams[:, :, 0], streams[:, :, 1]) + torch.testing.assert_close(streams[:, :, 1], streams[:, :, 2]) + + +def test_dspark_markov_head_returns_full_logits_and_embedding(): + + class _Head(torch.nn.Module): + + def forward(self, hidden_states, runtime_gather_output): + assert runtime_gather_output + return torch.cat((hidden_states, hidden_states + 10), dim=-1), None + + module = DeepseekV41DSparkMarkovHead.__new__(DeepseekV41DSparkMarkovHead) + torch.nn.Module.__init__(module) + module.embed = torch.nn.Embedding.from_pretrained(torch.tensor([[1., 2.], [3., 4.], [5., 6.]]), ) + module.head = _Head() + + logits, embedding = module(torch.tensor([0, 2])) + + torch.testing.assert_close(embedding, torch.tensor([[1., 2.], [5., 6.]])) + torch.testing.assert_close(logits, torch.tensor([[1., 2., 11., 12.], [5., 6., 15., 16.]])) + + +@pytest.fixture +def single_rank_model_parallel(tmp_path): + if dist.is_initialized() and dist.get_world_size() != 1: + pytest.skip('Single-rank model-parallel test.') + created_process_group = not dist.is_initialized() + created_model_parallel = not mpu.model_parallel_is_initialized() + try: + if created_process_group: + dist.init_process_group( + 'gloo', + init_method=f'file://{tmp_path}/single-rank-dist-init', + rank=0, + world_size=1, + ) + if created_model_parallel: + mpu.initialize_model_parallel(tensor_model_parallel_size=1) + yield + finally: + if created_model_parallel and mpu.model_parallel_is_initialized(): + mpu.destroy_model_parallel() + if created_process_group and dist.is_initialized(): + dist.destroy_process_group() + + +def test_dspark_tp_modules_construct_and_run_on_one_rank(single_rank_model_parallel): + config = TransformerConfig( + num_layers=1, + hidden_size=4, + num_attention_heads=1, + use_cpu_initialization=True, + params_dtype=torch.float32, + ) + config.padded_vocab_size = 8 + config.dspark_markov_rank = 2 + config.dspark_target_layer_ids = [0, 1] + config.dspark_block_size = 3 + config.dspark_noise_token_id = 7 + config.num_residual_streams = 2 + + markov = DeepseekV41DSparkMarkovHead(config) + logits, markov_embed = markov(torch.tensor([0, 7])) + assert logits.shape == (2, 8) + assert markov_embed.shape == (2, 2) + assert logits.dtype == torch.float32 + + dspark_input = DeepseekV41DSparkInput(config) + draft_embedding = VocabParallelEmbedding( + 8, + 4, + init_method=config.init_method, + config=config, + ) + hidden_states, main_x, draft_ids = dspark_input( + torch.randn(1, 2, 8), + torch.tensor([1, 2]), + draft_embedding, + ) + assert hidden_states.shape == (3, 2, 8) + assert main_x.shape == (1, 2, 4) + assert draft_ids.shape == (3, 2) + + output_layer = ColumnParallelLinear( + 4, + 8, + config=config, + init_method=config.init_method, + bias=False, + gather_output=False, + skip_bias_add=False, + ) + dspark_output = DeepseekV41DSparkOutput(config) + output_ids, logits, confidence = dspark_output( + hidden_states.view(3, 2, 2, 4).mean(dim=2), + torch.tensor([1, 2]), + output_layer, + ) + assert output_ids.shape == (2, 4) + assert logits.shape == (2, 3, 8) + assert confidence.shape == (2, 3) + + +def test_attach_dspark_freezes_the_draft_stack(single_rank_model_parallel, monkeypatch): + """The draft stack has to be attached frozen. + + It is deliberately kept out of the training forward -- it exists so the checkpoint's ``mtp.*`` + weights have somewhere to be loaded into and saved from -- so none of its parameters can ever + hold anything but a zero gradient. Leaving them trainable costs more than the optimizer state and + gradient buffers it needlessly allocates: Adam's weight decay is decoupled from the gradient, so + every step still multiplies them by ``1 - lr * wd`` with nothing pushing back, and the weights the + stack exists to carry erode over a long run. + """ + config = TransformerConfig( + num_layers=1, + hidden_size=4, + num_attention_heads=1, + use_cpu_initialization=True, + params_dtype=torch.float32, + ) + config.padded_vocab_size = 8 + config.dspark_num_layers = 1 + config.dspark_markov_rank = 2 + config.dspark_target_layer_ids = [0] + config.dspark_block_size = 3 + config.dspark_noise_token_id = 7 + config.num_residual_streams = 2 + config.mhc_single_pass = True + + class _Layer(torch.nn.Module): + """Stands in for a draft TransformerLayer: one parameter, and an ``mlp`` to look for a router on.""" + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 2)) + self.mlp = torch.nn.Module() + + class _LanguageModel(torch.nn.Module): + + def __init__(self): + super().__init__() + self.trainable = torch.nn.Parameter(torch.ones(2)) + self.config = config + self.vocab_size = 8 + self.pg_collection = SimpleNamespace(tp=None) + + loader = DeepseekV41Loader.__new__(DeepseekV41Loader) + loader.config = config + monkeypatch.setattr(loader, 'get_dspark_layer_spec', lambda: (config, [object()]), raising=False) + monkeypatch.setattr(loader, '_set_linear_is_expert', lambda module: None, raising=False) + monkeypatch.setattr(deepseek_v41_module, 'build_module', lambda *args, **kwargs: _Layer()) + + language_model = _LanguageModel() + loader._attach_dspark(language_model, post_process=True) + + trainable = [name for name, p in language_model.dspark.named_parameters() if p.requires_grad] + assert list(language_model.dspark.parameters()), 'the draft stack was attached with no parameters' + assert not trainable, trainable + # this ``language_model`` has no base ``embedding`` to seed drafts from, so the stack built its own + assert not language_model.dspark_word_embeddings.weight.requires_grad + # and nothing outside the draft stack was frozen along the way + assert language_model.trainable.requires_grad + + +def test_dspark_attention_sink_and_ring_cache(): + attention = DeepseekV41DSparkAttention.__new__(DeepseekV41DSparkAttention) + torch.nn.Module.__init__(attention) + attention.window_size = 3 + attention._dspark_window_kv_cache = None + attention.config = SimpleNamespace(v_head_dim=2) + attention.core_attention = SimpleNamespace(attn_sink=torch.nn.Parameter(torch.tensor([0.0]))) + + query = torch.tensor([[[[1.0, 0.0]]]]) + key_value = torch.tensor([[[[1.0, 0.0]]], [[[0.0, 1.0]]]]) + actual = attention._latent_attention(query, key_value) + scale = 2**-0.5 + probabilities = torch.softmax(torch.tensor([scale, 0.0, 0.0]), dim=0)[:2] + expected = torch.tensor([[[[probabilities[0], probabilities[1]]]]]) + torch.testing.assert_close(actual, expected) + + first = torch.arange(5 * 1 * 1 * 2, dtype=torch.float32).view(5, 1, 1, 2) + cache, valid_lengths = attention._write_main_cache(first, start_pos=0) + assert torch.equal(valid_lengths, torch.tensor([3])) + torch.testing.assert_close(cache[0], first[3].squeeze(-2)) + torch.testing.assert_close(cache[1], first[4].squeeze(-2)) + torch.testing.assert_close(cache[2], first[2].squeeze(-2)) + update = torch.tensor([[[[20.0, 21.0]]]]) + cache, valid_lengths = attention._write_main_cache(update, start_pos=5) + assert torch.equal(valid_lengths, torch.tensor([3])) + torch.testing.assert_close(cache[2], update[0].squeeze(-2)) + + +def test_dspark_verification_accepts_only_strict_matching_prefix(): + draft_ids = torch.tensor([ + [10, 11, 12, 13], + [20, 21, 22, 23], + [30, 31, 32, 33], + ]) + target_ids = torch.tensor([ + [11, 12, 13, 14], + [21, 99, 23, 24], + [98, 32, 33, 34], + ]) + + result = verify_dspark_draft(draft_ids, target_ids) + + assert torch.equal(result.accepted_lengths, torch.tensor([3, 1, 0])) + assert torch.equal(result.next_tokens, torch.tensor([14, 99, 98])) + assert torch.equal(result.accepted_mask, + torch.tensor([ + [True, True, True], + [True, False, False], + [False, False, False], + ])) + + confidence = torch.tensor([[10.0, -10.0, 10.0]] * 3) + result = verify_dspark_draft(draft_ids, target_ids, confidence, confidence_threshold=0.5) + assert torch.equal(result.accepted_lengths, torch.tensor([1, 1, 0])) + + +def test_dspark_state_preserves_recompute_inputs(): + main_hidden = torch.randn(2, 1, 4) + rotary = torch.randn(2, 1, 1, 2) + cache_slots = torch.tensor([3]) + state = DeepseekV41DSparkState(main_hidden, rotary, cache_slots) + + attention_kwargs = state.attention_kwargs() + assert attention_kwargs['dspark_main_hidden'] is main_hidden + assert attention_kwargs['dspark_main_rotary_pos_emb'] is rotary + assert attention_kwargs['dspark_cache_slots'] is cache_slots + tensors, restore = state.save_for_recompute() + restored = restore(tensors) + assert restored.main_hidden is main_hidden + assert restored.main_rotary_pos_emb is rotary + assert restored.cache_slots is cache_slots + + +def test_dspark_stack_prefill_and_decode_lifecycle(): + + class _Input(torch.nn.Module): + + def forward(self, main_hidden, input_ids, embedding): + hidden = embedding(input_ids).unsqueeze(0).expand(2, -1, -1) + return hidden, main_hidden[..., :2], None + + class _Attention(torch.nn.Module): + + def __init__(self): + super().__init__() + self.prefill_args = None + + def prefill_dspark( + self, + main_hidden, + rotary_pos_emb, + inference_context, + start_pos=0, + cache_slots=None, + ): + self.prefill_args = ( + main_hidden, + rotary_pos_emb, + inference_context, + start_pos, + cache_slots, + ) + + class _Layer(torch.nn.Module): + + def __init__(self): + super().__init__() + self.self_attention = _Attention() + self.received_main_hidden = None + + def forward(self, hidden_states, **kwargs): + state = kwargs['cross_layer_state'] + self.received_main_hidden = state.main_hidden + mhc_state = kwargs['mhc_state'] + mhc_state.pre_mix = hidden_states.new_full(hidden_states.shape[:2] + (2, ), 0.5) + return hidden_states + 1, None + + class _Output(torch.nn.Module): + + def forward(self, hidden_states, input_ids, output_layer, temperature, sample_fn=None): + return hidden_states, input_ids, temperature + + stack = DeepseekV41DSparkStack.__new__(DeepseekV41DSparkStack) + torch.nn.Module.__init__(stack) + stack.config = SimpleNamespace(num_residual_streams=2, use_fused_mhc=False) + stack.input = _Input() + stack.layers = torch.nn.ModuleList([_Layer(), _Layer()]) + stack.output = _Output() + embedding = torch.nn.Embedding.from_pretrained(torch.arange(10).float().unsqueeze(-1).expand(-1, 4)) + main_hidden = torch.randn(1, 2, 4) + rotary = torch.randn(1, 1, 1, 2) + + assert stack( + main_hidden, + torch.tensor([1, 2]), + embedding, + lambda *_args, **_kwargs: None, + start_pos=0, + main_rotary_pos_emb=rotary, + ) is None + for layer in stack.layers: + prefill_main, prefill_rotary, prefill_context, prefill_start, prefill_slots = ( + layer.self_attention.prefill_args) + torch.testing.assert_close(prefill_main, main_hidden[..., :2]) + assert prefill_rotary is rotary + assert prefill_context is None + assert prefill_start == 0 + assert prefill_slots is None + + contracted, returned_ids, temperature = stack( + main_hidden, + torch.tensor([1, 2]), + embedding, + lambda *_args, **_kwargs: None, + start_pos=1, + temperature=0.5, + ) + assert contracted.shape == (2, 2, 2) + assert torch.equal(returned_ids, torch.tensor([1, 2])) + assert temperature == 0.5 + for layer in stack.layers: + assert layer.received_main_hidden is not None + + +def test_dspark_confidence_head_uses_fp32_projection(): + config = SimpleNamespace( + hidden_size=3, + dspark_markov_rank=2, + use_cpu_initialization=True, + ) + module = DeepseekV41DSparkConfidenceHead(config) + module.proj.weight.data.copy_(torch.tensor([[1., 2., 3., 4., 5.]])) + hidden_states = torch.tensor([[[1., 2., 3.]]], dtype=torch.bfloat16) + markov_embed = torch.tensor([[[4., 5.]]], dtype=torch.bfloat16) + + actual = module(hidden_states, markov_embed) + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, torch.tensor([[55.]])) + + +def test_dspark_output_applies_markov_recurrence_in_block_order(): + + class _OutputLayer(torch.nn.Module): + + def forward(self, hidden_states, runtime_gather_output): + assert runtime_gather_output + return hidden_states.new_zeros((*hidden_states.shape[:-1], 5)), None + + class _MarkovHead(torch.nn.Module): + + def forward(self, token_ids): + logits = torch.nn.functional.one_hot((token_ids + 1) % 5, 5).float() * 10 + return logits, token_ids.float().unsqueeze(-1) + + class _ConfidenceHead(torch.nn.Module): + + def forward(self, hidden_states, markov_embed): + return hidden_states[..., 0].float() + markov_embed[..., 0] + + module = DeepseekV41DSparkOutput.__new__(DeepseekV41DSparkOutput) + torch.nn.Module.__init__(module) + module.block_size = 3 + module.norm = torch.nn.Identity() + module.markov_head = _MarkovHead() + module.confidence_head = _ConfidenceHead() + hidden_states = torch.tensor([ + [[1., 0.], [2., 0.]], + [[3., 0.], [4., 0.]], + [[5., 0.], [6., 0.]], + ]) + + output_ids, logits, confidence = module( + hidden_states, + torch.tensor([0, 2]), + _OutputLayer(), + ) + + assert torch.equal(output_ids, torch.tensor([[0, 1, 2, 3], [2, 3, 4, 0]])) + assert logits.shape == (2, 3, 5) + torch.testing.assert_close(confidence, torch.tensor([[1., 4., 7.], [4., 7., 10.]])) + assert torch.equal(dspark_sample(logits, temperature=0), output_ids[:, 1:]) + + +def test_engram_adapter_remaps_checkpoint_layers_to_megatron_layers(tmp_path): + if not engram_adapter.has_native_engram(): + pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') + artifact = tmp_path / 'tokenizer-map.json' + artifact.write_text( + json.dumps({ + 'format': 'megatron-engram-token-map', + 'version': 1, + 'source_vocab_size': 8, + 'compressed_vocab_size': 8, + 'pad_token_id': 0, + 'compressed_pad_token_id': 0, + 'max_ngram_order': 3, + 'hash_seed': 0, + 'layer_ids': [0, 2], + 'layer_multipliers': { + '0': [11, 13, 15], + '2': [17, 19, 21] + }, + 'remap': list(range(8)), + })) + config = engram_adapter.build_deepseek_v41_engram_config( + placement_layer_ids=(1, 3), + hash_layer_ids=(0, 2), + excluded_token_ids=(99, ), + global_vocab_sizes=(17, 19), + max_ngram_order=3, + num_hash_heads=1, + memory_dim=4, + kernel_size=1, + hash_seed=0, + boundary_token_id=0, + tokenizer_map_path=str(artifact), + ) + + assert config.layer_ids == (1, 3) + assert config.hash_layer_ids == (0, 2) + assert config.layer_multipliers == {1: (11, 13, 15), 3: (17, 19, 21)} + assert set(config.table_sizes_by_layer) == {1, 3} + assert config.excluded_token_ids == (99, ) + + +def _engram_config_for_validation(tmp_path): + artifact = tmp_path / 'tokenizer-map.json' + artifact.write_text( + json.dumps({ + 'format': 'megatron-engram-token-map', + 'version': 1, + 'source_vocab_size': 8, + 'compressed_vocab_size': 8, + 'pad_token_id': 0, + 'compressed_pad_token_id': 0, + 'max_ngram_order': 3, + 'hash_seed': 0, + 'layer_ids': [0], + 'layer_multipliers': { + '0': [11, 13, 15] + }, + 'remap': list(range(8)), + })) + return engram_adapter.build_deepseek_v41_engram_config( + placement_layer_ids=(1, ), + hash_layer_ids=(0, ), + global_vocab_sizes=(17, 19), + max_ngram_order=3, + num_hash_heads=1, + memory_dim=4, + kernel_size=1, + hash_seed=0, + boundary_token_id=0, + tokenizer_map_path=str(artifact), + ) + + +def test_engram_config_allows_context_and_virtual_pipeline_but_keeps_the_other_guards(tmp_path): + if not engram_adapter.has_native_engram(): + pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') + config = _engram_config_for_validation(tmp_path) + parallelism = dict( + context_parallel_size=2, + tensor_model_parallel_size=1, + expert_tensor_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + sequence_parallel=False, + ) + + # V4.1 hashes the full sequence locally and slices it, so CP no longer has to be 1. + config._validate_parallelism(SimpleNamespace(**parallelism), None) + # VPP is now allowed too: Engram.forward is self-contained and layer placement uses the + # vp_stage-aware global layer_number, so the upstream blanket VPP guard is dropped. + config._validate_parallelism(SimpleNamespace(**{**parallelism, 'virtual_pipeline_model_parallel_size': 2}), None) + # ... but only the CP and VPP guards are relaxed; every other parallelism check still fires. + with pytest.raises(ValueError, match='expert_tensor_parallel_size'): + config._validate_parallelism(SimpleNamespace(**{**parallelism, 'expert_tensor_parallel_size': 2}), None) + + +def test_engram_config_allows_packed_sequences_without_losing_the_pipeline_guard(tmp_path): + if not engram_adapter.has_native_engram(): + pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') + config = _engram_config_for_validation(tmp_path) + assert not config.variant_spec.supports_packed_sequences + + config._validate_packed_sequences(SimpleNamespace(pipeline_model_parallel_size=1), packed_sequences=True) + # The temporary variant override must not leak into the hashing path. + assert not config.variant_spec.supports_packed_sequences + + with pytest.raises(ValueError, match='pipeline_model_parallel_size > 2'): + config._validate_packed_sequences(SimpleNamespace(pipeline_model_parallel_size=4), packed_sequences=True) + + +def test_engram_hash_blocks_suffixes_after_excluded_token(): + hashes = engram_adapter._hash_token_windows( + token_windows=torch.tensor([[[5, -1, 7]]]), + tokenizer_remap=None, + multipliers=torch.tensor([1, 10, 100]), + table_sizes=torch.tensor([997, 991]), + max_ngram_order=3, + num_hash_heads=1, + boundary_token_id=0, + invalid_token_id=-1, + ) + + assert torch.equal(hashes, torch.tensor([[[5, 5]]])) + + +def _ngram_hash_kwargs(): + return dict( + tokenizer_remap=None, + multipliers=torch.tensor([11, 13, 15]), + table_sizes=torch.tensor([997, 991]), + max_ngram_order=3, + num_hash_heads=1, + boundary_token_id=0, + reset_at_boundary=False, + ) + + +def test_engram_packed_hashes_match_separately_hashed_documents(): + # The DeepSeek variant carries no boundary token in the stream, so cu_seqlens is the only + # thing that stops an n-gram window from reaching into the previous packed document. + packed_row = torch.tensor([[5, 6, 7, 8, 9]]) + kwargs = _ngram_hash_kwargs() + + packed = engram_adapter._build_ngram_hashes(packed_row, cu_seqlens=torch.tensor([0, 2, 5]), **kwargs) + separate = torch.cat( + ( + engram_adapter._build_ngram_hashes(packed_row[:, :2], **kwargs), + engram_adapter._build_ngram_hashes(packed_row[:, 2:], **kwargs), + ), + dim=1, + ) + + assert torch.equal(packed, separate) + # Without the reset the second document would mix in tokens 5 and 6. + assert not torch.equal(packed, engram_adapter._build_ngram_hashes(packed_row, **kwargs)) + + +def test_engram_hashes_are_unchanged_when_cu_seqlens_spans_one_document(): + packed_row = torch.tensor([[5, -1, 7, 8]]) + kwargs = _ngram_hash_kwargs() + + assert torch.equal( + engram_adapter._build_ngram_hashes(packed_row, cu_seqlens=torch.tensor([0, 4]), **kwargs), + engram_adapter._build_ngram_hashes(packed_row, **kwargs), + ) + + +def test_engram_rejects_cu_seqlens_that_does_not_cover_the_row(): + with pytest.raises(ValueError, match='cu_seqlens ends at'): + engram_adapter._build_ngram_hashes( + torch.tensor([[5, 6, 7]]), cu_seqlens=torch.tensor([0, 2]), **_ngram_hash_kwargs()) + + +def _bare_engram(context_parallel_size=1, sequence_parallel=False): + module = engram_adapter.DeepseekV41Engram.__new__(engram_adapter.DeepseekV41Engram) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace(context_parallel_size=context_parallel_size, sequence_parallel=sequence_parallel) + return module + + +def test_engram_context_parallel_slices_reassemble_the_global_hashes(monkeypatch): + from mcore_bridge.utils import megatron_utils + + global_hashes = torch.arange(2 * 8 * 3).view(2, 8, 3) + cp_size = 4 + monkeypatch.setattr(megatron_utils.mpu, 'get_context_parallel_world_size', lambda: cp_size) + + slices = [] + for cp_rank in range(cp_size): + monkeypatch.setattr(megatron_utils.mpu, 'get_context_parallel_rank', lambda rank=cp_rank: rank) + slices.append(_bare_engram(cp_size)._slice_for_context_parallel(global_hashes)) + + # Contiguous partitioning must hand rank r the block [r * local, (r + 1) * local). + assert all(item.shape == (2, 2, 3) for item in slices) + assert torch.equal(torch.cat(slices, dim=1), global_hashes) + + +def test_contiguous_cp_reconstruct_inverts_the_matching_split(monkeypatch): + """Multimodal V4.1 splits embeddings/input_ids itself, so both directions must + honour ``cp_partition_mode``: reconstructing a contiguous shard with the zigzag + layout silently reorders tokens away from what the DSv4 THD CP forward assumes.""" + from mcore_bridge.utils import megatron_utils + + global_ids = torch.arange(8).view(1, 8) + cp_size, cp_rank = 2, 1 + local_length = global_ids.shape[1] // cp_size + monkeypatch.setattr(megatron_utils.mpu, 'get_context_parallel_world_size', lambda: cp_size) + monkeypatch.setattr(megatron_utils.mpu, 'get_context_parallel_rank', lambda: cp_rank) + monkeypatch.setattr(megatron_utils.mpu, 'get_context_parallel_group', lambda: 'cp-group') + + shard = megatron_utils.split_cp_inputs(global_ids, None, 1, cp_partition_mode='contiguous') + + def fake_all_gather(output_list, tensor, group=None): + assert group == 'cp-group' + assert torch.equal(tensor, shard) + for rank, buffer in enumerate(output_list): + buffer.copy_(global_ids[:, rank * local_length:(rank + 1) * local_length]) + + monkeypatch.setattr(torch.distributed, 'all_gather', fake_all_gather) + assert torch.equal( + megatron_utils.reconstruct_tensor_cp(shard, None, dim=1, cp_partition_mode='contiguous'), + global_ids, + ) + # The default zigzag layout must not be applied to a contiguous shard. + assert not torch.equal(megatron_utils.reconstruct_tensor_cp(shard, None, dim=1), global_ids) + + +def test_engram_gathers_cp_sharded_input_ids_but_leaves_full_ones_alone(monkeypatch): + full_ids = torch.arange(8).view(1, 8) + cp_size, cp_rank = 4, 2 + local_length = 2 + + monkeypatch.setattr(engram_adapter.mpu, 'get_context_parallel_group', lambda: 'cp-group') + + def fake_all_gather(output_list, tensor, group=None): + assert group == 'cp-group' + assert torch.equal(tensor, full_ids[:, cp_rank * local_length:(cp_rank + 1) * local_length]) + for rank, buffer in enumerate(output_list): + buffer.copy_(full_ids[:, rank * local_length:(rank + 1) * local_length]) + + monkeypatch.setattr(torch.distributed, 'all_gather', fake_all_gather) + module = _bare_engram(cp_size) + + shard = full_ids[:, cp_rank * local_length:(cp_rank + 1) * local_length] + assert torch.equal(module._gather_input_ids_for_context_parallel(shard, local_length), full_ids) + # Multimodal models keep input_ids whole and split the embeddings instead. + assert module._gather_input_ids_for_context_parallel(full_ids, local_length) is full_ids + with pytest.raises(ValueError, match='matches neither'): + module._gather_input_ids_for_context_parallel(full_ids[:, :5], local_length) + + +def test_engram_cp_local_sequence_length_undoes_the_inner_sp_split(monkeypatch): + monkeypatch.setattr(engram_adapter, 'get_pg_size', lambda group: 2, raising=False) + hidden_states = torch.zeros(4, 1, 8) + + assert _bare_engram(2)._cp_local_sequence_length(hidden_states) == 4 + module = _bare_engram(2, sequence_parallel=True) + module.tp_group = None + assert module._cp_local_sequence_length(hidden_states) == 8 + + +class _Table: + + def __init__(self, global_rows, row_start, row_end, dim): + self.global_num_embeddings = global_rows + self.row_start = row_start + self.row_end = row_end + self.weight = torch.nn.Parameter(torch.empty(row_end - row_start, dim, dtype=torch.bfloat16)) + + @property + def local_num_embeddings(self): + return self.row_end - self.row_start + + +def test_safetensor_lazy_loader_reads_only_requested_rows(tmp_path): + path = tmp_path / 'model.safetensors' + tensor = torch.arange(40, dtype=torch.float32).view(10, 4) + save_file({'table': tensor}, path) + + with SafetensorLazyLoader(str(tmp_path)) as loader: + lazy = loader.get_state_dict()['table'] + sliced = lazy.load_slice(slice(3, 6)) + + torch.testing.assert_close(sliced, tensor[3:6]) + + +def test_dspark_bridge_uses_tp_layout_and_official_endpoint_names(): + bridge = DeepseekV41Bridge.__new__(DeepseekV41Bridge) + bridge.config = SimpleNamespace(task_type='causal_lm', dspark_num_layers=3) + assert bridge._get_tp_split_dim('input.main_proj.weight') == 1 + assert bridge._get_tp_split_dim('output.markov_head.embed.weight') == 0 + assert bridge._get_tp_split_dim('output.markov_head.head.weight') == 0 + assert bridge._get_tp_split_dim('output.confidence_head.proj.weight') is None + + calls = [] + + def record(_module, mg_key, _state, hf_key, to_mcore): + calls.append((mg_key, hf_key, to_mcore)) + + bridge._set_state_dict = record + result = bridge._set_dspark_endpoints(object(), {}, to_mcore=False) + + assert result == {} + assert calls == [ + ('input.main_proj.weight', 'main_proj.weight', False), + ('input.main_norm.weight', 'main_norm.weight', False), + ('output.norm.weight', 'norm.weight', False), + ('output.markov_head.embed.weight', 'markov_head.embed.weight', False), + ('output.markov_head.head.weight', 'markov_head.head.weight', False), + ('output.confidence_head.proj.weight', 'confidence_head.proj.weight', False), + ] + + +def test_dspark_bridge_loads_dedicated_embedding_with_padding_and_tp_shard(): + bridge = DeepseekV41Bridge.__new__(DeepseekV41Bridge) + bridge.hf_embed_key = 'model.embed.weight' + bridge.config = SimpleNamespace(padded_vocab_size=8) + bridge.tp_size = 2 + bridge.tp_rank = 1 + + class _Lazy: + + def __init__(self, tensor): + self.tensor = tensor + + def load(self): + return self.tensor + + # Six real vocab rows are padded up to padded_vocab_size=8, then split across TP. + hf_rows = torch.arange(6 * 4, dtype=torch.float32).view(6, 4) + padded = torch.nn.functional.pad(hf_rows, (0, 0, 0, 2)) + expected = padded.chunk(2, dim=0)[1] + + embedding = SimpleNamespace(weight=torch.zeros(4, 4)) + bridge._load_dspark_word_embeddings(embedding, {'model.embed.weight': _Lazy(hf_rows)}) + + torch.testing.assert_close(embedding.weight, expected) + + +def test_engram_flat_fp8_table_is_dequantized_into_local_prime_shards(): + tables = [_Table(5, 1, 4, 4), _Table(7, 4, 7, 4)] + engram = SimpleNamespace( + embedding=SimpleNamespace(tables=tables), + # HF layer 1 -> doubled-space attention layer_number ``2 * 1 + 1``. + layer_number=3, + ) + bridge = DeepseekV41Bridge.__new__(DeepseekV41Bridge) + bridge.config = SimpleNamespace(engram_num_embeddings=[12], engram_layer_ids=[1]) + bridge._ENGRAM_LOAD_CHUNK_ROWS = 2 + + raw = (torch.arange(48, dtype=torch.float32).view(12, 4) % 8).to(torch.float8_e4m3fn) + scale = torch.tensor([[1.0, 0.5]] * 12, dtype=torch.float32) + + class _Lazy: + + def __init__(self, tensor): + self.tensor = tensor + self.slices = [] + + def load(self): + raise AssertionError('full tensor loading is forbidden for Engram tables') + + def load_slice(self, slices): + self.slices.append(slices) + return self.tensor[slices] + + lazy_weight, lazy_scale = _Lazy(raw), _Lazy(scale) + bridge._load_engram_embedding( + engram, + { + 'engram.embed.weight': lazy_weight, + 'engram.embed.weight_scale_inv': lazy_scale, + }, + ) + + expected = DeepseekV41Bridge._dequantize_engram_rows(raw, scale).to(torch.bfloat16) + torch.testing.assert_close(tables[0].weight, expected[1:4]) + torch.testing.assert_close(tables[1].weight, expected[9:12]) + assert [(item.start, item.stop) for item in lazy_weight.slices] == [(1, 3), (3, 4), (9, 11), (11, 12)] + assert [(item.start, item.stop) for item in lazy_scale.slices] == [(1, 3), (3, 4), (9, 11), (11, 12)] + + +def test_engram_dense_weights_are_dequantized_and_split_like_official_wkv(): + + class _Lazy: + + def __init__(self, tensor): + self.tensor = tensor + + def load(self): + return self.tensor + + hidden_size, num_streams, memory_dim = 3, 2, 4 + engram = SimpleNamespace( + num_streams=num_streams, + hidden_size=hidden_size, + engram_config=SimpleNamespace(total_memory_dim=memory_dim), + key_projection=SimpleNamespace(weight=torch.nn.Parameter(torch.empty(num_streams * hidden_size, memory_dim))), + value_projection=SimpleNamespace(weight=torch.nn.Parameter(torch.empty(hidden_size, memory_dim))), + query_norm=SimpleNamespace(weight=torch.nn.Parameter(torch.empty(num_streams * hidden_size))), + key_norm=SimpleNamespace(weight=torch.nn.Parameter(torch.empty(num_streams * hidden_size))), + ) + bridge = DeepseekV41Bridge.__new__(DeepseekV41Bridge) + bridge._load_engram_embedding = lambda *_args: None + raw_wkv = (torch.arange(36, dtype=torch.float32).view(9, 4) % 8).to(torch.float8_e4m3fn) + scale = torch.tensor([[1.0, 0.5]] * 9, dtype=torch.float32) + q_weight = torch.arange(6, dtype=torch.float32).view(2, 3) + k_weight = q_weight + 10 + + bridge._set_layer_engram( + SimpleNamespace(engram=engram), + { + 'engram.wkv.weight': _Lazy(raw_wkv), + 'engram.wkv.weight_scale_inv': _Lazy(scale), + 'engram.q_weight': _Lazy(q_weight), + 'engram.k_weight': _Lazy(k_weight), + }, + to_mcore=True, + ) + + expected = bridge._dequantize_engram_rows(raw_wkv, scale) + torch.testing.assert_close(engram.key_projection.weight, expected[:6]) + torch.testing.assert_close(engram.value_projection.weight, expected[6:]) + torch.testing.assert_close(engram.query_norm.weight, q_weight.flatten()) + torch.testing.assert_close(engram.key_norm.weight, k_weight.flatten()) + + +def test_vision_and_aligner_match_official_equations(): + torch.manual_seed(7) + vision_config = SimpleNamespace( + hidden_size=8, + num_attention_heads=2, + intermediate_size=6, + num_hidden_layers=2, + patch_size=2, + rope_theta=10000.0, + downsample_ratio=2, + ) + vision = DeepseekV41VisionTransformer(vision_config) + aligner = DeepseekV41Aligner(vision_config, text_hidden_size=10) + patches = torch.randn(6, 3, 2, 2) + + actual_vision = vision(patches, n_h=2, n_w=3) + x = actual_vision.view(2, 3, -1).permute(2, 0, 1) + x = torch.nn.functional.pad(x, (0, 1, 0, 0)) + unfolded = torch.nn.functional.unfold(x.unsqueeze(0), 2, stride=2).squeeze(0).transpose(0, 1) + expected = aligner.w2(torch.nn.functional.gelu(aligner.w1(unfolded))) + + actual = aligner(actual_vision, n_h=2, n_w=3) + assert actual_vision.shape == (6, 8) + assert actual.shape == (2, 10) + torch.testing.assert_close(actual, expected) + assert set(vision.state_dict()) == { + 'patch_embed.proj.weight', + 'patch_embed.proj.bias', + 'blocks.0.norm1.weight', + 'blocks.0.attn.wqkv.weight', + 'blocks.0.attn.wqkv.bias', + 'blocks.0.attn.wo.weight', + 'blocks.0.attn.wo.bias', + 'blocks.0.norm2.weight', + 'blocks.0.mlp.w1.weight', + 'blocks.0.mlp.w2.weight', + 'blocks.1.norm1.weight', + 'blocks.1.attn.wqkv.weight', + 'blocks.1.attn.wqkv.bias', + 'blocks.1.attn.wo.weight', + 'blocks.1.attn.wo.bias', + 'blocks.1.norm2.weight', + 'blocks.1.mlp.w1.weight', + 'blocks.1.mlp.w2.weight', + 'norm.weight', + } + + +def test_vision_merges_official_image_span_layout(): + torch.manual_seed(11) + vision_config = SimpleNamespace( + hidden_size=8, + num_attention_heads=2, + intermediate_size=6, + num_hidden_layers=1, + patch_size=2, + rope_theta=10000.0, + downsample_ratio=2, + ) + config = SimpleNamespace( + hf_config=SimpleNamespace(image_token_id=42, vision_config=vision_config), + hidden_size=10, + language_model_only=False, + params_dtype=torch.float32, + ) + module = DeepseekV41Vision(config) + device = module.image_start.device + input_ids = torch.tensor([[5, 42, 42, 42, 42, 6]], device=device) + token_types = torch.tensor([[-1, 0, 1, 2, 3, -1]], device=device) + inputs_embeds = torch.zeros(1, 6, 10, device=device) + patches = torch.randn(4, 3, 2, 2, device=device) + grid = torch.tensor([[1, 2, 2]], device=device) + + image_features = module.encode_images(patches, grid) + actual = module.get_inputs_embeds( + inputs_embeds, + input_ids=input_ids, + pixel_values=patches, + image_grid_thw=grid, + image_token_types=token_types, + ) + + torch.testing.assert_close(actual[0, 0], inputs_embeds[0, 0]) + torch.testing.assert_close(actual[0, 1], module.image_start) + torch.testing.assert_close(actual[0, 2], image_features[0]) + torch.testing.assert_close(actual[0, 3], module.image_newline) + torch.testing.assert_close(actual[0, 4], module.image_end) + torch.testing.assert_close(actual[0, 5], inputs_embeds[0, 5]) diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py new file mode 100644 index 0000000..b21f3a0 --- /dev/null +++ b/tests/test_deepseek_v41_hybrid.py @@ -0,0 +1,601 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Unit tests for the DeepSeek-V4.1 GPT-space -> HybridStack config derivation (B1a). + +Pure logic, no GPU / distributed init required. +""" +from mcore_bridge.model.gpts.deepseek_v41 import HybridLayerConfig, derive_hybrid_layer_config + + +def test_tiny_all_moe_zero_ratio(): + # tiny checkpoint: 4 layers, sliding-window only (ratio 0), every layer MoE. + r = derive_hybrid_layer_config(4, [0, 0, 0, 0], [1, 1, 1, 1]) + assert isinstance(r, HybridLayerConfig) + assert r.hybrid_layer_pattern == 'DEDEDEDE' + assert r.num_layers == 8 + assert r.csa_compress_ratios == [0] * 8 + assert r.csa2_kv_source_layers == [] + assert r.csa2_index_source_layers == [] + assert r.csa2_candidate_source_layer is None + + +def test_matches_csa2_pipeline_reference_layout(): + # The upstream test_csa2_pipeline._config bakes the *hybrid* ratios directly as + # [0,0,2,0,2,0,1,0,1,0,1,0] (even slots = real per-attention ratio). Deriving from the + # GPT-space 6-layer config must reproduce that exact 2x layout, and re-map the source + # layers from GPT index j to hybrid attention index 2*j. + r = derive_hybrid_layer_config( + 6, + [0, 2, 2, 1, 1, 1], + [1] * 6, + csa2_kv_source_layers=[1, 3], + csa2_index_source_layers=[2], + csa2_candidate_source_layer=1, + ) + assert r.csa_compress_ratios == [0, 0, 2, 0, 2, 0, 1, 0, 1, 0, 1, 0] + assert r.hybrid_layer_pattern == 'DEDEDEDEDEDE' + assert r.csa2_kv_source_layers == [2, 6] + assert r.csa2_index_source_layers == [4] + assert r.csa2_candidate_source_layer == 2 + + +def test_dense_prefix_from_first_k_dense_replace(): + # first_k_dense_replace=1 -> moe_layer_freq [0,1,1,1]; the first MLP slot is dense '-'. + r = derive_hybrid_layer_config(4, [0, 0, 0, 0], [0, 1, 1, 1]) + assert r.hybrid_layer_pattern == 'D-DEDEDE' + + +def test_int_moe_layer_freq_convention(): + # int N -> layer i is MoE iff i % N == 0 (megatron-core moe_logging convention). + r = derive_hybrid_layer_config(4, [0] * 4, 2) + assert r.hybrid_layer_pattern == 'DED-DED-' + + +def test_no_experts_all_dense(): + r = derive_hybrid_layer_config(3, [0, 0, 0], None) + assert r.hybrid_layer_pattern == 'D-D-D-' + assert r.num_layers == 6 + + +def test_ratio_length_mismatch_raises(): + import pytest + with pytest.raises(ValueError): + derive_hybrid_layer_config(4, [0, 0, 0], [1, 1, 1, 1]) + with pytest.raises(ValueError): + derive_hybrid_layer_config(4, [0, 0, 0, 0], [1, 1, 1]) + + +def test_loader_build_hybrid_config_does_not_mutate_original(): + # The loader derives its own config copy so the caller's config is never mutated. + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Loader + + loader = object.__new__(DeepseekV41Loader) + loader.config = SimpleNamespace( + num_layers=4, + csa_compress_ratios=[0, 0, 0, 0], + moe_layer_freq=[1, 1, 1, 1], + csa2_kv_source_layers=[], + csa2_index_source_layers=[], + csa2_candidate_source_layer=None, + mtp_num_layers=1, + is_hybrid_model=False, + hybrid_layer_pattern=None, + ) + original = loader.config + cfg = loader._build_hybrid_config() + + assert cfg is not original + assert cfg.num_layers == 8 + assert cfg.hybrid_layer_pattern == 'DEDEDEDE' + assert cfg.csa_compress_ratios == [0] * 8 + assert cfg.moe_layer_freq == [0, 1, 0, 1, 0, 1, 0, 1] + assert cfg.is_hybrid_model is True + # MTP stays disabled (V4.1 ``mtp.*`` keys are DSpark). + assert cfg.mtp_num_layers is None + # caller config left untouched + assert original.num_layers == 4 + assert original.mtp_num_layers == 1 + assert original.is_hybrid_model is False + + +def _make_bridge(engram_layer_ids, enable_hyper_connections): + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Bridge + + bridge = object.__new__(DeepseekV41Bridge) + bridge.config = SimpleNamespace( + engram_layer_ids=engram_layer_ids, enable_hyper_connections=enable_hyper_connections) + return bridge + + +def test_hybrid_layer_state_fans_out_attn_and_mlp(): + # A single HF layer i must fan out to hybrid layer 2*i (attention half) and 2*i+1 (MLP half), + # each unwrapping ``inner_layer`` and routing the wrapper's single hyper-connection to the + # matching hc channel. Engram fires only on the attention half of an engram layer. + from types import SimpleNamespace + + bridge = _make_bridge(engram_layer_ids=[1], enable_hyper_connections=True) + calls = [] + bridge._set_layer_attn = lambda inner, local, hf_idx, to_mcore: calls.append(('attn', inner, hf_idx)) or {} + bridge._set_layer_mlp = lambda inner, local, hf_idx, to_mcore: calls.append(('mlp', inner, hf_idx)) or {} + bridge._set_one_hyper_connection = lambda hc, local, hf_key, to_mcore: calls.append(('hc', hf_key, hc)) + bridge._set_layer_engram = lambda mg_layer, local, to_mcore: calls.append(('engram', mg_layer)) + + wrappers = {idx: SimpleNamespace(inner_layer=f'inner{idx}', hyper_connection=f'hc{idx}') for idx in range(4)} + for idx in range(4): + res = bridge._set_hybrid_layer_state(wrappers[idx], {}, 'model.layers.', idx, to_mcore=False) + assert isinstance(res, dict) + + assert ('attn', 'inner0', 0) in calls + assert ('mlp', 'inner1', 0) in calls + assert ('attn', 'inner2', 1) in calls + assert ('mlp', 'inner3', 1) in calls + # each wrapper's single hyper-connection maps to attn (even) / ffn (odd) + assert ('hc', 'attn', 'hc0') in calls + assert ('hc', 'ffn', 'hc1') in calls + assert ('hc', 'attn', 'hc2') in calls + assert ('hc', 'ffn', 'hc3') in calls + # engram only on the attention half of HF layer 1 + engram_calls = [c for c in calls if c[0] == 'engram'] + assert len(engram_calls) == 1 + assert engram_calls[0][1] is wrappers[2] + + +def test_hybrid_layer_state_without_wrapper_uses_layer_directly(): + # enable_hyper_connections=False: no wrapper, so the payload is the layer itself and no + # hyper-connection channel is written. + bridge = _make_bridge(engram_layer_ids=[], enable_hyper_connections=False) + seen = {} + bridge._set_layer_attn = lambda inner, local, hf_idx, to_mcore: seen.update(attn_inner=inner) or {} + bridge._set_one_hyper_connection = lambda *a, **k: seen.update(hc=True) + + plain_layer = object() # no ``inner_layer`` attribute + bridge._set_hybrid_layer_state(plain_layer, {}, 'model.layers.', 0, to_mcore=False) + assert seen['attn_inner'] is plain_layer + assert 'hc' not in seen + + +def test_hybrid_layer_state_prefixes_by_hf_index_on_export(): + # Both halves of HF layer i write under ``model.layers.{i}.`` so the exported checkpoint + # keeps one layer per index. + bridge = _make_bridge(engram_layer_ids=[], enable_hyper_connections=False) + bridge._set_layer_mlp = lambda inner, local, hf_idx, to_mcore: {'ffn.w1.weight': 1} + + res = bridge._set_hybrid_layer_state(object(), {}, 'model.layers.', 5, to_mcore=False) + assert res == {'model.layers.2.ffn.w1.weight': 1} + + +def test_engram_placement_and_hf_layer_id_round_trip(): + # HF layer e -> attention hybrid layer_number 2*e + 1; the bridge maps it back to e so the + # engram_num_embeddings validation (keyed by 0-based HF ids) still resolves. + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Bridge, DeepseekV41Loader + + loader = object.__new__(DeepseekV41Loader) + assert loader._engram_placement_layer_ids([1, 3]) == (3, 7) + assert loader._engram_placement_layer_ids([0]) == (1, ) + + bridge = object.__new__(DeepseekV41Bridge) + for hf_id in (0, 1, 3, 10): + layer_number = 2 * hf_id + 1 + assert bridge._engram_hf_layer_id(SimpleNamespace(layer_number=layer_number)) == hf_id + + +def test_num_hybrid_layers_normalizes_both_layer_spaces(): + # ``_convert`` sees ``self.config`` in two layer spaces. On load it is the GPT-space config + # (num_layers == N, no pattern) whose built decoder holds 2*N layers; on export it is the + # doubled hybrid config (num_layers == 2*N, pattern populated). Both must yield 2*N so the + # loop matches the decoder's real layer count -- a regression guard for the export bug where + # ``2 * num_layers`` on the doubled config over-counted and dereferenced None layers. + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Bridge + + load_cfg = SimpleNamespace(num_layers=4, hybrid_layer_pattern=None) + export_cfg = SimpleNamespace(num_layers=8, hybrid_layer_pattern='DEDEDEDE') + assert DeepseekV41Bridge._num_hybrid_layers(load_cfg) == 8 + assert DeepseekV41Bridge._num_hybrid_layers(export_cfg) == 8 + # A config missing the attribute entirely is treated as GPT-space (load). + assert DeepseekV41Bridge._num_hybrid_layers(SimpleNamespace(num_layers=3)) == 6 + + +import pytest # noqa: E402 + +from mcore_bridge.model.gpts.deepseek_v41 import ( # noqa: E402 + DeepseekV41HyperConnectionHybridLayer, HyperConnectionHybridLayer) + +requires_hybrid = pytest.mark.skipif( + DeepseekV41HyperConnectionHybridLayer is None, reason='megatron hybrid stack not importable') + + +@requires_hybrid +def test_hc_wrapper_applies_engram_on_nstream_before_delegating(): + # The V4.1 wrapper subclass overrides ``forward``: for an Engram-carrying inner layer it adds + # the n-stream Engram delta to ``hidden_states`` BEFORE delegating to the base wrapper forward + # (aggregation + fast-path attention), so the delta lands on the pre-aggregation streams. + # A plain inner layer delegates unchanged with no Engram add. + from types import SimpleNamespace + from unittest.mock import patch + + import torch + + # Engram-carrying layer: a constant unit delta is added, so the tensor handed to the base + # forward is the input plus that delta. + def engram(hidden_states, input_ids, inference_context): + return torch.ones_like(hidden_states) + + engram_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) + engram_layer.inner_layer = SimpleNamespace(engram=engram) + h = torch.zeros(2, 1, 4) + ids = torch.zeros(2, 1, dtype=torch.long) + with patch.object(HyperConnectionHybridLayer, 'forward', return_value='OUT') as base_fwd: + assert engram_layer.forward(h, input_ids=ids) == 'OUT' + passed = base_fwd.call_args.args[0] + assert torch.equal(passed, torch.ones_like(h)) # delta added before delegating + + # Plain layer (no Engram): delegate unchanged, forwarding the original tensor untouched. + plain_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) + plain_layer.inner_layer = SimpleNamespace(engram=None) + with patch.object(HyperConnectionHybridLayer, 'forward', return_value='OUT') as base_fwd: + assert plain_layer.forward(h, input_ids=ids) == 'OUT' + assert base_fwd.call_args.args[0] is h + + +@requires_hybrid +def test_hc_wrapper_requires_input_ids_for_engram_layer(): + # An Engram layer cannot run without token IDs (needed for the n-gram hash), so forward + # raises rather than silently dropping the Engram contribution. + from types import SimpleNamespace + + import pytest + import torch + + engram_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) + engram_layer.inner_layer = SimpleNamespace(engram=lambda *a, **k: 0) + with pytest.raises(ValueError, match='input token IDs'): + engram_layer.forward(torch.zeros(2, 1, 4), input_ids=None) + + +@requires_hybrid +def test_rewrap_swaps_class_only_on_engram_wrappers(): + # Post-build retrofit: only wrappers whose inner layer built an Engram module get the + # subclass; every other wrapper keeps the base class (and its fast path). + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Loader + + engram_wrapper = object.__new__(HyperConnectionHybridLayer) + engram_wrapper.inner_layer = SimpleNamespace(engram=object()) + plain_wrapper = object.__new__(HyperConnectionHybridLayer) + plain_wrapper.inner_layer = SimpleNamespace(engram=None) + model = SimpleNamespace(decoder=SimpleNamespace(layers=[engram_wrapper, plain_wrapper])) + + loader = object.__new__(DeepseekV41Loader) + loader.config = SimpleNamespace(enable_hyper_connections=True) + loader._rewrap_engram_hyper_connection_layers(model) + + assert type(engram_wrapper) is DeepseekV41HyperConnectionHybridLayer + assert type(plain_wrapper) is HyperConnectionHybridLayer + + +@requires_hybrid +def test_rewrap_noop_without_hyper_connections(): + # No wrapping happens at all when hyper-connections are off, so nothing to retrofit. + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Loader + + engram_wrapper = object.__new__(HyperConnectionHybridLayer) + engram_wrapper.inner_layer = SimpleNamespace(engram=object()) + model = SimpleNamespace(decoder=SimpleNamespace(layers=[engram_wrapper])) + + loader = object.__new__(DeepseekV41Loader) + loader.config = SimpleNamespace(enable_hyper_connections=False) + loader._rewrap_engram_hyper_connection_layers(model) + + assert type(engram_wrapper) is HyperConnectionHybridLayer + + +from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41HybridStackModel # noqa: E402 + + +def _seg_config(pattern, **overrides): + from types import SimpleNamespace + cfg = dict( + hybrid_layer_pattern=pattern, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + num_layers_in_first_pipeline_stage=None, + num_layers_in_last_pipeline_stage=None, + pipeline_model_parallel_layout=None, + mtp_num_layers=None, + ) + cfg.update(overrides) + return SimpleNamespace(**cfg) + + +@requires_hybrid +def test_segment_main_pattern_block_aligned_even_split(): + # 4 blocks over pp=2 -> two whole blocks per stage. + cfg = _seg_config('DEDEDEDE', pipeline_model_parallel_size=2) + assert DeepseekV41HybridStackModel._segment_main_pattern(cfg) == 'DEDE|DEDE' + + +@requires_hybrid +def test_segment_main_pattern_uneven_split_front_loads_extra_blocks(): + # 4 blocks over pp=3 -> divmod(4,3)=(1,1): first stage gets 2 blocks, the rest 1 each. + cfg = _seg_config('DEDEDEDE', pipeline_model_parallel_size=3) + assert DeepseekV41HybridStackModel._segment_main_pattern(cfg) == 'DEDE|DE|DE' + + +@requires_hybrid +def test_segment_main_pattern_vpp_multiplies_stages(): + # pp=2 * vp=2 = 4 stages, consecutive segments ordered (vp0,pp0),(vp0,pp1),(vp1,pp0),(vp1,pp1) + # to match upstream segment_index = vp_stage * pp_size + pp_rank. + cfg = _seg_config('DEDEDEDE', pipeline_model_parallel_size=2, virtual_pipeline_model_parallel_size=2) + assert DeepseekV41HybridStackModel._segment_main_pattern(cfg) == 'DE|DE|DE|DE' + + +@requires_hybrid +def test_segment_main_pattern_pp1_is_noop(): + cfg = _seg_config('DEDEDEDE', pipeline_model_parallel_size=1) + assert DeepseekV41HybridStackModel._segment_main_pattern(cfg) == 'DEDEDEDE' + + +@requires_hybrid +def test_segment_main_pattern_respects_explicit_pipes_and_uneven_layout(): + # An explicit '|' layout or num_layers_in_first/last_pipeline_stage is passed through + # untouched; upstream + the post-build even-boundary guard validate it. + assert DeepseekV41HybridStackModel._segment_main_pattern(_seg_config('DEDE|DEDE', + pipeline_model_parallel_size=2)) == 'DEDE|DEDE' + assert DeepseekV41HybridStackModel._segment_main_pattern( + _seg_config('DEDEDEDE', pipeline_model_parallel_size=2, num_layers_in_first_pipeline_stage=2)) == 'DEDEDEDE' + + +@requires_hybrid +def test_segment_main_pattern_raises_when_stage_gets_no_block(): + import pytest + # 2 blocks cannot cover 4 stages. + with pytest.raises(ValueError, match='at least one attention'): + DeepseekV41HybridStackModel._segment_main_pattern(_seg_config('DEDE', pipeline_model_parallel_size=4)) + + +@requires_hybrid +def test_segment_main_pattern_rejects_pipeline_layout(): + import pytest + with pytest.raises(ValueError, match='pipeline_model_parallel_layout'): + DeepseekV41HybridStackModel._segment_main_pattern( + _seg_config('DEDEDEDE', pipeline_model_parallel_size=2, pipeline_model_parallel_layout=[[0], [1]])) + + +@requires_hybrid +def test_resolve_hybrid_layer_pattern_segments_then_defers_to_base(): + # With MTP off the resolver just returns the block-segmented main pattern; pp=1 returns the + # bare pattern (base resolver, no MTP suffix appended). + assert DeepseekV41HybridStackModel._resolve_hybrid_layer_pattern( + _seg_config('DEDEDEDE', pipeline_model_parallel_size=2)) == 'DEDE|DEDE' + assert DeepseekV41HybridStackModel._resolve_hybrid_layer_pattern( + _seg_config('DEDEDEDE', pipeline_model_parallel_size=1)) == 'DEDEDEDE' + + +from mcore_bridge.model.gpts.deepseek_v41 import ( # noqa: E402 + DeepseekV41Bridge, DeepseekV41Loader) + +# --- DSpark (``mtp.*``) draft stack ------------------------------------------------------------- + + +def _dspark_bridge(dspark_num_layers=1): + # object.__new__ so no distributed init; only the DSpark dispatch fields are needed. Record + # calls into the shared ``_convert_dspark_stack`` so we assert dispatch + stage guards without + # building a real stack (that is the GPU acceptance step). + from types import SimpleNamespace + + bridge = object.__new__(DeepseekV41Bridge) + bridge.config = SimpleNamespace(dspark_num_layers=dspark_num_layers) + calls = [] + bridge._convert_dspark_stack = (lambda language_model, dspark, hf_state_dict, hf_prefix, to_mcore: (calls.append( + (language_model, dspark)) or iter(['SENTINEL']))) + return bridge, calls + + +def test_hybrid_convert_additional_layers_maps_dspark_via_lm(): + # The hybrid model is text-only (no ``language_model`` wrapper), so ``_lm`` resolves the model + # itself; the DSpark stack attached there is mapped through the shared base helper. + from types import SimpleNamespace + + bridge, calls = _dspark_bridge() + dspark = object() + mg_model = SimpleNamespace(dspark=dspark) # no ``language_model`` -> _lm returns mg_model + out = list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=True, is_pp_last_stage=True)) + assert out == ['SENTINEL'] + assert calls == [(mg_model, dspark)] + + +def test_hybrid_convert_additional_layers_resolves_language_model_wrapper(): + # When a multimodal wrapper is present, ``_lm`` unwraps it and the + # DSpark stack is looked up on the nested language model. + from types import SimpleNamespace + + bridge, calls = _dspark_bridge() + dspark = object() + language_model = SimpleNamespace(dspark=dspark) + mg_model = SimpleNamespace(language_model=language_model) + out = list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=True, is_pp_last_stage=True)) + assert out == ['SENTINEL'] + assert calls == [(language_model, dspark)] + + +def test_hybrid_convert_additional_layers_skips_without_dspark(): + # No draft stack configured -> nothing to convert, base helper untouched. + from types import SimpleNamespace + + bridge, calls = _dspark_bridge(dspark_num_layers=0) + mg_model = SimpleNamespace(dspark=object()) + out = list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=True, is_pp_last_stage=True)) + assert out == [] + assert calls == [] + + +def test_hybrid_convert_additional_layers_load_skips_non_last_stage(): + # On load only the final stage owns the stack; earlier stages are guarded out. + from types import SimpleNamespace + + bridge, calls = _dspark_bridge() + mg_model = SimpleNamespace() # no ``dspark`` on this stage + out = list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=True, is_pp_last_stage=False)) + assert out == [] + assert calls == [] + + +def test_hybrid_convert_additional_layers_export_uses_proxy_on_non_last_stage(): + # On export a non-last PP stage has no draft stack, but it must still issue the same collective + # sequence as the final stage, so an empty structural proxy with the configured layer count + # stands in (skipping quietly would desynchronize the pp group). + from types import SimpleNamespace + + bridge, calls = _dspark_bridge() + mg_model = SimpleNamespace() # no ``dspark`` + out = list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=False, is_pp_last_stage=False)) + assert out == ['SENTINEL'] + (language_model, dspark), = calls + assert language_model is mg_model + assert dspark.layers == [None] * bridge.config.dspark_num_layers + + +def test_hybrid_convert_additional_layers_raises_on_last_stage_without_stack(): + import pytest + from types import SimpleNamespace + + bridge, _ = _dspark_bridge() + mg_model = SimpleNamespace() # last stage but stack missing -> misbuilt model + with pytest.raises(RuntimeError, match='DSpark weights require the draft stack'): + list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=False, is_pp_last_stage=True)) + + +# --- multimodal wrapper hosting the hybrid backbone --------------------------------------------- + + +def test_multimodal_wrapper_hosts_hybrid_backbone(): + # The wrapper is the stock multimodal model with the language-model class swapped for the + # PP-capable hybrid backbone; everything else (vision tower, image-embed injection) is inherited. + from mcore_bridge.model.gpts.deepseek_v41 import (DeepseekV41HybridStackModel, DeepseekV41MultimodalModel) + from mcore_bridge.model.mm_gpt_model import MultimodalGPTModel + assert issubclass(DeepseekV41MultimodalModel, MultimodalGPTModel) + assert DeepseekV41MultimodalModel.language_model_cls is DeepseekV41HybridStackModel + + +def test_hybrid_stack_exposes_extra_forward_keys(): + # ``MultimodalGPTModel.forward`` reads ``language_model.extra_forward_keys``; the hybrid backbone + # must expose the same empty default the GPT ``GPTModel`` carries. + assert DeepseekV41HybridStackModel.extra_forward_keys == [] + + +def test_hybrid_loader_model_cls_is_multimodal_wrapper(): + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41MultimodalModel + assert DeepseekV41Loader.model_cls is DeepseekV41MultimodalModel + + +def _pre_process_bridge(monkeypatch): + # Stub the base implementation (word embeddings + config-guarded vision/aligner block) and the + # marker writes so we observe delegation + ordering without a real model. + from mcore_bridge.bridge.gpt_bridge import GPTBridge + + called, markers = [], [] + monkeypatch.setattr(GPTBridge, '_convert_pre_process', lambda self, mg, sd, pfx, tm: called.append( + (mg, pfx, tm)) or {'SUPER': True}) + bridge = object.__new__(DeepseekV41Bridge) + bridge._set_state_dict = (lambda mod, mkey, sd, hkey, tm: markers.append((mod, mkey, sd, hkey, tm))) + return bridge, called, markers + + +def test_pre_process_delegates_to_base_and_maps_image_markers(monkeypatch): + # First PP stage of a multimodal model: delegate to the base bridge, then map the three + # ``image_*`` markers, which live on the wrapper (not the backbone). + from types import SimpleNamespace + + bridge, called, markers = _pre_process_bridge(monkeypatch) + mg_model = SimpleNamespace(visual=object()) + hf_state_dict = {'x': 1} + out = bridge._convert_pre_process(mg_model, hf_state_dict, '', to_mcore=True) + assert out == {'SUPER': True} + assert called == [(mg_model, '', True)] + # On load the markers are read out of the incoming HF dict. + assert markers == [(mg_model, f'visual.{name}', hf_state_dict, f'model.{name}', True) + for name in ('image_start', 'image_end', 'image_newline')] + + +def test_pre_process_runs_vision_block_even_without_visual(monkeypatch): + # A non-first PP stage builds ``visual=None``, but the vision block must still run: every rank + # has to issue the same pp-group collectives, otherwise the last stage's per-layer ``has_model`` + # all-reduce reads a stale value and ``_convert`` raises ``StopIteration``. + from types import SimpleNamespace + + bridge, called, markers = _pre_process_bridge(monkeypatch) + mg_model = SimpleNamespace(visual=None) + out = bridge._convert_pre_process(mg_model, {'x': 1}, '', to_mcore=False) + assert out == {'SUPER': True} + assert called == [(mg_model, '', False)] + # On export the markers are written into the dict the base call returned. + assert [(item[1], item[2] is out, item[4]) for item in markers + ] == [(f'visual.{name}', True, False) for name in ('image_start', 'image_end', 'image_newline')] + + +def test_hybrid_set_word_embeddings_resolves_via_lm(): + # ``_set_word_embeddings`` must resolve the LM through ``_lm`` so both the wrapper and a bare + # backbone map ``embedding.word_embeddings.weight`` onto the right module. + from types import SimpleNamespace + + bridge = object.__new__(DeepseekV41Bridge) + bridge.hf_embed_key = 'model.embed_tokens.weight' + recorded = [] + bridge._set_state_dict = lambda mod, mkey, sd, hkey, tm: recorded.append((mod, mkey, hkey, tm)) + + language_model = SimpleNamespace(tag='lm') + bridge._set_word_embeddings(SimpleNamespace(language_model=language_model), {}, to_mcore=True) + bare = SimpleNamespace() # no wrapper -> _lm returns the model itself + bridge._set_word_embeddings(bare, {}, to_mcore=False) + assert recorded == [ + (language_model, 'embedding.word_embeddings.weight', 'model.embed_tokens.weight', True), + (bare, 'embedding.word_embeddings.weight', 'model.embed_tokens.weight', False), + ] + + +@requires_hybrid +def test_hybrid_forward_unpacks_extra_block_kwargs(monkeypatch): + # ``MultimodalGPTModel.forward`` funnels the decoder's extra kwargs through ``extra_block_kwargs`` + # (the GPTModel calling convention), but upstream ``HybridModel.forward`` has no such parameter -- + # it threads ``input_ids`` itself. The hybrid stack must unpack that container before delegating, + # strip visual keys, and forward anything else, otherwise a text-only wrapper run raises + # ``HybridModel.forward() got an unexpected keyword argument 'extra_block_kwargs'``. + from mcore_bridge.model.gpts import deepseek_v41 as hyb + + received = {} + monkeypatch.setattr(hyb.HybridModel, 'forward', lambda self, *a, **k: received.update(args=a, kwargs=k) or 'OUT') + stack = object.__new__(DeepseekV41HybridStackModel) # no distributed init; forward is self-contained + out = DeepseekV41HybridStackModel.forward( + stack, input_ids=1, extra_block_kwargs={ + 'image_grid_thw': 7, + 'foo': 'bar' + }) + assert out == 'OUT' + kwargs = received['kwargs'] + assert 'extra_block_kwargs' not in kwargs # container unpacked, not forwarded verbatim + assert 'image_grid_thw' not in kwargs # visual key stripped + assert kwargs['foo'] == 'bar' # unknown extra kwarg still threaded through + assert kwargs['input_ids'] == 1 + + +@requires_hybrid +def test_hybrid_forward_rejects_pixel_values_from_extra_block_kwargs(monkeypatch): + # Defense in depth: a multimodal batch that smuggles ``pixel_values`` via ``extra_block_kwargs`` + # must still hit the text-only guard (the wrapper injects image embeds and clears them, so the + # backbone never legitimately sees pixels). + from mcore_bridge.model.gpts import deepseek_v41 as hyb + + monkeypatch.setattr(hyb.HybridModel, 'forward', lambda self, *a, **k: 'OUT') + stack = object.__new__(DeepseekV41HybridStackModel) + with pytest.raises(NotImplementedError, match='text-only'): + DeepseekV41HybridStackModel.forward(stack, input_ids=1, extra_block_kwargs={'pixel_values': 1}) diff --git a/tests/test_deepseek_v41_packing_cp.py b/tests/test_deepseek_v41_packing_cp.py new file mode 100644 index 0000000..85843c5 --- /dev/null +++ b/tests/test_deepseek_v41_packing_cp.py @@ -0,0 +1,195 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""GPU regression tests for DeepSeek-V4.1 packed and context-parallel backward. + +The fixture is deliberately checkpoint-free: it drives the real V4.1 Engram hashing, +packed-boundary, contiguous-CP gather/slice and differentiable lookup/projection path with a +tiny parameter set. The CP case launches two local NCCL workers from pytest, so ordinary +CPU or single-GPU jobs collect it safely and report a skip. +""" +import itertools +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from megatron.core import parallel_state +from megatron.core.packed_seq_params import PackedSeqParams + +from mcore_bridge.model.modules import engram as engram_adapter + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA') +requires_two_gpus = pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires two CUDA devices') +requires_nccl = pytest.mark.skipif(not dist.is_nccl_available(), reason='requires NCCL') +requires_native_engram = pytest.mark.skipif( + not engram_adapter.has_native_engram(), reason='requires Megatron-Core Engram support') + + +def _packed_params(lengths, device): + boundaries = torch.tensor([0, *itertools.accumulate(lengths)], dtype=torch.int32, device=device) + return PackedSeqParams( + qkv_format='thd', + cu_seqlens_q=boundaries, + cu_seqlens_kv=boundaries, + max_seqlen_q=max(lengths), + max_seqlen_kv=max(lengths), + cp_partition_mode='contiguous', + ) + + +def _tiny_engram(context_parallel_size, device): + """Construct the real adapter around tiny deterministic PyTorch projections.""" + module = engram_adapter.DeepseekV41Engram.__new__(engram_adapter.DeepseekV41Engram) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + context_parallel_size=context_parallel_size, + sequence_parallel=False, + cp_partition_mode='contiguous', + ) + module.engram_config = SimpleNamespace( + excluded_token_ids=(), + max_ngram_order=3, + num_hash_heads=1, + num_tables=2, + hash_boundary_token_id=0, + boundary_token_id=0, + variant_spec=SimpleNamespace(resets_windows_at_boundary_token=False), + ) + module.num_streams = 2 + module.hidden_size = 4 + module.tp_group = None + module.tokenizer_remap = None + module.register_buffer('hash_multipliers', torch.tensor([17, 31, 43], dtype=torch.int64, device=device)) + module.register_buffer('table_sizes', torch.tensor([17, 19], dtype=torch.int64, device=device)) + module.embedding = torch.nn.Embedding(19, 2, device=device) + module.value_projection = torch.nn.Linear(4, 4, bias=False, device=device) + module.key_projection = torch.nn.Linear(4, 8, bias=False, device=device) + module.key_norm = torch.nn.Identity() + module.query_norm = torch.nn.Identity() + return module + + +def _fixture_tensors(device): + input_ids = torch.tensor([[5, 6, 7, 11, 12, 13, 14, 15]], dtype=torch.long, device=device) + hidden = torch.linspace(-0.75, 0.75, steps=64, device=device).view(8, 1, 8) + probe = torch.linspace(0.5, -0.25, steps=64, device=device).view_as(hidden) + return input_ids, hidden, probe + + +def _parameter_grads(module): + return {name: parameter.grad.detach().clone() for name, parameter in module.named_parameters()} + + +def _assert_finite_nonzero(grads): + for name, grad in grads.items(): + assert torch.isfinite(grad).all(), f'{name} gradient is not finite' + assert grad.abs().sum() > 0, f'{name} gradient is all zero' + + +@requires_cuda +@requires_native_engram +def test_deepseek_v41_packed_backward_matches_separate_documents(): + """A packed document boundary must isolate both activations and parameter/input gradients.""" + if dist.is_initialized(): + pytest.skip('single-rank packing test') + device = torch.device('cuda', 0) + torch.cuda.set_device(device) + torch.manual_seed(2026) + packed_model = _tiny_engram(1, device) + separate_model = _tiny_engram(1, device) + separate_model.load_state_dict(packed_model.state_dict()) + input_ids, hidden, probe = _fixture_tensors(device) + + packed_hidden = hidden.detach().clone().requires_grad_(True) + packed_model._bridge_packed_seq_params = _packed_params([3, 5], device) + packed_output = packed_model(packed_hidden, input_ids) + (packed_output * probe).sum().backward() + packed_grads = _parameter_grads(packed_model) + + separate_hidden = hidden.detach().clone().requires_grad_(True) + separate_outputs = [] + offset = 0 + for length in (3, 5): + separate_outputs.append( + separate_model( + separate_hidden[offset:offset + length], + input_ids[:, offset:offset + length], + )) + offset += length + separate_output = torch.cat(separate_outputs, dim=0) + (separate_output * probe).sum().backward() + separate_grads = _parameter_grads(separate_model) + + torch.testing.assert_close(packed_output, separate_output, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(packed_hidden.grad, separate_hidden.grad, atol=1e-6, rtol=1e-6) + assert packed_grads.keys() == separate_grads.keys() + for name in packed_grads: + torch.testing.assert_close(packed_grads[name], separate_grads[name], atol=1e-6, rtol=1e-6) + _assert_finite_nonzero(packed_grads) + + +def _cp_backward_worker(rank, world_size, init_file): + try: + torch.cuda.set_device(rank) + device = torch.device('cuda', rank) + dist.init_process_group( + backend='nccl', + init_method=f'file://{init_file}', + rank=rank, + world_size=world_size, + ) + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world_size, + ) + torch.manual_seed(2026) + full_model = _tiny_engram(1, device) + cp_model = _tiny_engram(world_size, device) + cp_model.load_state_dict(full_model.state_dict()) + input_ids, hidden, probe = _fixture_tensors(device) + packed = _packed_params([3, 5], device) + + full_hidden = hidden.detach().clone().requires_grad_(True) + full_model._bridge_packed_seq_params = packed + full_output = full_model(full_hidden, input_ids) + (full_output * probe).sum().backward() + full_grads = _parameter_grads(full_model) + + local_length = hidden.shape[0] // world_size + start = rank * local_length + stop = start + local_length + local_hidden = hidden[start:stop].detach().clone().requires_grad_(True) + cp_model._bridge_packed_seq_params = packed + local_output = cp_model(local_hidden, input_ids[:, start:stop]) + (local_output * probe[start:stop]).sum().backward() + cp_grads = _parameter_grads(cp_model) + for grad in cp_grads.values(): + dist.all_reduce(grad) + + output_shards = [torch.empty_like(local_output) for _ in range(world_size)] + hidden_grad_shards = [torch.empty_like(local_hidden.grad) for _ in range(world_size)] + dist.all_gather(output_shards, local_output.detach()) + dist.all_gather(hidden_grad_shards, local_hidden.grad) + + torch.testing.assert_close(torch.cat(output_shards), full_output, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(torch.cat(hidden_grad_shards), full_hidden.grad, atol=1e-6, rtol=1e-6) + assert cp_grads.keys() == full_grads.keys() + for name in cp_grads: + torch.testing.assert_close(cp_grads[name], full_grads[name], atol=1e-5, rtol=1e-5) + _assert_finite_nonzero(cp_grads) + finally: + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@requires_cuda +@requires_two_gpus +@requires_nccl +@requires_native_engram +def test_deepseek_v41_packed_cp2_backward_matches_cp1(tmp_path): + """Two real NCCL CP ranks must reconstruct the packed forward and summed backward.""" + init_file = str(tmp_path / 'deepseek-v41-cp2-init') + mp.spawn(_cp_backward_worker, args=(2, init_file), nprocs=2, join=True) diff --git a/tests/test_hybrid_compat.py b/tests/test_hybrid_compat.py index 3fba665..9551dc9 100644 --- a/tests/test_hybrid_compat.py +++ b/tests/test_hybrid_compat.py @@ -9,6 +9,15 @@ from mcore_bridge.config.parser import hf_to_mcore_config from mcore_bridge.model.register import get_mcore_model +_TE_ATTN_ENV_VARS = ('NVTE_FLASH_ATTN', 'NVTE_FUSED_ATTN', 'NVTE_UNFUSED_ATTN') + + +@pytest.fixture(autouse=True) +def _isolate_te_attention_backend(monkeypatch): + """Keep Megatron's process-wide attention backend selection local to each test.""" + for variable in _TE_ATTN_ENV_VARS: + monkeypatch.delenv(variable, raising=False) + def _config(name): if name == 'qwen3_5': @@ -114,6 +123,10 @@ def test_other_model_forward_backward(name): from mcore_bridge.model.gpts.deepseek_v4 import McoreDSv4HybridSelfAttention if McoreDSv4HybridSelfAttention is object: pytest.skip('the installed release has no DSv4; that existing dev dependency is not a new requirement') + elif name == 'nemotron_h': + from megatron.core.ssm.mamba_mixer import HAVE_MAMBA_SSM + if not HAVE_MAMBA_SSM: + pytest.skip('requires the optional mamba-ssm dependency') with _parallel_context(): config = _config(name) model = get_mcore_model(config)[0].cuda().train() diff --git a/tests/test_save_missing_weights_dsv4.py b/tests/test_save_missing_weights_dsv4.py index 06589d9..81aecc5 100644 --- a/tests/test_save_missing_weights_dsv4.py +++ b/tests/test_save_missing_weights_dsv4.py @@ -11,21 +11,40 @@ * `mtp.*` keys never appear twice under two different naming schemes, which is what would happen if Megatron also exported its own `model.mtp.*` weights. """ +import json import os +import shutil +import tempfile + +import pytest +import torch +from megatron.core import parallel_state +from safetensors.torch import load_file, save_file + +_DIST_ENV_DEFAULTS = { + 'RANK': '0', + 'LOCAL_RANK': '0', + 'WORLD_SIZE': '1', + 'MASTER_ADDR': '127.0.0.1', + 'MASTER_PORT': '29901', +} +_TE_ATTN_ENV_VARS = ('NVTE_FLASH_ATTN', 'NVTE_FUSED_ATTN', 'NVTE_UNFUSED_ATTN') + + +@pytest.fixture(autouse=True) +def _isolate_megatron_runtime(monkeypatch): + """Provide a clean single-rank runtime and release all process-wide Megatron state.""" + for variable, value in _DIST_ENV_DEFAULTS.items(): + if variable not in os.environ: + monkeypatch.setenv(variable, value) + for variable in _TE_ATTN_ENV_VARS: + monkeypatch.delenv(variable, raising=False) + yield + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() -os.environ['CUDA_VISIBLE_DEVICES'] = '0' -# The megatron entrypoint initializes torch.distributed via env:// rendezvous. -os.environ.setdefault('RANK', '0') -os.environ.setdefault('LOCAL_RANK', '0') -os.environ.setdefault('WORLD_SIZE', '1') -os.environ.setdefault('MASTER_ADDR', '127.0.0.1') -os.environ.setdefault('MASTER_PORT', '29901') - -import json # noqa: E402 -import shutil # noqa: E402 -import tempfile # noqa: E402 -import torch # noqa: E402 -from safetensors.torch import load_file, save_file # noqa: E402 MODEL_TYPE = 'deepseek_v4' TEMPLATE = 'deepseek_v4_flash' @@ -276,32 +295,35 @@ def test_dsv4_mtp_weights_dropped_by_default(): def test_dsv4_no_duplicate_mtp_when_megatron_exports_it(): - """Guard against storing the same DSpark parameters under two naming schemes. + """Treat native `mtp.*` and Megatron `model.mtp.*` names as one weight identity.""" + from mcore_bridge.model.gpts.deepseek_v4 import DeepseekV4Bridge + + class RecordingSaver: + + def __init__(self): + self.tensors = {} + + def add_tensor(self, key, tensor): + self.tensors[key] = tensor - When Megatron does materialize MTP layers it writes them as `model.mtp.*`, - while the source checkpoint names them `mtp.*`. Both sets would then land in - the output, doubling the size and leaving it ambiguous which one is loaded. - """ with tempfile.TemporaryDirectory() as tmp_dir: - model_dir = _build_fake_checkpoint(os.path.join(tmp_dir, 'src')) - try: - output_dir = _export( - model_dir, os.path.join(tmp_dir, 'mtp'), save_missing_weights=True, mtp_num_layers=NUM_MTP_STAGES) - except Exception as e: # noqa: BLE001 - # Expected today: `_convert_mtp_extra` looks for the pre-0731 `enorm.weight` - # layout, so Megatron cannot build the DSpark stages at all. - print(f'SKIP: Megatron cannot load DSpark MTP layers yet ({type(e).__name__}: {e})') - return - exported = _load_exported(output_dir) + source = { + 'mtp.0.input.main_proj.weight': torch.ones(2, 2), + 'vision.encoder.weight': torch.full((2, 2), 2.0), + } + save_file(source, os.path.join(tmp_dir, 'model.safetensors')) + saver = RecordingSaver() + bridge = object.__new__(DeepseekV4Bridge) + + bridge._save_missing_weights(saver, {'model.mtp.0.input.main_proj.weight'}, tmp_dir) - megatron_mtp = {k for k in exported if k.startswith('model.mtp.')} - restored_mtp = {k for k in exported if k.startswith('mtp.')} - assert not (megatron_mtp - and restored_mtp), (f'DSpark weights stored twice: {len(megatron_mtp)} keys as `model.mtp.*` and ' - f'{len(restored_mtp)} keys as `mtp.*`') + assert 'mtp.0.input.main_proj.weight' not in saver.tensors + assert torch.equal(saver.tensors['vision.encoder.weight'], source['vision.encoder.weight']) if __name__ == '__main__': + for variable, value in _DIST_ENV_DEFAULTS.items(): + os.environ.setdefault(variable, value) test_dsv4_mtp_weights_restored() test_dsv4_mtp_weights_dropped_by_default() test_dsv4_no_duplicate_mtp_when_megatron_exports_it()