From fe8e498f07398260747fa7e03709888002eec884 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Tue, 15 Sep 2026 19:52:33 +0800 Subject: [PATCH 01/17] feat(deepseek-v41): add DeepSeek-V4.1-Flash bridge - Add deepseek_v41 GPT bridge: CSA2 compressor/indexer, single-pass mHC, native trainable Engram modules, Vision/Aligner, DSpark (MTP) draft stack - Add DSpark stack module (markov_head / confidence_head endpoints) - Config parser/model_config support for composite deepseek_v41 config - GPTBridge.export_weights: add skip_unsupported_export flag (RL weight sync skips engram export; save/checkpoint path unaffected) - FP8/FP4 dequant handling in safetensors util - Add engram load/export unit test --- src/mcore_bridge/__init__.py | 4 +- src/mcore_bridge/bridge/gpt_bridge.py | 13 + src/mcore_bridge/config/__init__.py | 2 +- src/mcore_bridge/config/model_config.py | 53 + src/mcore_bridge/config/parser.py | 50 +- src/mcore_bridge/model/constant.py | 1 + src/mcore_bridge/model/gpt_model.py | 3 +- src/mcore_bridge/model/gpts/__init__.py | 2 +- src/mcore_bridge/model/gpts/deepseek_v4.py | 9 + src/mcore_bridge/model/gpts/deepseek_v41.py | 1244 +++++++++++++++++++ src/mcore_bridge/model/mm_gpt_model.py | 2 + src/mcore_bridge/model/modules/__init__.py | 12 + src/mcore_bridge/model/modules/dspark.py | 456 +++++++ src/mcore_bridge/model/register.py | 1 + src/mcore_bridge/utils/safetensors.py | 21 +- tests/test_deepseek_v41_engram.py | 782 ++++++++++++ 16 files changed, 2646 insertions(+), 9 deletions(-) create mode 100644 src/mcore_bridge/model/gpts/deepseek_v41.py create mode 100644 src/mcore_bridge/model/modules/dspark.py create mode 100644 tests/test_deepseek_v41_engram.py diff --git a/src/mcore_bridge/__init__.py b/src/mcore_bridge/__init__.py index 50100df8..e0fbc16f 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 41caf7cb..3c79e170 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1758,6 +1758,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(): @@ -1838,6 +1842,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: @@ -1948,6 +1954,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. @@ -1965,6 +1972,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. @@ -1975,6 +1987,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() diff --git a/src/mcore_bridge/config/__init__.py b/src/mcore_bridge/config/__init__.py index 204ce711..09fa6b77 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 c5231404..bedfb830 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 @@ -242,6 +243,28 @@ 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) + 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 @@ -355,6 +378,31 @@ def __post_init__(self): self.mtp_num_layers = 1 else: self.mtp_unroll_steps = self.mtp_num_layers + if self.dspark_num_layers is not None or 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, + 'dspark_num_experts': self.dspark_num_experts, + 'dspark_router_topk': self.dspark_router_topk, + } + 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.') + 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: @@ -418,3 +466,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 ef554700..0a5ff912 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'], @@ -79,11 +80,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'], @@ -185,8 +209,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' @@ -204,6 +228,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 1296f342..ed193a7e 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 52946102..cf96a052 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 655415bc..40a1bfbb 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 2d8e6257..73897b14 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -162,6 +162,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 @@ -304,6 +310,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] @@ -378,6 +385,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, 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 00000000..baa1e4dc --- /dev/null +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -0,0 +1,1244 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1-Flash (text backbone) bridge for megatron-core. + +Wires the V4.1 *language* 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 text integration also attaches the native trainable Engram modules and +loads their EP-local table rows directly from the official flat FP8 tensors. +Vision, MTP and DSpark are integrated separately. + +The V4.1 loader deliberately selects megatron-core's native TransformerBlock, +whose forward owns the per-call CSA2State and SinglePassMHCState lifecycle. Other +mcore-bridge models keep the custom TransformerBlock path. +""" +import copy +import os +from contextlib import contextmanager + +import torch +import torch.nn.functional as F +import transformer_engine +from megatron.core import parallel_state +from megatron.core.models.engram.config import EngramConfig +from megatron.core.models.engram.layer_specs import apply_engram_to_layer_spec +from megatron.core.tensor_parallel.layers import VocabParallelEmbedding +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +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 typing import Optional + +from mcore_bridge.config import MLAModelConfig +from mcore_bridge.model.modules.dspark import DeepseekV41DSparkStack + +from ..constant import ModelType +from ..mm_gpt_model import MultimodalGPTModel +from ..register import ModelMeta, register_model +from .deepseek_v4 import ( + _apply_mla_rope, + DeepseekV4Bridge, + DeepseekV4GPTModel, + 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 + + +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) + + +class DeepseekV41GPTModel(DeepseekV4GPTModel): + """V4.1 language model with opt-in DSpark target-layer capture. + + DSpark consumes the attention inputs of its target layers. The official target + IDs are zero-based; Megatron layer numbers are one-based. Capturing is opt-in + so regular training does not retain three large activation graphs. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._capture_dspark_hidden = False + self._dspark_hidden_states = {} + self._dspark_hook_handles = [] + target_ids = tuple(self.config.dspark_target_layer_ids or ()) + for layer in self.decoder.layers: + layer_id = layer.layer_number - 1 + if layer_id not in target_ids: + continue + self._dspark_hook_handles.append( + layer.register_forward_pre_hook(self._make_dspark_capture_hook(layer_id), with_kwargs=True)) + + @staticmethod + def _contract_dspark_target_hidden(hidden_states: torch.Tensor, num_streams: int): + if hidden_states.ndim != 3: + raise ValueError(f'DSpark target hidden states must be [s, b, n*h], got {tuple(hidden_states.shape)}.') + if hidden_states.shape[-1] % num_streams: + raise ValueError( + f'DSpark target hidden width {hidden_states.shape[-1]} is not divisible by {num_streams} streams.') + return hidden_states.unflatten(-1, (num_streams, -1)).mean(dim=-2) + + def _make_dspark_capture_hook(self, layer_id): + + def _capture(_module, args, kwargs): + if not self._capture_dspark_hidden: + return + hidden_states = kwargs.get('hidden_states') + if hidden_states is None and args: + hidden_states = args[0] + if hidden_states is None: + raise ValueError(f'Could not capture the attention input for DSpark target layer {layer_id}.') + self._dspark_hidden_states[layer_id] = self._contract_dspark_target_hidden( + hidden_states, self.config.num_residual_streams) + + return _capture + + @contextmanager + def capture_dspark_hidden_states(self): + if not self.config.dspark_target_layer_ids: + raise ValueError('DSpark target-layer capture requested, but DSpark is not configured.') + self._dspark_hidden_states = {} + self._capture_dspark_hidden = True + try: + yield + finally: + self._capture_dspark_hidden = False + + def get_dspark_main_hidden(self, clear: bool = True): + target_ids = tuple(self.config.dspark_target_layer_ids or ()) + missing = [layer_id for layer_id in target_ids if layer_id not in self._dspark_hidden_states] + if missing: + raise RuntimeError( + f'DSpark target layers {missing} were not captured on this pipeline rank. ' + 'The DSpark draft stack must be colocated with all target layers.') + result = torch.cat([self._dspark_hidden_states[layer_id] for layer_id in target_ids], dim=-1) + if clear: + self._dspark_hidden_states = {} + return result + + def forward(self, *args, **kwargs): + inference_context = kwargs.get('inference_context') or kwargs.get('inference_params') + capture_dspark = ( + hasattr(self, 'dspark') + and not self.training + and inference_context is not None + and inference_context.is_dynamic_batching() + and inference_context.num_speculative_tokens > 0 + ) + if not capture_dspark: + return super().forward(*args, **kwargs) + if inference_context.using_cuda_graph_this_step(): + raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') + with self.capture_dspark_hidden_states(): + return super().forward(*args, **kwargs) + + def _dspark_rotary_for_positions(self, position_ids: torch.Tensor): + if self.position_embedding_type != 'rope' or self.rotary_pos_emb is None: + raise RuntimeError('DSpark requires RoPE position embeddings.') + position_ids = position_ids.to(dtype=torch.long) + max_position = int(position_ids.max().item()) + 1 + rotary_table = self.rotary_pos_emb(max_position) + if isinstance(rotary_table, dict): + rotary_table = rotary_table['main'] + selected = rotary_table.index_select(0, position_ids.reshape(-1)) + # Drop the table's singleton batch dimension. The requested position tensor + # supplies the batch axes for packed main tokens or the parallel draft block. + return selected.reshape(*position_ids.shape, *rotary_table.shape[2:]) + + def _dspark_word_embeddings(self): + """Return the token embedding DSpark uses to embed its draft seed. + + On a single stage (or a standard-MTP stage) the base input embedding is + colocated and reused. On a PP>1 last stage with untied embeddings the base + model has no ``embedding``; a dedicated replicated DSpark embedding built + in ``build_model`` and loaded from ``model.embed_tokens.weight`` is used. + """ + if hasattr(self, 'embedding'): + return self.embedding.word_embeddings + dspark_embedding = getattr(self, 'dspark_word_embeddings', None) + if dspark_embedding is None: + raise RuntimeError( + 'DSpark requires an input word embedding on its pipeline stage, but neither ' + 'the base embedding nor a dedicated DSpark embedding is present.') + return dspark_embedding + + def forward_dspark( + self, + main_hidden, + input_ids, + *, + start_pos, + rotary_pos_emb, + main_rotary_pos_emb, + inference_context=None, + temperature=0.0, + sample_fn=None, + cache_slots=None, + prefill_only=None, + ): + if not hasattr(self, 'dspark'): + raise RuntimeError('DSpark is not available on this pipeline stage.') + if not hasattr(self, 'output_layer'): + raise RuntimeError('DSpark requires the output head on its pipeline stage.') + return self.dspark( + main_hidden, + input_ids, + self._dspark_word_embeddings(), + self.output_layer, + start_pos=start_pos, + rotary_pos_emb=rotary_pos_emb, + main_rotary_pos_emb=main_rotary_pos_emb, + inference_context=inference_context, + temperature=temperature, + sample_fn=sample_fn, + cache_slots=cache_slots, + prefill_only=prefill_only, + ) + + def compute_dspark_speculative_tokens( + self, + next_token_ids, + accepted_token_counts, + last_accepted_seq_indices, + num_speculative_tokens, + inference_context, + sample_fn, + ): + """Commit verified target states and produce one parallel DSpark draft block.""" + if inference_context.using_cuda_graph_this_step(): + raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') + if num_speculative_tokens > self.config.dspark_block_size: + raise ValueError( + f'Requested {num_speculative_tokens} speculative tokens, but DSpark block size is ' + f'{self.config.dspark_block_size}.') + + main_hidden = self.get_dspark_main_hidden() + if self.config.sequence_parallel and parallel_state.get_tensor_model_parallel_world_size() > 1: + main_hidden = gather_from_sequence_parallel_region(main_hidden, group=self.tp_group) + if main_hidden.ndim != 3 or main_hidden.shape[1] != 1: + raise RuntimeError( + 'DSpark dynamic inference expects packed target states [tokens, 1, targets*h], ' + f'got {tuple(main_hidden.shape)}.') + + active_count = inference_context.total_request_count - inference_context.paused_request_count + active_slice = slice(inference_context.paused_request_count, inference_context.total_request_count) + query_lengths = inference_context.request_query_lengths[active_slice].to(dtype=torch.long) + active_token_count = int(query_lengths.sum().item()) + if main_hidden.shape[0] != active_token_count: + raise RuntimeError( + f'DSpark captured {main_hidden.shape[0]} target rows for {active_token_count} active tokens.') + + device = main_hidden.device + request_ids = inference_context.request_ids[active_slice].to(device=device, dtype=torch.long) + live_request_ids = inference_context.request_ids[:inference_context.total_request_count].to( + device=device, dtype=torch.long) + cache_slots = self.dspark.resolve_cache_slots(request_ids, live_request_ids) + token_positions = inference_context.token_to_position_in_request[:active_token_count].to( + device=device, dtype=torch.long) + accepted_token_counts = accepted_token_counts[:active_count].to(device='cpu', dtype=torch.long) + + offset = 0 + for request_index, query_length in enumerate(query_lengths.tolist()): + if request_index < inference_context.num_decode_requests: + accepted_length = min(query_length, int(accepted_token_counts[request_index].item()) + 1) + else: + accepted_length = query_length + if accepted_length: + token_slice = slice(offset, offset + accepted_length) + positions = token_positions[token_slice] + if positions.numel() > 1 and not torch.all(positions[1:] == positions[:-1] + 1): + raise RuntimeError('DSpark cache updates require contiguous per-request token positions.') + self.dspark.update_main_cache( + main_hidden[token_slice], + self._dspark_rotary_for_positions(positions), + start_pos=positions[0], + cache_slots=cache_slots[request_index:request_index + 1], + inference_context=inference_context, + ) + offset += query_length + + last_indices = last_accepted_seq_indices[:active_count].to(device=device, dtype=torch.long) + last_hidden = main_hidden.index_select(0, last_indices).transpose(0, 1).contiguous() + main_positions = token_positions.index_select(0, last_indices) + draft_positions = main_positions.unsqueeze(0) + 1 + torch.arange( + self.config.dspark_block_size, device=device).unsqueeze(1) + output_ids, _, _ = self.forward_dspark( + last_hidden, + next_token_ids[:active_count], + start_pos=main_positions, + rotary_pos_emb=self._dspark_rotary_for_positions(draft_positions), + main_rotary_pos_emb=self._dspark_rotary_for_positions(main_positions), + inference_context=inference_context, + sample_fn=sample_fn, + cache_slots=cache_slots, + prefill_only=False, + ) + return output_ids[:, 1:num_speculative_tokens + 1].transpose(0, 1).contiguous() + + +class DeepseekV41MultimodalGPTModel(MultimodalGPTModel): + language_model_cls = DeepseekV41GPTModel + + @property + def vocab_size(self): + return self.language_model.vocab_size + + def forward_with_dspark_hidden(self, *args, **kwargs): + with self.language_model.capture_dspark_hidden_states(): + output = self.forward(*args, **kwargs) + return output, self.language_model.get_dspark_main_hidden() + + def forward_dspark(self, *args, **kwargs): + return self.language_model.forward_dspark(*args, **kwargs) + + def compute_dspark_speculative_tokens(self, *args, **kwargs): + return self.language_model.compute_dspark_speculative_tokens(*args, **kwargs) + + +class DeepseekV41Loader(DeepseekV4Loader): + model_cls = DeepseekV41MultimodalGPTModel + # Native V4.1 forward owns CSA2State + SinglePassMHCState. Using it only for + # this loader avoids changing the custom bridge block used by V4/DSpark/MTP. + transformer_block = McoreTransformerBlock + + def _get_engram_config(self): + hf_layer_ids = tuple(self.config.engram_layer_ids or ()) + if not hf_layer_ids: + return None + 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 = EngramConfig( + global_vocab_sizes=(self.config.engram_vocab_size,) * (max_ngram_order - 1), + # TransformerLayer numbers are 1-based, while the official checkpoint and + # PCG64 multiplier seeds use the original 0-based HF layer IDs. + layer_ids=tuple(layer_id + 1 for layer_id in 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 + dspark_config.num_moe_experts = self.config.dspark_num_experts + 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 + return dspark_config, layer_specs + + def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): + model = super().build_model(pre_process, post_process, vp_stage) + if not self.config.dspark_num_layers or not post_process: + return model + language_model = model.language_model + 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, + ) + for index, layer_spec in enumerate(dspark_layer_specs) + ] + language_model.dspark = DeepseekV41DSparkStack(dspark_config, layers) + self._set_linear_is_expert(language_model.dspark) + # 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, + ) + return model + + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + from megatron.core.models.gpt.experimental_attention_variant_module_specs import \ + get_transformer_block_with_experimental_attention_variant_spec + transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec(self.config, vp_stage) + for layer_spec in transformer_layer_spec.layer_specs: + layer_spec.submodules.self_attention.module = DSv4HybridSelfAttention + core_attention_submodules = layer_spec.submodules.self_attention.submodules.core_attention.submodules + if getattr(core_attention_submodules, 'compressor', None) is not None: + core_attention_submodules.compressor.module = CSA2Compressor + if getattr(core_attention_submodules, 'indexer', None) is not None: + # CSA2 indexer is flat (no nested compressor). + core_attention_submodules.indexer.module = CSA2Indexer + engram_config = self._get_engram_config() + if engram_config is not None: + transformer_layer_spec = apply_engram_to_layer_spec(transformer_layer_spec, engram_config) + return transformer_layer_spec + + +class DeepseekV41Bridge(DeepseekV4Bridge): + _ENGRAM_LOAD_CHUNK_ROWS = 65536 + additional_dim0_keys = DeepseekV4Bridge.additional_dim0_keys | {'embed', 'head'} + additional_dim1_keys = DeepseekV4Bridge.additional_dim1_keys | {'main_proj'} + + def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): + 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 _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)) + expected_rows = self.config.engram_num_embeddings[ + self.config.engram_layer_ids.index(engram.layer_number - 1)] + if flat_offset != expected_rows: + raise ValueError( + f'Engram layer {engram.layer_number - 1} 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_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): + layer_prefix = f'{hf_prefix}{layer_idx}.' + local_state = self._remove_prefix(hf_state_dict, layer_prefix) if to_mcore else {} + result = super()._set_layer_state(mg_layer, hf_state_dict, hf_prefix, layer_idx, to_mcore) + if layer_idx in (self.config.engram_layer_ids or []): + self._set_layer_engram(mg_layer, local_state, to_mcore) + if not to_mcore and local_state: + result.update(self._add_prefix(local_state, layer_prefix)) + return result + + 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_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage): + if not self.config.dspark_num_layers or (to_mcore and not is_pp_last_stage): + return + language_model = mg_model.language_model if self.is_multimodal else mg_model + dspark = getattr(language_model, 'dspark', None) + if dspark is None: + raise RuntimeError('DSpark weights require the draft stack on the final 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 + 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 + + def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): + # V4.1 single-pass mHC has no learned hc_head_*; skip the V4 hc_head mapping + # and only handle the plain final layernorm (base GPTBridge behaviour). + super(DeepseekV4Bridge, self)._set_final_layernorm(lm_model, hf_state_dict, to_mcore) + + +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 ee610f31..1f1fd89e 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -82,6 +82,7 @@ def forward( packed_seq_params: PackedSeqParams = 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 @@ -107,6 +108,7 @@ def forward( attention_mask=attention_mask, decoder_input=decoder_input, labels=labels, + inference_context=inference_context, inference_params=inference_params, packed_seq_params=packed_seq_params, extra_block_kwargs=kwargs, diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py index 996a3bb2..9eb3331c 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 00000000..d70dec8d --- /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/register.py b/src/mcore_bridge/model/register.py index 847c9b0b..6548c07c 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/safetensors.py b/src/mcore_bridge/utils/safetensors.py index 2e22f0e9..0ad3ea3d 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 00000000..b1479a0b --- /dev/null +++ b/tests/test_deepseek_v41_engram.py @@ -0,0 +1,782 @@ +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from megatron.core import mpu +from megatron.core.inference.text_generation_controllers.text_generation_controller import TextGenerationController +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.deepseek_v41 import ( + DeepseekV41Aligner, + DeepseekV41Bridge, + DeepseekV41DSparkAttention, + DeepseekV41GPTModel, + DeepseekV41Vision, + DeepseekV41VisionTransformer, +) +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_target_hidden_averages_mhc_streams(): + hidden = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).view(2, 3, 20) + + actual = DeepseekV41GPTModel._contract_dspark_target_hidden(hidden, num_streams=4) + expected = hidden.view(2, 3, 4, 5).mean(dim=2) + + assert actual.shape == (2, 3, 5) + torch.testing.assert_close(actual, expected) + + +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.]])) + + +def test_dspark_tp_modules_construct_and_run_on_one_rank(tmp_path): + if dist.is_initialized() and dist.get_world_size() != 1: + pytest.skip('Single-rank DSpark TP smoke test.') + if not dist.is_initialized(): + dist.init_process_group( + 'gloo', + init_method=f'file://{tmp_path}/dspark-dist-init', + rank=0, + world_size=1, + ) + if not mpu.model_parallel_is_initialized(): + mpu.initialize_model_parallel(tensor_model_parallel_size=1) + + 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_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_dspark_model_commits_only_verified_states_before_proposal(): + class _DSpark: + + def __init__(self): + self.updates = [] + + def resolve_cache_slots(self, request_ids, live_request_ids): + assert torch.equal(request_ids.cpu(), torch.tensor([10, 20])) + assert torch.equal(live_request_ids.cpu(), torch.tensor([10, 20])) + return torch.tensor([2, 0], device=request_ids.device) + + def update_main_cache( + self, + main_hidden, + rotary_pos_emb, + *, + start_pos, + cache_slots, + inference_context, + ): + self.updates.append((main_hidden.clone(), start_pos.clone(), cache_slots.clone())) + + model = DeepseekV41GPTModel.__new__(DeepseekV41GPTModel) + torch.nn.Module.__init__(model) + model.config = SimpleNamespace(sequence_parallel=False, dspark_block_size=3) + model.dspark = _DSpark() + captured = torch.arange(5 * 4, dtype=torch.float32).view(5, 1, 4) + model.get_dspark_main_hidden = lambda: captured + model._dspark_rotary_for_positions = lambda positions: positions.float() + proposal_args = {} + + def forward_dspark(main_hidden, input_ids, **kwargs): + proposal_args.update(main_hidden=main_hidden, input_ids=input_ids, **kwargs) + output_ids = torch.tensor([[31, 32, 33, 34], [41, 42, 43, 44]]) + return output_ids, None, None + + model.forward_dspark = forward_dspark + context = SimpleNamespace( + total_request_count=2, + paused_request_count=0, + num_decode_requests=1, + request_query_lengths=torch.tensor([3, 2], dtype=torch.int32), + request_ids=torch.tensor([10, 20], dtype=torch.int32), + token_to_position_in_request=torch.tensor([5, 6, 7, 0, 1], dtype=torch.int32), + using_cuda_graph_this_step=lambda: False, + ) + + proposals = model.compute_dspark_speculative_tokens( + next_token_ids=torch.tensor([31, 41]), + accepted_token_counts=torch.tensor([1, 0]), + last_accepted_seq_indices=torch.tensor([1, 4]), + num_speculative_tokens=2, + inference_context=context, + sample_fn=lambda logits: logits.argmax(dim=-1), + ) + + assert len(model.dspark.updates) == 2 + torch.testing.assert_close(model.dspark.updates[0][0], captured[:2]) + torch.testing.assert_close(model.dspark.updates[1][0], captured[3:5]) + assert model.dspark.updates[0][1].item() == 5 + assert model.dspark.updates[1][1].item() == 0 + assert torch.equal(proposal_args['start_pos'], torch.tensor([6, 1])) + torch.testing.assert_close(proposal_args['main_hidden'], captured[[1, 4]].transpose(0, 1)) + assert torch.equal(proposals, torch.tensor([[32, 42], [33, 43]])) + + +def test_controller_routes_speculative_proposals_to_dspark_provider(): + calls = {} + + class _Model: + + def compute_dspark_speculative_tokens(self, **kwargs): + calls.update(kwargs) + return torch.tensor([[7, 8], [9, 10]]) + + context = SimpleNamespace( + total_request_count=2, + paused_request_count=0, + _nvls_dispatcher=None, + ) + controller = TextGenerationController.__new__(TextGenerationController) + controller.inference_wrapped_model = SimpleNamespace(inference_context=context) + controller._unwrapped_model = _Model() + controller._is_last_pp_stage = True + controller.model_is_pipeline_parallel = False + controller.model_config = SimpleNamespace(dspark_block_size=3) + controller.num_speculative_tokens = 2 + controller._sampled_tokens_cuda = torch.tensor([5, 6]) + controller._accepted_token_counts_per_request = torch.tensor([1, 0]) + controller._last_accepted_seq_indices = torch.tensor([1, 3]) + controller._sampled_mtp_tokens_cuda = torch.empty(2, 2, dtype=torch.long) + controller._sample_from_logits_2d = lambda logits: logits.argmax(dim=-1) + + controller._compute_dspark_and_sample() + + assert torch.equal(controller._sampled_mtp_tokens_cuda, torch.tensor([[7, 8], [9, 10]])) + assert calls['inference_context'] is context + assert calls['sample_fn'] is controller._sample_from_logits_2d + assert calls['num_speculative_tokens'] == 2 + + +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_word_embeddings_resolver_prefers_base_then_dedicated(): + resolver = DeepseekV41GPTModel._dspark_word_embeddings + model = DeepseekV41GPTModel.__new__(DeepseekV41GPTModel) + torch.nn.Module.__init__(model) + # Neither the base embedding nor a dedicated DSpark embedding is present. + with pytest.raises(RuntimeError): + resolver(model) + # A PP>1 last stage falls back to the dedicated DSpark embedding. + dedicated = object() + model.dspark_word_embeddings = dedicated + assert resolver(model) is dedicated + # When the base embedding is colocated it always takes priority. + base = object() + model.embedding = SimpleNamespace(word_embeddings=base) + assert resolver(model) is base + + +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), + layer_number=2, + ) + 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]) From 50987ddb82235763f2c35532b2198af57fd85079 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Tue, 15 Sep 2026 20:54:10 +0800 Subject: [PATCH 02/17] feat(deepseek-v41): add DSpark and Engram support --- src/mcore_bridge/config/model_config.py | 13 +- src/mcore_bridge/inference/__init__.py | 5 + src/mcore_bridge/inference/dspark.py | 125 ++++++ src/mcore_bridge/model/gpts/deepseek_v41.py | 59 ++- src/mcore_bridge/model/modules/engram.py | 358 ++++++++++++++++++ src/mcore_bridge/model/modules/topk_router.py | 54 ++- tests/test_deepseek_v41_engram.py | 125 +++++- 7 files changed, 716 insertions(+), 23 deletions(-) create mode 100644 src/mcore_bridge/inference/__init__.py create mode 100644 src/mcore_bridge/inference/dspark.py create mode 100644 src/mcore_bridge/model/modules/engram.py diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index bedfb830..acd4bbe8 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -173,6 +173,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' @@ -244,6 +248,9 @@ class ModelConfig(TransformerConfig): 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 @@ -302,7 +309,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 @@ -344,6 +350,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: diff --git a/src/mcore_bridge/inference/__init__.py b/src/mcore_bridge/inference/__init__.py new file mode 100644 index 00000000..428bc63b --- /dev/null +++ b/src/mcore_bridge/inference/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +from .dspark import DeepseekV41DynamicInferenceEngine, DeepseekV41TextGenerationController + +__all__ = ['DeepseekV41DynamicInferenceEngine', 'DeepseekV41TextGenerationController'] diff --git a/src/mcore_bridge/inference/dspark.py b/src/mcore_bridge/inference/dspark.py new file mode 100644 index 00000000..927245b7 --- /dev/null +++ b/src/mcore_bridge/inference/dspark.py @@ -0,0 +1,125 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1 DSpark adapters for Megatron's dynamic inference API.""" + +from contextlib import contextmanager + +import torch +from megatron.core.inference.communication_utils import broadcast_from_last_pipeline_stage +from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher + + +@contextmanager +def _standard_mtp_compatibility(config): + """Satisfy legacy MTP-only constructor validation without changing model semantics.""" + sentinel = object() + original_num_layers = getattr(config, 'mtp_num_layers', sentinel) + original_repeated = getattr(config, 'mtp_use_repeated_layer', sentinel) + config.mtp_num_layers = max(getattr(config, 'mtp_num_layers', 0) or 0, 1) + config.mtp_use_repeated_layer = True + try: + yield + finally: + if original_num_layers is sentinel: + delattr(config, 'mtp_num_layers') + else: + config.mtp_num_layers = original_num_layers + if original_repeated is sentinel: + delattr(config, 'mtp_use_repeated_layer') + else: + config.mtp_use_repeated_layer = original_repeated + + +def _validate_dspark_speculation(model_config, num_speculative_tokens): + if num_speculative_tokens <= 0: + return + block_size = getattr(model_config, 'dspark_block_size', 0) + if not getattr(model_config, 'dspark_num_layers', None): + raise ValueError('DSpark speculative decoding requires dspark_num_layers.') + if not block_size or num_speculative_tokens > block_size: + raise ValueError( + f'num_speculative_tokens={num_speculative_tokens} must not exceed ' + f'dspark_block_size={block_size}.') + if getattr(model_config, 'cuda_graph_impl', None) == 'local': + raise ValueError('DSpark speculative decoding does not support local CUDA graphs yet.') + + +class DeepseekV41TextGenerationController(TextGenerationController): + """Route Megatron's standard speculative loop through the parallel DSpark draft stack.""" + + def __init__(self, inference_wrapped_model, tokenizer): + model_config = inference_wrapped_model.model.config + inference_config = inference_wrapped_model.inference_context.config + self._uses_dspark = bool(getattr(model_config, 'dspark_num_layers', None)) + if self._uses_dspark: + _validate_dspark_speculation(model_config, inference_config.num_speculative_tokens) + # PR #7224 validates speculative decoding as standard serial MTP. DSpark is + # separate, so adapt only while the upstream constructor initializes buffers. + with _standard_mtp_compatibility(model_config): + super().__init__(inference_wrapped_model, tokenizer) + self._uses_dspark = True + self.num_mtp_depths = 0 + else: + super().__init__(inference_wrapped_model, tokenizer) + + def _compute_dspark_and_sample(self): + context = self.inference_wrapped_model.inference_context + active_request_count = context.total_request_count - context.paused_request_count + speculative_tokens = None + + if self._is_last_pp_stage: + compute_dspark = getattr(self._unwrapped_model, 'compute_dspark_speculative_tokens', None) + if compute_dspark is None: + raise RuntimeError( + 'DSpark speculative decoding requires compute_dspark_speculative_tokens() ' + 'on the last pipeline stage.') + if context._nvls_dispatcher: + NVLSAllGatherVDispatcher.modify_real_token_count_for_mtp( + active_request_count * self.model_config.dspark_block_size) + speculative_tokens = compute_dspark( + next_token_ids=self._sampled_tokens_cuda[:active_request_count], + accepted_token_counts=self._accepted_token_counts_per_request[:active_request_count], + last_accepted_seq_indices=self._last_accepted_seq_indices, + num_speculative_tokens=self.num_speculative_tokens, + inference_context=context, + sample_fn=self._sample_from_logits_2d, + ) + expected_shape = (self.num_speculative_tokens, active_request_count) + if tuple(speculative_tokens.shape) != expected_shape: + raise RuntimeError( + f'DSpark returned speculative tokens with shape {tuple(speculative_tokens.shape)}; ' + f'expected {expected_shape}.') + + if self.model_is_pipeline_parallel: + speculative_tokens = broadcast_from_last_pipeline_stage( + [self.num_speculative_tokens, active_request_count], + dtype=torch.int64, + tensor=speculative_tokens, + pp_group=self.pp_group, + ) + self._sampled_mtp_tokens_cuda[ + :self.num_speculative_tokens, :active_request_count + ].copy_(speculative_tokens) + + def _compute_serial_mtp_and_sample(self): + # The upstream event loop invokes this extension point after verification and KV + # rewind. Reusing it avoids copying Megatron's large scheduling loop. + if self._uses_dspark: + return self._compute_dspark_and_sample() + return super()._compute_serial_mtp_and_sample() + + +class DeepseekV41DynamicInferenceEngine(DynamicInferenceEngine): + """Dynamic engine adapter that validates DSpark instead of serial-MTP depth.""" + + def __init__(self, controller, context): + model_config = controller.inference_wrapped_model.model.config + if getattr(model_config, 'dspark_num_layers', None): + _validate_dspark_speculation(model_config, context.config.num_speculative_tokens) + with _standard_mtp_compatibility(model_config): + super().__init__(controller, context) + else: + super().__init__(controller, context) diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index baa1e4dc..4d7ffb17 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -30,8 +30,6 @@ import torch.nn.functional as F import transformer_engine from megatron.core import parallel_state -from megatron.core.models.engram.config import EngramConfig -from megatron.core.models.engram.layer_specs import apply_engram_to_layer_spec from megatron.core.tensor_parallel.layers import VocabParallelEmbedding from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region from megatron.core.transformer.module import MegatronModule, mark_keep_in_fp32 @@ -42,6 +40,12 @@ from mcore_bridge.config import MLAModelConfig from mcore_bridge.model.modules.dspark import DeepseekV41DSparkStack +from mcore_bridge.model.modules.engram import ( + adapt_deepseek_v41_layer_specs, + allow_engram_inference, + build_deepseek_v41_engram_config, + has_native_engram, +) from ..constant import ModelType from ..mm_gpt_model import MultimodalGPTModel @@ -648,20 +652,32 @@ def get_dspark_main_hidden(self, clear: bool = True): return result def forward(self, *args, **kwargs): + input_ids = kwargs.get('input_ids', args[0] if args else None) inference_context = kwargs.get('inference_context') or kwargs.get('inference_params') - capture_dspark = ( - hasattr(self, 'dspark') - and not self.training - and inference_context is not None - and inference_context.is_dynamic_batching() - and inference_context.num_speculative_tokens > 0 - ) - if not capture_dspark: - return super().forward(*args, **kwargs) - if inference_context.using_cuda_graph_this_step(): - raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') - with self.capture_dspark_hidden_states(): - return super().forward(*args, **kwargs) + if inference_context is None and len(args) > 5: + inference_context = args[5] + extra_block_kwargs = kwargs.get('extra_block_kwargs') + if extra_block_kwargs is None and len(args) > 7: + extra_block_kwargs = args[7] + + with allow_engram_inference(self.config, input_ids, extra_block_kwargs) as block_kwargs: + if len(args) > 7: + args = (*args[:7], block_kwargs, *args[8:]) + else: + kwargs['extra_block_kwargs'] = block_kwargs + capture_dspark = ( + hasattr(self, 'dspark') + and not self.training + and inference_context is not None + and inference_context.is_dynamic_batching() + and inference_context.num_speculative_tokens > 0 + ) + if not capture_dspark: + return super().forward(*args, **kwargs) + if inference_context.using_cuda_graph_this_step(): + raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') + with self.capture_dspark_hidden_states(): + return super().forward(*args, **kwargs) def _dspark_rotary_for_positions(self, position_ids: torch.Tensor): if self.position_embedding_type != 'rope' or self.rotary_pos_emb is None: @@ -836,6 +852,11 @@ 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', @@ -859,11 +880,11 @@ def _get_engram_config(self): max_ngram_order = self.config.engram_max_ngram_size image_token_id = getattr(self.config.hf_config, 'image_token_id', None) - engram_config = EngramConfig( + 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 - # PCG64 multiplier seeds use the original 0-based HF layer IDs. - layer_ids=tuple(layer_id + 1 for layer_id in hf_layer_ids), + # tokenizer artifact use the original 0-based HF layer IDs. + placement_layer_ids=tuple(layer_id + 1 for layer_id in hf_layer_ids), hash_layer_ids=hf_layer_ids, max_ngram_order=max_ngram_order, num_hash_heads=self.config.engram_n_heads, @@ -967,7 +988,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): core_attention_submodules.indexer.module = CSA2Indexer engram_config = self._get_engram_config() if engram_config is not None: - transformer_layer_spec = apply_engram_to_layer_spec(transformer_layer_spec, engram_config) + transformer_layer_spec = adapt_deepseek_v41_layer_specs(transformer_layer_spec, engram_config) return transformer_layer_spec diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py new file mode 100644 index 00000000..afeb2893 --- /dev/null +++ b/src/mcore_bridge/model/modules/engram.py @@ -0,0 +1,358 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1 adapters for NVIDIA Megatron-Core's optional Engram modules.""" + +from contextlib import contextmanager + +import torch +from torch import Tensor + +try: + 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 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 + + +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 +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 _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, +) -> 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: + suffixes = [compressed] + for shift in range(1, max_ngram_order): + suffixes.append(torch.nn.functional.pad( + compressed, (shift, 0), value=boundary_token_id)[:, :sequence_length]) + 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 _hash_windows(self, token_windows: Tensor) -> Tensor: + return _hash_token_windows( + token_windows, + 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, + invalid_token_id=-1, + ) + + def _static_inference_hashes(self, input_ids: Tensor, context) -> tuple[Tensor, Tensor]: + batch_size, sequence_length = input_ids.shape + shape = (context.max_batch_size, context.max_sequence_length) + cache = getattr(context, 'engram_token_cache', None) + if cache is None or cache.device != input_ids.device or cache.shape != shape: + cache = input_ids.new_full(shape, self.engram_config.boundary_token_id) + context.engram_token_cache = cache + + batch_start = context.batch_size_offset + batch_end = batch_start + batch_size + sequence_start = context.sequence_len_offset + sequence_end = sequence_start + sequence_length + if batch_end > shape[0] or sequence_end > shape[1]: + raise ValueError('Engram inference token cache is too small for the current batch/chunk.') + masked_ids, live = self._mask_excluded_tokens(input_ids) + cache[batch_start:batch_end, sequence_start:sequence_end] = masked_ids + shifts = torch.arange( + self.engram_config.max_ngram_order, device=input_ids.device, dtype=torch.long) + positions = torch.arange( + sequence_start, sequence_end, device=input_ids.device, dtype=torch.long).unsqueeze(-1) - shifts + gather_positions = positions.clamp_min(0).reshape(1, -1).expand(batch_size, -1) + windows = cache[batch_start:batch_end].gather(1, gather_positions).view( + batch_size, sequence_length, self.engram_config.max_ngram_order) + windows = torch.where( + positions.unsqueeze(0) >= 0, + windows, + windows.new_full((), self.engram_config.boundary_token_id), + ) + return self._hash_windows(windows), live + + def _dynamic_inference_hashes(self, input_ids: Tensor, context) -> tuple[Tensor, Tensor]: + if input_ids.shape[0] != 1: + raise ValueError('Dynamic Engram inference expects flattened input_ids with batch size 1.') + shape = (context.max_requests, context.max_sequence_length) + cache = getattr(context, 'engram_token_cache', None) + if cache is None or cache.device != input_ids.device or cache.shape != shape: + cache = input_ids.new_full(shape, self.engram_config.boundary_token_id) + context.engram_token_cache = cache + + total_tokens = input_ids.shape[1] + active_tokens = min(int(context.active_token_count), total_tokens) + live = torch.zeros_like(input_ids, dtype=torch.bool) + hashes = input_ids.new_zeros((1, total_tokens, self.engram_config.num_tables), dtype=torch.long) + if active_tokens == 0: + return hashes, live + request_indices = context.gpu_view.token_to_request_idx[:active_tokens].long() + token_positions = context.gpu_view.token_to_position_in_request[:active_tokens].long() + if request_indices.min() < 0 or request_indices.max() >= shape[0]: + raise ValueError('Dynamic Engram inference received an out-of-range request index.') + if token_positions.min() < 0 or token_positions.max() >= shape[1]: + raise ValueError('Dynamic Engram inference received an out-of-range token position.') + masked_ids, active_live = self._mask_excluded_tokens(input_ids[:, :active_tokens]) + cache[request_indices, token_positions] = masked_ids.squeeze(0) + shifts = torch.arange( + self.engram_config.max_ngram_order, device=input_ids.device, dtype=torch.long) + positions = token_positions.unsqueeze(-1) - shifts + windows = cache[request_indices.unsqueeze(-1).expand_as(positions), positions.clamp_min(0)] + windows = torch.where( + positions >= 0, + windows, + windows.new_full((), self.engram_config.boundary_token_id), + ) + hashes[:, :active_tokens] = self._hash_windows(windows.unsqueeze(0)) + live[:, :active_tokens] = active_live + return hashes, live + + def _build_hash_ids(self, input_ids: Tensor, inference_context=None) -> tuple[Tensor, Tensor]: + masked_ids, live = self._mask_excluded_tokens(input_ids) + if inference_context is None: + 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, + ) + return hashes, live + if inference_context.is_static_batching(): + return self._static_inference_hashes(input_ids, inference_context) + return self._dynamic_inference_hashes(input_ids, inference_context) + + 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 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]}.') + + nvtx_range_push('engram.hash') + try: + hash_ids, live_tokens = self._build_hash_ids(input_ids, inference_context) + 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.unsqueeze(-1), 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) + previous = getattr(engram, '_bridge_inference_context', None) + engram._bridge_inference_context = kwargs.get('inference_context') + try: + return super()._forward_attention(*args, **kwargs) + finally: + engram._bridge_inference_context = previous + + +if TransformerLayer is not None: + + class DeepseekV41TransformerLayer(_DeepseekV41EngramLayerMixin, TransformerLayer): + pass + + + class DeepseekV41HyperConnectionTransformerLayer( + _DeepseekV41EngramLayerMixin, HyperConnectionTransformerLayer): + pass +else: + DeepseekV41TransformerLayer = None + DeepseekV41HyperConnectionTransformerLayer = None + + +def adapt_deepseek_v41_layer_specs(transformer_layer_spec, engram_config): + """Attach the V4.1 Engram module and inference-aware layer subclasses.""" + if not has_native_engram(): + raise RuntimeError('The installed Megatron-Core does not provide Engram.') + from megatron.core.transformer.spec_utils import ModuleSpec + + engram_spec = ModuleSpec(module=DeepseekV41Engram, params={'engram_config': engram_config}) + for layer_spec in transformer_layer_spec.layer_specs: + if layer_spec.module is HyperConnectionTransformerLayer: + layer_spec.module = DeepseekV41HyperConnectionTransformerLayer + elif layer_spec.module is TransformerLayer: + layer_spec.module = DeepseekV41TransformerLayer + if not hasattr(layer_spec.submodules, 'engram'): + raise RuntimeError( + 'The installed Engram extension does not expose TransformerLayerSubmodules.engram.') + layer_spec.submodules.engram = engram_spec + return transformer_layer_spec + + +@contextmanager +def allow_engram_inference(model_config, input_ids, extra_block_kwargs): + """Bypass PR #7231's inference guard while preserving input_ids propagation.""" + if not getattr(model_config, 'engram_enabled', False): + yield extra_block_kwargs + return + if input_ids is None: + raise ValueError('Engram requires input token IDs on every pipeline stage.') + block_kwargs = dict(extra_block_kwargs or {}) + block_kwargs['input_ids'] = input_ids + model_config.engram_enabled = False + try: + yield block_kwargs + finally: + model_config.engram_enabled = True diff --git a/src/mcore_bridge/model/modules/topk_router.py b/src/mcore_bridge/model/modules/topk_router.py index d1c9400f..e666e7f1 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/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index b1479a0b..f3596d3e 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -1,15 +1,16 @@ +import json from types import SimpleNamespace import pytest import torch import torch.distributed as dist from megatron.core import mpu -from megatron.core.inference.text_generation_controllers.text_generation_controller import TextGenerationController 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.inference import DeepseekV41TextGenerationController from mcore_bridge.model.gpts.deepseek_v41 import ( DeepseekV41Aligner, DeepseekV41Bridge, @@ -18,6 +19,7 @@ DeepseekV41Vision, DeepseekV41VisionTransformer, ) +from mcore_bridge.model.modules import engram as engram_adapter from mcore_bridge.model.modules.dspark import ( DeepseekV41DSparkConfidenceHead, DeepseekV41DSparkInput, @@ -489,7 +491,7 @@ def compute_dspark_speculative_tokens(self, **kwargs): paused_request_count=0, _nvls_dispatcher=None, ) - controller = TextGenerationController.__new__(TextGenerationController) + controller = DeepseekV41TextGenerationController.__new__(DeepseekV41TextGenerationController) controller.inference_wrapped_model = SimpleNamespace(inference_context=context) controller._unwrapped_model = _Model() controller._is_last_pp_stage = True @@ -510,6 +512,125 @@ def compute_dspark_speculative_tokens(self, **kwargs): assert calls['num_speculative_tokens'] == 2 +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 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 test_engram_static_inference_cache_matches_full_sequence_hashing(): + module = engram_adapter.DeepseekV41Engram.__new__(engram_adapter.DeepseekV41Engram) + torch.nn.Module.__init__(module) + module.engram_config = SimpleNamespace( + excluded_token_ids=(99,), + max_ngram_order=3, + num_hash_heads=1, + hash_boundary_token_id=0, + boundary_token_id=0, + num_tables=2, + variant_spec=SimpleNamespace(resets_windows_at_boundary_token=False), + ) + module.tokenizer_remap = None + module.hash_multipliers = torch.tensor([11, 13, 15]) + module.table_sizes = torch.tensor([997, 991]) + full_hashes, full_live = module._build_hash_ids(torch.tensor([[1, 2, 3]])) + context = SimpleNamespace( + max_batch_size=1, + max_sequence_length=8, + batch_size_offset=0, + sequence_len_offset=0, + is_static_batching=lambda: True, + ) + + prefill_hashes, prefill_live = module._build_hash_ids(torch.tensor([[1, 2]]), context) + context.sequence_len_offset = 2 + decode_hashes, decode_live = module._build_hash_ids(torch.tensor([[3]]), context) + + assert torch.equal(torch.cat((prefill_hashes, decode_hashes), dim=1), full_hashes) + assert torch.equal(torch.cat((prefill_live, decode_live), dim=1), full_live) + + +def test_engram_layer_spec_uses_bridge_owned_module(): + if not engram_adapter.has_native_engram(): + pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_layer import ( + HyperConnectionTransformerLayer, + TransformerLayerSubmodules, + ) + + layer_spec = ModuleSpec( + module=HyperConnectionTransformerLayer, + submodules=TransformerLayerSubmodules(), + ) + block_spec = SimpleNamespace(layer_specs=[layer_spec]) + config = SimpleNamespace(layer_ids=(1,)) + + engram_adapter.adapt_deepseek_v41_layer_specs(block_spec, config) + + assert layer_spec.module is engram_adapter.DeepseekV41HyperConnectionTransformerLayer + assert layer_spec.submodules.engram.module is engram_adapter.DeepseekV41Engram + + +def test_allow_engram_inference_preserves_input_ids_and_restores_flag(): + config = SimpleNamespace(engram_enabled=True) + input_ids = torch.tensor([[1, 2]]) + + with engram_adapter.allow_engram_inference(config, input_ids, {'marker': 1}) as kwargs: + assert not config.engram_enabled + assert kwargs['marker'] == 1 + assert kwargs['input_ids'] is input_ids + + assert config.engram_enabled + + class _Table: def __init__(self, global_rows, row_start, row_end, dim): From 3a31a08d00a02f87eb835ff7bc7bbc232ea886d3 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 15:46:53 +0800 Subject: [PATCH 03/17] wip --- src/mcore_bridge/model/gpts/deepseek_v41.py | 29 +- .../model/gpts/deepseek_v41_hybrid.py | 492 ++++++++++++++++++ src/mcore_bridge/model/mm_gpt_model.py | 10 +- src/mcore_bridge/model/modules/engram.py | 155 +++++- src/mcore_bridge/utils/megatron_utils.py | 15 +- tests/test_deepseek_v41_engram.py | 208 ++++++++ tests/test_deepseek_v41_hybrid.py | 260 +++++++++ 7 files changed, 1153 insertions(+), 16 deletions(-) create mode 100644 src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py create mode 100644 tests/test_deepseek_v41_hybrid.py diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 4d7ffb17..54624bdf 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -25,6 +25,7 @@ import copy import os from contextlib import contextmanager +from types import SimpleNamespace import torch import torch.nn.functional as F @@ -848,6 +849,14 @@ class DeepseekV41Loader(DeepseekV4Loader): # this loader avoids changing the custom bridge block used by V4/DSpark/MTP. transformer_block = McoreTransformerBlock + def _engram_placement_layer_ids(self, hf_layer_ids): + """Map 0-based HF Engram layer IDs to 1-based ``TransformerLayer`` placement numbers. + + On the GPT stack HF layer ``e`` is one ``TransformerLayer`` numbered ``e + 1``. The + HybridStack loader overrides this because there each HF layer becomes two hybrid layers. + """ + return tuple(layer_id + 1 for layer_id in hf_layer_ids) + def _get_engram_config(self): hf_layer_ids = tuple(self.config.engram_layer_ids or ()) if not hf_layer_ids: @@ -884,7 +893,7 @@ def _get_engram_config(self): 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=tuple(layer_id + 1 for layer_id in 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, @@ -941,6 +950,11 @@ def get_dspark_layer_spec(self): 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 build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): @@ -1053,13 +1067,22 @@ def _load_engram_embedding(self, engram, hf_state_dict): 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(engram.layer_number - 1)] + self.config.engram_layer_ids.index(hf_layer_id)] if flat_offset != expected_rows: raise ValueError( - f'Engram layer {engram.layer_number - 1} expected {expected_rows} flat rows, ' + f'Engram layer {hf_layer_id} expected {expected_rows} flat rows, ' f'but its prime tables contain {flat_offset}.') + def _engram_hf_layer_id(self, engram): + """Recover the 0-based HF layer ID from a built Engram's 1-based ``layer_number``. + + On the GPT stack ``layer_number == hf_id + 1``. The HybridStack bridge overrides this + because its Engram sits on the doubled-space attention layer ``2 * hf_id + 1``. + """ + return engram.layer_number - 1 + def _set_layer_engram(self, mg_layer, hf_state_dict, to_mcore): engram = self._get_layer_engram(mg_layer) if to_mcore: diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py new file mode 100644 index 00000000..1694c5ea --- /dev/null +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -0,0 +1,492 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""DeepSeek-V4.1 on megatron-core's ``HybridModel`` (pipeline-parallel path). + +The default :class:`DeepseekV41Loader` builds a ``GPTModel`` whose custom +``TransformerBlock`` owns the CSA2 / single-pass-mHC forward. Upstream refuses to +run that block under pipeline parallelism:: + + # transformer_block.py:304-311 + if (config.pipeline_model_parallel_size > 1 + and config.experimental_attention_variant == "dsv4_hybrid" + and config.dsv4_version == "v4.1"): + raise ValueError("V4.1 pipeline parallelism requires HybridModel and its payload adapter") + +so PP>1 requires the native ``HybridModel`` + ``CSA2HybridAdapter`` typed-payload path. + +On ``HybridModel`` one *pattern symbol is one layer*: a GPT ``attn+mlp`` layer becomes +two hybrid layers -- an attention-only layer (symbol ``D``) followed by an MLP-only +layer (``E`` for MoE, ``-`` for dense). So a hybrid stack has ``2 * num_layers`` layers +and every per-layer config array that CSA2 indexes by ``layer_number - 1`` must be +re-expanded into this doubled index space (see :func:`derive_hybrid_layer_config`). + +This module keeps the GPT loader untouched (golden baseline) and adds the hybrid path +alongside it; both are validated to agree before the default is switched. +""" +import copy +from dataclasses import dataclass +from typing import List, Optional, Sequence, Union + +import torch +import torch.distributed as dist +from megatron.core import mpu +from tqdm import tqdm + +from mcore_bridge.utils import is_master + +from ..modules.engram import DeepseekV41Engram, DeepseekV41TransformerLayer +from .deepseek_v41 import (CSA2Compressor, CSA2Indexer, DeepseekV41Bridge, DeepseekV41Loader, + DSv4HybridSelfAttention) + +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 + + +if _HYBRID_MODEL_AVAILABLE: + + class DeepseekV41HyperConnectionHybridLayer(HyperConnectionHybridLayer): + """Hyper-connection wrapper that keeps Engram inside the mHC layer delta. + + With ``enable_hyper_connections=True`` (always set for V4.1, see parser.py) HybridStack + wraps every layer in :class:`HyperConnectionHybridLayer`. Its eager forward takes a + *fast path* (:meth:`_call_inner_transformer_layer_without_local_bda`) that calls the + inner layer's ``_forward_self_attention_output_with_bias`` directly. That method skips + ``_forward_attention`` -> ``_maybe_apply_engram`` entirely, so it (a) never adds the + Engram residual to the n-stream layer delta and (b) never forwards ``input_ids`` to the + attention branch. Both silently drop Engram on the PP path. + + For the (few) layers that actually carry an Engram module we therefore decline the fast + path by returning ``None``. :meth:`HyperConnectionHybridLayer.forward` then falls back to + ``_call_inner_layer``, which runs the inner ``DeepseekV41TransformerLayer``'s full + ``forward`` (the ``_DeepseekV41EngramLayerMixin`` stashes the inference context there) and + computes ``layer_output - aggregated`` -- Engram delta included -- reproducing the + GPTModel golden path exactly. Non-Engram layers keep the fast path untouched. + + This subclass adds no state and overrides one method, so it is applied by an in-place + ``__class__`` swap on the already-built wrappers (see + :meth:`DeepseekV41HybridLoader._rewrap_engram_hyper_connection_layers`) -- HybridStack + hard-codes the wrapper class with no spec hook. The fast path is also invoked by the + CUDA-graph capture body, which is out of scope for this change (plan: no CUDA Graph); + returning ``None`` there would raise rather than miscompute. + """ + + def _call_inner_transformer_layer_without_local_bda(self, *args, **kwargs): + if getattr(self.inner_layer, 'engram', None) is not None: + return None + return super()._call_inner_transformer_layer_without_local_bda(*args, **kwargs) +else: + DeepseekV41HyperConnectionHybridLayer = None + + +@dataclass +class HybridLayerConfig: + """Per-layer config re-expanded from GPT 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 a GPT-space V4.1 config into the doubled hybrid layer space. + + Each GPT 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)), + ) + + +class DeepseekV41HybridLoader(DeepseekV41Loader): + """Build DeepSeek-V4.1 on ``HybridModel`` (the PP-capable path). + + Reuses :class:`DeepseekV41Loader`'s Engram config resolution and MLA/CSA2 knowledge, but + swaps the model class to the native ``HybridModel`` and rewrites the layer config into the + doubled hybrid layer space (see :func:`derive_hybrid_layer_config`). The golden GPT + ``DeepseekV41Loader`` is left untouched; this loader derives its own config copy so both + paths can coexist in one process. + + B1 covers the text backbone only. MTP (B2), DSpark capture (B3) and the multimodal wrapper + (B4) are added on top; here MTP is disabled so the backbone can be aligned in isolation. + """ + + model_cls = HybridModel + + 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] + # MTP on HybridModel is B2 (its inner attention cannot use the CSA2 'D' symbol, which + # rejects is_mtp_layer). Disable it for the B1 backbone-only alignment. + 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 same fp8-parity module swaps the GPT loader uses, 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 like the GPT ``adapt_deepseek_v41_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}) + return spec + + 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 via ``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.""" + 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, + ) + self._rewrap_engram_hyper_connection_layers(model) + self._set_linear_is_expert(model) + return model + + +class DeepseekV41HybridBridge(DeepseekV41Bridge): + """Weight bridge for the ``HybridModel`` backbone (PP-capable path). + + The GPT bridge maps one HF layer onto one ``TransformerLayer`` that owns both attention + and MLP. On ``HybridModel`` that layer 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 (the GPT layer instead carried two: + ``self_attention_hyper_connection`` + ``mlp_hyper_connection``). The GPT + ``hc_{attn,ffn}_*`` HF keys therefore split across the two wrappers. + + B1 handles the text backbone only. It treats the model as its own language model (the + multimodal wrapper is B4) and skips MTP (B2). ``self.config`` stays in GPT layer space + (``num_layers == N``); the model's decoder holds ``2 * N`` layers. + """ + + @staticmethod + def _lm(mg_model): + """Resolve the language model. B1's HybridModel is text-only (no ``language_model`` + wrapper); B4 will nest it under a multimodal container.""" + language_model = getattr(mg_model, 'language_model', None) + return mg_model if language_model is None else language_model + + def _engram_hf_layer_id(self, engram): + # Engram lives on the doubled-space attention layer ``2 * hf_id + 1`` (see + # ``DeepseekV41HybridLoader._engram_placement_layer_ids``), so map it back to HF space. + return (engram.layer_number - 1) // 2 + + def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): + # Text-only word embeddings; visual embeds (image_start/end/newline) are B4. + if to_mcore: + hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) + else: + hf_state_dict = {} + lm_model = self._lm(mg_model) + self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key, to_mcore) + if to_mcore: + return {} + return self._add_prefix(hf_state_dict, hf_prefix) + + 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``); the block-level output hyper-connection head is unchanged. + self._set_state_dict(lm_model, 'decoder.final_norm.weight', hf_state_dict, self.hf_final_layernorm_key, + to_mcore) + for key in ['hc_head_base', 'hc_head_fn', 'hc_head_scale']: + self._set_state_dict(lm_model, f'decoder.{key}', hf_state_dict, f'model.{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) + if has_hyper_connection: + 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(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 (B2); the multimodal-wrapper indirection is B4. + """ + 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 = {} + # HybridStack layer_number spans the doubled space (i + 1 + pp_offset), matching this + # loop's hybrid index so the PP-availability window below stays correct. + num_hybrid_layers = 2 * self.config.num_layers + 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()) diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index ae37b8b0..2dfdd638 100644 --- a/src/mcore_bridge/model/mm_gpt_model.py +++ b/src/mcore_bridge/model/mm_gpt_model.py @@ -59,7 +59,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) @@ -91,11 +92,13 @@ def forward( 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) @@ -106,7 +109,8 @@ 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, diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py index afeb2893..d149466f 100644 --- a/src/mcore_bridge/model/modules/engram.py +++ b/src/mcore_bridge/model/modules/engram.py @@ -1,12 +1,16 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """DeepSeek-V4.1 adapters for NVIDIA Megatron-Core's optional Engram modules.""" +import dataclasses from contextlib import contextmanager 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 ( @@ -18,7 +22,7 @@ HyperConnectionTransformerLayer, TransformerLayer, ) - from megatron.core.utils import nvtx_range_pop, nvtx_range_push + from megatron.core.utils import get_pg_size, nvtx_range_pop, nvtx_range_push except ImportError: EngramConfig = None Engram = torch.nn.Module @@ -31,6 +35,22 @@ def has_native_engram() -> bool: return EngramConfig is not None +class _ContextParallelSizeOneView: + """Read-only view of a transformer config that reports ``context_parallel_size == 1``. + + Lets the Engram validators reuse every upstream parallelism check except the CP one, + without mutating the shared transformer config. + """ + + context_parallel_size = 1 + + 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): @@ -58,6 +78,30 @@ def _load_tokenizer_map(self): return super()._load_tokenizer_map() finally: self.layer_ids = placement_layer_ids + + def _validate_parallelism(self, transformer_config, sequence_length): + # DeepseekV41Engram hashes the full sequence locally and then selects its own CP + # slice, so no window ever has to cross a CP rank boundary. Drop only the upstream + # `context_parallel_size == 1` rejection and keep every other check. + 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( + _ContextParallelSizeOneView(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 @@ -108,6 +152,24 @@ def _hash_token_windows( 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, @@ -117,6 +179,7 @@ def _build_ngram_hashes( 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) @@ -127,10 +190,22 @@ def _build_ngram_hashes( 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): - suffixes.append(torch.nn.functional.pad( - compressed, (shift, 0), value=boundary_token_id)[:, :sequence_length]) + 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, @@ -241,7 +316,8 @@ def _dynamic_inference_hashes(self, input_ids: Tensor, context) -> tuple[Tensor, live[:, :active_tokens] = active_live return hashes, live - def _build_hash_ids(self, input_ids: Tensor, inference_context=None) -> tuple[Tensor, Tensor]: + def _build_hash_ids(self, input_ids: Tensor, inference_context=None, + cu_seqlens: Tensor | None = None) -> tuple[Tensor, Tensor]: masked_ids, live = self._mask_excluded_tokens(input_ids) if inference_context is None: hashes = _build_ngram_hashes( @@ -253,27 +329,88 @@ def _build_hash_ids(self, input_ids: Tensor, inference_context=None) -> tuple[Te 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 if inference_context.is_static_batching(): return self._static_inference_hashes(input_ids, inference_context) return self._dynamic_inference_hashes(input_ids, inference_context) + 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) + 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 + if context_parallel and inference_context is not None: + raise ValueError('Engram inference does not support context parallelism.') + + 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) + if context_parallel and getattr(packed_seq_params, 'cp_partition_mode', + 'zigzag') != 'contiguous': + raise ValueError( + "Engram with context parallelism requires cp_partition_mode='contiguous', " + 'matching the DSv4 THD CP forward.') nvtx_range_push('engram.hash') try: - hash_ids, live_tokens = self._build_hash_ids(input_ids, inference_context) + 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, inference_context, 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.unsqueeze(-1), hidden_states.shape[0], self.tp_group).squeeze(-1) + live_tokens, hidden_states.shape[0], self.tp_group).squeeze(-1) finally: nvtx_range_pop('engram.hash') @@ -300,12 +437,18 @@ 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: diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index 419ed539..5108100a 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/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index f3596d3e..4e126d78 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -550,6 +550,74 @@ def test_engram_adapter_remaps_checkpoint_layers_to_megatron_layers(tmp_path): 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_parallelism_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) + + with pytest.raises(ValueError, match='expert_tensor_parallel_size'): + config._validate_parallelism( + SimpleNamespace(**{**parallelism, 'expert_tensor_parallel_size': 2}), None) + with pytest.raises(ValueError, match='virtual pipeline'): + config._validate_parallelism( + SimpleNamespace(**{**parallelism, 'virtual_pipeline_model_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]]]), @@ -597,6 +665,146 @@ def test_engram_static_inference_cache_matches_full_sequence_hashing(): assert torch.equal(torch.cat((prefill_live, decode_live), dim=1), full_live) +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) + 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 + + def test_engram_layer_spec_uses_bridge_owned_module(): if not engram_adapter.has_native_engram(): pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py new file mode 100644 index 00000000..e132e99f --- /dev/null +++ b/tests/test_deepseek_v41_hybrid.py @@ -0,0 +1,260 @@ +# 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_hybrid 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 golden GPT path stays intact. + from types import SimpleNamespace + + from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41HybridLoader + + loader = object.__new__(DeepseekV41HybridLoader) + 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 disabled for the B1 backbone-only path. + assert cfg.mtp_num_layers is None + # golden GPT 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_hybrid import DeepseekV41HybridBridge + + bridge = object.__new__(DeepseekV41HybridBridge) + 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_hybrid import DeepseekV41HybridBridge, DeepseekV41HybridLoader + + loader = object.__new__(DeepseekV41HybridLoader) + assert loader._engram_placement_layer_ids([1, 3]) == (3, 7) + assert loader._engram_placement_layer_ids([0]) == (1,) + + bridge = object.__new__(DeepseekV41HybridBridge) + 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 + + +import pytest # noqa: E402 + +from mcore_bridge.model.gpts.deepseek_v41_hybrid 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_declines_fast_path_only_for_engram_layers(): + # The V4.1 wrapper subclass returns None (forcing the full-forward `_call_inner_layer` + # path, which applies Engram) iff the inner layer carries an Engram module; otherwise it + # must delegate unchanged to the base fast path. + from types import SimpleNamespace + from unittest.mock import patch + + engram_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) + engram_layer.inner_layer = SimpleNamespace(engram=object()) + assert engram_layer._call_inner_transformer_layer_without_local_bda('h', 'mask') is None + + plain_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) + plain_layer.inner_layer = SimpleNamespace(engram=None) + sentinel = object() + with patch.object( + HyperConnectionHybridLayer, + '_call_inner_transformer_layer_without_local_bda', + return_value=sentinel) as base_call: + assert plain_layer._call_inner_transformer_layer_without_local_bda('h', 'mask') is sentinel + base_call.assert_called_once() + + +@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_hybrid import DeepseekV41HybridLoader + + 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__(DeepseekV41HybridLoader) + 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_hybrid import DeepseekV41HybridLoader + + engram_wrapper = object.__new__(HyperConnectionHybridLayer) + engram_wrapper.inner_layer = SimpleNamespace(engram=object()) + model = SimpleNamespace(decoder=SimpleNamespace(layers=[engram_wrapper])) + + loader = object.__new__(DeepseekV41HybridLoader) + loader.config = SimpleNamespace(enable_hyper_connections=False) + loader._rewrap_engram_hyper_connection_layers(model) + + assert type(engram_wrapper) is HyperConnectionHybridLayer + From 4076d63e9a6bbb712ad19928c0355258c9fcc6fa Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 18:39:14 +0800 Subject: [PATCH 04/17] feat(deepseek-v41): support pipeline parallelism via HybridModel (B1 text backbone) Migrate the DeepSeek-V4.1 text backbone from GPTModel to the PP-capable HybridModel path, keeping GPTModel as the golden baseline. - Add DeepseekV41HybridLoader/Bridge: 1 GPT layer -> 2 hybrid layers (D attn + E/- mlp), 2x-derived CSA2 arrays/pattern, block-aligned PP/VPP segmentation, Engram placement at hybrid layer 2e+1, and an HC wrapper forward-override that applies the Engram delta on the n-stream residual to match the golden order. - Route to the hybrid path when pipeline_model_parallel_size > 1, with a deepseek_v41_hybrid config flag (via --megatron_extra_kwargs) to force it at PP1. - Fix DSv4HybridSelfAttention layer_type lookup for the doubled hybrid index space. - Normalize _convert layer count across the load (GPT-space) / export (doubled) config spaces via _num_hybrid_layers, fixing an export crash on out-of-range layers. GPU-validated on a tiny model: hybrid PP1/PP2/PP2xEP2/PP2xDP2 all align with the GPTModel baseline at iter-1 (<0.06% rel loss/grad) with clean weight exports. Unit tests: 62 passed (hybrid 27 + engram 35). --- src/mcore_bridge/config/model_config.py | 5 + src/mcore_bridge/model/gpts/deepseek_v4.py | 10 +- src/mcore_bridge/model/gpts/deepseek_v41.py | 36 ++ .../model/gpts/deepseek_v41_hybrid.py | 315 ++++++++++++++++-- tests/test_deepseek_v41_hybrid.py | 193 ++++++++++- 5 files changed, 509 insertions(+), 50 deletions(-) diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index 4e000947..09f96c98 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -255,6 +255,11 @@ class ModelConfig(TransformerConfig): mhc_sinkhorn_iterations: int = 20 mhc_init_gating_factor: float = 0.01 moe_n_hash_layers: int = 0 + # DeepSeek-V4.1 pipeline-parallel path selector (None = auto). The default GPTModel path is + # the golden baseline and cannot run PP>1, so the HybridModel loader/bridge are auto-selected + # when pipeline_model_parallel_size > 1. Set explicitly (e.g. via --megatron_extra_kwargs) to + # force the hybrid path on at PP1 (baseline alignment) or off. See deepseek_v41.py. + deepseek_v41_hybrid: Optional[bool] = None # deepseek-v4.1 engram (HF layer IDs are 0-based) # Declared here as well so the bridge remains importable on the PR #7224 baseline, diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index 5a6c1dce..ae49049a 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 diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 54624bdf..777a2b89 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -843,12 +843,38 @@ def compute_dspark_speculative_tokens(self, *args, **kwargs): return self.language_model.compute_dspark_speculative_tokens(*args, **kwargs) +def _deepseek_v41_use_hybrid(config) -> bool: + """Whether to build DeepSeek-V4.1 on the ``HybridModel`` (PP-capable) path. + + The default ``GPTModel`` path is the golden baseline and stays the default until the hybrid + path is fully aligned (plan step B5). Upstream refuses to run the V4.1 ``dsv4_hybrid`` block + under pipeline parallelism, so the hybrid loader/bridge are auto-selected whenever + ``pipeline_model_parallel_size > 1``. The ``deepseek_v41_hybrid`` config flag (settable via + ``--megatron_extra_kwargs``) overrides this: ``True`` forces the hybrid path on at PP1 (used to + align it against the GPTModel baseline), ``False`` keeps GPTModel even at PP>1. + """ + forced = getattr(config, 'deepseek_v41_hybrid', None) + if forced is not None: + return bool(forced) + return (getattr(config, 'pipeline_model_parallel_size', 1) or 1) > 1 + + class DeepseekV41Loader(DeepseekV4Loader): model_cls = DeepseekV41MultimodalGPTModel # Native V4.1 forward owns CSA2State + SinglePassMHCState. Using it only for # this loader avoids changing the custom bridge block used by V4/DSpark/MTP. transformer_block = McoreTransformerBlock + def __new__(cls, config=None, *args, **kwargs): + # Auto-route to the HybridModel loader on the PP path (or when forced); the subclass + # instantiates itself directly, so the ``cls is`` guard prevents re-dispatch. ``config`` + # is optional so ``__new__(cls)`` (used by tests to skip __init__) keeps working. + if cls is DeepseekV41Loader and config is not None and _deepseek_v41_use_hybrid(config): + from .deepseek_v41_hybrid import DeepseekV41HybridLoader + if DeepseekV41HybridLoader is not None: + return super().__new__(DeepseekV41HybridLoader) + return super().__new__(cls) + def _engram_placement_layer_ids(self, hf_layer_ids): """Map 0-based HF Engram layer IDs to 1-based ``TransformerLayer`` placement numbers. @@ -1011,6 +1037,16 @@ class DeepseekV41Bridge(DeepseekV4Bridge): additional_dim0_keys = DeepseekV4Bridge.additional_dim0_keys | {'embed', 'head'} additional_dim1_keys = DeepseekV4Bridge.additional_dim1_keys | {'main_proj'} + def __new__(cls, config=None, *args, **kwargs): + # Mirror the loader's routing so the bridge and the built model always agree on the path + # (this bridge is created in ``ModelConfig.__post_init__`` where PP size is already set). + # ``config`` is optional so ``__new__(cls)`` (used by tests to skip __init__) keeps working. + if cls is DeepseekV41Bridge and config is not None and _deepseek_v41_use_hybrid(config): + from .deepseek_v41_hybrid import DeepseekV41HybridBridge + if DeepseekV41HybridBridge is not None: + return super().__new__(DeepseekV41HybridBridge) + return super().__new__(cls) + def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): result = super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) target = hf_state_dict if to_mcore else result diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index 1694c5ea..3f5e5452 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -29,11 +29,13 @@ import torch import torch.distributed as dist from megatron.core import mpu +from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding from tqdm import tqdm from mcore_bridge.utils import is_master from ..modules.engram import DeepseekV41Engram, DeepseekV41TransformerLayer +from ..rope import get_rope_inv_freq from .deepseek_v41 import (CSA2Compressor, CSA2Indexer, DeepseekV41Bridge, DeepseekV41Loader, DSv4HybridSelfAttention) @@ -54,37 +56,225 @@ if _HYBRID_MODEL_AVAILABLE: class DeepseekV41HyperConnectionHybridLayer(HyperConnectionHybridLayer): - """Hyper-connection wrapper that keeps Engram inside the mHC layer delta. - - With ``enable_hyper_connections=True`` (always set for V4.1, see parser.py) HybridStack - wraps every layer in :class:`HyperConnectionHybridLayer`. Its eager forward takes a - *fast path* (:meth:`_call_inner_transformer_layer_without_local_bda`) that calls the - inner layer's ``_forward_self_attention_output_with_bias`` directly. That method skips - ``_forward_attention`` -> ``_maybe_apply_engram`` entirely, so it (a) never adds the - Engram residual to the n-stream layer delta and (b) never forwards ``input_ids`` to the - attention branch. Both silently drop Engram on the PP path. - - For the (few) layers that actually carry an Engram module we therefore decline the fast - path by returning ``None``. :meth:`HyperConnectionHybridLayer.forward` then falls back to - ``_call_inner_layer``, which runs the inner ``DeepseekV41TransformerLayer``'s full - ``forward`` (the ``_DeepseekV41EngramLayerMixin`` stashes the inference context there) and - computes ``layer_output - aggregated`` -- Engram delta included -- reproducing the - GPTModel golden path exactly. Non-Engram layers keep the fast path untouched. - - This subclass adds no state and overrides one method, so it is applied by an in-place - ``__class__`` swap on the already-built wrappers (see - :meth:`DeepseekV41HybridLoader._rewrap_engram_hyper_connection_layers`) -- HybridStack - hard-codes the wrapper class with no spec hook. The fast path is also invoked by the - CUDA-graph capture body, which is out of scope for this change (plan: no CUDA Graph); - returning ``None`` there would raise rather than miscompute. + """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). This reproduces + the GPTModel golden path exactly. 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:`DeepseekV41HybridLoader._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 _call_inner_transformer_layer_without_local_bda(self, *args, **kwargs): - if getattr(self.inner_layer, 'engram', None) is not None: + 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:`DeepseekV41HybridBridge`, 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 (B2) is still appended by the base resolver, so a + segmented main becomes ``seg0|seg1|.../mtp``. + """ + + # B1 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, + # like the GPT ``DeepseekV41MultimodalGPTModel``). Expose it so that guard short-circuits; + # the real vision tower arrives with the multimodal wrapper in B4. + visual = None + + @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 the GPT path) 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._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``).""" + 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' + return { + '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), + } + + def _inject_dsv4_rotary_pos_emb(self, module, args, kwargs): + if kwargs.get('rotary_pos_emb') is not None: return None - return super()._call_inner_transformer_layer_without_local_bda(*args, **kwargs) + 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 (B4) and then + # cleared before the language model runs; the base HybridModel.forward never accepts them. + # The B1 backbone is text-only, so strip them here. For a text batch the GPT wrapper's + # ``get_inputs_embeds`` is a numeric no-op (``_zero_parameter_dependency`` adds ``0 * + # vision_params``), so dropping them keeps parity with the GPTModel baseline. + _visual_forward_keys = ('pixel_values', 'image_grid_thw', 'image_token_types', 'token_types') + + def forward(self, *args, **kwargs): + if kwargs.get('pixel_values') is not None: + raise NotImplementedError( + 'DeepSeek-V4.1 hybrid (pipeline-parallel) path is text-only in B1; multimodal ' + 'inputs require the B4 multimodal wrapper.') + for key in self._visual_forward_keys: + kwargs.pop(key, None) + return super().forward(*args, **kwargs) else: DeepseekV41HyperConnectionHybridLayer = None + DeepseekV41HybridStackModel = None @dataclass @@ -185,7 +375,7 @@ class DeepseekV41HybridLoader(DeepseekV41Loader): (B4) are added on top; here MTP is disabled so the backbone can be aligned in isolation. """ - model_cls = HybridModel + model_cls = DeepseekV41HybridStackModel def _engram_placement_layer_ids(self, hf_layer_ids): """On HybridStack, HF layer ``e`` becomes the attention-only 'D' layer at hybrid index @@ -248,8 +438,32 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): # 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 the GPT router swap). + 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 @@ -309,8 +523,11 @@ class DeepseekV41HybridBridge(DeepseekV41Bridge): ``hc_{attn,ffn}_*`` HF keys therefore split across the two wrappers. B1 handles the text backbone only. It treats the model as its own language model (the - multimodal wrapper is B4) and skips MTP (B2). ``self.config`` stays in GPT layer space - (``num_layers == N``); the model's decoder holds ``2 * N`` layers. + multimodal wrapper is B4) and skips MTP (B2). ``self.config`` is seen in two layer spaces + depending on direction: on load it is the original GPT-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. """ @staticmethod @@ -320,6 +537,23 @@ def _lm(mg_model): 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 GPT-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 # ``DeepseekV41HybridLoader._engram_placement_layer_ids``), so map it back to HF space. @@ -359,11 +593,11 @@ def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcor 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``); the block-level output hyper-connection head is unchanged. + # ``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) - for key in ['hc_head_base', 'hc_head_fn', 'hc_head_scale']: - self._set_state_dict(lm_model, f'decoder.{key}', hf_state_dict, f'model.{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). @@ -417,6 +651,14 @@ def _set_hybrid_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, hybri 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): + """B1 hybrid backbone has no DSpark draft stack (that is B3), so there are no additional + layers to convert -- unlike the GPT bridge, whose base method walks ``mg_model.language_model`` + (absent on the text-only ``DeepseekV41HybridStackModel``). MTP (B2) is likewise skipped in + :meth:`_convert`. Yielding nothing keeps the backbone-only load/export intact.""" + return + yield # noqa: keep this an empty generator (matches the base method's protocol) + 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. @@ -442,9 +684,12 @@ def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqd 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 = {} - # HybridStack layer_number spans the doubled space (i + 1 + pp_offset), matching this - # loop's hybrid index so the PP-availability window below stays correct. - num_hybrid_layers = 2 * self.config.num_layers + # 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) diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py index e132e99f..fe8f9061 100644 --- a/tests/test_deepseek_v41_hybrid.py +++ b/tests/test_deepseek_v41_hybrid.py @@ -187,6 +187,24 @@ def test_engram_placement_and_hf_layer_id_round_trip(): 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_hybrid import DeepseekV41HybridBridge + + load_cfg = SimpleNamespace(num_layers=4, hybrid_layer_pattern=None) + export_cfg = SimpleNamespace(num_layers=8, hybrid_layer_pattern='DEDEDEDE') + assert DeepseekV41HybridBridge._num_hybrid_layers(load_cfg) == 8 + assert DeepseekV41HybridBridge._num_hybrid_layers(export_cfg) == 8 + # A config missing the attribute entirely is treated as GPT-space (load). + assert DeepseekV41HybridBridge._num_hybrid_layers(SimpleNamespace(num_layers=3)) == 6 + + import pytest # noqa: E402 from mcore_bridge.model.gpts.deepseek_v41_hybrid import ( # noqa: E402 @@ -197,26 +215,51 @@ def test_engram_placement_and_hf_layer_id_round_trip(): @requires_hybrid -def test_hc_wrapper_declines_fast_path_only_for_engram_layers(): - # The V4.1 wrapper subclass returns None (forcing the full-forward `_call_inner_layer` - # path, which applies Engram) iff the inner layer carries an Engram module; otherwise it - # must delegate unchanged to the base fast path. +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), reproducing the GPTModel golden order. A plain inner + # layer delegates unchanged with no Engram add. from types import SimpleNamespace from unittest.mock import patch - engram_layer = object.__new__(DeepseekV41HyperConnectionHybridLayer) - engram_layer.inner_layer = SimpleNamespace(engram=object()) - assert engram_layer._call_inner_transformer_layer_without_local_bda('h', 'mask') is None + 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) - sentinel = object() - with patch.object( - HyperConnectionHybridLayer, - '_call_inner_transformer_layer_without_local_bda', - return_value=sentinel) as base_call: - assert plain_layer._call_inner_transformer_layer_without_local_bda('h', 'mask') is sentinel - base_call.assert_called_once() + 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 @@ -258,3 +301,125 @@ def test_rewrap_noop_without_hyper_connections(): assert type(engram_wrapper) is HyperConnectionHybridLayer + +from mcore_bridge.model.gpts.deepseek_v41_hybrid 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, _deepseek_v41_use_hybrid) +from mcore_bridge.model.gpts.deepseek_v41_hybrid import ( # noqa: E402 + DeepseekV41HybridBridge, DeepseekV41HybridLoader) + + +def _route_config(pp, forced=None): + from types import SimpleNamespace + return SimpleNamespace(pipeline_model_parallel_size=pp, deepseek_v41_hybrid=forced) + + +def test_use_hybrid_auto_on_pp_and_forced_override(): + # Auto: GPTModel at PP1 (golden baseline), HybridModel once PP>1 (upstream blocks GPT there). + assert _deepseek_v41_use_hybrid(_route_config(1)) is False + assert _deepseek_v41_use_hybrid(_route_config(2)) is True + # Explicit flag wins either way (force-on aligns hybrid vs GPT at PP1; force-off stays GPT). + assert _deepseek_v41_use_hybrid(_route_config(1, forced=True)) is True + assert _deepseek_v41_use_hybrid(_route_config(2, forced=False)) is False + + +@requires_hybrid +def test_loader_new_dispatches_to_hybrid(): + # __new__ routing only (no __init__), so no distributed init is required. + assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(2))) is DeepseekV41HybridLoader + assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1, forced=True))) is DeepseekV41HybridLoader + assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1))) is DeepseekV41Loader + # A directly instantiated subclass must not re-dispatch (cls-is guard). + assert type(DeepseekV41HybridLoader.__new__(DeepseekV41HybridLoader, _route_config(1))) is DeepseekV41HybridLoader + + +@requires_hybrid +def test_bridge_new_dispatches_to_hybrid(): + assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(2))) is DeepseekV41HybridBridge + assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1, forced=True))) is DeepseekV41HybridBridge + assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1))) is DeepseekV41Bridge + assert type(DeepseekV41HybridBridge.__new__(DeepseekV41HybridBridge, _route_config(1))) is DeepseekV41HybridBridge + From e01303b84dbf3a54f9cf53149372f7c8500835e8 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 20:16:20 +0800 Subject: [PATCH 05/17] feat(deepseek-v41): map DSpark draft stack (mtp.*) under HybridModel (B3) Backbone-agnostic DSpark support for the PP-capable HybridModel path: - deepseek_v41.py: extract _attach_dspark() and _convert_dspark_stack() from the GPT build_model/_convert_additional_layers (behavior-preserving). - deepseek_v41_hybrid.py: build_model attaches DSpark on the text-only model itself; _convert_additional_layers maps mtp.* via _lm() with a non-last-stage guard, reusing _convert_dspark_stack. - Inference-time capture hook rehang deferred (documented): hybrid has no speculative-decoding forward, so a standalone hook would be dead code. Verified (tiny SFT): iter-1 loss/grad match the B1 hybrid baseline within noise; mtp.* key-ledger round-trips cleanly (52 DSpark keys, the sole delta vs B1 export, no unexpected keys vs golden GPT baseline). 33 unit tests pass. --- src/mcore_bridge/model/gpts/deepseek_v41.py | 25 +++++- .../model/gpts/deepseek_v41_hybrid.py | 46 +++++++--- tests/test_deepseek_v41_hybrid.py | 88 +++++++++++++++++++ 3 files changed, 146 insertions(+), 13 deletions(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 777a2b89..eb23daf5 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -985,9 +985,23 @@ def get_dspark_layer_spec(self): def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): model = super().build_model(pre_process, post_process, vp_stage) + self._attach_dspark(model.language_model, post_process) + return model + + def _attach_dspark(self, language_model, post_process): + """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 GPT path passes + ``model.language_model``; the hybrid path (which has no ``language_model`` wrapper) passes + the ``HybridModel`` itself -- both expose ``pg_collection`` / ``vocab_size`` / ``config``. + 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. + """ if not self.config.dspark_num_layers or not post_process: - return model - language_model = model.language_model + return dspark_config, dspark_layer_specs = self.get_dspark_layer_spec() layers = [ build_module( @@ -1012,7 +1026,6 @@ def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[in config=language_model.config, tp_group=language_model.pg_collection.tp, ) - return model def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): from megatron.core.models.gpt.experimental_attention_variant_module_specs import \ @@ -1225,7 +1238,13 @@ def _convert_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcor dspark = getattr(language_model, 'dspark', None) if dspark is None: raise RuntimeError('DSpark weights require the draft stack on the final pipeline stage.') + yield from self._convert_dspark_stack(language_model, dspark, hf_state_dict, hf_prefix, to_mcore) + 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), so + both the GPT and hybrid bridges reuse it; they differ only in how they locate ``dspark`` + and guard 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, diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index 3f5e5452..0516519a 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -371,8 +371,11 @@ class DeepseekV41HybridLoader(DeepseekV41Loader): ``DeepseekV41Loader`` is left untouched; this loader derives its own config copy so both paths can coexist in one process. - B1 covers the text backbone only. MTP (B2), DSpark capture (B3) and the multimodal wrapper - (B4) are added on top; here MTP is disabled so the backbone can be aligned in isolation. + B1 covers the text backbone; B3 adds the DSpark (``mtp.*``) draft stack on top (attached in + :meth:`build_model`, mapped in :meth:`DeepseekV41HybridBridge._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 + multimodal wrapper (B4) is still added separately. """ model_cls = DeepseekV41HybridStackModel @@ -408,8 +411,10 @@ def _build_hybrid_config(self): # 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] - # MTP on HybridModel is B2 (its inner attention cannot use the CSA2 'D' symbol, which - # rejects is_mtp_layer). Disable it for the B1 backbone-only alignment. + # 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``, B3). + # Keep it disabled so no ``MultiTokenPredictionBlock`` is built. cfg.mtp_num_layers = None return cfg @@ -500,6 +505,14 @@ def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[in ) self._rewrap_engram_hyper_connection_layers(model) self._set_linear_is_expert(model) + # DSpark (B3): the ``mtp.*`` draft stack is backbone-agnostic (plain + # experimental-attention layers), so reuse the GPT loader's builder. The hybrid model is + # text-only (no ``language_model`` wrapper, see :meth:`DeepseekV41HybridBridge._lm`), so + # the stack attaches to the model itself. Inference-time target-layer capture on + # HybridStack is deferred (it is not exercised by training / weight round-trip, mirroring + # B1's deferral of ``allow_engram_inference``); the stack only needs to exist so its + # parameters are loaded / saved via ``mtp.*``. + self._attach_dspark(model, post_process) return model @@ -652,12 +665,25 @@ def _set_hybrid_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, hybri 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): - """B1 hybrid backbone has no DSpark draft stack (that is B3), so there are no additional - layers to convert -- unlike the GPT bridge, whose base method walks ``mg_model.language_model`` - (absent on the text-only ``DeepseekV41HybridStackModel``). MTP (B2) is likewise skipped in - :meth:`_convert`. Yielding nothing keeps the backbone-only load/export intact.""" - return - yield # noqa: keep this an empty generator (matches the base method's protocol) + """Map the DSpark (``mtp.*``) draft stack (B3). + + The draft layers are plain experimental-attention ``TransformerLayer`` instances -- + identical in both paths -- so the base :meth:`DeepseekV41Bridge._convert_dspark_stack` + mapping is reused verbatim; only where the stack lives differs. On the hybrid path it is + attached to the model itself (no ``language_model`` wrapper, so use :meth:`_lm`) and only + on the final pipeline stage, so non-last stages have nothing to convert (on load the base + guard skips them; on export ``dspark`` is simply absent). 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 not is_pp_last_stage: + # Export from a non-last PP stage: the draft stack lives on the final stage only. + return + raise RuntimeError('DSpark weights require the draft stack on the final pipeline stage.') + 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. diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py index fe8f9061..288c690d 100644 --- a/tests/test_deepseek_v41_hybrid.py +++ b/tests/test_deepseek_v41_hybrid.py @@ -423,3 +423,91 @@ def test_bridge_new_dispatches_to_hybrid(): assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1))) is DeepseekV41Bridge assert type(DeepseekV41HybridBridge.__new__(DeepseekV41HybridBridge, _route_config(1))) is DeepseekV41HybridBridge + +# --- B3: DSpark (``mtp.*``) draft stack on the hybrid path --------------------------------------- + +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__(DeepseekV41HybridBridge) + 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(): + # Forward-compat with B4: 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_skips_non_last_stage_without_stack(): + # On export a non-last PP stage has no draft stack; it must skip quietly (not raise), unlike + # the final stage where a missing stack is a real error. + 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 == [] + assert calls == [] + + +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)) + From 02ac65194fc506a0907c516398365d453f6b3369 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 20:54:57 +0800 Subject: [PATCH 06/17] feat(deepseek-v41): host HybridModel backbone in multimodal wrapper (B4) Add DeepseekV41MultimodalHybridModel (composition over DeepseekV41Multimodal- GPTModel, swapping language_model_cls to the PP-capable HybridModel backbone) and route the hybrid loader/bridge through it, so the vision tower + image-embed injection reuse the golden GPT multimodal path verbatim. - DeepseekV41HybridStackModel: expose extra_forward_keys=[]; forward() now unpacks the wrapper's extra_block_kwargs container (upstream HybridModel.forward has no such param -- it threads input_ids itself) before stripping visual keys. - HybridLoader.build_model: resolve language_model via getattr and apply the engram/is_expert/dspark fix-ups on the nested backbone. - HybridBridge: _set_word_embeddings via _lm; _convert_pre_process dispatches to the GPT vision path when visual is present, else text-only word emb (guards the visual=None non-first PP stage on export). Tests: 41 passed (+2 forward extra_block_kwargs unpack / pixel_values guard). GPU accept: text-only hybrid_pp1_b4 iter-1 loss=12.89138508/grad=102.55239105 in the aligned family band (wrapper does not perturb text loss); export key ledger 285 keys, 0 missing/0 unexpected vs both source and golden GPT baseline (+42 vision/aligner/image_* over B3's 243). --- .../model/gpts/deepseek_v41_hybrid.py | 97 +++++++++++---- tests/test_deepseek_v41_hybrid.py | 112 ++++++++++++++++++ 2 files changed, 189 insertions(+), 20 deletions(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index 0516519a..2df6ed22 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -37,7 +37,7 @@ from ..modules.engram import DeepseekV41Engram, DeepseekV41TransformerLayer from ..rope import get_rope_inv_freq from .deepseek_v41 import (CSA2Compressor, CSA2Indexer, DeepseekV41Bridge, DeepseekV41Loader, - DSv4HybridSelfAttention) + DeepseekV41MultimodalGPTModel, DSv4HybridSelfAttention) try: from megatron.core.models.hybrid.hybrid_block import HyperConnectionHybridLayer @@ -146,6 +146,12 @@ class DeepseekV41HybridStackModel(HybridModel): # the real vision tower arrives with the multimodal wrapper in B4. visual = None + # ``MultimodalGPTModel.forward`` (B4 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 treats the hybrid backbone exactly like the GPT one. + extra_forward_keys: List[str] = [] + @staticmethod def _segment_main_pattern(config) -> Optional[str]: pattern = config.hybrid_layer_pattern @@ -265,6 +271,14 @@ def _inject_dsv4_rotary_pos_emb(self, module, args, kwargs): _visual_forward_keys = ('pixel_values', 'image_grid_thw', 'image_token_types', 'token_types') def forward(self, *args, **kwargs): + # The B4 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( 'DeepSeek-V4.1 hybrid (pipeline-parallel) path is text-only in B1; multimodal ' @@ -272,9 +286,32 @@ def forward(self, *args, **kwargs): for key in self._visual_forward_keys: kwargs.pop(key, None) return super().forward(*args, **kwargs) + + class DeepseekV41MultimodalHybridModel(DeepseekV41MultimodalGPTModel): + """Multimodal wrapper (B4) hosting the PP-capable ``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). So the only change from + the GPT :class:`DeepseekV41MultimodalGPTModel` is swapping the language-model class; 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). The DSpark + speculative-decoding helpers inherited from the GPT wrapper are inference-only and stay + deferred on the hybrid path (see :meth:`DeepseekV41HybridLoader.build_model`). + """ + + language_model_cls = DeepseekV41HybridStackModel else: DeepseekV41HyperConnectionHybridLayer = None DeepseekV41HybridStackModel = None + DeepseekV41MultimodalHybridModel = None @dataclass @@ -374,11 +411,12 @@ class DeepseekV41HybridLoader(DeepseekV41Loader): B1 covers the text backbone; B3 adds the DSpark (``mtp.*``) draft stack on top (attached in :meth:`build_model`, mapped in :meth:`DeepseekV41HybridBridge._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 - multimodal wrapper (B4) is still added separately. + V4.1 -- its ``mtp.*`` checkpoint keys *are* DSpark -- so it stays disabled here. B4 wraps the + backbone in :class:`DeepseekV41MultimodalHybridModel` (the vision tower + image-embed + injection), mirroring the GPT :class:`DeepseekV41MultimodalGPTModel`. """ - model_cls = DeepseekV41HybridStackModel + model_cls = DeepseekV41MultimodalHybridModel def _engram_placement_layer_ids(self, hf_layer_ids): """On HybridStack, HF layer ``e`` becomes the attention-only 'D' layer at hybrid index @@ -491,10 +529,15 @@ def _rewrap_engram_hyper_connection_layers(self, model): layer.__class__ = DeepseekV41HyperConnectionHybridLayer def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): - """Build via ``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.""" + """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:`DeepseekV41MultimodalHybridModel` (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 (matching the GPT wrapper, + where the same fix-ups and DSpark live under ``language_model``).""" self._hybrid_config = self._build_hybrid_config() model = self.model_cls( config=self._hybrid_config, @@ -503,16 +546,16 @@ def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[in post_process=post_process, vp_stage=vp_stage, ) - self._rewrap_engram_hyper_connection_layers(model) - self._set_linear_is_expert(model) + language_model = getattr(model, 'language_model', model) + self._rewrap_engram_hyper_connection_layers(language_model) + self._set_linear_is_expert(language_model) # DSpark (B3): the ``mtp.*`` draft stack is backbone-agnostic (plain - # experimental-attention layers), so reuse the GPT loader's builder. The hybrid model is - # text-only (no ``language_model`` wrapper, see :meth:`DeepseekV41HybridBridge._lm`), so - # the stack attaches to the model itself. Inference-time target-layer capture on - # HybridStack is deferred (it is not exercised by training / weight round-trip, mirroring - # B1's deferral of ``allow_engram_inference``); the stack only needs to exist so its - # parameters are loaded / saved via ``mtp.*``. - self._attach_dspark(model, post_process) + # experimental-attention layers), so reuse the GPT loader's builder. It attaches to the + # hybrid text backbone (``language_model.dspark``), mirroring the GPT wrapper. Inference-time + # target-layer capture on HybridStack is deferred (it is not exercised by training / weight + # round-trip, mirroring B1's deferral of ``allow_engram_inference``); the stack only needs + # to exist so its parameters are loaded / saved via ``mtp.*``. + self._attach_dspark(language_model, post_process) return model @@ -572,14 +615,28 @@ def _engram_hf_layer_id(self, engram): # ``DeepseekV41HybridLoader._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 (B4) 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): - # Text-only word embeddings; visual embeds (image_start/end/newline) are B4. + # First pipeline stage of a multimodal model: the vision tower + aligner + image_* markers + # live on the wrapper (``mg_model.visual``). Reuse the GPT ``DeepseekV41Bridge`` pre-process + # verbatim (``MultimodalGPTBridge`` word-embeddings + vision/aligner block, then the + # image_start/end/newline markers); ``_set_word_embeddings`` above resolves the LM via + # ``_lm``. ``super()`` here is ``DeepseekV41Bridge`` (MRO), matching the GPT path exactly. + if getattr(mg_model, 'visual', None) is not None: + return super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) + # No vision tower on this rank (text-only backbone, or a non-first PP stage where the + # wrapper built ``visual=None``): map only the word embeddings. if to_mcore: hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) else: hf_state_dict = {} - lm_model = self._lm(mg_model) - self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key, to_mcore) + self._set_word_embeddings(mg_model, hf_state_dict, to_mcore) if to_mcore: return {} return self._add_prefix(hf_state_dict, hf_prefix) diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py index 288c690d..1051f85b 100644 --- a/tests/test_deepseek_v41_hybrid.py +++ b/tests/test_deepseek_v41_hybrid.py @@ -511,3 +511,115 @@ def test_hybrid_convert_additional_layers_raises_on_last_stage_without_stack(): 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)) + +# --- B4: multimodal wrapper hosting the hybrid backbone ----------------------------------------- + +def test_multimodal_hybrid_wrapper_hosts_hybrid_backbone(): + # The B4 wrapper is just the GPT 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 DeepseekV41MultimodalGPTModel + from mcore_bridge.model.gpts.deepseek_v41_hybrid import (DeepseekV41HybridStackModel, + DeepseekV41MultimodalHybridModel) + assert issubclass(DeepseekV41MultimodalHybridModel, DeepseekV41MultimodalGPTModel) + assert DeepseekV41MultimodalHybridModel.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_hybrid import DeepseekV41MultimodalHybridModel + assert DeepseekV41HybridLoader.model_cls is DeepseekV41MultimodalHybridModel + + +def test_hybrid_pre_process_delegates_to_gpt_when_visual_present(monkeypatch): + # First PP stage of a multimodal model: the wrapper carries a vision tower, so pre-process must + # reuse the GPT DeepseekV41Bridge path (word embeddings + vision/aligner + image_* markers). + from types import SimpleNamespace + + bridge = object.__new__(DeepseekV41HybridBridge) + called = [] + monkeypatch.setattr(DeepseekV41Bridge, '_convert_pre_process', + lambda self, mg, sd, pfx, tm: called.append((mg, pfx, tm)) or {'SUPER': True}) + mg_model = SimpleNamespace(visual=object()) + out = bridge._convert_pre_process(mg_model, {}, '', to_mcore=True) + assert out == {'SUPER': True} + assert called == [(mg_model, '', True)] + + +def test_hybrid_pre_process_text_only_when_no_visual(monkeypatch): + # No vision tower on this rank (text backbone, or a non-first PP stage where ``visual=None``): + # only the word embeddings are mapped, and the GPT vision path is never entered. + from types import SimpleNamespace + + monkeypatch.setattr(DeepseekV41Bridge, '_convert_pre_process', + lambda *a, **k: (_ for _ in ()).throw(AssertionError('vision path must not run'))) + bridge = object.__new__(DeepseekV41HybridBridge) + calls = [] + bridge._set_word_embeddings = lambda mg, sd, tm: calls.append((mg, tm)) + bridge._remove_prefix = lambda sd, pfx: sd + bridge._add_prefix = lambda sd, pfx: sd + mg_model = SimpleNamespace(visual=None) + assert bridge._convert_pre_process(mg_model, {'x': 1}, '', to_mcore=True) == {} + assert calls == [(mg_model, True)] + + +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__(DeepseekV41HybridBridge) + 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_hybrid 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_hybrid 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}) + + From 33c5b804e38667e1e28f19ebaee8c5d4ebd94594 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 21:07:10 +0800 Subject: [PATCH 07/17] feat(deepseek-v41): default to HybridModel path for all layouts (B5) Switch _deepseek_v41_use_hybrid to return True by default so every layout (including PP1/single-GPU) builds on the PP-capable HybridModel path, now that B1-B4 validated it against the GPTModel golden baseline (iter-1 loss/grad within the bf16/MoE non-determinism band, clean weight key ledger). The GPTModel path is retained as a force-off regression baseline via deepseek_v41_hybrid=False. Update routing docstrings/module header and loader/bridge/use_hybrid unit tests (41 passed). GPU smoke at PP1 with no force flag confirms default routing builds DeepseekV41MultimodalHybridModel and aligns (iter-1 loss 12.8874/grad 102.589). VPP remains an optional follow-up (needs Engram VPP validation coverage). --- src/mcore_bridge/model/gpts/deepseek_v41.py | 15 +++++++------- .../model/gpts/deepseek_v41_hybrid.py | 7 ++++--- tests/test_deepseek_v41_hybrid.py | 20 ++++++++++++------- 3 files changed, 25 insertions(+), 17 deletions(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index eb23daf5..1bc21ff6 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -846,17 +846,18 @@ def compute_dspark_speculative_tokens(self, *args, **kwargs): def _deepseek_v41_use_hybrid(config) -> bool: """Whether to build DeepSeek-V4.1 on the ``HybridModel`` (PP-capable) path. - The default ``GPTModel`` path is the golden baseline and stays the default until the hybrid - path is fully aligned (plan step B5). Upstream refuses to run the V4.1 ``dsv4_hybrid`` block - under pipeline parallelism, so the hybrid loader/bridge are auto-selected whenever - ``pipeline_model_parallel_size > 1``. The ``deepseek_v41_hybrid`` config flag (settable via - ``--megatron_extra_kwargs``) overrides this: ``True`` forces the hybrid path on at PP1 (used to - align it against the GPTModel baseline), ``False`` keeps GPTModel even at PP>1. + The ``HybridModel`` path is now the default (plan step B5): B1-B4 validated it against the + ``GPTModel`` golden baseline (iter-1 loss/grad within the bf16/MoE non-determinism band and a + clean weight key ledger), and it is the only path that supports pipeline parallelism, so it is + selected for every layout. The ``deepseek_v41_hybrid`` config flag (settable via + ``--megatron_extra_kwargs``) overrides this: ``False`` drops back to the ``GPTModel`` golden + baseline (kept as a regression path; note upstream refuses ``GPTModel`` at ``PP>1``), ``True`` + is redundant but still forces hybrid. """ forced = getattr(config, 'deepseek_v41_hybrid', None) if forced is not None: return bool(forced) - return (getattr(config, 'pipeline_model_parallel_size', 1) or 1) > 1 + return True class DeepseekV41Loader(DeepseekV4Loader): diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index 2df6ed22..2c94dbae 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """DeepSeek-V4.1 on megatron-core's ``HybridModel`` (pipeline-parallel path). -The default :class:`DeepseekV41Loader` builds a ``GPTModel`` whose custom +The legacy :class:`DeepseekV41Loader` builds a ``GPTModel`` whose custom ``TransformerBlock`` owns the CSA2 / single-pass-mHC forward. Upstream refuses to run that block under pipeline parallelism:: @@ -19,8 +19,9 @@ and every per-layer config array that CSA2 indexes by ``layer_number - 1`` must be re-expanded into this doubled index space (see :func:`derive_hybrid_layer_config`). -This module keeps the GPT loader untouched (golden baseline) and adds the hybrid path -alongside it; both are validated to agree before the default is switched. +This module keeps the GPT loader available as a force-off regression baseline; the +hybrid path is now the default for every layout (plan step B5) after both were +validated to agree at iter-1 loss/grad and on the weight key ledger. """ import copy from dataclasses import dataclass diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py index 1051f85b..d68b67bd 100644 --- a/tests/test_deepseek_v41_hybrid.py +++ b/tests/test_deepseek_v41_hybrid.py @@ -397,30 +397,36 @@ def _route_config(pp, forced=None): return SimpleNamespace(pipeline_model_parallel_size=pp, deepseek_v41_hybrid=forced) -def test_use_hybrid_auto_on_pp_and_forced_override(): - # Auto: GPTModel at PP1 (golden baseline), HybridModel once PP>1 (upstream blocks GPT there). - assert _deepseek_v41_use_hybrid(_route_config(1)) is False +def test_use_hybrid_default_on_and_forced_override(): + # Default (B5 switch): HybridModel for every layout now that B1-B4 align with the GPT baseline. + assert _deepseek_v41_use_hybrid(_route_config(1)) is True assert _deepseek_v41_use_hybrid(_route_config(2)) is True - # Explicit flag wins either way (force-on aligns hybrid vs GPT at PP1; force-off stays GPT). - assert _deepseek_v41_use_hybrid(_route_config(1, forced=True)) is True + # Explicit force-off drops back to the GPTModel golden baseline (kept as a regression path). + assert _deepseek_v41_use_hybrid(_route_config(1, forced=False)) is False assert _deepseek_v41_use_hybrid(_route_config(2, forced=False)) is False + # Force-on is redundant now but must still route to hybrid. + assert _deepseek_v41_use_hybrid(_route_config(1, forced=True)) is True @requires_hybrid def test_loader_new_dispatches_to_hybrid(): # __new__ routing only (no __init__), so no distributed init is required. + assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1))) is DeepseekV41HybridLoader assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(2))) is DeepseekV41HybridLoader assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1, forced=True))) is DeepseekV41HybridLoader - assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1))) is DeepseekV41Loader + # Force-off keeps the GPTModel golden baseline. + assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1, forced=False))) is DeepseekV41Loader # A directly instantiated subclass must not re-dispatch (cls-is guard). assert type(DeepseekV41HybridLoader.__new__(DeepseekV41HybridLoader, _route_config(1))) is DeepseekV41HybridLoader @requires_hybrid def test_bridge_new_dispatches_to_hybrid(): + assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1))) is DeepseekV41HybridBridge assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(2))) is DeepseekV41HybridBridge assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1, forced=True))) is DeepseekV41HybridBridge - assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1))) is DeepseekV41Bridge + # Force-off keeps the GPTModel golden baseline. + assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1, forced=False))) is DeepseekV41Bridge assert type(DeepseekV41HybridBridge.__new__(DeepseekV41HybridBridge, _route_config(1))) is DeepseekV41HybridBridge From 012ce0f6b6ff45202ffe00d571a4635681cf28a5 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Wed, 16 Sep 2026 21:56:21 +0800 Subject: [PATCH 08/17] feat(deepseek-v41): support virtual pipeline parallelism (VPP) Enable PP + VPP (interleaved schedule) for the DeepSeek-V4.1 HybridModel path. Layer segmentation was already vp_stage-aware from the B1 migration; this lands the remaining three fixes, all in mcore-bridge: - engram: relax the upstream Engram parallelism guard to also allow VPP (rename _ContextParallelSizeOneView -> _RelaxedParallelismView, hiding both context_parallel_size and virtual_pipeline_model_parallel_size while keeping every other check, e.g. etp!=tp). Engram.forward is self-contained (local n-gram hashing + slicing, no cross-stage state); layer placement is vp_stage-aware and the whole-block guard rejects half-block stages. - deepseek_v41: thread vp_stage into _attach_dspark's build_module so the DSpark draft layers satisfy get_transformer_layer_offset's VPP assertion (vp_stage is not None). Offset is 0 for the tiny draft stack, placement unchanged. - mm_gpt_model: surface the language model's typed-pipeline payload interface (pipeline_payload_factory / pipeline_payload_spec) on MultimodalGPTModel. get_attr_wrapped_model only descends via .module and cannot reach self.language_model, so the CSA2 custom cross-stage payload fell back to the shape-based P2PCommunicator, which calls .size() on the payload and crashes under both 1F1B and interleaved schedules. Backbones without a payload (GPTModel, GLM) expose None and keep the legacy path. Tests: test_deepseek_v41_engram.py + test_deepseek_v41_hybrid.py -> 76 passed. GPU smoke (tiny, multimodal default): PP2xVPP2 iter-1 loss 12.89163 / grad 102.519 aligns with same-model PP2 (1F1B) baseline 12.89124 / 102.558. --- src/mcore_bridge/model/gpts/deepseek_v41.py | 12 ++++++-- .../model/gpts/deepseek_v41_hybrid.py | 2 +- src/mcore_bridge/model/mm_gpt_model.py | 10 +++++++ src/mcore_bridge/model/modules/engram.py | 28 +++++++++++++------ tests/test_deepseek_v41_engram.py | 11 ++++---- 5 files changed, 46 insertions(+), 17 deletions(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 1bc21ff6..70c17d91 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -986,10 +986,10 @@ def get_dspark_layer_spec(self): def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): model = super().build_model(pre_process, post_process, vp_stage) - self._attach_dspark(model.language_model, post_process) + self._attach_dspark(model.language_model, post_process, vp_stage=vp_stage) return model - def _attach_dspark(self, language_model, post_process): + 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. @@ -1000,6 +1000,13 @@ def _attach_dspark(self, language_model, post_process): the ``HybridModel`` itself -- both expose ``pg_collection`` / ``vocab_size`` / ``config``. 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. + + ``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 @@ -1010,6 +1017,7 @@ def _attach_dspark(self, language_model, post_process): 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) ] diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index 2c94dbae..f93141f9 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -556,7 +556,7 @@ def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[in # target-layer capture on HybridStack is deferred (it is not exercised by training / weight # round-trip, mirroring B1's deferral of ``allow_engram_inference``); the stack only needs # to exist so its parameters are loaded / saved via ``mtp.*``. - self._attach_dspark(language_model, post_process) + self._attach_dspark(language_model, post_process, vp_stage=vp_stage) return model diff --git a/src/mcore_bridge/model/mm_gpt_model.py b/src/mcore_bridge/model/mm_gpt_model.py index 2dfdd638..356c1634 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: diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py index d149466f..711e906c 100644 --- a/src/mcore_bridge/model/modules/engram.py +++ b/src/mcore_bridge/model/modules/engram.py @@ -35,14 +35,23 @@ def has_native_engram() -> bool: return EngramConfig is not None -class _ContextParallelSizeOneView: - """Read-only view of a transformer config that reports ``context_parallel_size == 1``. - - Lets the Engram validators reuse every upstream parallelism check except the CP one, - without mutating the shared transformer config. +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 @@ -80,15 +89,16 @@ def _load_tokenizer_map(self): self.layer_ids = placement_layer_ids def _validate_parallelism(self, transformer_config, sequence_length): - # DeepseekV41Engram hashes the full sequence locally and then selects its own CP - # slice, so no window ever has to cross a CP rank boundary. Drop only the upstream - # `context_parallel_size == 1` rejection and keep every other check. + # 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( - _ContextParallelSizeOneView(transformer_config), sequence_length) + _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 diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index 4e126d78..c29cf2bc 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -579,7 +579,7 @@ def _engram_config_for_validation(tmp_path): ) -def test_engram_config_allows_context_parallelism_but_keeps_the_other_guards(tmp_path): +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) @@ -593,13 +593,14 @@ def test_engram_config_allows_context_parallelism_but_keeps_the_other_guards(tmp # 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) - with pytest.raises(ValueError, match='virtual pipeline'): - config._validate_parallelism( - SimpleNamespace(**{**parallelism, 'virtual_pipeline_model_parallel_size': 2}), None) def test_engram_config_allows_packed_sequences_without_losing_the_pipeline_guard(tmp_path): From 1461d8b0df010a3d339efe8db4e00b357e5a58b9 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Thu, 17 Sep 2026 11:02:50 +0800 Subject: [PATCH 09/17] fix(deepseek-v41): correct HybridModel export symmetry, packed MLA rotary and weight fidelity Found while validating PP / CP / baseline loss-grad alignment and the RL scripts on a 4-layer tiny model. * Keep the PP export collectives symmetric: `_convert_pre_process` no longer branches on this rank's `visual`, and non-last stages convert the DSpark stack through an empty structural proxy. A per-rank guard desynchronized the pp-group all-reduces so the last stage read a stale `has_model` and raised StopIteration. * Pre-index the MLA rotary table by `position_ids` under `thd` packing, mirroring what `GPTModel.forward` does. The DSv4 attention consumes per-token frequencies, so without this the packed and CP paths hit `freqs.shape[0] != tokens`. `HybridModel.forward` never threads `position_ids` into the decoder, hence the stash for the rotary pre-hook. * Reject non-contiguous CP partitioning on the non-packed Engram path too: it reads the layout from the transformer config instead of packed_seq_params, so zigzag used to slip through and silently mis-align hashes with the local hidden states. * Mark the DSpark router `expert_bias` keep-in-fp32. The draft stack never runs in the training forward, so mcore's lazy `_maintain_float32_expert_bias` never fired and the checkpoint's fp32 `mtp.*.ffn.gate.bias` round-tripped through bf16. * Guard the hyper-connection `alpha_*` export with `_peft_format` (both the GPT and hybrid copies). These are frozen base weights written outside `_set_state_dict`, so a LoRA run used to write `hc_*_scale` into adapter_model.safetensors and demand it back on load. Validation on the tiny model (iter-1 relative difference, the clean signal; a same-config rerun measures 0.0141% loss / 0.145% grad as the noise floor): PP2 0.0083% / 0.108%, CP2 0.0134% / 0.027%, PP2+packing 0.0068% / 0.031%. Full-param export stays at 285 keys matching the source; the LoRA adapter now holds 155 lora_A + 155 lora_B and nothing else. --- src/mcore_bridge/bridge/gpt_bridge.py | 6 +- src/mcore_bridge/model/gpts/deepseek_v41.py | 10 +++ .../model/gpts/deepseek_v41_hybrid.py | 89 ++++++++++++------- src/mcore_bridge/model/modules/engram.py | 19 ++-- 4 files changed, 87 insertions(+), 37 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 8afc4d70..7ea300cf 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']): diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 70c17d91..9b98062b 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -1023,6 +1023,16 @@ def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = ] 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 diff --git a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py index f93141f9..b1abfaa4 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py @@ -19,12 +19,12 @@ and every per-layer config array that CSA2 indexes by ``layer_number - 1`` must be re-expanded into this doubled index space (see :func:`derive_hybrid_layer_config`). -This module keeps the GPT loader available as a force-off regression baseline; the -hybrid path is now the default for every layout (plan step B5) after both were -validated to agree at iter-1 loss/grad and on the weight key ledger. +The HybridModel path is the sole maintained DeepSeek-V4.1 implementation. The legacy +GPTModel path is deprecated and scheduled for removal. """ import copy from dataclasses import dataclass +from types import SimpleNamespace from typing import List, Optional, Sequence, Union import torch @@ -217,6 +217,7 @@ def __init__(self, config, transformer_layer_spec, pre_process=True, post_proces # 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) @@ -245,14 +246,33 @@ def _build_dsv4_rotary_tables(self): 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``).""" + ``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' - return { + 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: @@ -286,7 +306,16 @@ def forward(self, *args, **kwargs): 'inputs require the B4 multimodal wrapper.') for key in self._visual_forward_keys: kwargs.pop(key, None) - return super().forward(*args, **kwargs) + # 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 DeepseekV41MultimodalHybridModel(DeepseekV41MultimodalGPTModel): """Multimodal wrapper (B4) hosting the PP-capable ``HybridModel`` backbone. @@ -624,23 +653,19 @@ def _set_word_embeddings(self, mg_model, hf_state_dict, to_mcore): to_mcore) def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): - # First pipeline stage of a multimodal model: the vision tower + aligner + image_* markers - # live on the wrapper (``mg_model.visual``). Reuse the GPT ``DeepseekV41Bridge`` pre-process - # verbatim (``MultimodalGPTBridge`` word-embeddings + vision/aligner block, then the - # image_start/end/newline markers); ``_set_word_embeddings`` above resolves the LM via - # ``_lm``. ``super()`` here is ``DeepseekV41Bridge`` (MRO), matching the GPT path exactly. - if getattr(mg_model, 'visual', None) is not None: - return super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) - # No vision tower on this rank (text-only backbone, or a non-first PP stage where the - # wrapper built ``visual=None``): map only the word embeddings. - if to_mcore: - hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) - else: - hf_state_dict = {} - self._set_word_embeddings(mg_model, hf_state_dict, to_mcore) - if to_mcore: - return {} - return self._add_prefix(hf_state_dict, hf_prefix) + # Delegate to ``DeepseekV41Bridge._convert_pre_process`` (``super()`` via MRO) unconditionally + # instead of branching on *this* rank's ``mg_model.visual``. On export every pipeline stage runs + # ``_convert`` -> ``_convert_pre_process``, and the base path issues the *same* pp-group collective + # sequence on all ranks: word-embeddings (routed through ``_lm`` by the ``_set_word_embeddings`` + # override), 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), + # so the vision tower is always present there and the base path behaves exactly as before. + return super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore: bool): if to_mcore: @@ -681,7 +706,10 @@ def _set_one_hyper_connection(self, hyper_connection, hf_state_dict, hf_key, to_ 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_*`` 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']): @@ -729,18 +757,17 @@ def _convert_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcor identical in both paths -- so the base :meth:`DeepseekV41Bridge._convert_dspark_stack` mapping is reused verbatim; only where the stack lives differs. On the hybrid path it is attached to the model itself (no ``language_model`` wrapper, so use :meth:`_lm`) and only - on the final pipeline stage, so non-last stages have nothing to convert (on load the base - guard skips them; on export ``dspark`` is simply absent). MTP (``mtp_num_layers``) is - skipped in :meth:`_convert` and does not apply to V4.1.""" + 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 not is_pp_last_stage: - # Export from a non-last PP stage: the draft stack lives on the final stage only. - return - raise RuntimeError('DSpark weights require the draft stack on the final pipeline stage.') + 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: '): diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py index 711e906c..9fedfa41 100644 --- a/src/mcore_bridge/model/modules/engram.py +++ b/src/mcore_bridge/model/modules/engram.py @@ -402,11 +402,20 @@ def forward(self, hidden_states: Tensor, input_ids: Tensor, inference_context=No 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) - if context_parallel and getattr(packed_seq_params, 'cp_partition_mode', - 'zigzag') != 'contiguous': - raise ValueError( - "Engram with context parallelism requires cp_partition_mode='contiguous', " - 'matching the DSv4 THD CP forward.') + 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: From 3da2c2ac4148d72732ecf6a535e612b85d712464 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Thu, 17 Sep 2026 12:13:51 +0800 Subject: [PATCH 10/17] refactor(deepseek-v41): drop the GPTModel path, single HybridModel implementation The GPTModel-based V4.1 path was only kept as the golden baseline while the HybridModel backbone was being brought up. Now that HybridModel is the default for every layout (TP/PP/VPP/CP/EP/DP + packing, all verified against the PP1 baseline), the two-file / two-class-layer split is pure overhead: every load or export walked DeepseekV41Hybrid{Loader,Bridge} -> DeepseekV41{Loader,Bridge}, where the parent half of each pair was mostly dead code. * merge deepseek_v41_hybrid.py into deepseek_v41.py (registration must live in an eagerly imported module) and collapse the two inheritance layers into one DeepseekV41Loader / DeepseekV41Bridge; the only cross-layer super() call (_convert_pre_process) is merged by hand * guard register_model with _HYBRID_MODEL_AVAILABLE: V4.1 now requires a megatron with megatron.core.models.hybrid * drop the ModelConfig.deepseek_v41_hybrid selector and both __new__ routers, which no longer select anything Removes (irreversible, recover from 1461d8b or earlier): DeepseekV41GPTModel and its DSpark speculative-decoding methods (forward_dspark / compute_dspark_speculative_tokens / the capture hooks, ~240 lines) plus their 3 unit tests. That code was already unreachable on the HybridModel backbone -- the wrapper delegated to methods HybridModel does not define, so any call raised AttributeError. The shared DeepseekV41DSparkAttention (prefill_dspark / reset_dspark_cache) and the mtp.* draft-stack weight conversion are untouched. Tests: 70 passed (test_deepseek_v41_engram.py, test_deepseek_v41_hybrid.py). Four stale expectations updated -- two were already failing before this commit (they still asserted the pre-PP-fix "skip the vision block / draft stack on non-last stages" behaviour), and two asserted GPT-path semantics that the merge replaced (monkeypatching the parent bridge; the non-doubled Engram layer_number). Tiny 4-layer regression, iteration-1 loss / grad_norm: PP1 baseline reproduced across the refactor : 0.0125% / 0.0166% packed CP1 baseline reproduced : 0.0003% / 0.0017% PP2 vs PP1 (was 0.0487% / 0.0484%) : 0.0147% / 0.0050% CP2 vs CP1 (was 0.0210% / 0.2688%) : 0.0179% / 0.2429% Export: full-parameter PP2 round-trips all 285 HF keys (no missing/unexpected); LoRA exports 310 adapter tensors (155 lora_A + 155 lora_B) and no base weights. --- src/mcore_bridge/config/model_config.py | 5 - src/mcore_bridge/model/gpts/deepseek_v41.py | 1345 +++++++++++------ .../model/gpts/deepseek_v41_hybrid.py | 848 ----------- tests/test_deepseek_v41_engram.py | 220 +-- tests/test_deepseek_v41_hybrid.py | 224 ++- 5 files changed, 1044 insertions(+), 1598 deletions(-) delete mode 100644 src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index 09f96c98..4e000947 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -255,11 +255,6 @@ class ModelConfig(TransformerConfig): mhc_sinkhorn_iterations: int = 20 mhc_init_gating_factor: float = 0.01 moe_n_hash_layers: int = 0 - # DeepSeek-V4.1 pipeline-parallel path selector (None = auto). The default GPTModel path is - # the golden baseline and cannot run PP>1, so the HybridModel loader/bridge are auto-selected - # when pipeline_model_parallel_size > 1. Set explicitly (e.g. via --megatron_extra_kwargs) to - # force the hybrid path on at PP1 (baseline alignment) or off. See deepseek_v41.py. - deepseek_v41_hybrid: Optional[bool] = None # deepseek-v4.1 engram (HF layer IDs are 0-based) # Declared here as well so the bridge remains importable on the PR #7224 baseline, diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 9b98062b..4c07cddc 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -1,10 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""DeepSeek-V4.1-Flash (text backbone) bridge for megatron-core. +"""DeepSeek-V4.1-Flash bridge for megatron-core. -Wires the V4.1 *language* 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: +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. @@ -14,50 +14,54 @@ ``hc_head_*`` params; per-layer ``hc_attn_*``/``hc_ffn_*`` are mapped by the base ``GPTBridge``. -The text integration also attaches the native trainable Engram modules and -loads their EP-local table rows directly from the official flat FP8 tensors. -Vision, MTP and DSpark are integrated separately. +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 V4.1 loader deliberately selects megatron-core's native TransformerBlock, -whose forward owns the per-call CSA2State and SinglePassMHCState lifecycle. Other -mcore-bridge models keep the custom TransformerBlock path. +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 contextlib import contextmanager +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 parallel_state +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.tensor_parallel.mappings import gather_from_sequence_parallel_region 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 typing import Optional +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 ( - adapt_deepseek_v41_layer_specs, - allow_engram_inference, 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 .deepseek_v4 import ( - _apply_mla_rope, - DeepseekV4Bridge, - DeepseekV4GPTModel, - DeepseekV4Loader, - DSv4HybridSelfAttention, -) +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 @@ -66,6 +70,19 @@ 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( @@ -110,8 +127,8 @@ def __init__(self, config, *args, **kwargs): 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))) + 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): @@ -127,7 +144,7 @@ def __init__(self, config, submodules, *args, **kwargs): 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) + 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) @@ -171,8 +188,7 @@ def _project_query(self, hidden_states, rotary_pos_emb): 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_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, @@ -202,22 +218,19 @@ def _normalize_cache_inputs(main_kv, start_pos, cache_slots): 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)}.') + 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)}.') + 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) + 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], @@ -266,7 +279,7 @@ def _latent_attention(self, query, key_value, valid_main_lengths=None): 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), + torch.cat((scores, sink.expand(scores.shape[:-1] + (1, ))), dim=-1), dim=-1, dtype=torch.float32, )[..., :-1] @@ -318,8 +331,7 @@ def forward( 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, 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) @@ -327,8 +339,7 @@ def forward( 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_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, @@ -379,9 +390,8 @@ def __init__(self, patch_size: int, hidden_size: int): 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)}.') + 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)) @@ -398,10 +408,7 @@ def __init__(self, hidden_size: int, num_heads: int): 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, 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)) @@ -451,8 +458,7 @@ def __init__(self, vision_config): 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]}.') + 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: @@ -539,7 +545,8 @@ def encode_images(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor 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]}.') + 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) @@ -559,20 +566,19 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): 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)}.') + 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.') + 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.') + 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) @@ -584,318 +590,562 @@ def get_inputs_embeds(self, inputs_embeds, **kwargs): return inputs_embeds.masked_scatter(expanded_mask, replacements) -class DeepseekV41GPTModel(DeepseekV4GPTModel): - """V4.1 language model with opt-in DSpark target-layer capture. +@dataclass +class HybridLayerConfig: + """Per-layer config re-expanded from HF layer space into hybrid (2x) layer space. - DSpark consumes the attention inputs of its target layers. The official target - IDs are zero-based; Megatron layer numbers are one-based. Capturing is opt-in - so regular training does not retain three large activation graphs. + 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``. """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._capture_dspark_hidden = False - self._dspark_hidden_states = {} - self._dspark_hook_handles = [] - target_ids = tuple(self.config.dspark_target_layer_ids or ()) - for layer in self.decoder.layers: - layer_id = layer.layer_number - 1 - if layer_id not in target_ids: - continue - self._dspark_hook_handles.append( - layer.register_forward_pre_hook(self._make_dspark_capture_hook(layer_id), with_kwargs=True)) + 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] - @staticmethod - def _contract_dspark_target_hidden(hidden_states: torch.Tensor, num_streams: int): - if hidden_states.ndim != 3: - raise ValueError(f'DSpark target hidden states must be [s, b, n*h], got {tuple(hidden_states.shape)}.') - if hidden_states.shape[-1] % num_streams: - raise ValueError( - f'DSpark target hidden width {hidden_states.shape[-1]} is not divisible by {num_streams} streams.') - return hidden_states.unflatten(-1, (num_streams, -1)).mean(dim=-2) - def _make_dspark_capture_hook(self, layer_id): +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. - def _capture(_module, args, kwargs): - if not self._capture_dspark_hidden: - return - hidden_states = kwargs.get('hidden_states') - if hidden_states is None and args: - hidden_states = args[0] - if hidden_states is None: - raise ValueError(f'Could not capture the attention input for DSpark target layer {layer_id}.') - self._dspark_hidden_states[layer_id] = self._contract_dspark_target_hidden( - hidden_states, self.config.num_residual_streams) - - return _capture - - @contextmanager - def capture_dspark_hidden_states(self): - if not self.config.dspark_target_layer_ids: - raise ValueError('DSpark target-layer capture requested, but DSpark is not configured.') - self._dspark_hidden_states = {} - self._capture_dspark_hidden = True - try: - yield - finally: - self._capture_dspark_hidden = False + 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)), + ) - def get_dspark_main_hidden(self, clear: bool = True): - target_ids = tuple(self.config.dspark_target_layer_ids or ()) - missing = [layer_id for layer_id in target_ids if layer_id not in self._dspark_hidden_states] - if missing: - raise RuntimeError( - f'DSpark target layers {missing} were not captured on this pipeline rank. ' - 'The DSpark draft stack must be colocated with all target layers.') - result = torch.cat([self._dspark_hidden_states[layer_id] for layer_id in target_ids], dim=-1) - if clear: - self._dspark_hidden_states = {} - return result - def forward(self, *args, **kwargs): - input_ids = kwargs.get('input_ids', args[0] if args else None) - inference_context = kwargs.get('inference_context') or kwargs.get('inference_params') - if inference_context is None and len(args) > 5: - inference_context = args[5] - extra_block_kwargs = kwargs.get('extra_block_kwargs') - if extra_block_kwargs is None and len(args) > 7: - extra_block_kwargs = args[7] - - with allow_engram_inference(self.config, input_ids, extra_block_kwargs) as block_kwargs: - if len(args) > 7: - args = (*args[:7], block_kwargs, *args[8:]) - else: - kwargs['extra_block_kwargs'] = block_kwargs - capture_dspark = ( - hasattr(self, 'dspark') - and not self.training - and inference_context is not None - and inference_context.is_dynamic_batching() - and inference_context.num_speculative_tokens > 0 - ) - if not capture_dspark: - return super().forward(*args, **kwargs) - if inference_context.using_cuda_graph_this_step(): - raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') - with self.capture_dspark_hidden_states(): - return super().forward(*args, **kwargs) +if _HYBRID_MODEL_AVAILABLE: - def _dspark_rotary_for_positions(self, position_ids: torch.Tensor): - if self.position_embedding_type != 'rope' or self.rotary_pos_emb is None: - raise RuntimeError('DSpark requires RoPE position embeddings.') - position_ids = position_ids.to(dtype=torch.long) - max_position = int(position_ids.max().item()) + 1 - rotary_table = self.rotary_pos_emb(max_position) - if isinstance(rotary_table, dict): - rotary_table = rotary_table['main'] - selected = rotary_table.index_select(0, position_ids.reshape(-1)) - # Drop the table's singleton batch dimension. The requested position tensor - # supplies the batch axes for packed main tokens or the parallel draft block. - return selected.reshape(*position_ids.shape, *rotary_table.shape[2:]) - - def _dspark_word_embeddings(self): - """Return the token embedding DSpark uses to embed its draft seed. - - On a single stage (or a standard-MTP stage) the base input embedding is - colocated and reused. On a PP>1 last stage with untied embeddings the base - model has no ``embedding``; a dedicated replicated DSpark embedding built - in ``build_model`` and loaded from ``model.embed_tokens.weight`` is used. - """ - if hasattr(self, 'embedding'): - return self.embedding.word_embeddings - dspark_embedding = getattr(self, 'dspark_word_embeddings', None) - if dspark_embedding is None: - raise RuntimeError( - 'DSpark requires an input word embedding on its pipeline stage, but neither ' - 'the base embedding nor a dedicated DSpark embedding is present.') - return dspark_embedding - - def forward_dspark( - self, - main_hidden, - input_ids, - *, - start_pos, - rotary_pos_emb, - main_rotary_pos_emb, - inference_context=None, - temperature=0.0, - sample_fn=None, - cache_slots=None, - prefill_only=None, - ): - if not hasattr(self, 'dspark'): - raise RuntimeError('DSpark is not available on this pipeline stage.') - if not hasattr(self, 'output_layer'): - raise RuntimeError('DSpark requires the output head on its pipeline stage.') - return self.dspark( - main_hidden, - input_ids, - self._dspark_word_embeddings(), - self.output_layer, - start_pos=start_pos, - rotary_pos_emb=rotary_pos_emb, - main_rotary_pos_emb=main_rotary_pos_emb, - inference_context=inference_context, - temperature=temperature, - sample_fn=sample_fn, - cache_slots=cache_slots, - prefill_only=prefill_only, - ) + class DeepseekV41HyperConnectionHybridLayer(HyperConnectionHybridLayer): + """Hyper-connection wrapper that applies Engram on the n-stream residual, matching GPT. - def compute_dspark_speculative_tokens( - self, - next_token_ids, - accepted_token_counts, - last_accepted_seq_indices, - num_speculative_tokens, - inference_context, - sample_fn, - ): - """Commit verified target states and produce one parallel DSpark draft block.""" - if inference_context.using_cuda_graph_this_step(): - raise RuntimeError('DSpark speculative decoding does not support CUDA graph replay yet.') - if num_speculative_tokens > self.config.dspark_block_size: - raise ValueError( - f'Requested {num_speculative_tokens} speculative tokens, but DSpark block size is ' - f'{self.config.dspark_block_size}.') - - main_hidden = self.get_dspark_main_hidden() - if self.config.sequence_parallel and parallel_state.get_tensor_model_parallel_world_size() > 1: - main_hidden = gather_from_sequence_parallel_region(main_hidden, group=self.tp_group) - if main_hidden.ndim != 3 or main_hidden.shape[1] != 1: - raise RuntimeError( - 'DSpark dynamic inference expects packed target states [tokens, 1, targets*h], ' - f'got {tuple(main_hidden.shape)}.') - - active_count = inference_context.total_request_count - inference_context.paused_request_count - active_slice = slice(inference_context.paused_request_count, inference_context.total_request_count) - query_lengths = inference_context.request_query_lengths[active_slice].to(dtype=torch.long) - active_token_count = int(query_lengths.sum().item()) - if main_hidden.shape[0] != active_token_count: - raise RuntimeError( - f'DSpark captured {main_hidden.shape[0]} target rows for {active_token_count} active tokens.') - - device = main_hidden.device - request_ids = inference_context.request_ids[active_slice].to(device=device, dtype=torch.long) - live_request_ids = inference_context.request_ids[:inference_context.total_request_count].to( - device=device, dtype=torch.long) - cache_slots = self.dspark.resolve_cache_slots(request_ids, live_request_ids) - token_positions = inference_context.token_to_position_in_request[:active_token_count].to( - device=device, dtype=torch.long) - accepted_token_counts = accepted_token_counts[:active_count].to(device='cpu', dtype=torch.long) - - offset = 0 - for request_index, query_length in enumerate(query_lengths.tolist()): - if request_index < inference_context.num_decode_requests: - accepted_length = min(query_length, int(accepted_token_counts[request_index].item()) + 1) - else: - accepted_length = query_length - if accepted_length: - token_slice = slice(offset, offset + accepted_length) - positions = token_positions[token_slice] - if positions.numel() > 1 and not torch.all(positions[1:] == positions[:-1] + 1): - raise RuntimeError('DSpark cache updates require contiguous per-request token positions.') - self.dspark.update_main_cache( - main_hidden[token_slice], - self._dspark_rotary_for_positions(positions), - start_pos=positions[0], - cache_slots=cache_slots[request_index:request_index + 1], - inference_context=inference_context, - ) - offset += query_length - - last_indices = last_accepted_seq_indices[:active_count].to(device=device, dtype=torch.long) - last_hidden = main_hidden.index_select(0, last_indices).transpose(0, 1).contiguous() - main_positions = token_positions.index_select(0, last_indices) - draft_positions = main_positions.unsqueeze(0) + 1 + torch.arange( - self.config.dspark_block_size, device=device).unsqueeze(1) - output_ids, _, _ = self.forward_dspark( - last_hidden, - next_token_ids[:active_count], - start_pos=main_positions, - rotary_pos_emb=self._dspark_rotary_for_positions(draft_positions), - main_rotary_pos_emb=self._dspark_rotary_for_positions(main_positions), - inference_context=inference_context, - sample_fn=sample_fn, - cache_slots=cache_slots, - prefill_only=False, - ) - return output_ids[:, 1:num_speculative_tokens + 1].transpose(0, 1).contiguous() + 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 -class DeepseekV41MultimodalGPTModel(MultimodalGPTModel): - language_model_cls = DeepseekV41GPTModel + 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. - @property - def vocab_size(self): - return self.language_model.vocab_size + 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`). - def forward_with_dspark_hidden(self, *args, **kwargs): - with self.language_model.capture_dspark_hidden_states(): - output = self.forward(*args, **kwargs) - return output, self.language_model.get_dspark_main_hidden() + 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_dspark(self, *args, **kwargs): - return self.language_model.forward_dspark(*args, **kwargs) + 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, + ) - def compute_dspark_speculative_tokens(self, *args, **kwargs): - return self.language_model.compute_dspark_speculative_tokens(*args, **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). + """ -def _deepseek_v41_use_hybrid(config) -> bool: - """Whether to build DeepSeek-V4.1 on the ``HybridModel`` (PP-capable) path. + language_model_cls = DeepseekV41HybridStackModel - The ``HybridModel`` path is now the default (plan step B5): B1-B4 validated it against the - ``GPTModel`` golden baseline (iter-1 loss/grad within the bf16/MoE non-determinism band and a - clean weight key ledger), and it is the only path that supports pipeline parallelism, so it is - selected for every layout. The ``deepseek_v41_hybrid`` config flag (settable via - ``--megatron_extra_kwargs``) overrides this: ``False`` drops back to the ``GPTModel`` golden - baseline (kept as a regression path; note upstream refuses ``GPTModel`` at ``PP>1``), ``True`` - is redundant but still forces hybrid. - """ - forced = getattr(config, 'deepseek_v41_hybrid', None) - if forced is not None: - return bool(forced) - return True + @property + def vocab_size(self): + return self.language_model.vocab_size +else: + DeepseekV41HyperConnectionHybridLayer = None + DeepseekV41HybridStackModel = None + DeepseekV41MultimodalModel = None class DeepseekV41Loader(DeepseekV4Loader): - model_cls = DeepseekV41MultimodalGPTModel - # Native V4.1 forward owns CSA2State + SinglePassMHCState. Using it only for - # this loader avoids changing the custom bridge block used by V4/DSpark/MTP. + """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 - def __new__(cls, config=None, *args, **kwargs): - # Auto-route to the HybridModel loader on the PP path (or when forced); the subclass - # instantiates itself directly, so the ``cls is`` guard prevents re-dispatch. ``config`` - # is optional so ``__new__(cls)`` (used by tests to skip __init__) keeps working. - if cls is DeepseekV41Loader and config is not None and _deepseek_v41_use_hybrid(config): - from .deepseek_v41_hybrid import DeepseekV41HybridLoader - if DeepseekV41HybridLoader is not None: - return super().__new__(DeepseekV41HybridLoader) - return super().__new__(cls) + model_cls = DeepseekV41MultimodalModel def _engram_placement_layer_ids(self, hf_layer_ids): - """Map 0-based HF Engram layer IDs to 1-based ``TransformerLayer`` placement numbers. + """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 - On the GPT stack HF layer ``e`` is one ``TransformerLayer`` numbered ``e + 1``. The - HybridStack loader overrides this because there each HF layer becomes two hybrid layers. + 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. """ - return tuple(layer_id + 1 for layer_id in hf_layer_ids) + 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.') + 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', + '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: @@ -908,16 +1158,14 @@ def _get_engram_config(self): 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)}.' - ) + 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), + 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), @@ -929,24 +1177,19 @@ def _get_engram_config(self): 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,)), + 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}.' - ) + 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) + 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): @@ -984,20 +1227,15 @@ def get_dspark_layer_spec(self): self._replace_router(SimpleNamespace(layer_specs=layer_specs)) return dspark_config, layer_specs - def build_model(self, pre_process=True, post_process=True, vp_stage: Optional[int] = None): - model = super().build_model(pre_process, post_process, vp_stage) - self._attach_dspark(model.language_model, post_process, vp_stage=vp_stage) - return model - 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 GPT path passes - ``model.language_model``; the hybrid path (which has no ``language_model`` wrapper) passes - the ``HybridModel`` itself -- both expose ``pg_collection`` / ``vocab_size`` / ``config``. + 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. @@ -1018,8 +1256,7 @@ def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = layer_number=index + 1, pg_collection=language_model.pg_collection, vp_stage=vp_stage, - ) - for index, layer_spec in enumerate(dspark_layer_specs) + ) for index, layer_spec in enumerate(dspark_layer_specs) ] language_model.dspark = DeepseekV41DSparkStack(dspark_config, layers) self._set_linear_is_expert(language_model.dspark) @@ -1046,46 +1283,270 @@ def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = tp_group=language_model.pg_collection.tp, ) - def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): - from megatron.core.models.gpt.experimental_attention_variant_module_specs import \ - get_transformer_block_with_experimental_attention_variant_spec - transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec(self.config, vp_stage) - for layer_spec in transformer_layer_spec.layer_specs: - layer_spec.submodules.self_attention.module = DSv4HybridSelfAttention - core_attention_submodules = layer_spec.submodules.self_attention.submodules.core_attention.submodules - if getattr(core_attention_submodules, 'compressor', None) is not None: - core_attention_submodules.compressor.module = CSA2Compressor - if getattr(core_attention_submodules, 'indexer', None) is not None: - # CSA2 indexer is flat (no nested compressor). - core_attention_submodules.indexer.module = CSA2Indexer - engram_config = self._get_engram_config() - if engram_config is not None: - transformer_layer_spec = adapt_deepseek_v41_layer_specs(transformer_layer_spec, engram_config) - return transformer_layer_spec - 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'} - def __new__(cls, config=None, *args, **kwargs): - # Mirror the loader's routing so the bridge and the built model always agree on the path - # (this bridge is created in ``ModelConfig.__post_init__`` where PP size is already set). - # ``config`` is optional so ``__new__(cls)`` (used by tests to skip __init__) keeps working. - if cls is DeepseekV41Bridge and config is not None and _deepseek_v41_use_hybrid(config): - from .deepseek_v41_hybrid import DeepseekV41HybridBridge - if DeepseekV41HybridBridge is not None: - return super().__new__(DeepseekV41HybridBridge) - return super().__new__(cls) + @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: @@ -1111,11 +1572,9 @@ 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)}.') + 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]}.') + 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) @@ -1136,20 +1595,10 @@ def _load_engram_embedding(self, engram, hf_state_dict): 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)] + 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 _engram_hf_layer_id(self, engram): - """Recover the 0-based HF layer ID from a built Engram's 1-based ``layer_number``. - - On the GPT stack ``layer_number == hf_id + 1``. The HybridStack bridge overrides this - because its Engram sits on the doubled-space attention layer ``2 * hf_id + 1``. - """ - return engram.layer_number - 1 + 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) @@ -1165,14 +1614,12 @@ def _set_layer_engram(self, mg_layer, hf_state_dict, to_mcore): 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)) + 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 @@ -1187,24 +1634,14 @@ def _set_layer_engram(self, mg_layer, hf_state_dict, to_mcore): 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] + 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_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): - layer_prefix = f'{hf_prefix}{layer_idx}.' - local_state = self._remove_prefix(hf_state_dict, layer_prefix) if to_mcore else {} - result = super()._set_layer_state(mg_layer, hf_state_dict, hf_prefix, layer_idx, to_mcore) - if layer_idx in (self.config.engram_layer_ids or []): - self._set_layer_engram(mg_layer, local_state, to_mcore) - if not to_mcore and local_state: - result.update(self._add_prefix(local_state, layer_prefix)) - return result - 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 {} @@ -1223,13 +1660,11 @@ def _set_dspark_endpoints(self, dspark, hf_state_dict, to_mcore): 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) + 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) @@ -1250,20 +1685,10 @@ def _load_dspark_word_embeddings(self, embedding, hf_state_dict): 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_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage): - if not self.config.dspark_num_layers or (to_mcore and not is_pp_last_stage): - return - language_model = mg_model.language_model if self.is_multimodal else mg_model - dspark = getattr(language_model, 'dspark', None) - if dspark is None: - raise RuntimeError('DSpark weights require the draft stack on the final pipeline stage.') - yield from self._convert_dspark_stack(language_model, dspark, hf_state_dict, hf_prefix, to_mcore) - 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), so - both the GPT and hybrid bridges reuse it; they differ only in how they locate ``dspark`` - and guard the pipeline stage.""" + 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, @@ -1345,18 +1770,14 @@ def _set_mla_attn_state(self, mg_attn, hf_state_dict, hf_prefix, layer_idx, to_m hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) return hf_state_dict - def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): - # V4.1 single-pass mHC has no learned hc_head_*; skip the V4 hc_head mapping - # and only handle the plain final layernorm (base GPTBridge behaviour). - super(DeepseekV4Bridge, self)._set_final_layernorm(lm_model, hf_state_dict, to_mcore) - - -register_model( - ModelMeta( - ModelType.deepseek_v41, - ['deepseek_v41'], - bridge_cls=DeepseekV41Bridge, - visual_cls=DeepseekV41Vision, - loader=DeepseekV41Loader, - config_cls=MLAModelConfig, - )) + +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/gpts/deepseek_v41_hybrid.py b/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py deleted file mode 100644 index b1abfaa4..00000000 --- a/src/mcore_bridge/model/gpts/deepseek_v41_hybrid.py +++ /dev/null @@ -1,848 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""DeepSeek-V4.1 on megatron-core's ``HybridModel`` (pipeline-parallel path). - -The legacy :class:`DeepseekV41Loader` builds a ``GPTModel`` whose custom -``TransformerBlock`` owns the CSA2 / single-pass-mHC forward. Upstream refuses to -run that block under pipeline parallelism:: - - # transformer_block.py:304-311 - if (config.pipeline_model_parallel_size > 1 - and config.experimental_attention_variant == "dsv4_hybrid" - and config.dsv4_version == "v4.1"): - raise ValueError("V4.1 pipeline parallelism requires HybridModel and its payload adapter") - -so PP>1 requires the native ``HybridModel`` + ``CSA2HybridAdapter`` typed-payload path. - -On ``HybridModel`` one *pattern symbol is one layer*: a GPT ``attn+mlp`` layer becomes -two hybrid layers -- an attention-only layer (symbol ``D``) followed by an MLP-only -layer (``E`` for MoE, ``-`` for dense). So a hybrid stack has ``2 * num_layers`` layers -and every per-layer config array that CSA2 indexes by ``layer_number - 1`` must be -re-expanded into this doubled index space (see :func:`derive_hybrid_layer_config`). - -The HybridModel path is the sole maintained DeepSeek-V4.1 implementation. The legacy -GPTModel path is deprecated and scheduled for removal. -""" -import copy -from dataclasses import dataclass -from types import SimpleNamespace -from typing import List, Optional, Sequence, Union - -import torch -import torch.distributed as dist -from megatron.core import mpu -from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding -from tqdm import tqdm - -from mcore_bridge.utils import is_master - -from ..modules.engram import DeepseekV41Engram, DeepseekV41TransformerLayer -from ..rope import get_rope_inv_freq -from .deepseek_v41 import (CSA2Compressor, CSA2Indexer, DeepseekV41Bridge, DeepseekV41Loader, - DeepseekV41MultimodalGPTModel, DSv4HybridSelfAttention) - -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 - - -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). This reproduces - the GPTModel golden path exactly. 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:`DeepseekV41HybridLoader._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:`DeepseekV41HybridBridge`, 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 (B2) is still appended by the base resolver, so a - segmented main becomes ``seg0|seg1|.../mtp``. - """ - - # B1 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, - # like the GPT ``DeepseekV41MultimodalGPTModel``). Expose it so that guard short-circuits; - # the real vision tower arrives with the multimodal wrapper in B4. - visual = None - - # ``MultimodalGPTModel.forward`` (B4 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 treats the hybrid backbone exactly like the GPT one. - 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 the GPT path) 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 (B4) and then - # cleared before the language model runs; the base HybridModel.forward never accepts them. - # The B1 backbone is text-only, so strip them here. For a text batch the GPT wrapper's - # ``get_inputs_embeds`` is a numeric no-op (``_zero_parameter_dependency`` adds ``0 * - # vision_params``), so dropping them keeps parity with the GPTModel baseline. - _visual_forward_keys = ('pixel_values', 'image_grid_thw', 'image_token_types', 'token_types') - - def forward(self, *args, **kwargs): - # The B4 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( - 'DeepSeek-V4.1 hybrid (pipeline-parallel) path is text-only in B1; multimodal ' - 'inputs require the B4 multimodal wrapper.') - 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 DeepseekV41MultimodalHybridModel(DeepseekV41MultimodalGPTModel): - """Multimodal wrapper (B4) hosting the PP-capable ``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). So the only change from - the GPT :class:`DeepseekV41MultimodalGPTModel` is swapping the language-model class; 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). The DSpark - speculative-decoding helpers inherited from the GPT wrapper are inference-only and stay - deferred on the hybrid path (see :meth:`DeepseekV41HybridLoader.build_model`). - """ - - language_model_cls = DeepseekV41HybridStackModel -else: - DeepseekV41HyperConnectionHybridLayer = None - DeepseekV41HybridStackModel = None - DeepseekV41MultimodalHybridModel = None - - -@dataclass -class HybridLayerConfig: - """Per-layer config re-expanded from GPT 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 a GPT-space V4.1 config into the doubled hybrid layer space. - - Each GPT 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)), - ) - - -class DeepseekV41HybridLoader(DeepseekV41Loader): - """Build DeepSeek-V4.1 on ``HybridModel`` (the PP-capable path). - - Reuses :class:`DeepseekV41Loader`'s Engram config resolution and MLA/CSA2 knowledge, but - swaps the model class to the native ``HybridModel`` and rewrites the layer config into the - doubled hybrid layer space (see :func:`derive_hybrid_layer_config`). The golden GPT - ``DeepseekV41Loader`` is left untouched; this loader derives its own config copy so both - paths can coexist in one process. - - B1 covers the text backbone; B3 adds the DSpark (``mtp.*``) draft stack on top (attached in - :meth:`build_model`, mapped in :meth:`DeepseekV41HybridBridge._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. B4 wraps the - backbone in :class:`DeepseekV41MultimodalHybridModel` (the vision tower + image-embed - injection), mirroring the GPT :class:`DeepseekV41MultimodalGPTModel`. - """ - - model_cls = DeepseekV41MultimodalHybridModel - - 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``, B3). - # 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 same fp8-parity module swaps the GPT loader uses, 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 like the GPT ``adapt_deepseek_v41_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 the GPT router swap). - 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:`DeepseekV41MultimodalHybridModel` (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 (matching the GPT wrapper, - where the same fix-ups and DSpark live under ``language_model``).""" - 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 (B3): the ``mtp.*`` draft stack is backbone-agnostic (plain - # experimental-attention layers), so reuse the GPT loader's builder. It attaches to the - # hybrid text backbone (``language_model.dspark``), mirroring the GPT wrapper. Inference-time - # target-layer capture on HybridStack is deferred (it is not exercised by training / weight - # round-trip, mirroring B1's deferral of ``allow_engram_inference``); 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 - - -class DeepseekV41HybridBridge(DeepseekV41Bridge): - """Weight bridge for the ``HybridModel`` backbone (PP-capable path). - - The GPT bridge maps one HF layer onto one ``TransformerLayer`` that owns both attention - and MLP. On ``HybridModel`` that layer 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 (the GPT layer instead carried two: - ``self_attention_hyper_connection`` + ``mlp_hyper_connection``). The GPT - ``hc_{attn,ffn}_*`` HF keys therefore split across the two wrappers. - - B1 handles the text backbone only. It treats the model as its own language model (the - multimodal wrapper is B4) and skips MTP (B2). ``self.config`` is seen in two layer spaces - depending on direction: on load it is the original GPT-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. - """ - - @staticmethod - def _lm(mg_model): - """Resolve the language model. B1's HybridModel is text-only (no ``language_model`` - wrapper); B4 will nest 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 GPT-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 - # ``DeepseekV41HybridLoader._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 (B4) 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): - # Delegate to ``DeepseekV41Bridge._convert_pre_process`` (``super()`` via MRO) unconditionally - # instead of branching on *this* rank's ``mg_model.visual``. On export every pipeline stage runs - # ``_convert`` -> ``_convert_pre_process``, and the base path issues the *same* pp-group collective - # sequence on all ranks: word-embeddings (routed through ``_lm`` by the ``_set_word_embeddings`` - # override), 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), - # so the vision tower is always present there and the base path behaves exactly as before. - return super()._convert_pre_process(mg_model, hf_state_dict, hf_prefix, to_mcore) - - 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 (B3). - - The draft layers are plain experimental-attention ``TransformerLayer`` instances -- - identical in both paths -- so the base :meth:`DeepseekV41Bridge._convert_dspark_stack` - mapping is reused verbatim; only where the stack lives differs. On the hybrid path it is - attached to the model itself (no ``language_model`` wrapper, 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 (B2); the multimodal-wrapper indirection is B4. - """ - 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()) diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index c29cf2bc..9be4d878 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -15,7 +15,6 @@ DeepseekV41Aligner, DeepseekV41Bridge, DeepseekV41DSparkAttention, - DeepseekV41GPTModel, DeepseekV41Vision, DeepseekV41VisionTransformer, ) @@ -33,16 +32,6 @@ from mcore_bridge.utils.safetensors import SafetensorLazyLoader -def test_dspark_target_hidden_averages_mhc_streams(): - hidden = torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).view(2, 3, 20) - - actual = DeepseekV41GPTModel._contract_dspark_target_hidden(hidden, num_streams=4) - expected = hidden.view(2, 3, 4, 5).mean(dim=2) - - assert actual.shape == (2, 3, 5) - torch.testing.assert_close(actual, expected) - - def test_dspark_config_is_kept_separate_from_standard_mtp(): text_config = SimpleNamespace( model_type='deepseek_v41_text', @@ -68,6 +57,7 @@ def test_dspark_config_is_kept_separate_from_standard_mtp(): def test_dspark_input_builds_parallel_noise_block(): + class _Projection(torch.nn.Module): def forward(self, hidden_states): @@ -101,6 +91,7 @@ def embedding(token_ids): def test_dspark_markov_head_returns_full_logits_and_embedding(): + class _Head(torch.nn.Module): def forward(self, hidden_states, runtime_gather_output): @@ -109,9 +100,7 @@ def forward(self, hidden_states, runtime_gather_output): 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.embed = torch.nn.Embedding.from_pretrained(torch.tensor([[1., 2.], [3., 4.], [5., 6.]]), ) module.head = _Head() logits, embedding = module(torch.tensor([0, 2])) @@ -233,11 +222,12 @@ def test_dspark_verification_accepts_only_strict_matching_prefix(): 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], - ])) + 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) @@ -262,6 +252,7 @@ def test_dspark_state_preserves_recompute_inputs(): def test_dspark_stack_prefill_and_decode_lifecycle(): + class _Input(torch.nn.Module): def forward(self, main_hidden, input_ids, embedding): @@ -301,7 +292,7 @@ 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) + 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): @@ -329,8 +320,7 @@ def forward(self, hidden_states, input_ids, output_layer, temperature, sample_fn ) is None for layer in stack.layers: prefill_main, prefill_rotary, prefill_context, prefill_start, prefill_slots = ( - layer.self_attention.prefill_args - ) + layer.self_attention.prefill_args) torch.testing.assert_close(prefill_main, main_hidden[..., :2]) assert prefill_rotary is rotary assert prefill_context is None @@ -370,6 +360,7 @@ def test_dspark_confidence_head_uses_fp32_projection(): def test_dspark_output_applies_markov_recurrence_in_block_order(): + class _OutputLayer(torch.nn.Module): def forward(self, hidden_states, runtime_gather_output): @@ -411,72 +402,6 @@ def forward(self, hidden_states, markov_embed): assert torch.equal(dspark_sample(logits, temperature=0), output_ids[:, 1:]) -def test_dspark_model_commits_only_verified_states_before_proposal(): - class _DSpark: - - def __init__(self): - self.updates = [] - - def resolve_cache_slots(self, request_ids, live_request_ids): - assert torch.equal(request_ids.cpu(), torch.tensor([10, 20])) - assert torch.equal(live_request_ids.cpu(), torch.tensor([10, 20])) - return torch.tensor([2, 0], device=request_ids.device) - - def update_main_cache( - self, - main_hidden, - rotary_pos_emb, - *, - start_pos, - cache_slots, - inference_context, - ): - self.updates.append((main_hidden.clone(), start_pos.clone(), cache_slots.clone())) - - model = DeepseekV41GPTModel.__new__(DeepseekV41GPTModel) - torch.nn.Module.__init__(model) - model.config = SimpleNamespace(sequence_parallel=False, dspark_block_size=3) - model.dspark = _DSpark() - captured = torch.arange(5 * 4, dtype=torch.float32).view(5, 1, 4) - model.get_dspark_main_hidden = lambda: captured - model._dspark_rotary_for_positions = lambda positions: positions.float() - proposal_args = {} - - def forward_dspark(main_hidden, input_ids, **kwargs): - proposal_args.update(main_hidden=main_hidden, input_ids=input_ids, **kwargs) - output_ids = torch.tensor([[31, 32, 33, 34], [41, 42, 43, 44]]) - return output_ids, None, None - - model.forward_dspark = forward_dspark - context = SimpleNamespace( - total_request_count=2, - paused_request_count=0, - num_decode_requests=1, - request_query_lengths=torch.tensor([3, 2], dtype=torch.int32), - request_ids=torch.tensor([10, 20], dtype=torch.int32), - token_to_position_in_request=torch.tensor([5, 6, 7, 0, 1], dtype=torch.int32), - using_cuda_graph_this_step=lambda: False, - ) - - proposals = model.compute_dspark_speculative_tokens( - next_token_ids=torch.tensor([31, 41]), - accepted_token_counts=torch.tensor([1, 0]), - last_accepted_seq_indices=torch.tensor([1, 4]), - num_speculative_tokens=2, - inference_context=context, - sample_fn=lambda logits: logits.argmax(dim=-1), - ) - - assert len(model.dspark.updates) == 2 - torch.testing.assert_close(model.dspark.updates[0][0], captured[:2]) - torch.testing.assert_close(model.dspark.updates[1][0], captured[3:5]) - assert model.dspark.updates[0][1].item() == 5 - assert model.dspark.updates[1][1].item() == 0 - assert torch.equal(proposal_args['start_pos'], torch.tensor([6, 1])) - torch.testing.assert_close(proposal_args['main_hidden'], captured[[1, 4]].transpose(0, 1)) - assert torch.equal(proposals, torch.tensor([[32, 42], [33, 43]])) - - def test_controller_routes_speculative_proposals_to_dspark_provider(): calls = {} @@ -516,23 +441,27 @@ 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)), - })) + 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,), + excluded_token_ids=(99, ), global_vocab_sizes=(17, 19), max_ngram_order=3, num_hash_heads=1, @@ -547,27 +476,30 @@ def test_engram_adapter_remaps_checkpoint_layers_to_megatron_layers(tmp_path): 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,) + 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)), - })) + 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,), + placement_layer_ids=(1, ), + hash_layer_ids=(0, ), global_vocab_sizes=(17, 19), max_ngram_order=3, num_hash_heads=1, @@ -595,12 +527,10 @@ def test_engram_config_allows_context_and_virtual_pipeline_but_keeps_the_other_g 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) + 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) + 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): @@ -609,14 +539,12 @@ def test_engram_config_allows_packed_sequences_without_losing_the_pipeline_guard 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) + 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) + config._validate_packed_sequences(SimpleNamespace(pipeline_model_parallel_size=4), packed_sequences=True) def test_engram_hash_blocks_suffixes_after_excluded_token(): @@ -638,7 +566,7 @@ def test_engram_static_inference_cache_matches_full_sequence_hashing(): module = engram_adapter.DeepseekV41Engram.__new__(engram_adapter.DeepseekV41Engram) torch.nn.Module.__init__(module) module.engram_config = SimpleNamespace( - excluded_token_ids=(99,), + excluded_token_ids=(99, ), max_ngram_order=3, num_hash_heads=1, hash_boundary_token_id=0, @@ -684,8 +612,7 @@ def test_engram_packed_hashes_match_separately_hashed_documents(): 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) + 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), @@ -718,8 +645,7 @@ def test_engram_rejects_cu_seqlens_that_does_not_cover_the_row(): 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) + module.config = SimpleNamespace(context_parallel_size=context_parallel_size, sequence_parallel=sequence_parallel) return module @@ -767,8 +693,7 @@ def fake_all_gather(output_list, tensor, group=None): 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) + 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): @@ -788,8 +713,7 @@ def fake_all_gather(output_list, tensor, group=None): 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) + 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'): @@ -820,7 +744,7 @@ def test_engram_layer_spec_uses_bridge_owned_module(): submodules=TransformerLayerSubmodules(), ) block_spec = SimpleNamespace(layer_specs=[layer_spec]) - config = SimpleNamespace(layer_ids=(1,)) + config = SimpleNamespace(layer_ids=(1, )) engram_adapter.adapt_deepseek_v41_layer_specs(block_spec, config) @@ -892,23 +816,6 @@ def record(_module, mg_key, _state, hf_key, to_mcore): ] -def test_dspark_word_embeddings_resolver_prefers_base_then_dedicated(): - resolver = DeepseekV41GPTModel._dspark_word_embeddings - model = DeepseekV41GPTModel.__new__(DeepseekV41GPTModel) - torch.nn.Module.__init__(model) - # Neither the base embedding nor a dedicated DSpark embedding is present. - with pytest.raises(RuntimeError): - resolver(model) - # A PP>1 last stage falls back to the dedicated DSpark embedding. - dedicated = object() - model.dspark_word_embeddings = dedicated - assert resolver(model) is dedicated - # When the base embedding is colocated it always takes priority. - base = object() - model.embedding = SimpleNamespace(word_embeddings=base) - assert resolver(model) is base - - def test_dspark_bridge_loads_dedicated_embedding_with_padding_and_tp_shard(): bridge = DeepseekV41Bridge.__new__(DeepseekV41Bridge) bridge.hf_embed_key = 'model.embed.weight' @@ -939,7 +846,8 @@ 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), - layer_number=2, + # 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]) @@ -978,6 +886,7 @@ def load_slice(self, slices): def test_engram_dense_weights_are_dequantized_and_split_like_official_wkv(): + class _Lazy: def __init__(self, tensor): @@ -991,8 +900,7 @@ def load(self): 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))), + 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))), diff --git a/tests/test_deepseek_v41_hybrid.py b/tests/test_deepseek_v41_hybrid.py index d68b67bd..b21f3a01 100644 --- a/tests/test_deepseek_v41_hybrid.py +++ b/tests/test_deepseek_v41_hybrid.py @@ -3,7 +3,7 @@ Pure logic, no GPU / distributed init required. """ -from mcore_bridge.model.gpts.deepseek_v41_hybrid import HybridLayerConfig, derive_hybrid_layer_config +from mcore_bridge.model.gpts.deepseek_v41 import HybridLayerConfig, derive_hybrid_layer_config def test_tiny_all_moe_zero_ratio(): @@ -65,12 +65,12 @@ def test_ratio_length_mismatch_raises(): def test_loader_build_hybrid_config_does_not_mutate_original(): - # The loader derives its own config copy so the golden GPT path stays intact. + # 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_hybrid import DeepseekV41HybridLoader + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Loader - loader = object.__new__(DeepseekV41HybridLoader) + loader = object.__new__(DeepseekV41Loader) loader.config = SimpleNamespace( num_layers=4, csa_compress_ratios=[0, 0, 0, 0], @@ -91,9 +91,9 @@ def test_loader_build_hybrid_config_does_not_mutate_original(): 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 disabled for the B1 backbone-only path. + # MTP stays disabled (V4.1 ``mtp.*`` keys are DSpark). assert cfg.mtp_num_layers is None - # golden GPT config left untouched + # caller config left untouched assert original.num_layers == 4 assert original.mtp_num_layers == 1 assert original.is_hybrid_model is False @@ -102,9 +102,9 @@ def test_loader_build_hybrid_config_does_not_mutate_original(): def _make_bridge(engram_layer_ids, enable_hyper_connections): from types import SimpleNamespace - from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41HybridBridge + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Bridge - bridge = object.__new__(DeepseekV41HybridBridge) + bridge = object.__new__(DeepseekV41Bridge) bridge.config = SimpleNamespace( engram_layer_ids=engram_layer_ids, enable_hyper_connections=enable_hyper_connections) return bridge @@ -123,10 +123,7 @@ def test_hybrid_layer_state_fans_out_attn_and_mlp(): 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) - } + 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) @@ -175,13 +172,13 @@ def test_engram_placement_and_hf_layer_id_round_trip(): # engram_num_embeddings validation (keyed by 0-based HF ids) still resolves. from types import SimpleNamespace - from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41HybridBridge, DeepseekV41HybridLoader + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Bridge, DeepseekV41Loader - loader = object.__new__(DeepseekV41HybridLoader) + loader = object.__new__(DeepseekV41Loader) assert loader._engram_placement_layer_ids([1, 3]) == (3, 7) - assert loader._engram_placement_layer_ids([0]) == (1,) + assert loader._engram_placement_layer_ids([0]) == (1, ) - bridge = object.__new__(DeepseekV41HybridBridge) + 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 @@ -195,19 +192,19 @@ def test_num_hybrid_layers_normalizes_both_layer_spaces(): # ``2 * num_layers`` on the doubled config over-counted and dereferenced None layers. from types import SimpleNamespace - from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41HybridBridge + 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 DeepseekV41HybridBridge._num_hybrid_layers(load_cfg) == 8 - assert DeepseekV41HybridBridge._num_hybrid_layers(export_cfg) == 8 + 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 DeepseekV41HybridBridge._num_hybrid_layers(SimpleNamespace(num_layers=3)) == 6 + assert DeepseekV41Bridge._num_hybrid_layers(SimpleNamespace(num_layers=3)) == 6 import pytest # noqa: E402 -from mcore_bridge.model.gpts.deepseek_v41_hybrid import ( # noqa: E402 +from mcore_bridge.model.gpts.deepseek_v41 import ( # noqa: E402 DeepseekV41HyperConnectionHybridLayer, HyperConnectionHybridLayer) requires_hybrid = pytest.mark.skipif( @@ -218,8 +215,8 @@ def test_num_hybrid_layers_normalizes_both_layer_spaces(): 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), reproducing the GPTModel golden order. A plain inner - # layer delegates unchanged with no Engram add. + # (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 @@ -268,7 +265,7 @@ def test_rewrap_swaps_class_only_on_engram_wrappers(): # subclass; every other wrapper keeps the base class (and its fast path). from types import SimpleNamespace - from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41HybridLoader + from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41Loader engram_wrapper = object.__new__(HyperConnectionHybridLayer) engram_wrapper.inner_layer = SimpleNamespace(engram=object()) @@ -276,7 +273,7 @@ def test_rewrap_swaps_class_only_on_engram_wrappers(): plain_wrapper.inner_layer = SimpleNamespace(engram=None) model = SimpleNamespace(decoder=SimpleNamespace(layers=[engram_wrapper, plain_wrapper])) - loader = object.__new__(DeepseekV41HybridLoader) + loader = object.__new__(DeepseekV41Loader) loader.config = SimpleNamespace(enable_hyper_connections=True) loader._rewrap_engram_hyper_connection_layers(model) @@ -289,20 +286,20 @@ 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_hybrid import DeepseekV41HybridLoader + 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__(DeepseekV41HybridLoader) + 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_hybrid import DeepseekV41HybridStackModel # noqa: E402 +from mcore_bridge.model.gpts.deepseek_v41 import DeepseekV41HybridStackModel # noqa: E402 def _seg_config(pattern, **overrides): @@ -352,11 +349,10 @@ def test_segment_main_pattern_pp1_is_noop(): 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('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' + _seg_config('DEDEDEDE', pipeline_model_parallel_size=2, num_layers_in_first_pipeline_stage=2)) == 'DEDEDEDE' @requires_hybrid @@ -364,8 +360,7 @@ 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)) + DeepseekV41HybridStackModel._segment_main_pattern(_seg_config('DEDE', pipeline_model_parallel_size=4)) @requires_hybrid @@ -387,63 +382,22 @@ def test_resolve_hybrid_layer_pattern_segments_then_defers_to_base(): from mcore_bridge.model.gpts.deepseek_v41 import ( # noqa: E402 - DeepseekV41Bridge, DeepseekV41Loader, _deepseek_v41_use_hybrid) -from mcore_bridge.model.gpts.deepseek_v41_hybrid import ( # noqa: E402 - DeepseekV41HybridBridge, DeepseekV41HybridLoader) - - -def _route_config(pp, forced=None): - from types import SimpleNamespace - return SimpleNamespace(pipeline_model_parallel_size=pp, deepseek_v41_hybrid=forced) - + DeepseekV41Bridge, DeepseekV41Loader) -def test_use_hybrid_default_on_and_forced_override(): - # Default (B5 switch): HybridModel for every layout now that B1-B4 align with the GPT baseline. - assert _deepseek_v41_use_hybrid(_route_config(1)) is True - assert _deepseek_v41_use_hybrid(_route_config(2)) is True - # Explicit force-off drops back to the GPTModel golden baseline (kept as a regression path). - assert _deepseek_v41_use_hybrid(_route_config(1, forced=False)) is False - assert _deepseek_v41_use_hybrid(_route_config(2, forced=False)) is False - # Force-on is redundant now but must still route to hybrid. - assert _deepseek_v41_use_hybrid(_route_config(1, forced=True)) is True - - -@requires_hybrid -def test_loader_new_dispatches_to_hybrid(): - # __new__ routing only (no __init__), so no distributed init is required. - assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1))) is DeepseekV41HybridLoader - assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(2))) is DeepseekV41HybridLoader - assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1, forced=True))) is DeepseekV41HybridLoader - # Force-off keeps the GPTModel golden baseline. - assert type(DeepseekV41Loader.__new__(DeepseekV41Loader, _route_config(1, forced=False))) is DeepseekV41Loader - # A directly instantiated subclass must not re-dispatch (cls-is guard). - assert type(DeepseekV41HybridLoader.__new__(DeepseekV41HybridLoader, _route_config(1))) is DeepseekV41HybridLoader - - -@requires_hybrid -def test_bridge_new_dispatches_to_hybrid(): - assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1))) is DeepseekV41HybridBridge - assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(2))) is DeepseekV41HybridBridge - assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1, forced=True))) is DeepseekV41HybridBridge - # Force-off keeps the GPTModel golden baseline. - assert type(DeepseekV41Bridge.__new__(DeepseekV41Bridge, _route_config(1, forced=False))) is DeepseekV41Bridge - assert type(DeepseekV41HybridBridge.__new__(DeepseekV41HybridBridge, _route_config(1))) is DeepseekV41HybridBridge +# --- DSpark (``mtp.*``) draft stack ------------------------------------------------------------- -# --- B3: DSpark (``mtp.*``) draft stack on the hybrid path --------------------------------------- - 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__(DeepseekV41HybridBridge) + 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']))) + 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 @@ -461,7 +415,7 @@ def test_hybrid_convert_additional_layers_maps_dspark_via_lm(): def test_hybrid_convert_additional_layers_resolves_language_model_wrapper(): - # Forward-compat with B4: when a multimodal wrapper is present, ``_lm`` unwraps it and the + # 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 @@ -496,16 +450,19 @@ def test_hybrid_convert_additional_layers_load_skips_non_last_stage(): assert calls == [] -def test_hybrid_convert_additional_layers_export_skips_non_last_stage_without_stack(): - # On export a non-last PP stage has no draft stack; it must skip quietly (not raise), unlike - # the final stage where a missing stack is a real error. +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 == [] - assert calls == [] + 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(): @@ -518,16 +475,16 @@ def test_hybrid_convert_additional_layers_raises_on_last_stage_without_stack(): list(bridge._convert_additional_layers(mg_model, {}, 'prefix.', to_mcore=False, is_pp_last_stage=True)) -# --- B4: multimodal wrapper hosting the hybrid backbone ----------------------------------------- +# --- multimodal wrapper hosting the hybrid backbone --------------------------------------------- + -def test_multimodal_hybrid_wrapper_hosts_hybrid_backbone(): - # The B4 wrapper is just the GPT multimodal model with the language-model class swapped for the +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 DeepseekV41MultimodalGPTModel - from mcore_bridge.model.gpts.deepseek_v41_hybrid import (DeepseekV41HybridStackModel, - DeepseekV41MultimodalHybridModel) - assert issubclass(DeepseekV41MultimodalHybridModel, DeepseekV41MultimodalGPTModel) - assert DeepseekV41MultimodalHybridModel.language_model_cls is DeepseekV41HybridStackModel + 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(): @@ -537,40 +494,53 @@ def test_hybrid_stack_exposes_extra_forward_keys(): def test_hybrid_loader_model_cls_is_multimodal_wrapper(): - from mcore_bridge.model.gpts.deepseek_v41_hybrid import DeepseekV41MultimodalHybridModel - assert DeepseekV41HybridLoader.model_cls is DeepseekV41MultimodalHybridModel + 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 -def test_hybrid_pre_process_delegates_to_gpt_when_visual_present(monkeypatch): - # First PP stage of a multimodal model: the wrapper carries a vision tower, so pre-process must - # reuse the GPT DeepseekV41Bridge path (word embeddings + vision/aligner + image_* markers). + 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 = object.__new__(DeepseekV41HybridBridge) - called = [] - monkeypatch.setattr(DeepseekV41Bridge, '_convert_pre_process', - lambda self, mg, sd, pfx, tm: called.append((mg, pfx, tm)) or {'SUPER': True}) + bridge, called, markers = _pre_process_bridge(monkeypatch) mg_model = SimpleNamespace(visual=object()) - out = bridge._convert_pre_process(mg_model, {}, '', to_mcore=True) + 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_hybrid_pre_process_text_only_when_no_visual(monkeypatch): - # No vision tower on this rank (text backbone, or a non-first PP stage where ``visual=None``): - # only the word embeddings are mapped, and the GPT vision path is never entered. +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 - monkeypatch.setattr(DeepseekV41Bridge, '_convert_pre_process', - lambda *a, **k: (_ for _ in ()).throw(AssertionError('vision path must not run'))) - bridge = object.__new__(DeepseekV41HybridBridge) - calls = [] - bridge._set_word_embeddings = lambda mg, sd, tm: calls.append((mg, tm)) - bridge._remove_prefix = lambda sd, pfx: sd - bridge._add_prefix = lambda sd, pfx: sd + bridge, called, markers = _pre_process_bridge(monkeypatch) mg_model = SimpleNamespace(visual=None) - assert bridge._convert_pre_process(mg_model, {'x': 1}, '', to_mcore=True) == {} - assert calls == [(mg_model, True)] + 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(): @@ -578,7 +548,7 @@ def test_hybrid_set_word_embeddings_resolves_via_lm(): # backbone map ``embedding.word_embeddings.weight`` onto the right module. from types import SimpleNamespace - bridge = object.__new__(DeepseekV41HybridBridge) + 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)) @@ -600,19 +570,21 @@ def test_hybrid_forward_unpacks_extra_block_kwargs(monkeypatch): # 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_hybrid as hyb + 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') + 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'}) + 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 '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 @@ -621,11 +593,9 @@ 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_hybrid as hyb + 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}) - - From a34fcc49d30e1eaa60338ae4e2665ce0223a5353 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 18 Sep 2026 10:29:30 +0800 Subject: [PATCH 11/17] fix(deepseek-v41): gate per-head query RMS norm to V4 only The shared DSv4 attention applied a per-head query RMS norm unconditionally. V4.1 already normalizes the query latent via q_layernorm and feeds wq_b's output straight into RoPE, so re-normalizing every head here rescaled the attention scores. Gate it on dsv4_version == 'v4'. --- src/mcore_bridge/model/gpts/deepseek_v4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index ae49049a..7484f0a7 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -215,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: From 313c87c4de9c97987736ee3bcbba6f0dfc340aff Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 18 Sep 2026 10:29:40 +0800 Subject: [PATCH 12/17] fix(deepseek-v41): build the engram in the bridge's TransformerLayer fork The bridge's TransformerLayer replaces upstream's __init__ rather than extending it, and had not mirrored upstream's engram composition point, so self.engram was never set. The inherited _forward_attention reads self.engram on every layer, so every model assembled through this fork raised AttributeError on the first forward. Mirror the upstream engram block here, guarding submodules.engram with getattr for Megatron-Core builds that predate the composition point. --- .../model/modules/transformer_layer.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/mcore_bridge/model/modules/transformer_layer.py b/src/mcore_bridge/model/modules/transformer_layer.py index 68d35c16..7dd0796a 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( From 870e32a1f6980c63150874848b941a8fd2f44622 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 18 Sep 2026 10:29:51 +0800 Subject: [PATCH 13/17] fix(deepseek-v41): key DSpark off block_size and freeze the draft stack num_nextn_predict_layers is the standard MTP key carried by every DeepSeek V3/V4 config, so reading it as the DSpark layer count made plain V3/V4 checkpoints fail DSpark validation. Mark a DSpark checkpoint by dspark_block_size instead (V4.1-Flash and the V4 Vision experiment carry it, plain V4-Flash does not) and clear dspark_num_layers otherwise. The draft stack's own expert counts are optional: a checkpoint omitting them reuses the backbone's MoE shape, so the builder and weight mapping fall back to it. Freeze the draft stack after attaching it: it never runs in the training forward, so its parameters only ever hold a zero gradient, but Adam's decoupled weight decay would still erode the mtp.* weights every step. The bridge reads and writes param.data directly, so the checkpoint still round-trips unchanged. --- src/mcore_bridge/config/model_config.py | 18 +++-- src/mcore_bridge/model/gpts/deepseek_v41.py | 29 ++++++-- tests/test_deepseek_v41_engram.py | 75 +++++++++++++++++++++ 3 files changed, 113 insertions(+), 9 deletions(-) diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py index 4e000947..942f2cf2 100644 --- a/src/mcore_bridge/config/model_config.py +++ b/src/mcore_bridge/config/model_config.py @@ -401,15 +401,19 @@ def __post_init__(self): self.mtp_num_layers = 1 else: self.mtp_unroll_steps = self.mtp_num_layers - if self.dspark_num_layers is not None or self.dspark_block_size: + # ``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, - 'dspark_num_experts': self.dspark_num_experts, - 'dspark_router_topk': self.dspark_router_topk, } missing_dspark = [name for name, value in required_dspark.items() if value is None] if missing_dspark: @@ -424,8 +428,12 @@ def __post_init__(self): 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.') - 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.') + # 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: diff --git a/src/mcore_bridge/model/gpts/deepseek_v41.py b/src/mcore_bridge/model/gpts/deepseek_v41.py index 4c07cddc..f82002d7 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v41.py +++ b/src/mcore_bridge/model/gpts/deepseek_v41.py @@ -1201,8 +1201,12 @@ def get_dspark_layer_spec(self): 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 - dspark_config.num_moe_experts = self.config.dspark_num_experts - dspark_config.moe_router_topk = self.config.dspark_router_topk + # 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 @@ -1237,7 +1241,8 @@ def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = 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. + 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 @@ -1282,6 +1287,19 @@ def _attach_dspark(self, language_model, post_process, vp_stage: Optional[int] = 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): @@ -1698,7 +1716,10 @@ def _convert_dspark_stack(self, language_model, dspark, hf_state_dict, hf_prefix yield original_num_experts = self.config.num_moe_experts - self.config.num_moe_experts = self.config.dspark_num_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) diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index 9be4d878..b0f527f1 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -11,10 +11,12 @@ from mcore_bridge.config.parser import _convert_config from mcore_bridge.inference import DeepseekV41TextGenerationController +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, ) @@ -178,6 +180,79 @@ def test_dspark_tp_modules_construct_and_run_on_one_rank(tmp_path): assert confidence.shape == (2, 3) +def test_attach_dspark_freezes_the_draft_stack(tmp_path, 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. + """ + if dist.is_initialized() and dist.get_world_size() != 1: + pytest.skip('Single-rank DSpark construction test.') + if not dist.is_initialized(): + dist.init_process_group( + 'gloo', + init_method=f'file://{tmp_path}/dspark-freeze-init', + rank=0, + world_size=1, + ) + if not mpu.model_parallel_is_initialized(): + mpu.initialize_model_parallel(tensor_model_parallel_size=1) + + 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) From d1a8d74a505d4a8fba75bc7f4acc9d488229398c Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 18 Sep 2026 16:26:03 +0800 Subject: [PATCH 14/17] test(deepseek-v41): cover packed and CP backward --- tests/test_deepseek_v41_engram.py | 48 +++--- tests/test_deepseek_v41_packing_cp.py | 195 ++++++++++++++++++++++++ tests/test_hybrid_compat.py | 13 ++ tests/test_save_missing_weights_dsv4.py | 47 ++++-- 4 files changed, 266 insertions(+), 37 deletions(-) create mode 100644 tests/test_deepseek_v41_packing_cp.py diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index b0f527f1..458d82e0 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -111,19 +111,31 @@ def forward(self, hidden_states, runtime_gather_output): torch.testing.assert_close(logits, torch.tensor([[1., 2., 11., 12.], [5., 6., 15., 16.]])) -def test_dspark_tp_modules_construct_and_run_on_one_rank(tmp_path): +@pytest.fixture +def single_rank_model_parallel(tmp_path): if dist.is_initialized() and dist.get_world_size() != 1: - pytest.skip('Single-rank DSpark TP smoke test.') - if not dist.is_initialized(): - dist.init_process_group( - 'gloo', - init_method=f'file://{tmp_path}/dspark-dist-init', - rank=0, - world_size=1, - ) - if not mpu.model_parallel_is_initialized(): - mpu.initialize_model_parallel(tensor_model_parallel_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, @@ -180,7 +192,7 @@ def test_dspark_tp_modules_construct_and_run_on_one_rank(tmp_path): assert confidence.shape == (2, 3) -def test_attach_dspark_freezes_the_draft_stack(tmp_path, monkeypatch): +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.*`` @@ -190,18 +202,6 @@ def test_attach_dspark_freezes_the_draft_stack(tmp_path, monkeypatch): 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. """ - if dist.is_initialized() and dist.get_world_size() != 1: - pytest.skip('Single-rank DSpark construction test.') - if not dist.is_initialized(): - dist.init_process_group( - 'gloo', - init_method=f'file://{tmp_path}/dspark-freeze-init', - rank=0, - world_size=1, - ) - if not mpu.model_parallel_is_initialized(): - mpu.initialize_model_parallel(tensor_model_parallel_size=1) - config = TransformerConfig( num_layers=1, hidden_size=4, diff --git a/tests/test_deepseek_v41_packing_cp.py b/tests/test_deepseek_v41_packing_cp.py new file mode 100644 index 00000000..85843c5c --- /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 3fba665c..9551dc9d 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 06589d9f..17b102c4 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' @@ -302,6 +321,8 @@ def test_dsv4_no_duplicate_mtp_when_megatron_exports_it(): 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() From 506c622b448696ba2099331cae91075665a245a0 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Fri, 18 Sep 2026 17:46:10 +0800 Subject: [PATCH 15/17] fix(deepseek-v41): harden missing weight export --- src/mcore_bridge/bridge/gpt_bridge.py | 8 +++- src/mcore_bridge/model/gpts/deepseek_v4.py | 6 +++ tests/test_deepseek_v41_engram.py | 2 +- tests/test_save_missing_weights_dsv4.py | 43 +++++++++++----------- 4 files changed, 36 insertions(+), 23 deletions(-) diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 7ea300cf..72d9a22d 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -2163,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. @@ -2179,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/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index 7484f0a7..c58f0aa3 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -552,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/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index 458d82e0..06848b98 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -796,7 +796,7 @@ def fake_all_gather(output_list, tensor, group=None): def test_engram_cp_local_sequence_length_undoes_the_inner_sp_split(monkeypatch): - monkeypatch.setattr(engram_adapter, 'get_pg_size', lambda group: 2) + 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 diff --git a/tests/test_save_missing_weights_dsv4.py b/tests/test_save_missing_weights_dsv4.py index 17b102c4..81aecc5a 100644 --- a/tests/test_save_missing_weights_dsv4.py +++ b/tests/test_save_missing_weights_dsv4.py @@ -295,29 +295,30 @@ 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 - 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) + class RecordingSaver: - 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.*`') + def __init__(self): + self.tensors = {} + + def add_tensor(self, key, tensor): + self.tensors[key] = tensor + + with tempfile.TemporaryDirectory() as tmp_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) + + 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__': From 14e42a4e973dce62e164570287ec3333ddd8a55e Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Sat, 19 Sep 2026 15:06:26 +0800 Subject: [PATCH 16/17] refactor(deepseek-v41): drop unused DSpark inference controller The mcore_bridge.inference package (DeepseekV41TextGenerationController / DeepseekV41DynamicInferenceEngine) was scaffolding for a Megatron-native DSpark speculative-decoding path that is not wired into either SFT or RL: swift trains via the forward/backward path and serves V4.1 rollout through vLLM, so nothing imported it beyond one isolated unit test. Remove the package and its test; the training-side DSpark stack (model/modules/dspark.py) is untouched. --- src/mcore_bridge/inference/__init__.py | 5 - src/mcore_bridge/inference/dspark.py | 125 ------------------------- tests/test_deepseek_v41_engram.py | 36 ------- 3 files changed, 166 deletions(-) delete mode 100644 src/mcore_bridge/inference/__init__.py delete mode 100644 src/mcore_bridge/inference/dspark.py diff --git a/src/mcore_bridge/inference/__init__.py b/src/mcore_bridge/inference/__init__.py deleted file mode 100644 index 428bc63b..00000000 --- a/src/mcore_bridge/inference/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. - -from .dspark import DeepseekV41DynamicInferenceEngine, DeepseekV41TextGenerationController - -__all__ = ['DeepseekV41DynamicInferenceEngine', 'DeepseekV41TextGenerationController'] diff --git a/src/mcore_bridge/inference/dspark.py b/src/mcore_bridge/inference/dspark.py deleted file mode 100644 index 927245b7..00000000 --- a/src/mcore_bridge/inference/dspark.py +++ /dev/null @@ -1,125 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""DeepSeek-V4.1 DSpark adapters for Megatron's dynamic inference API.""" - -from contextlib import contextmanager - -import torch -from megatron.core.inference.communication_utils import broadcast_from_last_pipeline_stage -from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine -from megatron.core.inference.text_generation_controllers.text_generation_controller import ( - TextGenerationController, -) -from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher - - -@contextmanager -def _standard_mtp_compatibility(config): - """Satisfy legacy MTP-only constructor validation without changing model semantics.""" - sentinel = object() - original_num_layers = getattr(config, 'mtp_num_layers', sentinel) - original_repeated = getattr(config, 'mtp_use_repeated_layer', sentinel) - config.mtp_num_layers = max(getattr(config, 'mtp_num_layers', 0) or 0, 1) - config.mtp_use_repeated_layer = True - try: - yield - finally: - if original_num_layers is sentinel: - delattr(config, 'mtp_num_layers') - else: - config.mtp_num_layers = original_num_layers - if original_repeated is sentinel: - delattr(config, 'mtp_use_repeated_layer') - else: - config.mtp_use_repeated_layer = original_repeated - - -def _validate_dspark_speculation(model_config, num_speculative_tokens): - if num_speculative_tokens <= 0: - return - block_size = getattr(model_config, 'dspark_block_size', 0) - if not getattr(model_config, 'dspark_num_layers', None): - raise ValueError('DSpark speculative decoding requires dspark_num_layers.') - if not block_size or num_speculative_tokens > block_size: - raise ValueError( - f'num_speculative_tokens={num_speculative_tokens} must not exceed ' - f'dspark_block_size={block_size}.') - if getattr(model_config, 'cuda_graph_impl', None) == 'local': - raise ValueError('DSpark speculative decoding does not support local CUDA graphs yet.') - - -class DeepseekV41TextGenerationController(TextGenerationController): - """Route Megatron's standard speculative loop through the parallel DSpark draft stack.""" - - def __init__(self, inference_wrapped_model, tokenizer): - model_config = inference_wrapped_model.model.config - inference_config = inference_wrapped_model.inference_context.config - self._uses_dspark = bool(getattr(model_config, 'dspark_num_layers', None)) - if self._uses_dspark: - _validate_dspark_speculation(model_config, inference_config.num_speculative_tokens) - # PR #7224 validates speculative decoding as standard serial MTP. DSpark is - # separate, so adapt only while the upstream constructor initializes buffers. - with _standard_mtp_compatibility(model_config): - super().__init__(inference_wrapped_model, tokenizer) - self._uses_dspark = True - self.num_mtp_depths = 0 - else: - super().__init__(inference_wrapped_model, tokenizer) - - def _compute_dspark_and_sample(self): - context = self.inference_wrapped_model.inference_context - active_request_count = context.total_request_count - context.paused_request_count - speculative_tokens = None - - if self._is_last_pp_stage: - compute_dspark = getattr(self._unwrapped_model, 'compute_dspark_speculative_tokens', None) - if compute_dspark is None: - raise RuntimeError( - 'DSpark speculative decoding requires compute_dspark_speculative_tokens() ' - 'on the last pipeline stage.') - if context._nvls_dispatcher: - NVLSAllGatherVDispatcher.modify_real_token_count_for_mtp( - active_request_count * self.model_config.dspark_block_size) - speculative_tokens = compute_dspark( - next_token_ids=self._sampled_tokens_cuda[:active_request_count], - accepted_token_counts=self._accepted_token_counts_per_request[:active_request_count], - last_accepted_seq_indices=self._last_accepted_seq_indices, - num_speculative_tokens=self.num_speculative_tokens, - inference_context=context, - sample_fn=self._sample_from_logits_2d, - ) - expected_shape = (self.num_speculative_tokens, active_request_count) - if tuple(speculative_tokens.shape) != expected_shape: - raise RuntimeError( - f'DSpark returned speculative tokens with shape {tuple(speculative_tokens.shape)}; ' - f'expected {expected_shape}.') - - if self.model_is_pipeline_parallel: - speculative_tokens = broadcast_from_last_pipeline_stage( - [self.num_speculative_tokens, active_request_count], - dtype=torch.int64, - tensor=speculative_tokens, - pp_group=self.pp_group, - ) - self._sampled_mtp_tokens_cuda[ - :self.num_speculative_tokens, :active_request_count - ].copy_(speculative_tokens) - - def _compute_serial_mtp_and_sample(self): - # The upstream event loop invokes this extension point after verification and KV - # rewind. Reusing it avoids copying Megatron's large scheduling loop. - if self._uses_dspark: - return self._compute_dspark_and_sample() - return super()._compute_serial_mtp_and_sample() - - -class DeepseekV41DynamicInferenceEngine(DynamicInferenceEngine): - """Dynamic engine adapter that validates DSpark instead of serial-MTP depth.""" - - def __init__(self, controller, context): - model_config = controller.inference_wrapped_model.model.config - if getattr(model_config, 'dspark_num_layers', None): - _validate_dspark_speculation(model_config, context.config.num_speculative_tokens) - with _standard_mtp_compatibility(model_config): - super().__init__(controller, context) - else: - super().__init__(controller, context) diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index 06848b98..17de7744 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -10,7 +10,6 @@ from safetensors.torch import save_file from mcore_bridge.config.parser import _convert_config -from mcore_bridge.inference import DeepseekV41TextGenerationController from mcore_bridge.model.gpts import deepseek_v41 as deepseek_v41_module from mcore_bridge.model.gpts.deepseek_v41 import ( DeepseekV41Aligner, @@ -477,41 +476,6 @@ def forward(self, hidden_states, markov_embed): assert torch.equal(dspark_sample(logits, temperature=0), output_ids[:, 1:]) -def test_controller_routes_speculative_proposals_to_dspark_provider(): - calls = {} - - class _Model: - - def compute_dspark_speculative_tokens(self, **kwargs): - calls.update(kwargs) - return torch.tensor([[7, 8], [9, 10]]) - - context = SimpleNamespace( - total_request_count=2, - paused_request_count=0, - _nvls_dispatcher=None, - ) - controller = DeepseekV41TextGenerationController.__new__(DeepseekV41TextGenerationController) - controller.inference_wrapped_model = SimpleNamespace(inference_context=context) - controller._unwrapped_model = _Model() - controller._is_last_pp_stage = True - controller.model_is_pipeline_parallel = False - controller.model_config = SimpleNamespace(dspark_block_size=3) - controller.num_speculative_tokens = 2 - controller._sampled_tokens_cuda = torch.tensor([5, 6]) - controller._accepted_token_counts_per_request = torch.tensor([1, 0]) - controller._last_accepted_seq_indices = torch.tensor([1, 3]) - controller._sampled_mtp_tokens_cuda = torch.empty(2, 2, dtype=torch.long) - controller._sample_from_logits_2d = lambda logits: logits.argmax(dim=-1) - - controller._compute_dspark_and_sample() - - assert torch.equal(controller._sampled_mtp_tokens_cuda, torch.tensor([[7, 8], [9, 10]])) - assert calls['inference_context'] is context - assert calls['sample_fn'] is controller._sample_from_logits_2d - assert calls['num_speculative_tokens'] == 2 - - 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.') From dc0ee2a219217e6600b734ff0bbd64d6d9abb636 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Sat, 19 Sep 2026 15:55:09 +0800 Subject: [PATCH 17/17] refactor(deepseek-v41): drop unused Engram inference path The Engram module was carrying a full Megatron-native inference machinery that swift never exercises: rollout/serving goes through vLLM and training only runs the forward, so inference_context is always None in production. Remove the inference-only hashing (_static_inference_hashes, _dynamic_inference_hashes and their private _hash_windows helper), collapse _build_hash_ids to the training path, and delete the two dead standalone helpers adapt_deepseek_v41_layer_specs (the loader builds the D-layer spec inline) and allow_engram_inference (bypass for a native-inference guard that is never entered). forward now keeps its inference_context parameter -- the bridge wrapper still passes it positionally -- but fails loud instead of silently running inference, so the training forward is unchanged. Drop the now-unused contextmanager import and the three tests that covered the removed code. --- src/mcore_bridge/model/modules/engram.py | 140 ++--------------------- tests/test_deepseek_v41_engram.py | 66 ----------- 2 files changed, 11 insertions(+), 195 deletions(-) diff --git a/src/mcore_bridge/model/modules/engram.py b/src/mcore_bridge/model/modules/engram.py index 9fedfa41..8fc9b850 100644 --- a/src/mcore_bridge/model/modules/engram.py +++ b/src/mcore_bridge/model/modules/engram.py @@ -2,7 +2,6 @@ """DeepSeek-V4.1 adapters for NVIDIA Megatron-Core's optional Engram modules.""" import dataclasses -from contextlib import contextmanager import torch from torch import Tensor @@ -248,104 +247,21 @@ def _mask_excluded_tokens(self, input_ids: Tensor) -> tuple[Tensor, Tensor]: live = live & (input_ids != token_id) return torch.where(live, input_ids, input_ids.new_full((), -1)), live - def _hash_windows(self, token_windows: Tensor) -> Tensor: - return _hash_token_windows( - token_windows, + 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, - invalid_token_id=-1, - ) - - def _static_inference_hashes(self, input_ids: Tensor, context) -> tuple[Tensor, Tensor]: - batch_size, sequence_length = input_ids.shape - shape = (context.max_batch_size, context.max_sequence_length) - cache = getattr(context, 'engram_token_cache', None) - if cache is None or cache.device != input_ids.device or cache.shape != shape: - cache = input_ids.new_full(shape, self.engram_config.boundary_token_id) - context.engram_token_cache = cache - - batch_start = context.batch_size_offset - batch_end = batch_start + batch_size - sequence_start = context.sequence_len_offset - sequence_end = sequence_start + sequence_length - if batch_end > shape[0] or sequence_end > shape[1]: - raise ValueError('Engram inference token cache is too small for the current batch/chunk.') - masked_ids, live = self._mask_excluded_tokens(input_ids) - cache[batch_start:batch_end, sequence_start:sequence_end] = masked_ids - shifts = torch.arange( - self.engram_config.max_ngram_order, device=input_ids.device, dtype=torch.long) - positions = torch.arange( - sequence_start, sequence_end, device=input_ids.device, dtype=torch.long).unsqueeze(-1) - shifts - gather_positions = positions.clamp_min(0).reshape(1, -1).expand(batch_size, -1) - windows = cache[batch_start:batch_end].gather(1, gather_positions).view( - batch_size, sequence_length, self.engram_config.max_ngram_order) - windows = torch.where( - positions.unsqueeze(0) >= 0, - windows, - windows.new_full((), self.engram_config.boundary_token_id), + self.engram_config.variant_spec.resets_windows_at_boundary_token, + cu_seqlens=cu_seqlens, ) - return self._hash_windows(windows), live - - def _dynamic_inference_hashes(self, input_ids: Tensor, context) -> tuple[Tensor, Tensor]: - if input_ids.shape[0] != 1: - raise ValueError('Dynamic Engram inference expects flattened input_ids with batch size 1.') - shape = (context.max_requests, context.max_sequence_length) - cache = getattr(context, 'engram_token_cache', None) - if cache is None or cache.device != input_ids.device or cache.shape != shape: - cache = input_ids.new_full(shape, self.engram_config.boundary_token_id) - context.engram_token_cache = cache - - total_tokens = input_ids.shape[1] - active_tokens = min(int(context.active_token_count), total_tokens) - live = torch.zeros_like(input_ids, dtype=torch.bool) - hashes = input_ids.new_zeros((1, total_tokens, self.engram_config.num_tables), dtype=torch.long) - if active_tokens == 0: - return hashes, live - request_indices = context.gpu_view.token_to_request_idx[:active_tokens].long() - token_positions = context.gpu_view.token_to_position_in_request[:active_tokens].long() - if request_indices.min() < 0 or request_indices.max() >= shape[0]: - raise ValueError('Dynamic Engram inference received an out-of-range request index.') - if token_positions.min() < 0 or token_positions.max() >= shape[1]: - raise ValueError('Dynamic Engram inference received an out-of-range token position.') - masked_ids, active_live = self._mask_excluded_tokens(input_ids[:, :active_tokens]) - cache[request_indices, token_positions] = masked_ids.squeeze(0) - shifts = torch.arange( - self.engram_config.max_ngram_order, device=input_ids.device, dtype=torch.long) - positions = token_positions.unsqueeze(-1) - shifts - windows = cache[request_indices.unsqueeze(-1).expand_as(positions), positions.clamp_min(0)] - windows = torch.where( - positions >= 0, - windows, - windows.new_full((), self.engram_config.boundary_token_id), - ) - hashes[:, :active_tokens] = self._hash_windows(windows.unsqueeze(0)) - live[:, :active_tokens] = active_live return hashes, live - def _build_hash_ids(self, input_ids: Tensor, inference_context=None, - cu_seqlens: Tensor | None = None) -> tuple[Tensor, Tensor]: - masked_ids, live = self._mask_excluded_tokens(input_ids) - if inference_context is None: - 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 - if inference_context.is_static_batching(): - return self._static_inference_hashes(input_ids, inference_context) - return self._dynamic_inference_hashes(input_ids, inference_context) - 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] @@ -388,6 +304,10 @@ def _slice_for_context_parallel(self, hashes: Tensor) -> Tensor: 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}.') @@ -395,8 +315,6 @@ def forward(self, hidden_states: Tensor, input_ids: Tensor, inference_context=No 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 - if context_parallel and inference_context is not None: - raise ValueError('Engram inference does not support context parallelism.') cu_seqlens = None if packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', None) == 'thd': @@ -422,7 +340,7 @@ def forward(self, hidden_states: Tensor, input_ids: Tensor, inference_context=No 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, inference_context, cu_seqlens) + 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) @@ -482,39 +400,3 @@ class DeepseekV41HyperConnectionTransformerLayer( else: DeepseekV41TransformerLayer = None DeepseekV41HyperConnectionTransformerLayer = None - - -def adapt_deepseek_v41_layer_specs(transformer_layer_spec, engram_config): - """Attach the V4.1 Engram module and inference-aware layer subclasses.""" - if not has_native_engram(): - raise RuntimeError('The installed Megatron-Core does not provide Engram.') - from megatron.core.transformer.spec_utils import ModuleSpec - - engram_spec = ModuleSpec(module=DeepseekV41Engram, params={'engram_config': engram_config}) - for layer_spec in transformer_layer_spec.layer_specs: - if layer_spec.module is HyperConnectionTransformerLayer: - layer_spec.module = DeepseekV41HyperConnectionTransformerLayer - elif layer_spec.module is TransformerLayer: - layer_spec.module = DeepseekV41TransformerLayer - if not hasattr(layer_spec.submodules, 'engram'): - raise RuntimeError( - 'The installed Engram extension does not expose TransformerLayerSubmodules.engram.') - layer_spec.submodules.engram = engram_spec - return transformer_layer_spec - - -@contextmanager -def allow_engram_inference(model_config, input_ids, extra_block_kwargs): - """Bypass PR #7231's inference guard while preserving input_ids propagation.""" - if not getattr(model_config, 'engram_enabled', False): - yield extra_block_kwargs - return - if input_ids is None: - raise ValueError('Engram requires input token IDs on every pipeline stage.') - block_kwargs = dict(extra_block_kwargs or {}) - block_kwargs['input_ids'] = input_ids - model_config.engram_enabled = False - try: - yield block_kwargs - finally: - model_config.engram_enabled = True diff --git a/tests/test_deepseek_v41_engram.py b/tests/test_deepseek_v41_engram.py index 17de7744..84356a07 100644 --- a/tests/test_deepseek_v41_engram.py +++ b/tests/test_deepseek_v41_engram.py @@ -601,38 +601,6 @@ def test_engram_hash_blocks_suffixes_after_excluded_token(): assert torch.equal(hashes, torch.tensor([[[5, 5]]])) -def test_engram_static_inference_cache_matches_full_sequence_hashing(): - module = engram_adapter.DeepseekV41Engram.__new__(engram_adapter.DeepseekV41Engram) - torch.nn.Module.__init__(module) - module.engram_config = SimpleNamespace( - excluded_token_ids=(99, ), - max_ngram_order=3, - num_hash_heads=1, - hash_boundary_token_id=0, - boundary_token_id=0, - num_tables=2, - variant_spec=SimpleNamespace(resets_windows_at_boundary_token=False), - ) - module.tokenizer_remap = None - module.hash_multipliers = torch.tensor([11, 13, 15]) - module.table_sizes = torch.tensor([997, 991]) - full_hashes, full_live = module._build_hash_ids(torch.tensor([[1, 2, 3]])) - context = SimpleNamespace( - max_batch_size=1, - max_sequence_length=8, - batch_size_offset=0, - sequence_len_offset=0, - is_static_batching=lambda: True, - ) - - prefill_hashes, prefill_live = module._build_hash_ids(torch.tensor([[1, 2]]), context) - context.sequence_len_offset = 2 - decode_hashes, decode_live = module._build_hash_ids(torch.tensor([[3]]), context) - - assert torch.equal(torch.cat((prefill_hashes, decode_hashes), dim=1), full_hashes) - assert torch.equal(torch.cat((prefill_live, decode_live), dim=1), full_live) - - def _ngram_hash_kwargs(): return dict( tokenizer_remap=None, @@ -769,40 +737,6 @@ def test_engram_cp_local_sequence_length_undoes_the_inner_sp_split(monkeypatch): assert module._cp_local_sequence_length(hidden_states) == 8 -def test_engram_layer_spec_uses_bridge_owned_module(): - if not engram_adapter.has_native_engram(): - pytest.skip('The PR #7224 baseline intentionally has no Engram extension.') - from megatron.core.transformer.spec_utils import ModuleSpec - from megatron.core.transformer.transformer_layer import ( - HyperConnectionTransformerLayer, - TransformerLayerSubmodules, - ) - - layer_spec = ModuleSpec( - module=HyperConnectionTransformerLayer, - submodules=TransformerLayerSubmodules(), - ) - block_spec = SimpleNamespace(layer_specs=[layer_spec]) - config = SimpleNamespace(layer_ids=(1, )) - - engram_adapter.adapt_deepseek_v41_layer_specs(block_spec, config) - - assert layer_spec.module is engram_adapter.DeepseekV41HyperConnectionTransformerLayer - assert layer_spec.submodules.engram.module is engram_adapter.DeepseekV41Engram - - -def test_allow_engram_inference_preserves_input_ids_and_restores_flag(): - config = SimpleNamespace(engram_enabled=True) - input_ids = torch.tensor([[1, 2]]) - - with engram_adapter.allow_engram_inference(config, input_ids, {'marker': 1}) as kwargs: - assert not config.engram_enabled - assert kwargs['marker'] == 1 - assert kwargs['input_ids'] is input_ids - - assert config.engram_enabled - - class _Table: def __init__(self, global_rows, row_start, row_end, dim):