From 92ec17f2c3db4541dc667dcba6818645a8655c72 Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Tue, 4 Aug 2026 16:59:19 +0800 Subject: [PATCH 1/2] lint --- src/mcore_bridge/bridge/gpt_bridge.py | 38 +++ tests/test_save_missing_weights.py | 134 +++++++++++ tests/test_save_missing_weights_dsv4.py | 307 ++++++++++++++++++++++++ 3 files changed, 479 insertions(+) create mode 100644 tests/test_save_missing_weights.py create mode 100644 tests/test_save_missing_weights_dsv4.py diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 1acbfa6..41caf7c 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math +import os import re import torch import torch.distributed as dist @@ -61,6 +62,7 @@ def __init__(self, config: ModelConfig): self._peft_format = False self._adapter_name = 'default' self._is_saving = False + self._source_model_dir = None self.model_type = config.hf_model_type self.llm_model_type = config.llm_model_type self.is_multimodal = config.is_multimodal @@ -1915,6 +1917,8 @@ def load_weights( """ self._peft_format = peft_format self._adapter_name = adapter_name + if not peft_format: + self._source_model_dir = hf_model_dir mg_models = unwrap_model(mg_models) self._disable_tqdm = False self._is_saving = False @@ -1997,6 +2001,7 @@ def save_weights( adapter_name: str = 'default', converter: Optional[Callable] = None, max_shard_size: str = '5GB', + save_missing_weights: Union[bool, str] = False, ) -> None: """Save Megatron model checkpoint in safetensors (HuggingFace) format. @@ -2013,10 +2018,15 @@ def save_weights( adapter_name: Name of the adapter for PEFT models. Defaults to 'default'. converter: Used to perform key-value conversion on the newly exported state_dict. max_shard_size: Maximum size of a single storage file, default is '5GB'. + save_missing_weights: Whether to copy tensors that exist in the source checkpoint but are + absent from the exported weights, such as submodules Megatron does not support. Pass a + path to specify the source checkpoint, otherwise the one recorded by `load_weights` is + used. Ignored when `peft_format` is True. """ gc_collect() saver = StreamingSafetensorSaver(save_dir=output_dir, max_shard_size=max_shard_size, peft_format=peft_format) mg_models = unwrap_model(mg_models) + saved_keys = set() for k, v in self.export_weights( mg_models, target_device='cpu', @@ -2028,9 +2038,37 @@ def save_weights( disable_tqdm=False, _is_saving=True): saver.add_tensor(k, v) + saved_keys.add(k) + if save_missing_weights and not peft_format: + source_model_dir = save_missing_weights if isinstance(save_missing_weights, str) else None + self._save_missing_weights(saver, saved_keys, source_model_dir) saver.finalize() dist.barrier() # Ensure all weights are saved completely + def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None: + """Copy tensors present in the source checkpoint but absent from the exported ones. + + Megatron only materializes the modules it knows about, so weights of unsupported + submodules (for instance the DSpark stages under `mtp.*`) would silently vanish + from the exported checkpoint. Restoring them verbatim keeps the saved model + functionally complete. + """ + source_model_dir = source_model_dir or self._source_model_dir + if source_model_dir is None or not is_master(): + return + if not os.path.isdir(source_model_dir): + logger.warning(f'Source model dir does not exist, skip restoring missing weights: {source_model_dir}') + return + with SafetensorLazyLoader(source_model_dir) as loader: + state_dict = loader.get_state_dict() + missing_keys = sorted(set(state_dict.keys()) - saved_keys) + if not missing_keys: + return + logger.info(f'Restoring {len(missing_keys)} weights from the source checkpoint ' + f'that were not exported by Megatron, e.g. {missing_keys[:3]}.') + for key in missing_keys: + saver.add_tensor(key, state_dict[key].load()) + @contextmanager def _patch_hf_initialize_weight(self): diff --git a/tests/test_save_missing_weights.py b/tests/test_save_missing_weights.py new file mode 100644 index 0000000..1fe4e71 --- /dev/null +++ b/tests/test_save_missing_weights.py @@ -0,0 +1,134 @@ +"""Verify that `save_weights(save_missing_weights=True)` restores tensors that +Megatron never materializes. + +The scenario mirrors DeepSeek-V4-Flash-0731, whose `mtp.*` (DSpark) weights are +not supported by Megatron: without the restore step they silently disappear from +the exported checkpoint. +""" +import os + +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', '29513') + +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_ID = 'Qwen/Qwen2-0.5B-Instruct' +MODEL_TYPE = 'qwen2' # a copied dir loses the model id, so type/template cannot be inferred +TEMPLATE = 'qwen' +# Stand-ins for weights of a submodule Megatron does not know about. +EXTRA_WEIGHTS = { + 'mtp.0.main_proj.weight': torch.randn(8, 24, dtype=torch.bfloat16), + 'mtp.0.main_norm.weight': torch.randn(8, dtype=torch.bfloat16), + 'mtp.2.markov_head.markov_w1.weight': torch.randn(16, 4, dtype=torch.bfloat16), +} + + +def _build_checkpoint_with_extra_weights(dst_dir: str) -> str: + """Copy the source model and inject weights that the bridge cannot consume.""" + from swift import safe_snapshot_download + src_dir = safe_snapshot_download(MODEL_ID) + shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True) + + index_path = os.path.join(dst_dir, 'model.safetensors.index.json') + if os.path.exists(index_path): + with open(index_path) as f: + index = json.load(f) + shard_name = sorted(set(index['weight_map'].values()))[0] + else: + index = None + shard_name = 'model.safetensors' + + shard_path = os.path.join(dst_dir, shard_name) + state_dict = load_file(shard_path) + state_dict.update(EXTRA_WEIGHTS) + save_file(state_dict, shard_path, metadata={'format': 'pt'}) + + if index is not None: + for key in EXTRA_WEIGHTS: + index['weight_map'][key] = shard_name + with open(index_path, 'w') as f: + json.dump(index, f) + return dst_dir + + +def _export(model_dir: str, output_dir: str, save_missing_weights: bool): + """Round-trip HF -> mcore -> HF through the real bridge.""" + from swift.megatron import MegatronExportArguments, megatron_export_main + mcore_dir = f'{output_dir}-mcore' + megatron_export_main( + MegatronExportArguments( + model=model_dir, + model_type=MODEL_TYPE, + template=TEMPLATE, + to_mcore=True, + output_dir=mcore_dir, + exist_ok=True, + torch_dtype='bfloat16', + )) + megatron_export_main( + MegatronExportArguments( + mcore_model=mcore_dir, + # `model` is only used as the source of the weights Megatron cannot export. + model=model_dir, + model_type=MODEL_TYPE, + template=TEMPLATE, + to_hf=True, + output_dir=output_dir, + exist_ok=True, + torch_dtype='bfloat16', + save_missing_weights=save_missing_weights, + )) + return output_dir + + +def _load_exported(output_dir: str): + index_path = os.path.join(output_dir, 'model.safetensors.index.json') + state_dict = {} + if os.path.exists(index_path): + with open(index_path) as f: + shards = sorted(set(json.load(f)['weight_map'].values())) + else: + shards = ['model.safetensors'] + for shard in shards: + state_dict.update(load_file(os.path.join(output_dir, shard))) + return state_dict + + +def test_save_missing_weights(): + """The injected weights must reappear byte-for-byte when the flag is on.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = _build_checkpoint_with_extra_weights(os.path.join(tmp_dir, 'src')) + output_dir = _export(model_dir, os.path.join(tmp_dir, 'restored'), save_missing_weights=True) + state_dict = _load_exported(output_dir) + + for key, expected in EXTRA_WEIGHTS.items(): + assert key in state_dict, f'{key} was not restored' + assert torch.equal(state_dict[key], expected), f'{key} was altered during restore' + # The regular weights must still be exported by Megatron, not copied blindly. + assert 'model.layers.0.self_attn.q_proj.weight' in state_dict + + +def test_save_missing_weights_disabled(): + """With the flag off the behaviour is unchanged: the extra weights are dropped.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = _build_checkpoint_with_extra_weights(os.path.join(tmp_dir, 'src')) + output_dir = _export(model_dir, os.path.join(tmp_dir, 'dropped'), save_missing_weights=False) + state_dict = _load_exported(output_dir) + + for key in EXTRA_WEIGHTS: + assert key not in state_dict, f'{key} leaked into the export' + assert 'model.layers.0.self_attn.q_proj.weight' in state_dict + + +if __name__ == '__main__': + test_save_missing_weights() + test_save_missing_weights_disabled() diff --git a/tests/test_save_missing_weights_dsv4.py b/tests/test_save_missing_weights_dsv4.py new file mode 100644 index 0000000..06589d9 --- /dev/null +++ b/tests/test_save_missing_weights_dsv4.py @@ -0,0 +1,307 @@ +"""Verify `save_missing_weights` against a synthetic DeepSeek-V4 checkpoint. + +The real DeepSeek-V4-Flash-0731 is far too large to test with, so this builds a +4-layer model from its config and fills every tensor with random values. The +`mtp.*` (DSpark) weights follow the 3-stage layout of the real checkpoint, +including the stage-specific extras (`main_proj`, `confidence_head`, +`markov_head`) that Megatron has no module for. + +Two properties are checked: + * with the flag on, the unsupported `mtp.*` weights survive the round-trip; + * `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 os + +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' +# Only the tokenizer files are read from here; the weights are generated locally. +# GIT_LFS_SKIP_SMUDGE=1 git clone https://www.modelscope.cn/deepseek-ai/DeepSeek-V4-Flash-0731.git +REFERENCE_MODEL_ID = 'deepseek-ai/DeepSeek-V4-Flash-0731' +TOKENIZER_FILES = ['tokenizer.json', 'tokenizer_config.json', 'generation_config.json'] + +NUM_LAYERS = 4 +NUM_MTP_STAGES = 3 +HIDDEN = 256 +VOCAB = 512 +N_EXPERTS = 4 +MOE_INTERMEDIATE = 128 +HC_MULT = 2 +Q_LORA_RANK = 64 +O_LORA_RANK = 32 +O_GROUPS = 2 +HEAD_DIM = 32 +NUM_HEADS = 4 +QK_ROPE_HEAD_DIM = 16 +MARKOV_RANK = 32 +DSPARK_TARGET_LAYERS = [1, 2, 3] + +# Shapes derived the same way the real checkpoint does (verified against 0731). +QK_HEAD_DIM = HEAD_DIM # wq_b rows are NUM_HEADS * head_dim +KV_DIM = HEAD_DIM # wkv rows; the rope part is not stored separately here +# wo_a is [o_groups * o_lora_rank, per_group_dim]; wo_b maps that back to hidden. +O_GROUP_DIM = NUM_HEADS * HEAD_DIM // O_GROUPS +O_A_ROWS = O_GROUPS * O_LORA_RANK +HC_STREAM = HC_MULT * HIDDEN # hc_*_fn columns +HC_ROWS = HC_MULT * (HC_MULT + 2) # hc_attn/ffn rows: width + depth + scale terms +HC_ALPHAS = 3 # the bridge reads alpha_pre / alpha_post / alpha_res from hc_*_scale + +# All layers stay dense: a ratio of 0 avoids the CSA compressor weights, which are +# irrelevant to weight persistence and would only add noise to this test. +COMPRESS_RATIOS = [0] * NUM_LAYERS + + +def _config() -> dict: + """A miniature version of the DeepSeek-V4-Flash-0731 config.""" + return { + 'architectures': ['DeepseekV4ForCausalLM'], + 'attention_bias': False, + 'attention_dropout': 0.0, + 'bos_token_id': 0, + 'eos_token_id': 1, + 'hc_eps': 1e-06, + 'hc_mult': HC_MULT, + 'hc_sinkhorn_iters': 20, + 'head_dim': HEAD_DIM, + 'hidden_act': 'silu', + 'hidden_size': HIDDEN, + 'index_head_dim': 32, + 'index_n_heads': 4, + 'index_topk': 32, + 'initializer_range': 0.02, + 'max_position_embeddings': 4096, + 'model_type': 'deepseek_v4', + 'moe_intermediate_size': MOE_INTERMEDIATE, + 'n_routed_experts': N_EXPERTS, + 'n_shared_experts': 1, + 'norm_topk_prob': True, + 'num_attention_heads': NUM_HEADS, + 'num_experts_per_tok': 2, + 'num_hidden_layers': NUM_LAYERS, + 'num_hash_layers': 0, + 'num_key_value_heads': 1, + 'num_nextn_predict_layers': NUM_MTP_STAGES, + 'o_groups': O_GROUPS, + 'o_lora_rank': O_LORA_RANK, + 'q_lora_rank': Q_LORA_RANK, + 'qk_rope_head_dim': QK_ROPE_HEAD_DIM, + 'rms_norm_eps': 1e-06, + 'rope_theta': 10000, + 'routed_scaling_factor': 1.5, + 'scoring_func': 'sqrtsoftplus', + 'sliding_window': 32, + 'swiglu_limit': 10.0, + 'tie_word_embeddings': False, + 'topk_method': 'noaux_tc', + 'torch_dtype': 'bfloat16', + 'use_cache': True, + 'vocab_size': VOCAB, + 'compress_rope_theta': 160000, + 'compress_ratios': COMPRESS_RATIOS, + 'dspark_block_size': 5, + 'dspark_noise_token_id': VOCAB - 1, + 'dspark_target_layer_ids': DSPARK_TARGET_LAYERS, + 'dspark_markov_rank': MARKOV_RANK, + } + + +def _rand(*shape) -> torch.Tensor: + return torch.randn(*shape, dtype=torch.bfloat16) * 0.02 + + +def _attn_and_ffn_weights(prefix: str) -> dict: + """Weights shared by every transformer block, main trunk and DSpark alike.""" + sd = { + f'{prefix}attn_norm.weight': _rand(HIDDEN), + f'{prefix}ffn_norm.weight': _rand(HIDDEN), + f'{prefix}attn.q_norm.weight': _rand(Q_LORA_RANK), + f'{prefix}attn.kv_norm.weight': _rand(KV_DIM), + f'{prefix}attn.attn_sink': _rand(NUM_HEADS), + f'{prefix}attn.wq_a.weight': _rand(Q_LORA_RANK, HIDDEN), + f'{prefix}attn.wq_b.weight': _rand(NUM_HEADS * QK_HEAD_DIM, Q_LORA_RANK), + f'{prefix}attn.wkv.weight': _rand(KV_DIM, HIDDEN), + f'{prefix}attn.wo_a.weight': _rand(O_A_ROWS, O_GROUP_DIM), + f'{prefix}attn.wo_b.weight': _rand(HIDDEN, O_A_ROWS), + f'{prefix}ffn.gate.weight': _rand(N_EXPERTS, HIDDEN), + f'{prefix}ffn.gate.bias': _rand(N_EXPERTS), + } + for name, shape in [('w1', (MOE_INTERMEDIATE, HIDDEN)), ('w2', (HIDDEN, MOE_INTERMEDIATE)), + ('w3', (MOE_INTERMEDIATE, HIDDEN))]: + sd[f'{prefix}ffn.shared_experts.{name}.weight'] = _rand(*shape) + for e in range(N_EXPERTS): + sd[f'{prefix}ffn.experts.{e}.{name}.weight'] = _rand(*shape) + # Hyper-connection parameters, present on every block of the real model. + # `base` holds hc_mult entries per residual stream, `scale` one per stream. + for tag in ['attn', 'ffn']: + sd[f'{prefix}hc_{tag}_base'] = _rand(HC_ROWS) + sd[f'{prefix}hc_{tag}_fn'] = _rand(HC_ROWS, HC_STREAM) + sd[f'{prefix}hc_{tag}_scale'] = _rand(HC_ALPHAS) + return sd + + +def _hc_head_weights(prefix: str) -> dict: + """Output-side hyper-connection: one row per stream, a single scale.""" + return { + f'{prefix}hc_head_base': _rand(HC_MULT), + f'{prefix}hc_head_fn': _rand(HC_MULT, HC_STREAM), + f'{prefix}hc_head_scale': _rand(1), + } + + +def _mtp_weights() -> dict: + """The 3 asymmetric DSpark stages, mirroring the real key layout.""" + sd = {} + for stage in range(NUM_MTP_STAGES): + sd.update(_attn_and_ffn_weights(f'mtp.{stage}.')) + if stage == 0: + # Stage 0 consumes the concatenated hidden states of the target layers. + sd['mtp.0.main_norm.weight'] = _rand(HIDDEN) + sd['mtp.0.main_proj.weight'] = _rand(HIDDEN, HIDDEN * len(DSPARK_TARGET_LAYERS)) + if stage == NUM_MTP_STAGES - 1: + # The last stage owns the output-side heads. + sd['mtp.2.norm.weight'] = _rand(HIDDEN) + sd['mtp.2.confidence_head.proj.weight'] = _rand(1, HIDDEN + QK_ROPE_HEAD_DIM) + sd['mtp.2.markov_head.markov_w1.weight'] = _rand(VOCAB, MARKOV_RANK) + sd['mtp.2.markov_head.markov_w2.weight'] = _rand(VOCAB, MARKOV_RANK) + sd.update(_hc_head_weights('mtp.2.')) + return sd + + +def _build_fake_checkpoint(dst_dir: str) -> str: + os.makedirs(dst_dir, exist_ok=True) + config = _config() + with open(os.path.join(dst_dir, 'config.json'), 'w') as f: + json.dump(config, f, indent=2) + # Reuse the real tokenizer; `download_model=False` keeps the multi-GB weights out. + from swift import safe_snapshot_download + ref_dir = safe_snapshot_download(REFERENCE_MODEL_ID, download_model=False) + for fname in TOKENIZER_FILES: + src = os.path.join(ref_dir, fname) + if os.path.exists(src): + shutil.copy(src, os.path.join(dst_dir, fname)) + + state_dict = { + 'embed.weight': _rand(VOCAB, HIDDEN), + 'head.weight': _rand(VOCAB, HIDDEN), + 'norm.weight': _rand(HIDDEN), + } + state_dict.update(_hc_head_weights('')) + for layer in range(NUM_LAYERS): + state_dict.update(_attn_and_ffn_weights(f'layers.{layer}.')) + state_dict.update(_mtp_weights()) + + save_file(state_dict, os.path.join(dst_dir, 'model.safetensors'), metadata={'format': 'pt'}) + return dst_dir + + +def _load_exported(output_dir: str) -> dict: + index_path = os.path.join(output_dir, 'model.safetensors.index.json') + if os.path.exists(index_path): + with open(index_path) as f: + shards = sorted(set(json.load(f)['weight_map'].values())) + else: + shards = ['model.safetensors'] + state_dict = {} + for shard in shards: + state_dict.update(load_file(os.path.join(output_dir, shard))) + return state_dict + + +def _trunk_keys(state_dict) -> list: + """Keys of layer 0 of the main trunk, whatever prefix the bridge chose.""" + return [k for k in state_dict if 'layers.0.' in k and not k.startswith('mtp.')] + + +def _export(model_dir: str, output_dir: str, save_missing_weights: bool, mtp_num_layers=None): + """Round-trip HF -> mcore -> HF through the real bridge.""" + from swift.megatron import MegatronExportArguments, megatron_export_main + mcore_dir = f'{output_dir}-mcore' + common = dict(model_type=MODEL_TYPE, template=TEMPLATE, exist_ok=True, torch_dtype='bfloat16') + if mtp_num_layers is not None: + common['mtp_num_layers'] = mtp_num_layers + megatron_export_main(MegatronExportArguments(model=model_dir, to_mcore=True, output_dir=mcore_dir, **common)) + megatron_export_main( + MegatronExportArguments( + mcore_model=mcore_dir, + # `model` is only used as the source of the weights Megatron cannot export. + model=model_dir, + to_hf=True, + output_dir=output_dir, + save_missing_weights=save_missing_weights, + **common)) + return output_dir + + +def test_dsv4_mtp_weights_restored(): + """Without Megatron MTP, every `mtp.*` tensor must come back byte-for-byte.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = _build_fake_checkpoint(os.path.join(tmp_dir, 'src')) + source = load_file(os.path.join(model_dir, 'model.safetensors')) + mtp_keys = {k for k in source if k.startswith('mtp.')} + assert len(mtp_keys) > 100, f'the fake checkpoint should have many mtp keys, got {len(mtp_keys)}' + + output_dir = _export(model_dir, os.path.join(tmp_dir, 'restored'), save_missing_weights=True) + exported = _load_exported(output_dir) + + for key in sorted(mtp_keys): + assert key in exported, f'{key} was not restored' + assert torch.equal(exported[key], source[key]), f'{key} was altered during restore' + # The trunk must still be produced by Megatron rather than copied over. + assert _trunk_keys(exported), 'trunk weights missing from export' + + +def test_dsv4_mtp_weights_dropped_by_default(): + """With the flag off the `mtp.*` weights are lost, as they are today.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = _build_fake_checkpoint(os.path.join(tmp_dir, 'src')) + output_dir = _export(model_dir, os.path.join(tmp_dir, 'dropped'), save_missing_weights=False) + exported = _load_exported(output_dir) + + assert not [k for k in exported if k.startswith('mtp.')], 'mtp weights leaked into the export' + assert _trunk_keys(exported), 'trunk weights missing from export' + + +def test_dsv4_no_duplicate_mtp_when_megatron_exports_it(): + """Guard against storing the same DSpark parameters under two naming schemes. + + 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) + + 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.*`') + + +if __name__ == '__main__': + test_dsv4_mtp_weights_restored() + test_dsv4_mtp_weights_dropped_by_default() + test_dsv4_no_duplicate_mtp_when_megatron_exports_it() From 91396180a3d3267a6daf8d2bab6402b943a9d04a Mon Sep 17 00:00:00 2001 From: z0o0ey Date: Sun, 9 Aug 2026 10:42:18 +0800 Subject: [PATCH 2/2] fix --- src/mcore_bridge/model/gpts/deepseek_v4.py | 57 ++++++++++++++++++---- src/mcore_bridge/patcher.py | 19 ++++---- 2 files changed, 57 insertions(+), 19 deletions(-) diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index 03d0641..2d8e625 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -60,6 +60,38 @@ def _patch_YarnRotaryEmbedding(config): delattr(config, attr) +def _apply_mla_rope(t, freqs, *, config, cu_seqlens, cp_group, inverse=False): + """Apply DSv4's MLA RoPE to a tensor whose frequencies are already expanded per token. + + `GPTModel` pre-indexes the rotary table by `position_ids`, so `freqs` is row-aligned with + `t`: row i holds the frequency of token i. That holds for every layout DSv4 supports -- + unpacked, packed (thd), and packed CP: swift CP-splits `position_ids` with the same + partition mode as the hidden states, so the pre-indexed frequencies come out rank-local + while still carrying absolute positions. The multiply is therefore purely elementwise and + needs no `cu_seqlens`-based segment alignment, under any `cp_partition_mode`. + + Enforcing that invariant here matters: when it does not hold, the generic + `apply_rotary_pos_emb` thd path re-derives positions from `cu_seqlens` assuming a *zigzag* + CP split, which is wrong for DSv4's contiguous split and would corrupt positions silently. + Asserting row alignment turns any future layout change into an immediate, explicit failure. + """ + assert freqs.shape[0] == t.shape[0], ( + f'DSv4 MLA RoPE expects per-token frequencies row-aligned with the input, got ' + f'freqs.shape[0]={freqs.shape[0]} vs tokens={t.shape[0]}. `GPTModel` must pre-index the ' + 'rotary table by `position_ids` (requires `apply_rope_fusion=False`), and under CP the ' + '`position_ids` must be split with the same partition mode as the hidden states.') + return apply_rotary_pos_emb( + t, + freqs, + config=config, + cu_seqlens=cu_seqlens, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + inverse=inverse, + ) + + class DSv4HybridSelfAttention(McoreDSv4HybridSelfAttention): def __init__(self, config, *args, **kwargs): @@ -153,6 +185,15 @@ def qkv_up_proj_and_rope_apply(q_compressed, When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; otherwise, they maintain the unpacked shape [s, b, ...]. In subsequent code comments, we uniformly use [num_tokens, ...] to denote [s, b, ...] or [t, ...] for two cases. + + RoPE frequency layout: `GPTModel` pre-indexes the rotary table by `position_ids` + (see gpt_model.py, the `not apply_rope_fusion` branch), so `rotary_pos_emb` here is + already expanded per token -- row i belongs to token i -- rather than being a + position->frequency lookup table. Every RoPE call below is therefore an elementwise + multiply; see `_apply_mla_rope` for why that invariant is asserted. Under CP the + frequencies arrive rank-local because `position_ids` is split alongside the hidden + states, which is also why the boundary rows carry their own frequencies + (`boundary_rotary_pos_emb`) instead of being re-derived from positions. """ # q_compressed: [num_tokens, q_lora_rank] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] @@ -166,6 +207,8 @@ def qkv_up_proj_and_rope_apply(q_compressed, if boundary_kv_compressed is not None: boundary_rows = boundary_kv_compressed.shape[0] kv_projection_input = torch.cat([boundary_kv_compressed, kv_compressed], dim=0) + # The boundary rows precede this rank's block, so their frequencies must precede + # too -- keeping kv_rotary_pos_emb row-aligned with kv_projection_input. kv_rotary_pos_emb = torch.cat([boundary_rotary_pos_emb, rotary_pos_emb], dim=0) else: kv_projection_input = kv_compressed @@ -182,14 +225,12 @@ def qkv_up_proj_and_rope_apply(q_compressed, # RoPE and query (shared for wkv and latent) # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] - q_pos_emb = apply_rotary_pos_emb( + q_pos_emb = _apply_mla_rope( q_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q, cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, - mla_output_remove_interleaving=True, ) # query: [num_tokens, n, (qk_head_dim + v_head_dim)] query = torch.cat([q_no_pe, q_pos_emb], dim=-1) @@ -197,14 +238,12 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] - k_pos_emb = apply_rotary_pos_emb( + k_pos_emb = _apply_mla_rope( k_pos_emb, kv_rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, - mla_output_remove_interleaving=True, ) # Single head: key = value = [num_tokens, 1, v_head_dim] @@ -384,15 +423,13 @@ def forward( rot_part_in = rot_part.squeeze(1) else: rot_part_in = rot_part - rot_part_out = apply_rotary_pos_emb( + rot_part_out = _apply_mla_rope( rot_part_in, rotary_pos_emb, - self.config, + config=self.config, cu_seqlens=cu_seqlens_kv, cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, inverse=True, - mla_output_remove_interleaving=True, ) if packed_seq: rot_part = rot_part_out.unsqueeze(1) diff --git a/src/mcore_bridge/patcher.py b/src/mcore_bridge/patcher.py index abf8ad0..df5347e 100644 --- a/src/mcore_bridge/patcher.py +++ b/src/mcore_bridge/patcher.py @@ -197,16 +197,17 @@ def forward(self, position_ids, mrope_section: List[int], mrope_interleaved: boo def _apply_rotary_pos_emb_thd(t: torch.Tensor, cu_seqlens: torch.Tensor, freqs: torch.Tensor, *args, **kwargs) -> torch.Tensor: cp_group = kwargs.pop('cp_group', None) - if cp_group is not None: - cp_size = cp_group.size() - else: - cp_size = mpu.get_context_parallel_world_size() + if cp_group is None: cp_group = mpu.get_context_parallel_group() - cu_seqlens_for_batched = cu_seqlens // cp_size - use_batched_rope = (freqs.dim() >= 1 and freqs.shape[0] == cu_seqlens_for_batched[-1]).item() - # The determination of mla_output_remove_interleaving: a quick solution for identifying deepseek_v4 - # (TODO: refactor) - if not use_batched_rope and not kwargs.get('mla_output_remove_interleaving', False): + # The fast path below reinterprets the thd tensor as bshd and multiplies it by `freqs` + # directly, which is only valid when `freqs` is already expanded per token. Compare against + # the token count of `t` itself rather than deriving one from `cu_seqlens`: callers that + # change the sequence length -- e.g. deepseek_v4's CSA compressor, which shortens kv and + # passes the compressed cu_seqlens -- must fall back to the upstream thd kernel, which + # aligns freqs per segment. For an ordinary packed batch both are equal, so the fast path + # is preserved. + use_batched_rope = freqs.dim() >= 1 and freqs.shape[0] == t.shape[0] + if not use_batched_rope: logger.warning_once('Using non-batched RoPE, which may affect performance.') return _origin_apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, *args, cp_group=cp_group, **kwargs)