From d557d10d09bd31b50cc161266144ea1f823ba5e3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 12 May 2026 17:05:05 -0700 Subject: [PATCH 1/8] [model, ckpt] feat: Add on-the-fly dequantization utilities for HF checkpoints (#3778) Signed-off-by: Chen Cui --- .../models/conversion/quantization_utils.py | 229 ++++++++++++++++++ .../models/deepseek/deepseek_v3_bridge.py | 45 +++- .../bridge/models/gpt_oss/gpt_oss_bridge.py | 59 +---- .../models/kimi_vl/kimi_k25_vl_bridge.py | 8 +- .../models/minimax_m2/minimax_m2_bridge.py | 23 +- .../models/ministral3/ministral3_bridge.py | 8 +- .../models/deepseek/test_deepseek_bridges.py | 100 +++++++- .../models/test_quantization_utils.py | 80 ++++++ 8 files changed, 463 insertions(+), 89 deletions(-) create mode 100644 src/megatron/bridge/models/conversion/quantization_utils.py create mode 100644 tests/unit_tests/models/test_quantization_utils.py diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py new file mode 100644 index 0000000000..7ad1e5802c --- /dev/null +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch + + +FP8_BLOCK_SIZE = 128 +FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + + +def is_fp8_tensor(tensor: torch.Tensor) -> bool: + """Return whether *tensor* uses one of PyTorch's FP8 dtypes.""" + return tensor.dtype in FP8_DTYPES + + +def dequantize_fp8_blockwise( + weight: torch.Tensor, + scale_inv: torch.Tensor, + *, + block_size: int = FP8_BLOCK_SIZE, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize FP8 weights with one scale per 2D block. + + DeepSeek-V3 and MiniMax-M2 store linear weights as FP8 tensors with a + separate ``*_scale_inv`` tensor. Each scale applies to one 128x128 weight + block by default. + """ + M, N = weight.shape + w = weight.float() + out = torch.empty_like(w) + sM, sN = scale_inv.shape + for bi in range(sM): + for bj in range(sN): + r0, r1 = bi * block_size, min((bi + 1) * block_size, M) + c0, c1 = bj * block_size, min((bj + 1) * block_size, N) + out[r0:r1, c0:c1] = w[r0:r1, c0:c1] * scale_inv[bi, bj] + return out.to(dtype) + + +def maybe_dequantize_fp8_blockwise( + weight: torch.Tensor, + scale_inv: torch.Tensor | None = None, + *, + block_size: int = FP8_BLOCK_SIZE, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize FP8 block-scaled weights, falling back to a plain cast.""" + if not is_fp8_tensor(weight): + return weight + if weight.ndim == 2 and scale_inv is not None: + return dequantize_fp8_blockwise(weight, scale_inv, block_size=block_size, dtype=dtype) + return weight.float().to(dtype) + + +def maybe_dequantize_fp8( + weight: torch.Tensor, + scale_inv: torch.Tensor | None = None, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Dequantize FP8 weights with a scalar or broadcastable scale tensor.""" + if not is_fp8_tensor(weight): + return weight + if scale_inv is None: + return weight.to(dtype) + return weight.to(dtype) * scale_inv.to(dtype) + + +def dequantize_mxfp4( + blocks: torch.Tensor, + scales: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, + rows_per_chunk: int = 32768 * 1024, +) -> torch.Tensor: + """Dequantize GPT-OSS MXFP4 block/scales tensors.""" + assert blocks.shape[:-1] == scales.shape, f"{blocks.shape=} does not match {scales.shape=}" + fp4_values = [ + +0.0, + +0.5, + +1.0, + +1.5, + +2.0, + +3.0, + +4.0, + +6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ] + scales = scales.to(torch.int32) - 127 + lut = torch.tensor(fp4_values, dtype=dtype, device=blocks.device) + + *prefix_shape, G, B = blocks.shape + rows_total = math.prod(prefix_shape) * G + + blocks = blocks.reshape(rows_total, B) + scales = scales.reshape(rows_total, 1) + + out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device) + + for r0 in range(0, rows_total, rows_per_chunk): + r1 = min(r0 + rows_per_chunk, rows_total) + + blk = blocks[r0:r1] + exp = scales[r0:r1] + + idx_lo = (blk & 0x0F).to(torch.long) + idx_hi = (blk >> 4).to(torch.long) + + sub = out[r0:r1] + sub[:, 0::2] = lut[idx_lo] + sub[:, 1::2] = lut[idx_hi] + + torch.ldexp(sub, exp, out=sub) + del idx_lo, idx_hi, blk, exp + + return out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2) + + +def dequantize_int4( + weight_packed: torch.Tensor, + weight_scale: torch.Tensor, + weight_shape: torch.Tensor, + group_size: int = 32, + device: str | torch.device | None = None, +) -> torch.Tensor: + """Dequantize Kimi INT4 packed weights to bfloat16. + + The checkpoint stores eight offset-binary INT4 values in each int32 slot and + carries per-group scales beside the packed tensor. + """ + del weight_shape, group_size + + local_out, local_packed_in = weight_packed.shape + local_in = local_packed_in * 8 + + target_device = weight_packed.device if device is None else torch.device(device) + use_cuda = target_device.type == "cuda" and torch.cuda.is_available() + + if use_cuda: + weight_packed = weight_packed.to(target_device) + weight_scale = weight_scale.to(target_device) + + shifts = torch.arange(8, device=weight_packed.device) * 4 + + packed_unsqueezed = weight_packed.unsqueeze(-1) + unpacked = ((packed_unsqueezed >> shifts) & 0xF).float() + unpacked = unpacked.reshape(local_out, local_in) + + unpacked = unpacked - 8 + + scale = weight_scale.float() + if scale.ndim == 1: + local_num_groups = scale.numel() // local_out + scale = scale.view(local_out, local_num_groups) + else: + scale = scale.view(local_out, -1) + + local_num_groups = scale.shape[1] + elements_per_group = local_in // local_num_groups + + scale_expanded = scale.repeat_interleave(elements_per_group, dim=1) + + if scale_expanded.shape[1] < local_in: + scale_expanded = torch.nn.functional.pad( + scale_expanded, (0, local_in - scale_expanded.shape[1]), value=scale_expanded[:, -1:].mean() + ) + scale_expanded = scale_expanded[:, :local_in] + result = unpacked * scale_expanded + + return result.to(torch.bfloat16) + + +def quantize_to_int4( + weight: torch.Tensor, + group_size: int = 32, + scale_dtype: torch.dtype = torch.bfloat16, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize bfloat16/float16 weights to Kimi INT4 packed format.""" + out_features, in_features = weight.shape + weight_shape = torch.tensor([out_features, in_features], dtype=torch.int32) + + w = weight.float() + + num_groups = (in_features + group_size - 1) // group_size + w_grouped = w.view(out_features, num_groups, -1) + + group_max = w_grouped.abs().amax(dim=-1) + scale = group_max / 7.0 + scale = scale.clamp(min=1e-10) + + scale_expanded = scale.unsqueeze(-1).expand_as(w_grouped) + w_q = (w_grouped / scale_expanded).round().clamp(-8, 7) + + w_q = w_q.view(out_features, -1)[:, :in_features] + w_q = (w_q + 8).to(torch.uint8) + + assert in_features % 8 == 0, f"in_features must be divisible by 8, got {in_features}" + + w_q_grouped = w_q.view(out_features, in_features // 8, 8).to(torch.int32) + + packed = torch.zeros(out_features, in_features // 8, dtype=torch.int32, device=weight.device) + for i in range(8): + packed |= (w_q_grouped[:, :, i] & 0xF) << (i * 4) + + weight_packed = packed + weight_scale = scale.to(scale_dtype) + + return weight_packed, weight_scale, weight_shape diff --git a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py index 30efe1c2da..1333aee9d2 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v3_bridge.py @@ -13,12 +13,13 @@ # limitations under the License. from functools import partial -from typing import Dict, Mapping +from typing import Dict, Mapping, Union import torch from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.bridge.models.conversion import quantization_utils from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask from megatron.bridge.models.conversion.param_mapping import AutoMapping @@ -36,6 +37,12 @@ HAVE_TE = False +__all__ = ["DeepSeekV3Bridge", "_dequant_fp8_blockwise"] + + +_dequant_fp8_blockwise = quantization_utils.dequantize_fp8_blockwise + + @MegatronModelBridge.register_bridge( source="DeepseekV3ForCausalLM", target=GPTModel, @@ -130,6 +137,42 @@ def mapping_registry(self) -> MegatronMappingRegistry: ) return MegatronMappingRegistry(*mapping_list) + def maybe_modify_loaded_hf_weight( + self, + hf_param: Union[str, dict[str, str]], + hf_state_dict: Mapping[str, torch.Tensor], + ) -> Union[torch.Tensor, dict[str, torch.Tensor]]: + """Load HF weights and dequantize FP8 tensors on the fly. + + DeepSeek-V3 ships linear weights as ``float8_e4m3fn`` with per-block scale + factors stored in ``_scale_inv`` (128x128 blocks). The true bf16 weight is:: + + w_bf16 = fp8_weight.float() * scale_inv_block + + Without this override the bridge would do a bare ``.to(bf16)`` cast in + ``ColumnParallelMapping.hf_to_megatron`` (param_mapping.py:905), discarding the + per-block scales — the resulting model produces random-looking logits. + """ + hf_weights = super().maybe_modify_loaded_hf_weight(hf_param, hf_state_dict) + + if isinstance(hf_weights, dict): + # Compound params (QKV / GatedMLP): dequantize each component individually. + return { + key: self._maybe_dequantize_fp8(tensor, hf_param[key], hf_state_dict) + for key, tensor in hf_weights.items() + } + return self._maybe_dequantize_fp8(hf_weights, hf_param, hf_state_dict) + + @staticmethod + def _maybe_dequantize_fp8( + weight: torch.Tensor, + param_name: str, + hf_state_dict: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Dequantize ``weight`` if it is stored as FP8 with a matching ``*_scale_inv``.""" + scale_key = param_name + "_scale_inv" + return quantization_utils.maybe_dequantize_fp8_blockwise(weight, hf_state_dict.get(scale_key)) + def maybe_modify_converted_hf_weight( self, task: WeightConversionTask, diff --git a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py index 29f38eae93..38cb9564ea 100644 --- a/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py +++ b/src/megatron/bridge/models/gpt_oss/gpt_oss_bridge.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math from typing import Dict, Mapping, Optional, Tuple, Union import torch @@ -27,6 +26,7 @@ QKVMapping, _align_expert_weight_to_shape, ) +from megatron.bridge.models.conversion.quantization_utils import dequantize_mxfp4 as _dequantize_mxfp4 from megatron.bridge.models.conversion.utils import get_module_and_param_from_name from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM @@ -302,60 +302,3 @@ def megatron_to_hf(self, megatron_weights: torch.Tensor, megatron_module: nn.Mod if len(megatron_weights.shape) == 2: megatron_weights = megatron_weights.transpose(0, 1) return super().megatron_to_hf(megatron_weights.contiguous(), megatron_module) - - -def _dequantize_mxfp4( - blocks: torch.Tensor, - scales: torch.Tensor, - *, - dtype: torch.dtype = torch.bfloat16, - rows_per_chunk: int = 32768 * 1024, -) -> torch.Tensor: - assert blocks.shape[:-1] == scales.shape, f"{blocks.shape=} does not match {scales.shape=}" - FP4_VALUES = [ - +0.0, - +0.5, - +1.0, - +1.5, - +2.0, - +3.0, - +4.0, - +6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ] - scales = scales.to(torch.int32) - 127 - lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device) - - *prefix_shape, G, B = blocks.shape - rows_total = math.prod(prefix_shape) * G - - blocks = blocks.reshape(rows_total, B) - scales = scales.reshape(rows_total, 1) - - out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device) - - for r0 in range(0, rows_total, rows_per_chunk): - r1 = min(r0 + rows_per_chunk, rows_total) - - blk = blocks[r0:r1] - exp = scales[r0:r1] - - # nibble indices -> int64 - idx_lo = (blk & 0x0F).to(torch.long) - idx_hi = (blk >> 4).to(torch.long) - - sub = out[r0:r1] - sub[:, 0::2] = lut[idx_lo] - sub[:, 1::2] = lut[idx_hi] - - torch.ldexp(sub, exp, out=sub) - del idx_lo, idx_hi, blk, exp - - return out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2) diff --git a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py index 7c73cc1abd..6087fbbba6 100644 --- a/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +++ b/src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py @@ -25,14 +25,14 @@ GatedMLPMapping, ReplicatedMapping, ) +from megatron.bridge.models.conversion.quantization_utils import ( + dequantize_int4, + quantize_to_int4, +) from megatron.bridge.models.deepseek.common import get_common_mapping_list from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM from megatron.bridge.models.kimi_vl.kimi_k25_vl_provider import KimiK25VLModelProvider from megatron.bridge.models.kimi_vl.modeling_kimi_k25_vl import KimiK25VLModel -from megatron.bridge.models.kimi_vl.utils import ( - dequantize_int4, - quantize_to_int4, -) try: diff --git a/src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py b/src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py index 05677f1c71..53068aaf4f 100644 --- a/src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py +++ b/src/megatron/bridge/models/minimax_m2/minimax_m2_bridge.py @@ -19,6 +19,7 @@ import torch.nn as nn from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.bridge.models.conversion import quantization_utils from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.param_mapping import ( @@ -30,22 +31,10 @@ from megatron.bridge.models.minimax_m2.minimax_m2_provider import minimax_m2_layer_spec -_FP8_BLOCK_SIZE = 128 +__all__ = ["MiniMaxM2Bridge", "_FullDimQKNormMapping", "_dequant_fp8_blockwise"] -def _dequant_fp8_blockwise(weight: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: - """Block-wise FP8 dequantization: out = fp8_val * scale_inv per 128x128 block.""" - M, N = weight.shape - B = _FP8_BLOCK_SIZE - w = weight.float() - out = torch.empty_like(w) - sM, sN = scale_inv.shape - for bi in range(sM): - for bj in range(sN): - r0, r1 = bi * B, min((bi + 1) * B, M) - c0, c1 = bj * B, min((bj + 1) * B, N) - out[r0:r1, c0:c1] = w[r0:r1, c0:c1] * scale_inv[bi, bj] - return out.to(torch.bfloat16) +_dequant_fp8_blockwise = quantization_utils.dequantize_fp8_blockwise class _FullDimQKNormMapping(MegatronParamMapping[torch.Tensor]): @@ -194,12 +183,8 @@ def maybe_modify_loaded_hf_weight( def _load_and_dequant(self, key: str, hf_state_dict: Mapping[str, torch.Tensor]) -> torch.Tensor: w = hf_state_dict[key] - if w.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2): - return w sinv_key = key + "_scale_inv" - if w.ndim == 2 and sinv_key in hf_state_dict: - return _dequant_fp8_blockwise(w, hf_state_dict[sinv_key]) - return w.float().to(torch.bfloat16) + return quantization_utils.maybe_dequantize_fp8_blockwise(w, hf_state_dict.get(sinv_key)) def mapping_registry(self) -> MegatronMappingRegistry: param_mappings = { diff --git a/src/megatron/bridge/models/ministral3/ministral3_bridge.py b/src/megatron/bridge/models/ministral3/ministral3_bridge.py index 031905e33c..271678c06a 100644 --- a/src/megatron/bridge/models/ministral3/ministral3_bridge.py +++ b/src/megatron/bridge/models/ministral3/ministral3_bridge.py @@ -44,6 +44,7 @@ QKVMapping, ReplicatedMapping, ) +from megatron.bridge.models.conversion.quantization_utils import maybe_dequantize_fp8 from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM from megatron.bridge.models.ministral3.ministral3_provider import Ministral3ModelProvider @@ -211,13 +212,8 @@ def _maybe_dequantize_fp8( w_bf16 = weight.to(bfloat16) * scale_inv """ - if weight.dtype != torch.float8_e4m3fn: - return weight scale_key = param_name + "_scale_inv" - if scale_key not in hf_state_dict: - return weight.to(torch.bfloat16) - scale_inv = hf_state_dict[scale_key].to(torch.bfloat16) - return weight.to(torch.bfloat16) * scale_inv + return maybe_dequantize_fp8(weight, hf_state_dict.get(scale_key)) # Register the bridge if Mistral3ForConditionalGeneration is available diff --git a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py index 7a8ed6dbdd..c5f81a274c 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_bridges.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_bridges.py @@ -24,7 +24,7 @@ from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask from megatron.bridge.models.deepseek.deepseek_v2_bridge import DeepSeekV2Bridge -from megatron.bridge.models.deepseek.deepseek_v3_bridge import DeepSeekV3Bridge +from megatron.bridge.models.deepseek.deepseek_v3_bridge import DeepSeekV3Bridge, _dequant_fp8_blockwise from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM from megatron.bridge.models.mla_provider import MLAModelProvider @@ -373,3 +373,101 @@ def test_export_skips_inv_freq_when_not_expected(self, mock_pretrained_v3): inv_key = "model.layers.0.self_attn.rotary_emb.inv_freq" assert inv_key not in result + + +class TestDeepSeekV3DequantFP8Blockwise: + """Unit tests for the standalone _dequant_fp8_blockwise helper.""" + + def test_identity_scale_inv(self): + """With scale_inv=1 the output equals the input cast to bfloat16.""" + weight = torch.ones(128, 128, dtype=torch.float8_e4m3fn) + scale_inv = torch.ones(1, 1) + result = _dequant_fp8_blockwise(weight, scale_inv) + + assert result.dtype == torch.bfloat16 + assert result.shape == (128, 128) + assert torch.all(result == 1.0) + + def test_scale_inv_applied_per_block(self): + """scale_inv value is multiplied block-wise across all 128x128 blocks.""" + weight = torch.ones(256, 256, dtype=torch.float8_e4m3fn) + scale_inv = torch.full((2, 2), 2.0) + result = _dequant_fp8_blockwise(weight, scale_inv) + + assert result.dtype == torch.bfloat16 + assert torch.all(result == 2.0) + + def test_distinct_scale_per_block(self): + """Each 128x128 block uses its own scale value.""" + weight = torch.ones(256, 256, dtype=torch.float8_e4m3fn) + scale_inv = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + result = _dequant_fp8_blockwise(weight, scale_inv).float() + + assert torch.all(result[:128, :128] == 1.0) + assert torch.all(result[:128, 128:] == 2.0) + assert torch.all(result[128:, :128] == 3.0) + assert torch.all(result[128:, 128:] == 4.0) + + def test_non_multiple_dim(self): + """Trailing partial block (dim not divisible by 128) is handled.""" + weight = torch.zeros(100, 70, dtype=torch.float8_e4m3fn) + scale_inv = torch.ones(1, 1) + result = _dequant_fp8_blockwise(weight, scale_inv) + + assert result.shape == (100, 70) + assert result.dtype == torch.bfloat16 + + +class TestDeepSeekV3MaybeModifyLoadedHFWeight: + """Unit tests for DeepSeekV3Bridge.maybe_modify_loaded_hf_weight (FP8 dequant on import).""" + + def test_passthrough_bfloat16(self): + """Non-FP8 weights are returned unchanged.""" + bridge = DeepSeekV3Bridge() + w = torch.randn(4, 4, dtype=torch.bfloat16) + state = {"layer.weight": w} + result = bridge.maybe_modify_loaded_hf_weight("layer.weight", state) + assert result is w + + def test_passthrough_float32(self): + """Non-FP8 (float32) weights pass through unchanged.""" + bridge = DeepSeekV3Bridge() + w = torch.randn(4, 4, dtype=torch.float32) + state = {"layer.weight": w} + result = bridge.maybe_modify_loaded_hf_weight("layer.weight", state) + assert result is w + + def test_dequants_fp8_when_scale_inv_present(self): + """FP8 weight with a ``*_scale_inv`` key is block-wise dequantized.""" + bridge = DeepSeekV3Bridge() + w = torch.ones(128, 128, dtype=torch.float8_e4m3fn) + sinv = torch.full((1, 1), 3.0) + state = {"layer.weight": w, "layer.weight_scale_inv": sinv} + result = bridge.maybe_modify_loaded_hf_weight("layer.weight", state) + + assert result.dtype == torch.bfloat16 + assert torch.all(result == 3.0) + + def test_fp8_without_scale_inv_cast_to_bfloat16(self): + """FP8 weight without ``*_scale_inv`` falls back to a plain float cast.""" + bridge = DeepSeekV3Bridge() + w = torch.ones(4, 4, dtype=torch.float8_e4m3fn) + state = {"layer.weight": w} + result = bridge.maybe_modify_loaded_hf_weight("layer.weight", state) + + assert result.dtype == torch.bfloat16 + + def test_dict_hf_param_each_key_processed(self): + """Compound (dict) hf_param dequantizes every sub-key independently.""" + bridge = DeepSeekV3Bridge() + w1 = torch.ones(128, 128, dtype=torch.float8_e4m3fn) + w2 = torch.ones(64, 64, dtype=torch.bfloat16) + sinv = torch.full((1, 1), 2.0) + state = {"key1": w1, "key1_scale_inv": sinv, "key2": w2} + result = bridge.maybe_modify_loaded_hf_weight({"gate": "key1", "up": "key2"}, state) + + assert isinstance(result, dict) + assert result["gate"].dtype == torch.bfloat16 + assert torch.all(result["gate"] == 2.0) + # Non-FP8 entries pass through unchanged. + assert result["up"] is w2 diff --git a/tests/unit_tests/models/test_quantization_utils.py b/tests/unit_tests/models/test_quantization_utils.py new file mode 100644 index 0000000000..b47bbe067b --- /dev/null +++ b/tests/unit_tests/models/test_quantization_utils.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from megatron.bridge.models.conversion.quantization_utils import ( + dequantize_fp8_blockwise, + dequantize_int4, + dequantize_mxfp4, + maybe_dequantize_fp8, + maybe_dequantize_fp8_blockwise, + quantize_to_int4, +) + + +def test_dequantize_fp8_blockwise_applies_distinct_scales(): + weight = torch.ones(256, 256, dtype=torch.float8_e4m3fn) + scale_inv = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + result = dequantize_fp8_blockwise(weight, scale_inv).float() + + assert result.dtype == torch.float32 + assert torch.all(result[:128, :128] == 1.0) + assert torch.all(result[:128, 128:] == 2.0) + assert torch.all(result[128:, :128] == 3.0) + assert torch.all(result[128:, 128:] == 4.0) + + +def test_maybe_dequantize_fp8_blockwise_passthrough_and_fallback_cast(): + bf16_weight = torch.ones(4, 4, dtype=torch.bfloat16) + assert maybe_dequantize_fp8_blockwise(bf16_weight) is bf16_weight + + fp8_weight = torch.ones(4, 4, dtype=torch.float8_e4m3fn) + result = maybe_dequantize_fp8_blockwise(fp8_weight) + + assert result.dtype == torch.bfloat16 + assert torch.all(result == 1.0) + + +def test_maybe_dequantize_fp8_applies_broadcastable_scale(): + fp8_weight = torch.ones(2, 2, dtype=torch.float8_e4m3fn) + scale_inv = torch.tensor([2.0]) + + result = maybe_dequantize_fp8(fp8_weight, scale_inv) + + assert result.dtype == torch.bfloat16 + assert torch.all(result == 2.0) + + +def test_dequantize_mxfp4_uses_low_then_high_nibbles(): + blocks = torch.tensor([[[0x21]]], dtype=torch.uint8) + scales = torch.tensor([[127]], dtype=torch.uint8) + + result = dequantize_mxfp4(blocks, scales, dtype=torch.float32) + + assert result.shape == (1, 2) + assert torch.equal(result, torch.tensor([[0.5, 1.0]])) + + +def test_quantize_dequantize_int4_preserves_shape_and_dtype(): + weight = torch.linspace(-1.0, 1.0, steps=32).view(1, 32).to(torch.bfloat16) + + packed, scale, shape = quantize_to_int4(weight) + result = dequantize_int4(packed, scale, shape) + + assert packed.shape == (1, 4) + assert shape.tolist() == [1, 32] + assert result.shape == weight.shape + assert result.dtype == torch.bfloat16 From 233036821c50ffb824c17a7577abbaac515098a8 Mon Sep 17 00:00:00 2001 From: weijiac0619 Date: Tue, 19 May 2026 22:28:00 -0700 Subject: [PATCH 2/8] DeepSeek V4 Bridge (#3562) Signed-off-by: weijiac Signed-off-by: weijiac Signed-off-by: chcui Signed-off-by: Chen Cui Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: chcui --- docs/models/llm/deepseek-v4.md | 34 + docs/models/llm/index.md | 1 + .../convert_checkpoints_multi_gpu.py | 56 +- .../hf_to_megatron_generate_text.py | 17 + examples/models/deepseek_v4/README.md | 59 ++ examples/models/deepseek_v4/conversion.sh | 107 ++ examples/models/deepseek_v4/inference.sh | 69 ++ .../bridge/models/deepseek/__init__.py | 2 + .../models/deepseek/deepseek_v4_bridge.py | 952 ++++++++++++++++++ .../bridge/models/hf_pretrained/state.py | 38 +- .../bridge/training/utils/config_utils.py | 4 + .../deepseek/test_deepseek_v4_conversion.py | 264 +++++ .../deepseek/test_deepseek_v4_bridge.py | 162 +++ 13 files changed, 1749 insertions(+), 16 deletions(-) create mode 100644 docs/models/llm/deepseek-v4.md create mode 100644 examples/models/deepseek_v4/README.md create mode 100755 examples/models/deepseek_v4/conversion.sh create mode 100755 examples/models/deepseek_v4/inference.sh create mode 100644 src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py create mode 100644 tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py create mode 100644 tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py diff --git a/docs/models/llm/deepseek-v4.md b/docs/models/llm/deepseek-v4.md new file mode 100644 index 0000000000..88b7dd8676 --- /dev/null +++ b/docs/models/llm/deepseek-v4.md @@ -0,0 +1,34 @@ +# DeepSeek V4 + +[DeepSeek-V4](https://github.com/deepseek-ai/DeepSeek-V4) is the next-generation Mixture-of-Experts language model from DeepSeek-AI. It extends the V3 design with **Hyper-Connections (mHC)** for multi-stream residual mixing, **Compressed Sparse Attention (CSA)** with a learned token-importance indexer (DSA), **hash-routed MoE layers** for the first few decoder blocks, and a refined **Multi-Token Prediction (MTP)** head with separate `e_proj` / `h_proj` projections. + +DeepSeek V4 models are supported via the Bridge system with auto-detected configuration and weight mapping. + +## Model Architecture Features + +- **Hybrid Attention (DSv4HybridSelfAttention)**: Per-layer mix of dense MLA and Compressed Sparse Attention selected by `compress_ratios` +- **Compressed Sparse Attention (CSA)** with **DSA Indexer**: Top-k token selection over windowed keys; `index_n_heads`, `index_head_dim`, `index_topk` control the indexer +- **Hyper-Connections (mHC)**: 4-stream residual mixing per layer (`hc_mult = 4`) with sinkhorn-iterated attention; per-MTP-layer `hc_head_*` learns output contraction +- **Hash-Routed MoE**: First few decoder layers use a deterministic vocab → expert mapping (`tid2eid`) instead of softmax routing +- **Multi-Token Prediction (MTP)**: One MTP layer with separate `e_proj` and `h_proj` projections (post-MCore #4518) +- **YaRN RoPE**: `rotary_scaling_factor=16`, `original_max_position_embeddings=65536`; `mscale=mscale_all_dim=1.0` for V4 +- **Sigmoid Gating with Expert Bias**: `noaux_tc` load balancing, `sqrtsoftplus` scoring, expert bias enabled +- **`o_groups` Output Projection**: `o_lora_rank` low-rank output projection split into `o_groups` parallel groups + +## Examples, Parallelism, and Limitations + +For checkpoint conversion and inference scripts, recommended parallelism settings, and current known limitations, see the [DeepSeek V4 examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4). + +## Hugging Face Model Cards & References + +### Hugging Face Model Cards +- DeepSeek-V4-Flash: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash +- DeepSeek-V4-Flash-Base: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-Base +- DeepSeek-V4-Pro: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro +- DeepSeek-V4-Pro-Base: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro-Base + +### Additional Resources +- GitHub Repository: https://github.com/deepseek-ai/DeepSeek-V4 + +## Related Docs +- DeepSeek V4 examples: [examples/models/deepseek_v4/README.md](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4) diff --git a/docs/models/llm/index.md b/docs/models/llm/index.md index 2e3de95454..ac69ca076e 100644 --- a/docs/models/llm/index.md +++ b/docs/models/llm/index.md @@ -7,6 +7,7 @@ This section documents Large Language Models supported by Megatron Bridge, with deepseek-v2.md deepseek-v3.md +deepseek-v4.md gemma2.md gemma3.md glm45.md diff --git a/examples/conversion/convert_checkpoints_multi_gpu.py b/examples/conversion/convert_checkpoints_multi_gpu.py index 9a64c0d24e..377d7a0987 100644 --- a/examples/conversion/convert_checkpoints_multi_gpu.py +++ b/examples/conversion/convert_checkpoints_multi_gpu.py @@ -53,6 +53,7 @@ """ import argparse +import datetime import os import sys @@ -84,6 +85,20 @@ def _check_distributed(): sys.exit(1) +def _ensure_distributed_initialized(timeout_minutes: int | None): + _check_distributed() + if timeout_minutes is None: + return + if torch.distributed.is_initialized(): + return + + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + torch.distributed.init_process_group( + "nccl", + timeout=datetime.timedelta(minutes=timeout_minutes), + ) + + @torchrun_main def import_hf_to_megatron( hf_model: str, @@ -94,9 +109,10 @@ def import_hf_to_megatron( etp: int = 1, torch_dtype: str = "bfloat16", trust_remote_code: bool = False, + distributed_timeout_minutes: int | None = None, ) -> None: """Import a HuggingFace model and save it as a distributed Megatron checkpoint.""" - _check_distributed() + _ensure_distributed_initialized(distributed_timeout_minutes) dtype = _parse_dtype(torch_dtype) print_rank_0(f"Importing: {hf_model} -> {megatron_path}") @@ -115,6 +131,15 @@ def import_hf_to_megatron( model_provider.expert_tensor_parallel_size = etp model_provider.pipeline_dtype = dtype model_provider.params_dtype = dtype + # Auto-generate pipeline layout for models that need it (e.g. DSv4 hash MoE) + if pp > 1 and hasattr(bridge._model_bridge, "generate_pipeline_layout"): + hf_config = bridge.hf_pretrained.config + num_layers = hf_config.num_hidden_layers + mtp = getattr(hf_config, "num_nextn_predict_layers", 0) or 0 + model_provider.pipeline_model_parallel_layout = bridge._model_bridge.generate_pipeline_layout( + num_layers, pp, mtp + ) + print_rank_0(f" Auto-generated pipeline layout for PP={pp} ({num_layers} layers, {mtp} MTP)") model_provider.finalize() model_provider.initialize_model_parallel(seed=0) @@ -151,9 +176,10 @@ def export_megatron_to_hf( show_progress: bool = True, distributed_save: bool = False, save_every_n_ranks: int = 1, + distributed_timeout_minutes: int | None = None, ) -> None: """Export a distributed Megatron checkpoint to HuggingFace format.""" - _check_distributed() + _ensure_distributed_initialized(distributed_timeout_minutes) dtype = _parse_dtype(torch_dtype) print_rank_0(f"Exporting: {megatron_path} -> {hf_path}") @@ -173,6 +199,23 @@ def export_megatron_to_hf( model_provider.expert_tensor_parallel_size = etp model_provider.pipeline_dtype = dtype model_provider.params_dtype = dtype + # For PP > 1 export, read pipeline layout from checkpoint + if pp > 1: + from pathlib import Path + + import yaml + + ckpt_path = Path(megatron_path) + for candidate in [ckpt_path, *ckpt_path.glob("iter_*")]: + rc = candidate / "run_config.yaml" + if rc.exists(): + with open(rc) as f: + cfg = yaml.safe_load(f) + saved_layout = cfg.get("model", {}).get("pipeline_model_parallel_layout") + if isinstance(saved_layout, list): + model_provider.pipeline_model_parallel_layout = saved_layout + print_rank_0(f" Read pipeline layout from checkpoint ({len(saved_layout)} stages)") + break model_provider.finalize() model_provider.initialize_model_parallel(seed=0) @@ -218,6 +261,12 @@ def _add_common_args(parser: argparse.ArgumentParser) -> None: help="Model precision (default: bfloat16)", ) parser.add_argument("--trust-remote-code", action="store_true", help="Allow custom model code execution") + parser.add_argument( + "--distributed-timeout-minutes", + type=int, + default=None, + help="Initialize the distributed process group with this timeout before model setup", + ) def main(): @@ -255,7 +304,6 @@ def main(): default=1, help="Only every N-th rank writes files (reduces I/O, only with --distributed-save)", ) - args = parser.parse_args() if not args.command: @@ -272,6 +320,7 @@ def main(): etp=args.etp, torch_dtype=args.torch_dtype, trust_remote_code=args.trust_remote_code, + distributed_timeout_minutes=args.distributed_timeout_minutes, ) elif args.command == "export": export_megatron_to_hf( @@ -288,6 +337,7 @@ def main(): show_progress=not args.no_progress, distributed_save=args.distributed_save, save_every_n_ranks=args.save_every_n_ranks, + distributed_timeout_minutes=args.distributed_timeout_minutes, ) diff --git a/examples/conversion/hf_to_megatron_generate_text.py b/examples/conversion/hf_to_megatron_generate_text.py index dd9aee56f6..b06e86e129 100644 --- a/examples/conversion/hf_to_megatron_generate_text.py +++ b/examples/conversion/hf_to_megatron_generate_text.py @@ -128,6 +128,23 @@ def main(args) -> None: model_provider.expert_tensor_parallel_size = etp model_provider.pipeline_dtype = torch.bfloat16 + # Read pipeline layout from checkpoint for PP > 1 + if pp > 1: + from pathlib import Path + + import yaml + + ckpt_path = Path(args.megatron_model_path) + for candidate in [ckpt_path, *ckpt_path.glob("iter_*")]: + rc = candidate / "run_config.yaml" + if rc.exists(): + with open(rc) as f: + cfg = yaml.safe_load(f) + saved_layout = cfg.get("model", {}).get("pipeline_model_parallel_layout") + if isinstance(saved_layout, list): + model_provider.pipeline_model_parallel_layout = saved_layout + break + # Once all overrides are set, finalize the model provider to ensure the post initialization logic is run model_provider.finalize() model_provider.initialize_model_parallel(seed=0) diff --git a/examples/models/deepseek_v4/README.md b/examples/models/deepseek_v4/README.md new file mode 100644 index 0000000000..ccbfc4f65f --- /dev/null +++ b/examples/models/deepseek_v4/README.md @@ -0,0 +1,59 @@ +# DeepSeek V4 + +End-to-end conversion and inference scripts for the DeepSeek V4 family on Megatron Bridge. + +The bridge supports four published variants out of the same code path. The on-disk quantisation differs between post-trained (Flash, Pro) and pretrained-only (Flash-Base, Pro-Base) models — see [`docs/models/llm/deepseek-v4.md`](../../../docs/models/llm/deepseek-v4.md) for the per-variant scheme. + +## MCore Dev Branch Requirement + +DSv4 imports require MCore changes that are not yet on a tagged release: PR [#3430](https://github.com/NVIDIA/Megatron-LM/pull/3430), PR [#4458](https://github.com/NVIDIA/Megatron-LM/pull/4458), PR [#4481](https://github.com/NVIDIA/Megatron-LM/pull/4481), and PR [#4518](https://github.com/NVIDIA/Megatron-LM/pull/4518), and PR [#4839](https://github.com/NVIDIA/Megatron-LM/pull/4839). Until these merge to Megatron-LM `main` and the bridge submodule pin advances, point `3rdparty/Megatron-LM` at the Megatron-LM `dev` branch: + +```bash +./scripts/switch_mcore.sh dev +uv sync +``` + +Use `./scripts/switch_mcore.sh main` and `uv sync --locked` to return to the pinned main-branch submodule. + +| Variant | HF path | Quant scheme | Validation | +|---------|---------|--------------|------------| +| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | FP8 attn + MXFP4 experts | Verified on GB200, last-token logit cosine 0.96-0.99 (short prompts ~0.98, long prompts >1024 tokens ~0.96-0.99) vs official inference | +| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | uniform FP8 (F32 scales) | Verified on GB200, last-real-token logit cosine 0.9866-0.9930, mean 0.9907 vs official inference | +| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | FP8 attn + MXFP4 experts | Import, export, inference verified on GB200 (PP=4 EP=8) and H100 (PP=16 EP=8) | +| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | uniform FP8 (F32 scales) | Same bridge code as Pro; end-to-end untested | + +## Examples + +- `conversion.sh` imports HF weights into Megatron Bridge and exports Megatron checkpoints back to HF format. +- `inference.sh` runs text generation against an HF or Megatron checkpoint. + +Run `bash conversion.sh` after setting `WORKSPACE` and `MODEL_VARIANT`. See each script's header comments for the expected environment variables and `#SBATCH` directives to edit before submitting. + +The bridge's `maybe_modify_loaded_hf_weight` hook dispatches dequantisation by tensor dtype: + +- `int8` -> MXFP4 packed nibbles -> `bfloat16` via the E2M1 lookup table and per-row 16-K-tile E8M0 scales +- `float8_e4m3fn` with companion `.scale` -> `bfloat16` via 128x128 block-scale expansion, handling both E8M0 and F32 scale dtypes + +No external dequantisation script is required. + +## Parallelism Configurations + +DSv4 currently requires **TP=1** because MLA tensor parallelism is not supported alongside the DSv4 hybrid attention path. Scale via expert and pipeline parallelism instead. + +| Model | TP | PP | EP | GPUs | GPU | Verified | +|-------|---:|---:|---:|-----:|-----|----------| +| DeepSeek-V4-Flash | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | +| DeepSeek-V4-Flash | 1 | 1 | 8 | 8 | H100 80GB | Import, export, inference | +| DeepSeek-V4-Flash-Base | 1 | 1 | 4 | 4 | GB200 192GB | Import, export, inference | +| DeepSeek-V4-Pro | 1 | 4 | 8 | 32 | GB200 192GB | Import, export, inference | +| DeepSeek-V4-Pro | 1 | 16 | 8 | 128 | H100 80GB | Import, export, inference | + +## Known Limitations + +- **MTP is disabled for inference** via `disable_mtp_for_inference()`. MTP weights are mapped end-to-end and loaded into the Megatron model. + +- **Fused mHC is not supported on H100.** Set `use_fused_mhc=False` in the bridge config when running on Hopper GPUs. Fused mHC is enabled by default and works on GB200. + +- **`fast_hadamard_transform` is optional.** When unavailable, DSA falls back to a PyTorch hadamard implementation. Throughput is lower but numerical behavior is unchanged. + +- **Logit parity is verified for Flash and Flash-Base** against the official inference stack at last-real-token logits. The remaining gap is structural, from different attention/HC kernel decompositions and accumulation precisions between MCore and official inference. diff --git a/examples/models/deepseek_v4/conversion.sh b/examples/models/deepseek_v4/conversion.sh new file mode 100755 index 0000000000..ac83ddca17 --- /dev/null +++ b/examples/models/deepseek_v4/conversion.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# DeepSeek-V4 import + export with the Bridge. +# +# DSv4 currently requires TP=1; scale via expert and pipeline parallelism (EP, PP). +# The Bridge dispatches FP8 / MXFP4 dequantisation by tensor dtype, so the +# same script works for Flash, Flash-Base, Pro, and Pro-Base. +# +# Override defaults by exporting environment variables before running: +# WORKSPACE: directory for converted Megatron checkpoints (default: /workspace) +# MODEL_VARIANT: one of DeepSeek-V4-Flash, DeepSeek-V4-Flash-Base, +# DeepSeek-V4-Pro, DeepSeek-V4-Pro-Base +# (default: DeepSeek-V4-Flash-Base) +# EP: expert-parallel size (default: 4 for Flash, 8 for Pro) +# PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro) +# +# Defaults below are for GB200 (192 GB). For H100 (80 GB) configs, see README.md. + +set -xeuo pipefail + +WORKSPACE=${WORKSPACE:-/workspace} +MODEL_VARIANT=${MODEL_VARIANT:-DeepSeek-V4-Flash-Base} +HF_MODEL_ID="deepseek-ai/${MODEL_VARIANT}" + +if [[ -z "${EP:-}" ]]; then + case "${MODEL_VARIANT}" in + DeepSeek-V4-Pro*) EP=8 ;; + *) EP=4 ;; + esac +fi +if [[ -z "${PP:-}" ]]; then + case "${MODEL_VARIANT}" in + DeepSeek-V4-Pro*) PP=4 ;; + *) PP=1 ;; + esac +fi +TP=1 + +MEGATRON_DIR="${WORKSPACE}/models/${MODEL_VARIANT}" +EXPORT_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-hf-export" +ITER=iter_0000000 + +# 1) Import HF -> Megatron (FP8 / MXFP4 dequantised to bfloat16 in-flight) +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/convert_checkpoints_multi_gpu.py import \ + --hf-model "${HF_MODEL_ID}" \ + --megatron-path "${MEGATRON_DIR}" \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --torch-dtype bfloat16 \ + --trust-remote-code + +# 2) Compare HF and Megatron logits on a short prompt +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/compare_hf_and_megatron/compare.py \ + --hf_model_path "${HF_MODEL_ID}" \ + --megatron_model_path "${MEGATRON_DIR}" \ + --prompt "Hello, how are you?" \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --trust-remote-code + +# 3) Export Megatron -> HF (round-trip) +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/convert_checkpoints_multi_gpu.py export \ + --hf-model "${HF_MODEL_ID}" \ + --megatron-path "${MEGATRON_DIR}/${ITER}" \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --torch-dtype bfloat16 \ + --hf-path "${EXPORT_DIR}" \ + --distributed-save \ + --trust-remote-code + +# 4) Round-trip validation (bf16 -> Megatron -> bf16) +# DSv4 HF weights are quantized (FP8/MXFP4), so the first import dequantises +# to bfloat16. A true lossless roundtrip re-imports the exported bf16 checkpoint +# and compares against the first export. +ROUNDTRIP_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-roundtrip" +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/convert_checkpoints_multi_gpu.py import \ + --hf-model "${EXPORT_DIR}" \ + --megatron-path "${ROUNDTRIP_DIR}" \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --torch-dtype bfloat16 \ + --trust-remote-code + +ROUNDTRIP_EXPORT_DIR="${WORKSPACE}/models/${MODEL_VARIANT}-roundtrip-export" +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/convert_checkpoints_multi_gpu.py export \ + --hf-model "${EXPORT_DIR}" \ + --megatron-path "${ROUNDTRIP_DIR}" \ + --hf-path "${ROUNDTRIP_EXPORT_DIR}" \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --torch-dtype bfloat16 \ + --distributed-save \ + --trust-remote-code diff --git a/examples/models/deepseek_v4/inference.sh b/examples/models/deepseek_v4/inference.sh new file mode 100755 index 0000000000..0e660a54c6 --- /dev/null +++ b/examples/models/deepseek_v4/inference.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# DeepSeek-V4 text generation with the Bridge. +# +# Override defaults by exporting environment variables: +# WORKSPACE: directory holding the imported Megatron checkpoint (default: /workspace) +# MODEL_VARIANT: one of DeepSeek-V4-Flash, DeepSeek-V4-Flash-Base, +# DeepSeek-V4-Pro, DeepSeek-V4-Pro-Base +# EP: expert-parallel size (default: 4 for Flash, 8 for Pro) +# PP: pipeline-parallel size (default: 1 for Flash, 4 for Pro) +# PROMPT: prompt string (default: "Explain hyper-connections in transformer models.") +# +# Defaults below are for GB200 (192 GB). For H100 (80 GB) configs, see README.md. + +set -xeuo pipefail + +WORKSPACE=${WORKSPACE:-/workspace} +MODEL_VARIANT=${MODEL_VARIANT:-DeepSeek-V4-Flash-Base} +HF_MODEL_ID="deepseek-ai/${MODEL_VARIANT}" +PROMPT=${PROMPT:-"Explain hyper-connections in transformer models."} + +if [[ -z "${EP:-}" ]]; then + case "${MODEL_VARIANT}" in + DeepSeek-V4-Pro*) EP=8 ;; + *) EP=4 ;; + esac +fi +if [[ -z "${PP:-}" ]]; then + case "${MODEL_VARIANT}" in + DeepSeek-V4-Pro*) PP=4 ;; + *) PP=1 ;; + esac +fi +TP=1 + +# Inference directly from the HF checkpoint (Bridge dequantises in-flight). +uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path "${HF_MODEL_ID}" \ + --prompt "${PROMPT}" \ + --max_new_tokens 100 \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --trust-remote-code + +# Inference from a previously-imported Megatron checkpoint (faster cold start). +MEGATRON_DIR="${WORKSPACE}/models/${MODEL_VARIANT}" +if [[ -d "${MEGATRON_DIR}/iter_0000000" ]]; then + uv run python -m torch.distributed.run --nproc_per_node=$((PP * EP)) \ + examples/conversion/hf_to_megatron_generate_text.py \ + --hf_model_path "${HF_MODEL_ID}" \ + --megatron_model_path "${MEGATRON_DIR}" \ + --prompt "${PROMPT}" \ + --max_new_tokens 100 \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --trust-remote-code +fi diff --git a/src/megatron/bridge/models/deepseek/__init__.py b/src/megatron/bridge/models/deepseek/__init__.py index af9620a389..968868c28d 100644 --- a/src/megatron/bridge/models/deepseek/__init__.py +++ b/src/megatron/bridge/models/deepseek/__init__.py @@ -14,9 +14,11 @@ from megatron.bridge.models.deepseek.deepseek_v2_bridge import DeepSeekV2Bridge # noqa: F401 from megatron.bridge.models.deepseek.deepseek_v3_bridge import DeepSeekV3Bridge # noqa: F401 +from megatron.bridge.models.deepseek.deepseek_v4_bridge import DeepSeekV4Bridge # noqa: F401 __all__ = [ "DeepSeekV2Bridge", "DeepSeekV3Bridge", + "DeepSeekV4Bridge", ] diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py new file mode 100644 index 0000000000..768aa30295 --- /dev/null +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -0,0 +1,952 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bridge for the DeepSeek-V4 model family. + +The bridge covers DeepSeek-V4 variants that share the ``deepseek_v4`` HF config +schema. It derives dimension- and layer-dependent fields from the HF config and +dispatches checkpoint import by tensor dtype so FP8 and FP8+MXFP4 formats can +share the same conversion path. + +Checkpoint format notes: DeepSeek-V4 uses a custom serialisation format that +differs from standard HuggingFace Transformers naming conventions: + + - embed.weight (not model.embed_tokens.weight) + - head.weight (not lm_head.weight) + - norm.weight (not model.norm.weight) + - layers.N.attn_norm.weight / layers.N.ffn_norm.weight + - layers.N.attn.wq_a / wq_b / wkv / wo_a / wo_b … + - layers.N.ffn.gate / experts / shared_experts … + - layers.N.hc_attn_fn / hc_attn_base / hc_attn_scale (Hyper-Connections) + - layers.N.hc_ffn_fn / hc_ffn_base / hc_ffn_scale + - hc_head_fn / hc_head_base / hc_head_scale (global HC head, learned output contraction) + - mtp.N.* (MTP layers) + +Quantisation schemes: Two on-disk formats coexist in this family. The bridge +dispatches purely on tensor dtype, so the same code path handles both: + + Released variant Attn / shared experts Routed experts + ------------------- ------------------------ ---------------------------- + Flash (post-trained) FP8_E4M3 + F8_E8M0 (...) MXFP4 packed I8 + F8_E8M0 + Flash-Base / Pro / FP8_E4M3 + F32 (...) FP8_E4M3 + F32 (...) + Pro-Base (raw) + +All scale tensors are 128x128 block-tile geometry (scale.shape[i] == ceil(weight.shape[i]/128)) +except the MXFP4 expert path, where scale is per-row over 32-element K-tiles. +``maybe_modify_loaded_hf_weight`` flattens both F8_E8M0 and F32 scales to +F32 via ``.to(torch.float32)`` and selects the tile expansion automatically. +All weights are dequantised to bfloat16 during import. + +MoE router note: Hash-routing layers (layer_number <= moe_n_hash_layers) +contain a `tid2eid` buffer (int32 vocab→expert lookup table). Buffers are not +parameters, so Megatron does not expose them via `named_parameters()`. +The bridge handles `tid2eid` via `maybe_modify_loaded_hf_weight()` and +a dedicated `_Tid2EidMapping` that writes it into `state_dict` directly. + +Megatron-Core prerequisites: + - HyperConnectionModule + - DSv4HybridSelfAttention / CompressedSparseAttention / CSAIndexer / Compressor + - Hash-routing tid2eid support and SwiGLU clamp + - Separate MTP e_proj / h_proj modules with hyper-connections +""" + +from typing import Dict, Mapping + +import torch +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_transformer_block_with_experimental_attention_variant_spec as _get_exp_attn_spec, +) +from megatron.core.models.gpt.gpt_model import GPTModel + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge, WeightConversionTask +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + ColumnParallelMapping, + GatedMLPMapping, + MegatronParamMapping, + ReplicatedMapping, +) +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.mla_provider import MLAModelProvider + + +try: + import transformer_engine # noqa: F401 + + HAVE_TE = True +except (ImportError, ModuleNotFoundError): + HAVE_TE = False + + +_DSV4_LAYER_TYPE_TO_COMPRESS_RATIO = { + "sliding_attention": 0, + "compressed_sparse_attention": 4, + "heavily_compressed_attention": 128, +} + +_DSV4_COMPRESS_RATIO_TO_LAYER_TYPE = { + ratio: layer_type for layer_type, ratio in _DSV4_LAYER_TYPE_TO_COMPRESS_RATIO.items() +} + + +def _dsv4_num_hash_layers(hf_config) -> int: + num_hash_layers = getattr(hf_config, "num_hash_layers", None) + if num_hash_layers is not None: + return int(num_hash_layers) + + mlp_layer_types = getattr(hf_config, "mlp_layer_types", None) + if mlp_layer_types is None: + return 0 + + n_hash = 0 + for layer_type in mlp_layer_types: + if layer_type != "hash_moe": + break + n_hash += 1 + + if any(layer_type == "hash_moe" for layer_type in mlp_layer_types[n_hash:]): + raise ValueError("DeepSeek-V4 hash MoE layers must be a contiguous prefix.") + + return n_hash + + +def _dsv4_compress_ratios(hf_config) -> list[int]: + num_hidden_layers = int(hf_config.num_hidden_layers) + num_mtp_layers = int(getattr(hf_config, "num_nextn_predict_layers", 0) or 0) + expected_len = num_hidden_layers + num_mtp_layers + + compress_ratios = getattr(hf_config, "compress_ratios", None) + if compress_ratios is not None: + ratios = [int(ratio) for ratio in compress_ratios] + else: + layer_types = getattr(hf_config, "layer_types", None) + compress_rates = getattr(hf_config, "compress_rates", None) + if layer_types is None or compress_rates is None: + raise ValueError( + "HF config missing 'compress_ratios' and native 'layer_types'/'compress_rates'. " + "DeepSeek-V4 requires per-layer compression ratios." + ) + + ratios = [] + for layer_type in layer_types: + if layer_type == "sliding_attention": + ratios.append(0) + elif layer_type in compress_rates: + ratios.append(int(compress_rates[layer_type])) + elif layer_type in _DSV4_LAYER_TYPE_TO_COMPRESS_RATIO: + ratios.append(_DSV4_LAYER_TYPE_TO_COMPRESS_RATIO[layer_type]) + else: + raise ValueError(f"Unsupported DeepSeek-V4 attention layer type: {layer_type!r}") + + if len(ratios) == num_hidden_layers and num_mtp_layers: + ratios.extend([0] * num_mtp_layers) + + if len(ratios) < expected_len: + raise ValueError( + f"DeepSeek-V4 compression ratios length ({len(ratios)}) is shorter than " + f"num_hidden_layers + num_nextn_predict_layers ({expected_len})." + ) + + return ratios[:expected_len] + + +# --------------------------------------------------------------------------- +# Custom mapping helpers +# --------------------------------------------------------------------------- + + +class _HCAlphaMapping(MegatronParamMapping): + """Map Megatron's three scalar HC alpha parameters to/from the V4 checkpoint's + 3-element hc_*_scale tensor. + + V4 checkpoint : layers.N.hc_attn_scale shape [3] = [alpha_pre, alpha_post, alpha_res] + Megatron : three separate nn.Parameter([1]) tensors + """ + + def __init__(self, megatron_pre: str, megatron_post: str, megatron_res: str, hf_param: str): + # We register under the alpha_pre path; the others are handled inside hf_to_megatron. + super().__init__(megatron_param=megatron_pre, hf_param=hf_param) + self._megatron_post = megatron_post + self._megatron_res = megatron_res + + @staticmethod + def _resolve_single(pattern: str, captures) -> str: + result = pattern + ci = 0 + while "**" in result and ci < len(captures): + result = result.replace("**", captures[ci], 1) + ci += 1 + ci = 0 + while "*" in result and ci < len(captures): + result = result.replace("*", captures[ci], 1) + ci += 1 + return result + + def resolve(self, captures): + resolved_mg, resolved_hf = self._resolve_names(captures) + resolved_post = self._resolve_single(self._megatron_post, captures) + resolved_res = self._resolve_single(self._megatron_res, captures) + return _HCAlphaMapping( + megatron_pre=resolved_mg, + megatron_post=resolved_post, + megatron_res=resolved_res, + hf_param=resolved_hf, + ) + + def hf_to_megatron(self, hf_weights, megatron_module): + # hf_weights is hc_*_scale [3]; we write alpha_pre here (index 0). + # alpha_post and alpha_res are handled by their own mappings when registered. + target = hf_weights.to(megatron_module.alpha_pre.device) + return target[0:1] + + def megatron_to_hf(self, megatron_weights, megatron_module): + # megatron_weights is alpha_pre [1]; gather all 3 from the same module. + # With PP > 1, megatron_module may be None on non-owning ranks, + # so we broadcast alpha_post and alpha_res alongside alpha_pre. + post_tensor = megatron_module.alpha_post.detach() if megatron_module is not None else None + res_tensor = megatron_module.alpha_res.detach() if megatron_module is not None else None + megatron_weights = self.broadcast_from_pp_rank(megatron_weights, cache_key=str(self.hf_param)) + post = self.broadcast_from_pp_rank(post_tensor, cache_key=str(self.hf_param) + "_post") + res = self.broadcast_from_pp_rank(res_tensor, cache_key=str(self.hf_param) + "_res") + if megatron_weights is None: + return {} + megatron_weights = self.maybe_dequantize(megatron_weights) + return {self.hf_param: torch.cat([megatron_weights.float(), post.float(), res.float()])} + + +class _HCAlphaSecondaryMapping(MegatronParamMapping): + """Secondary mapping for alpha_post (index=1) or alpha_res (index=2). + + Import: extracts element [index] from the 3-element hc_*_scale tensor. + Export: returns {} because the primary _HCAlphaMapping (alpha_pre) already + exports all three alpha values together. This mapping just suppresses the + "No mapping found" warning for the secondary Megatron params during export. + """ + + def __init__(self, megatron_param: str, hf_scale_param: str, index: int): + super().__init__(megatron_param=megatron_param, hf_param=hf_scale_param) + self._index = index + self.allow_hf_name_mismatch = True # export is no-op; skip hf_keys check + + def hf_to_megatron(self, hf_weights, megatron_module): + attr = "alpha_post" if self._index == 1 else "alpha_res" + target = hf_weights.to(getattr(megatron_module, attr).device) + return target[self._index : self._index + 1] + + def resolve(self, captures): + resolved_mg, resolved_hf = self._resolve_names(captures) + return _HCAlphaSecondaryMapping(resolved_mg, resolved_hf, self._index) + + def megatron_to_hf(self, megatron_weights, megatron_module): + # Already handled by the primary alpha_pre _HCAlphaMapping + return {} + + +# MXFP4 E2M1 lookup table: nibble value 0-15 → float32 +# FP4 E2M1 (sign=1, exponent=2, mantissa=1, exp_bias=1): +# positive: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +# negative: [-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] +_FP4_E2M1_TABLE = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) + + +class _ReplicatedOptional(ReplicatedMapping): + """ReplicatedMapping for CSA-optional weights (compressor / indexer). + + Sets allow_hf_name_mismatch=True so the export path does not validate + the HF key against the real checkpoint's key set. Compressor and indexer + weights only exist on non-hash layers; when we build a tiny smoke-test + model whose layer indices don't match the production compress_ratios, a + strict hf_keys check would wrongly skip those weights. + + resolve_wildcards() uses type(self)(...) which preserves this subclass, + so allow_hf_name_mismatch stays True after wildcard expansion. + """ + + def __init__(self, megatron_param: str, hf_param: str) -> None: + super().__init__(megatron_param, hf_param) + self.allow_hf_name_mismatch = True + + +# --------------------------------------------------------------------------- +# Bridge registration +# --------------------------------------------------------------------------- + + +@MegatronModelBridge.register_bridge( + source="DeepseekV4ForCausalLM", + target=GPTModel, + provider=MLAModelProvider, + model_type="deepseek_v4", +) +class DeepSeekV4Bridge(MegatronModelBridge): + """Megatron Bridge implementation for DeepSeek-V4 causal language models.""" + + # ------------------------------------------------------------------ + # Provider configuration + # ------------------------------------------------------------------ + + @staticmethod + def generate_pipeline_layout(num_layers: int, pp: int, mtp_layers: int = 1) -> list[list[str]]: + """Generate a pipeline-parallel layout for DSv4 models. + + DSv4 with hash MoE routing requires an explicit pipeline layout when PP > 1. + The layout distributes decoder layers across PP stages, placing the embedding + on the first stage and MTP + loss on the last stage. + + Args: + num_layers: Number of decoder layers (e.g. 43 for Flash, 61 for Pro). + pp: Pipeline parallel size. + mtp_layers: Number of MTP layers (default 1). + + Returns: + List of lists, where each inner list describes one pipeline stage. + """ + base, rem = num_layers // pp, num_layers % pp + layout = [] + for i in range(pp): + n = base + (1 if i < rem else 0) + stage = ["decoder"] * n + if i == 0: + stage = ["embedding"] + stage + if i == pp - 1: + stage = stage + ["mtp"] * mtp_layers + ["loss"] + layout.append(stage) + return layout + + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider: + provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config + + # ---- Attention ---- + provider.experimental_attention_variant = "dsv4_hybrid" + provider.multi_latent_attention = True + # V4 uses a heterogeneous per-layer spec (hash vs MLA layers differ); + # override the default transformer_engine_layer_spec with the experimental + # attention variant block spec builder. + # GPTModelProvider IS the TransformerConfig (not cfg.transformer) + provider.transformer_layer_spec = _get_exp_attn_spec + provider.qk_layernorm = True + provider.normalization = "RMSNorm" + provider.add_bias_linear = False + + # V4 MLA geometry + # head_dim = 512 (nope_dim + rope_dim = 448 + 64) + provider.v_head_dim = hf_config.head_dim # 512 + provider.qk_pos_emb_head_dim = hf_config.qk_rope_head_dim # 64 + # qk_head_dim and kv_lora_rank derived automatically in DSv4HybridConfig + provider.q_lora_rank = hf_config.q_lora_rank # 1024 + provider.o_groups = hf_config.o_groups # 8 + provider.o_lora_rank = hf_config.o_lora_rank # 1024 + + # ---- Rotary embeddings (YaRN) ---- + # Two separate RoPE bases in V4: + # - compress_rope_theta for compressed-KV layers + # - rope_theta for pure sliding-window layers (layers 0,1) + # Megatron keeps the regular and compressed CSA RoPE bases separately. + provider.apply_rope_fusion = True + provider.rope_type = "yarn" + rope_params = getattr(hf_config, "rope_scaling", None) or getattr(hf_config, "rope_parameters", None) or {} + if "compress" in rope_params: + main_rope_params = rope_params.get("main", {}) + compress_rope_params = rope_params["compress"] + else: + main_rope_params = rope_params + compress_rope_params = rope_params + provider.rotary_base = float(main_rope_params.get("rope_theta", hf_config.rope_theta)) # 10000 + provider.csa_compress_rotary_base = float( + compress_rope_params.get("rope_theta", getattr(hf_config, "compress_rope_theta", provider.rotary_base)) + ) # 160000 + provider.rotary_scaling_factor = float(compress_rope_params["factor"]) # 16 + provider.original_max_position_embeddings = int( + compress_rope_params["original_max_position_embeddings"] + ) # 65536 + provider.beta_fast = float(compress_rope_params.get("beta_fast", 32)) + provider.beta_slow = float(compress_rope_params.get("beta_slow", 1)) + # DSv4 has no mscale in HF config; Set both equal to cancel out (like DSv3). + provider.mscale = 1.0 + provider.mscale_all_dim = 1.0 + + # ---- CSA (Compressed Sparse Attention) ---- + # Legacy configs ship compress_ratios, while native Transformers configs + # expose layer_types + compress_rates. MCore consumes the flattened list. + _cr = _dsv4_compress_ratios(hf_config) + _mtp = getattr(hf_config, "num_nextn_predict_layers", None) + if _mtp is None: + import logging + + logging.warning( + "HF config missing 'num_nextn_predict_layers'; defaulting to 0. " + "DeepSeek-V4-Flash uses num_nextn_predict_layers=1." + ) + _mtp = 0 + _expected = hf_config.num_hidden_layers + _mtp + provider.csa_compress_ratios = _cr[:_expected] + provider.csa_window_size = hf_config.sliding_window # 128 + + # DSA indexer geometry (matches index_n_heads / index_head_dim / index_topk in config) + provider.dsa_indexer_n_heads = hf_config.index_n_heads # 64 + provider.dsa_indexer_head_dim = hf_config.index_head_dim # 128 + provider.dsa_indexer_topk = hf_config.index_topk # 512 + + # ---- Hyper-Connections (mHC) ---- + provider.enable_hyper_connections = True + provider.use_fused_mhc = True + provider.num_residual_streams = hf_config.hc_mult # 4 + provider.mhc_sinkhorn_iterations = hf_config.hc_sinkhorn_iters # 20 + + # ---- MoE ---- + provider.gated_linear_unit = True + provider.moe_grouped_gemm = True + provider.moe_router_pre_softmax = False # V4 uses post-topk normalisation + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_router_load_balancing_type = "noaux_tc" + provider.moe_shared_expert_overlap = True + provider.moe_router_score_function = hf_config.scoring_func # "sqrtsoftplus" + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + provider.moe_permute_fusion = True + provider.moe_aux_loss_coeff = 0.0 + provider.moe_router_topk = hf_config.num_experts_per_tok # 6 + provider.norm_topk_prob = hf_config.norm_topk_prob + provider.moe_router_topk_scaling_factor = hf_config.routed_scaling_factor # 1.5 + + # Hash routing + provider.moe_n_hash_layers = _dsv4_num_hash_layers(hf_config) # 3 for DSv4 Flash + provider.actual_vocab_size = hf_config.vocab_size # 129280 + + # SwiGLU activation clamp + provider.activation_func_clamp_value = hf_config.swiglu_limit # 10.0 + + # All 43 layers are MoE (no dense prefix unlike V3) + provider.moe_layer_freq = [1] * hf_config.num_hidden_layers + provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size * hf_config.n_shared_experts + + # ---- MTP ---- + provider.mtp_num_layers = getattr(hf_config, "num_nextn_predict_layers", 0) or None + + # ---- Misc ---- + provider.share_embeddings_and_output_weights = bool(hf_config.tie_word_embeddings) + provider.gradient_accumulation_fusion = True + provider.bias_dropout_fusion = True + provider.cross_entropy_fusion_impl = "te" + provider.cross_entropy_loss_fusion = True + provider.masked_softmax_fusion = True + provider.persist_layer_norm = True + provider.hidden_dropout = 0.0 + provider.attention_softmax_in_fp32 = False + provider.make_vocab_size_divisible_by = 1280 + provider.seq_length = 4096 + + return provider + + # ------------------------------------------------------------------ + # Export: HF config reconstruction + # ------------------------------------------------------------------ + + @classmethod + def megatron_to_hf_config(cls, provider: MLAModelProvider) -> dict: + hf_cfg = super(DeepSeekV4Bridge, cls).megatron_to_hf_config(provider) + + hf_cfg["num_nextn_predict_layers"] = getattr(provider, "mtp_num_layers", None) or 0 + num_hidden_layers = hf_cfg.get("num_hidden_layers", getattr(provider, "num_layers", 0)) + num_hash_layers = getattr(provider, "moe_n_hash_layers", 0) + hf_cfg["num_hash_layers"] = num_hash_layers + hf_cfg["mlp_layer_types"] = ["hash_moe"] * min(num_hidden_layers, num_hash_layers) + ["moe"] * max( + 0, num_hidden_layers - num_hash_layers + ) + hf_cfg["swiglu_limit"] = getattr(provider, "activation_func_clamp_value", 0.0) + + compress_ratios = getattr(provider, "csa_compress_ratios", None) + if compress_ratios is not None: + num_mtp = hf_cfg.get("num_nextn_predict_layers", 0) + expected_len = num_hidden_layers + num_mtp + compress_ratios = list(compress_ratios) + if len(compress_ratios) == num_hidden_layers and num_mtp: + compress_ratios = compress_ratios + [0] * num_mtp + hf_cfg["compress_ratios"] = compress_ratios[:expected_len] + hf_cfg["layer_types"] = [ + _DSV4_COMPRESS_RATIO_TO_LAYER_TYPE[ratio] for ratio in hf_cfg["compress_ratios"][:num_hidden_layers] + ] + hf_cfg["compress_rates"] = { + "compressed_sparse_attention": _DSV4_LAYER_TYPE_TO_COMPRESS_RATIO["compressed_sparse_attention"], + "heavily_compressed_attention": _DSV4_LAYER_TYPE_TO_COMPRESS_RATIO["heavily_compressed_attention"], + } + + hf_cfg["sliding_window"] = getattr(provider, "csa_window_size", 128) + hf_cfg["hc_mult"] = getattr(provider, "num_residual_streams", 4) + hf_cfg["hc_sinkhorn_iters"] = getattr(provider, "mhc_sinkhorn_iterations", 20) + hf_cfg["n_shared_experts"] = getattr(provider, "moe_shared_expert_intermediate_size", 0) // hf_cfg.get( + "moe_intermediate_size", 1 + ) + + return hf_cfg + + # ------------------------------------------------------------------ + # FP8 / MXFP4 dequantisation on import + # ------------------------------------------------------------------ + + def _dequant_mxfp4( + self, + hf_param: str, + weight_i8: torch.Tensor, + hf_state_dict: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Dequantize MXFP4-packed expert weights (I8 body + F8_E8M0 block scale). + + Each I8 byte stores 2 FP4 E2M1 nibbles (low nibble = first element, + high nibble = second element, i.e. [lo0, hi0, lo1, hi1, ...]). + One E8M0 scale covers 32 consecutive FP4 elements along K. + + Args: + hf_param: Weight tensor name (used to derive scale key). + weight_i8: Raw I8 tensor of shape (M, K//2). + hf_state_dict: Full HF state dict for scale lookup. + + Returns: + Dequantised weight as bfloat16 of shape (M, K). + """ + scale_key = hf_param[: -len(".weight")] + ".scale" if hf_param.endswith(".weight") else None + if scale_key is None or scale_key not in hf_state_dict: + return weight_i8.to(torch.bfloat16) + + scale = hf_state_dict[scale_key] # F8_E8M0, shape (M, K_scale) + + # Reinterpret int8 as uint8 so bitwise ops give correct 0-15 nibbles + w_u8 = weight_i8.view(torch.uint8) # (M, K_packed) + lo = (w_u8 & 0xF).to(torch.int64) # (M, K_packed) + hi = (w_u8 >> 4).to(torch.int64) # (M, K_packed) + + table = _FP4_E2M1_TABLE.to(weight_i8.device) + # stack [lo_row, hi_row] along last dim then flatten → interleaved (M, K_logical) + # This creates a contiguous output avoiding slow strided writes. + logical = torch.stack([table[lo], table[hi]], dim=-1).reshape( + weight_i8.shape[0], -1 + ) # (M, K_logical), float32 + + # E8M0 scale: value = 2^(e - 127), block_size = K_logical / K_scale = 32 + scale_f32 = scale.to(torch.float32) # E8M0 .to(f32) already gives 2^(e-127) + block_size = logical.shape[1] // scale_f32.shape[1] + # repeat_interleave along dim=1 expands (M, K_scale) → (M, K_logical) + scale_exp = scale_f32.repeat_interleave(block_size, dim=1) + + return (logical * scale_exp).to(torch.bfloat16) + + def maybe_modify_loaded_hf_weight( + self, + hf_param, + hf_state_dict: Mapping[str, torch.Tensor], + ): + """Dequantise quantized weights using their accompanying block-scale tensor. + + V4 stores attention/embedding weights as float8_e4m3fn with 128x128-block + scales, and expert FFN weights as MXFP4 packed (I8, 2 nibbles/byte) with + F8_E8M0 per-32-element scales. For dict hf_param (GatedMLPMapping etc.), + dequantizes each key individually so expert gate/up weights are also handled. + """ + if isinstance(hf_param, dict): + # Recurse for each key so each string key gets dequantized individually + return {k: self.maybe_modify_loaded_hf_weight(v, hf_state_dict) for k, v in hf_param.items()} + + weight = hf_state_dict[hf_param] + + # MXFP4 packed (I8): expert FFN gate/up/down weights + if weight.dtype == torch.int8: + return self._dequant_mxfp4(hf_param, weight, hf_state_dict) + + if weight.dtype != torch.float8_e4m3fn: + return weight + + if not hf_param.endswith(".weight"): + return weight.to(torch.bfloat16) + + scale_key = hf_param[: -len(".weight")] + ".scale" + if scale_key not in hf_state_dict: + return weight.to(torch.bfloat16) + + scale = hf_state_dict[scale_key] # [ceil(out/128), ceil(in/128)], float8_e8m0fnu + weight_bf16 = weight.to(torch.bfloat16) + + # Detect scale format: 128-tile (attn) vs per-row/16-tile (expert FFN) + # Attention: scale = (ceil(M/128), ceil(K/128)) + # Expert FFN: scale = (M, ceil(K/16)) [per-row, 16-element K-tiles] + scale_f32 = scale.to(torch.float32) + if scale_f32.dim() == 1: + # 1-D scale: single value per row (broadcast over K) + if scale_f32.shape[0] == weight_bf16.shape[0]: + scale_exp = scale_f32.unsqueeze(1) + else: + scale_exp = scale_f32.repeat_interleave(128)[: weight_bf16.shape[0]].unsqueeze(1) + elif scale_f32.shape[0] == weight_bf16.shape[0]: + # Per-row format: scale[i, j] covers weight[i, j*16:(j+1)*16] + tile_k = weight_bf16.shape[1] // scale_f32.shape[1] + scale_exp = scale_f32.repeat_interleave(tile_k, dim=1)[:, : weight_bf16.shape[1]] + else: + # 128-tile format: scale[i, j] covers weight[i*128:(i+1)*128, j*128:(j+1)*128] + scale_exp = scale_f32.repeat_interleave(128, dim=0)[: weight_bf16.shape[0]] + scale_exp = scale_exp.repeat_interleave(128, dim=1)[:, : weight_bf16.shape[1]] + if scale_exp.shape != weight_bf16.shape: + raise RuntimeError( + f"FP8 dequant shape mismatch for {hf_param!r}: " + f"weight={tuple(weight_bf16.shape)} scale={tuple(scale.shape)} " + f"scale_exp={tuple(scale_exp.shape)}" + ) + return (weight_bf16.to(torch.float32) * scale_exp).to(torch.bfloat16) + + # ------------------------------------------------------------------ + # Weight mapping registry + # ------------------------------------------------------------------ + + def mapping_registry(self) -> MegatronMappingRegistry: # noqa: C901 + hf_config = self.hf_config + num_mtp = getattr(hf_config, "num_nextn_predict_layers", 0) # 1 + + mappings = [] + + # ------ Embeddings / LM head / final norm ------ + mappings += [ + AutoMapping("embedding.word_embeddings.weight", "embed.weight"), + AutoMapping("output_layer.weight", "head.weight"), + AutoMapping("decoder.final_layernorm.weight", "norm.weight"), + # Global HC head (lives on TransformerBlock, not a parallel module → replicated) + ReplicatedMapping("decoder.hc_head_fn", "hc_head_fn"), + ReplicatedMapping("decoder.hc_head_base", "hc_head_base"), + ReplicatedMapping("decoder.hc_head_scale", "hc_head_scale"), + ] + + # ------ Per-layer mappings ------ + mappings += [ + # Layer norms + AutoMapping( + "decoder.layers.*.input_layernorm.weight", + "layers.*.attn_norm.weight", + ), + AutoMapping( + "decoder.layers.*.pre_mlp_layernorm.weight", + "layers.*.ffn_norm.weight", + ), + # Q down / Q norm / Q up (MLA) + AutoMapping( + "decoder.layers.*.self_attention.linear_q_down_proj.weight", + "layers.*.attn.wq_a.weight", + ), + AutoMapping( + "decoder.layers.*.self_attention.q_layernorm.weight", + "layers.*.attn.q_norm.weight", + ), + AutoMapping( + "decoder.layers.*.self_attention.linear_q_up_proj.weight", + "layers.*.attn.wq_b.weight", + ), + # KV (single projection) / KV norm + AutoMapping( + "decoder.layers.*.self_attention.linear_kv_proj.weight", + "layers.*.attn.wkv.weight", + ), + AutoMapping( + "decoder.layers.*.self_attention.kv_layernorm.weight", + "layers.*.attn.kv_norm.weight", + ), + # Factored output projection: wo_a (group param) + wo_b (row-parallel linear) + # linear_o_group_proj is a plain nn.Parameter (all o_groups on every TP rank) + ReplicatedMapping( + "decoder.layers.*.self_attention.linear_o_group_proj", + "layers.*.attn.wo_a.weight", + ), + AutoMapping( + "decoder.layers.*.self_attention.linear_proj.weight", + "layers.*.attn.wo_b.weight", + ), + # Attention sink: split by TP (size = num_heads // TP on each rank) + ColumnParallelMapping( + "decoder.layers.*.self_attention.core_attention.attn_sink", + "layers.*.attn.attn_sink", + ), + # Compressor (compress_ratio > 1 layers: 128x and 4x) + # All compressor linears use parallel_mode="duplicated" -> ReplicatedMapping + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.compressor.linear_wkv.weight", + "layers.*.attn.compressor.wkv.weight", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.compressor.linear_wgate.weight", + "layers.*.attn.compressor.wgate.weight", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.compressor.ape", + "layers.*.attn.compressor.ape", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.compressor.norm.weight", + "layers.*.attn.compressor.norm.weight", + ), + # Indexer (compress_ratio == 4 layers only) + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.linear_wq_b.weight", + "layers.*.attn.indexer.wq_b.weight", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.linear_weights_proj.weight", + "layers.*.attn.indexer.weights_proj.weight", + ), + # Indexer sub-compressor (each indexer has its own compressor) + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.compressor.linear_wkv.weight", + "layers.*.attn.indexer.compressor.wkv.weight", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.compressor.linear_wgate.weight", + "layers.*.attn.indexer.compressor.wgate.weight", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.compressor.ape", + "layers.*.attn.indexer.compressor.ape", + ), + _ReplicatedOptional( + "decoder.layers.*.self_attention.core_attention.indexer.compressor.norm.weight", + "layers.*.attn.indexer.compressor.norm.weight", + ), + # MoE router weight and expert bias + AutoMapping( + "decoder.layers.*.mlp.router.weight", + "layers.*.ffn.gate.weight", + ), + AutoMapping( + "decoder.layers.*.mlp.router.expert_bias", + "layers.*.ffn.gate.bias", + ), + # Hash-routing lookup table (buffer, not a parameter) + AutoMapping( + "decoder.layers.*.mlp.router.tid2eid", + "layers.*.ffn.gate.tid2eid", + ), + # Routed expert MLP (w1=gate, w3=up, w2=down in V4 naming) + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="layers.*.ffn.experts.*.w1.weight", + up="layers.*.ffn.experts.*.w3.weight", + ), + AutoMapping( + "decoder.layers.*.mlp.experts.linear_fc2.weight*", + "layers.*.ffn.experts.*.w2.weight", + ), + # Shared expert MLP + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate="layers.*.ffn.shared_experts.w1.weight", + up="layers.*.ffn.shared_experts.w3.weight", + ), + AutoMapping( + "decoder.layers.*.mlp.shared_experts.linear_fc2.weight", + "layers.*.ffn.shared_experts.w2.weight", + ), + # Hyper-Connections: attn HC (HyperConnectionModule not in AutoMapping registry → replicated) + ReplicatedMapping( + "decoder.layers.*.self_attention_hyper_connection.mapping_proj.weight", + "layers.*.hc_attn_fn", + ), + ReplicatedMapping( + "decoder.layers.*.self_attention_hyper_connection.bias", + "layers.*.hc_attn_base", + ), + # Hyper-Connections: FFN HC + ReplicatedMapping( + "decoder.layers.*.mlp_hyper_connection.mapping_proj.weight", + "layers.*.hc_ffn_fn", + ), + ReplicatedMapping( + "decoder.layers.*.mlp_hyper_connection.bias", + "layers.*.hc_ffn_base", + ), + ] + + # HC alpha scalars need custom concatenation mapping (per-layer, both attn and ffn) + # These are wildcarded across all layers. + mappings += [ + _HCAlphaMapping( + megatron_pre="decoder.layers.*.self_attention_hyper_connection.alpha_pre", + megatron_post="decoder.layers.*.self_attention_hyper_connection.alpha_post", + megatron_res="decoder.layers.*.self_attention_hyper_connection.alpha_res", + hf_param="layers.*.hc_attn_scale", + ), + _HCAlphaMapping( + megatron_pre="decoder.layers.*.mlp_hyper_connection.alpha_pre", + megatron_post="decoder.layers.*.mlp_hyper_connection.alpha_post", + megatron_res="decoder.layers.*.mlp_hyper_connection.alpha_res", + hf_param="layers.*.hc_ffn_scale", + ), + ] + + # HC alpha secondary: register alpha_post and alpha_res to suppress export warnings + mappings += [ + _HCAlphaSecondaryMapping( + "decoder.layers.*.self_attention_hyper_connection.alpha_post", + "layers.*.hc_attn_scale", + 1, + ), + _HCAlphaSecondaryMapping( + "decoder.layers.*.self_attention_hyper_connection.alpha_res", + "layers.*.hc_attn_scale", + 2, + ), + _HCAlphaSecondaryMapping( + "decoder.layers.*.mlp_hyper_connection.alpha_post", + "layers.*.hc_ffn_scale", + 1, + ), + _HCAlphaSecondaryMapping( + "decoder.layers.*.mlp_hyper_connection.alpha_res", + "layers.*.hc_ffn_scale", + 2, + ), + ] + + # ------ MTP layer mappings ------ + # MTP layers mirror the main layer structure under mtp.layers.N.* + for mtp_idx in range(num_mtp): + ck_pfx = f"mtp.{mtp_idx}" # checkpoint prefix + mg_pfx = f"mtp.layers.{mtp_idx}" # Megatron prefix + + # Standard transformer weights (shared pattern with main layers) + _mtp_plain = [ + (f"{mg_pfx}.mtp_model_layer.input_layernorm.weight", f"{ck_pfx}.attn_norm.weight"), + (f"{mg_pfx}.mtp_model_layer.pre_mlp_layernorm.weight", f"{ck_pfx}.ffn_norm.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.linear_q_down_proj.weight", f"{ck_pfx}.attn.wq_a.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.q_layernorm.weight", f"{ck_pfx}.attn.q_norm.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.linear_q_up_proj.weight", f"{ck_pfx}.attn.wq_b.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.linear_kv_proj.weight", f"{ck_pfx}.attn.wkv.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.kv_layernorm.weight", f"{ck_pfx}.attn.kv_norm.weight"), + (f"{mg_pfx}.mtp_model_layer.self_attention.linear_proj.weight", f"{ck_pfx}.attn.wo_b.weight"), + (f"{mg_pfx}.mtp_model_layer.mlp.router.weight", f"{ck_pfx}.ffn.gate.weight"), + (f"{mg_pfx}.mtp_model_layer.mlp.router.expert_bias", f"{ck_pfx}.ffn.gate.bias"), + (f"{mg_pfx}.mtp_model_layer.mlp.router.tid2eid", f"{ck_pfx}.ffn.gate.tid2eid"), + ( + f"{mg_pfx}.mtp_model_layer.mlp.shared_experts.linear_fc2.weight", + f"{ck_pfx}.ffn.shared_experts.w2.weight", + ), + # MTP-specific norms / projections + (f"{mg_pfx}.enorm.weight", f"{ck_pfx}.enorm.weight"), + (f"{mg_pfx}.hnorm.weight", f"{ck_pfx}.hnorm.weight"), + (f"{mg_pfx}.final_layernorm.weight", f"{ck_pfx}.norm.weight"), + ] + # MTP HC params use ReplicatedMapping (HyperConnectionModule not in AutoMapping registry) + _mtp_hc_plain = [ + ( + f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.mapping_proj.weight", + f"{ck_pfx}.hc_attn_fn", + ), + (f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.bias", f"{ck_pfx}.hc_attn_base"), + (f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.mapping_proj.weight", f"{ck_pfx}.hc_ffn_fn"), + (f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.bias", f"{ck_pfx}.hc_ffn_base"), + # Per-MTP-layer HC head (output contraction); mirrors decoder.hc_head_* mappings. + (f"{mg_pfx}.hc_head_fn", f"{ck_pfx}.hc_head_fn"), + (f"{mg_pfx}.hc_head_base", f"{ck_pfx}.hc_head_base"), + (f"{mg_pfx}.hc_head_scale", f"{ck_pfx}.hc_head_scale"), + ] + for mg, hf in _mtp_plain: + mappings.append(AutoMapping(mg, hf)) + for mg, hf in _mtp_hc_plain: + mappings.append(ReplicatedMapping(mg, hf)) + # MTP attn_sink: TP-split like the main model attn_sink + mappings.append( + ColumnParallelMapping( + f"{mg_pfx}.mtp_model_layer.self_attention.core_attention.attn_sink", + f"{ck_pfx}.attn.attn_sink", + ) + ) + # linear_o_group_proj is a plain nn.Parameter (all o_groups on every TP rank) + mappings.append( + ReplicatedMapping( + f"{mg_pfx}.mtp_model_layer.self_attention.linear_o_group_proj", + f"{ck_pfx}.attn.wo_a.weight", + ) + ) + + # MTP e_proj + h_proj are separate ColumnParallelLinear projections + # when the MTP layer uses hyper-connections. + # AutoMapping auto-detects ColumnParallelLinear and shards along dim 0. + mappings += [ + AutoMapping(f"{mg_pfx}.e_proj.weight", f"{ck_pfx}.e_proj.weight"), + AutoMapping(f"{mg_pfx}.h_proj.weight", f"{ck_pfx}.h_proj.weight"), + ] + + # MTP gated MLP (routed experts + shared expert) + mappings += [ + GatedMLPMapping( + megatron_param=f"{mg_pfx}.mtp_model_layer.mlp.experts.linear_fc1.weight*", + gate=f"{ck_pfx}.ffn.experts.*.w1.weight", + up=f"{ck_pfx}.ffn.experts.*.w3.weight", + ), + AutoMapping( + f"{mg_pfx}.mtp_model_layer.mlp.experts.linear_fc2.weight*", + f"{ck_pfx}.ffn.experts.*.w2.weight", + ), + GatedMLPMapping( + megatron_param=f"{mg_pfx}.mtp_model_layer.mlp.shared_experts.linear_fc1.weight", + gate=f"{ck_pfx}.ffn.shared_experts.w1.weight", + up=f"{ck_pfx}.ffn.shared_experts.w3.weight", + ), + ] + + # MTP HC alpha scalars + mappings += [ + _HCAlphaMapping( + megatron_pre=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_pre", + megatron_post=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_post", + megatron_res=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_res", + hf_param=f"{ck_pfx}.hc_attn_scale", + ), + _HCAlphaMapping( + megatron_pre=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_pre", + megatron_post=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_post", + megatron_res=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_res", + hf_param=f"{ck_pfx}.hc_ffn_scale", + ), + ] + + # MTP HC alpha secondary: suppress export warnings for post/res + for _hc_mg_sub, _hc_hf_key in [ + ("self_attention_hyper_connection", "hc_attn_scale"), + ("mlp_hyper_connection", "hc_ffn_scale"), + ]: + mappings += [ + _HCAlphaSecondaryMapping( + f"{mg_pfx}.mtp_model_layer.{_hc_mg_sub}.alpha_post", + f"{ck_pfx}.{_hc_hf_key}", + 1, + ), + _HCAlphaSecondaryMapping( + f"{mg_pfx}.mtp_model_layer.{_hc_mg_sub}.alpha_res", + f"{ck_pfx}.{_hc_hf_key}", + 2, + ), + ] + + return MegatronMappingRegistry(*mappings) + + # ------------------------------------------------------------------ + # Export: synthesise inv_freq (keeps roundtrip HF compat) + # ------------------------------------------------------------------ + + def maybe_modify_converted_hf_weight( + self, + task: WeightConversionTask, + converted_weights_dict: Dict[str, torch.Tensor], + hf_state_dict: Mapping[str, torch.Tensor], + ) -> Dict[str, torch.Tensor]: + """No-op for V4: the checkpoint does not contain inv_freq tensors.""" + return converted_weights_dict diff --git a/src/megatron/bridge/models/hf_pretrained/state.py b/src/megatron/bridge/models/hf_pretrained/state.py index 238e007d4e..b42dcc9657 100644 --- a/src/megatron/bridge/models/hf_pretrained/state.py +++ b/src/megatron/bridge/models/hf_pretrained/state.py @@ -979,8 +979,16 @@ def _save_generator_distributed( if is_saver_rank: missing_keys = assigned_expected_keys - set(buffered_tensors.keys()) if missing_keys: - missing_str = ", ".join(sorted(missing_keys)) - print(f"Rank {rank}: Missing tensors for keys: {missing_str}", flush=True) + missing_keys_sorted = sorted(missing_keys) + missing_preview = ", ".join(missing_keys_sorted[:20]) + missing_suffix = "" + if len(missing_keys_sorted) > 20: + missing_suffix = f", ... (+{len(missing_keys_sorted) - 20} more)" + print( + f"Rank {rank}: Missing {len(missing_keys_sorted)} tensors for keys: " + f"{missing_preview}{missing_suffix}", + flush=True, + ) for fname in assigned_filenames: keys_for_file = filename_to_keys_map[fname] @@ -990,23 +998,24 @@ def _save_generator_distributed( save_file(tensors_to_save, output_path / fname) actually_saved_keys.update(tensors_to_save.keys()) - # Gather all saved keys from all ranks to rank 0 + # Rank 0 builds the index from the files that were actually written. + # This avoids all_gather_object on very large key lists, which can + # allocate excessive CUDA memory in distributed runs. if is_distributed: - # Convert set to list for gathering - local_saved_keys_list = list(actually_saved_keys) if is_saver_rank else [] - gathered_keys = [None] * world_size - torch.distributed.all_gather_object(gathered_keys, local_saved_keys_list) + torch.distributed.barrier() if rank == 0: - # Aggregate all saved keys from all ranks + from safetensors import safe_open + all_saved_keys_aggregated = set() - for keys_list in gathered_keys: - if keys_list: - all_saved_keys_aggregated.update(keys_list) + for fname in all_filenames: + file_path = output_path / fname + if not file_path.exists(): + continue + with safe_open(file_path, framework="pt", device="cpu") as f: + all_saved_keys_aggregated.update(f.keys()) else: all_saved_keys_aggregated = set() - - torch.distributed.barrier() else: all_saved_keys_aggregated = actually_saved_keys @@ -1028,3 +1037,6 @@ def _save_generator_distributed( output_index_file = output_path / "model.safetensors.index.json" with open(output_index_file, "w") as f: json.dump(new_index_data, f, indent=4) + + if is_distributed: + torch.distributed.barrier() diff --git a/src/megatron/bridge/training/utils/config_utils.py b/src/megatron/bridge/training/utils/config_utils.py index e5baf36247..0a9c033ce4 100644 --- a/src/megatron/bridge/training/utils/config_utils.py +++ b/src/megatron/bridge/training/utils/config_utils.py @@ -283,6 +283,10 @@ def _convert_value_to_dict(cls, value: Any) -> Any: elif hasattr(value, "to_cfg_dict"): # Allow non-Container classes to implement own custom method return value.to_cfg_dict() + elif hasattr(value, "input_data") and type(value).__module__.startswith("megatron.core"): + # PipelineParallelLayerLayout: serialize as the original plain list + # so it can be deserialized without special instantiation logic. + return value.input_data elif is_dataclass(value) and not isinstance(value, type): # Handle regular dataclasses result = {} diff --git a/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py new file mode 100644 index 0000000000..0350e520e2 --- /dev/null +++ b/tests/functional_tests/test_groups/models/deepseek/test_deepseek_v4_conversion.py @@ -0,0 +1,264 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Functional toy-model conversion tests for DeepSeek V4.""" + +import importlib.util +import json +import subprocess +from pathlib import Path + +import pytest + + +def _has_dsv4_in_transformers() -> bool: + try: + from transformers import DeepseekV4Config, DeepseekV4ForCausalLM # noqa: F401 + + return True + except Exception: + return False + + +def _has_dsv4_in_mcore() -> bool: + try: + return all( + importlib.util.find_spec(mod) is not None + for mod in ( + "megatron.core.transformer.hyper_connection", + "megatron.core.transformer.experimental_attention_variant.csa", + "megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention", + ) + ) + except ModuleNotFoundError: + return False + + +pytestmark = [ + pytest.mark.skipif( + not _has_dsv4_in_transformers(), + reason="transformers does not yet ship DeepseekV4ForCausalLM (HF hub only via trust_remote_code).", + ), + pytest.mark.skipif( + not _has_dsv4_in_mcore(), + reason="megatron-core does not yet ship DSv4 prerequisites (PRs #3430 / #4458 / #4481 / #4518).", + ), +] + + +# Toy config tuned to satisfy DSv4 invariants at minimum size: +# - len(compress_ratios) == num_hidden_layers + num_nextn_predict_layers +# - sliding_window <= max_position_embeddings +# - vocab_size large enough for hash routing and divisible by the DSv4 provider's 1280 vocab padding +# - n_routed_experts divisible by num_experts_per_tok and the EP sizes we test +HF_DEEPSEEK_V4_TOY_MODEL_CONFIG = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 1024, + "head_dim": 256, + "qk_rope_head_dim": 32, + "intermediate_size": 2048, + "max_position_embeddings": 4096, + "moe_intermediate_size": 512, + "n_routed_experts": 8, + "n_shared_experts": 1, + "num_attention_heads": 16, + "num_experts_per_tok": 4, + "num_hidden_layers": 4, + "num_key_value_heads": 4, + "num_nextn_predict_layers": 0, # disable MTP for the toy + "q_lora_rank": 256, + "o_lora_rank": 256, + "o_groups": 4, + "compress_ratios": [0, 4, 4, 4], # 4 entries == num_hidden_layers (mtp=0) + "sliding_window": 64, + "index_n_heads": 4, + "index_head_dim": 32, + "index_topk": 32, + "hc_mult": 4, + "hc_sinkhorn_iters": 4, + "norm_topk_prob": True, + "scoring_func": "sqrtsoftplus", + "routed_scaling_factor": 1.0, + "rope_theta": 10000, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 16, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + "vocab_size": 12800, + "torch_dtype": "bfloat16", +} + + +def _hf_to_bridge_state_dict(hf_state_dict: dict, num_layers: int) -> dict: + """Translate native Transformers DSv4 parameter names to the released checkpoint layout.""" + bridge_state = {} + + def copy(src: str, dst: str) -> None: + if src in hf_state_dict: + bridge_state[dst] = hf_state_dict[src].detach().cpu().contiguous() + + copy("model.embed_tokens.weight", "embed.weight") + copy("lm_head.weight", "head.weight") + copy("model.norm.weight", "norm.weight") + copy("model.hc_head.hc_fn", "hc_head_fn") + copy("model.hc_head.hc_base", "hc_head_base") + copy("model.hc_head.hc_scale", "hc_head_scale") + + for layer_idx in range(num_layers): + hf_prefix = f"model.layers.{layer_idx}" + ckpt_prefix = f"layers.{layer_idx}" + + copy(f"{hf_prefix}.input_layernorm.weight", f"{ckpt_prefix}.attn_norm.weight") + copy(f"{hf_prefix}.post_attention_layernorm.weight", f"{ckpt_prefix}.ffn_norm.weight") + copy(f"{hf_prefix}.self_attn.q_a_proj.weight", f"{ckpt_prefix}.attn.wq_a.weight") + copy(f"{hf_prefix}.self_attn.q_a_norm.weight", f"{ckpt_prefix}.attn.q_norm.weight") + copy(f"{hf_prefix}.self_attn.q_b_proj.weight", f"{ckpt_prefix}.attn.wq_b.weight") + copy(f"{hf_prefix}.self_attn.kv_proj.weight", f"{ckpt_prefix}.attn.wkv.weight") + copy(f"{hf_prefix}.self_attn.kv_norm.weight", f"{ckpt_prefix}.attn.kv_norm.weight") + copy(f"{hf_prefix}.self_attn.o_a_proj.weight", f"{ckpt_prefix}.attn.wo_a.weight") + copy(f"{hf_prefix}.self_attn.o_b_proj.weight", f"{ckpt_prefix}.attn.wo_b.weight") + copy(f"{hf_prefix}.self_attn.sinks", f"{ckpt_prefix}.attn.attn_sink") + + compressor_prefix = f"{hf_prefix}.self_attn.compressor" + copy(f"{compressor_prefix}.kv_proj.weight", f"{ckpt_prefix}.attn.compressor.wkv.weight") + copy(f"{compressor_prefix}.gate_proj.weight", f"{ckpt_prefix}.attn.compressor.wgate.weight") + copy(f"{compressor_prefix}.position_bias", f"{ckpt_prefix}.attn.compressor.ape") + copy(f"{compressor_prefix}.kv_norm.weight", f"{ckpt_prefix}.attn.compressor.norm.weight") + + indexer_prefix = f"{compressor_prefix}.indexer" + copy(f"{indexer_prefix}.q_b_proj.weight", f"{ckpt_prefix}.attn.indexer.wq_b.weight") + copy(f"{indexer_prefix}.weights_proj.weight", f"{ckpt_prefix}.attn.indexer.weights_proj.weight") + copy(f"{indexer_prefix}.kv_proj.weight", f"{ckpt_prefix}.attn.indexer.compressor.wkv.weight") + copy(f"{indexer_prefix}.gate_proj.weight", f"{ckpt_prefix}.attn.indexer.compressor.wgate.weight") + copy(f"{indexer_prefix}.position_bias", f"{ckpt_prefix}.attn.indexer.compressor.ape") + copy(f"{indexer_prefix}.kv_norm.weight", f"{ckpt_prefix}.attn.indexer.compressor.norm.weight") + + copy(f"{hf_prefix}.mlp.gate.weight", f"{ckpt_prefix}.ffn.gate.weight") + if f"{hf_prefix}.mlp.gate.e_score_correction_bias" in hf_state_dict: + copy(f"{hf_prefix}.mlp.gate.e_score_correction_bias", f"{ckpt_prefix}.ffn.gate.bias") + copy(f"{hf_prefix}.mlp.gate.tid2eid", f"{ckpt_prefix}.ffn.gate.tid2eid") + + gate_up = hf_state_dict[f"{hf_prefix}.mlp.experts.gate_up_proj"].detach().cpu() + gate, up = gate_up.chunk(2, dim=1) + down = hf_state_dict[f"{hf_prefix}.mlp.experts.down_proj"].detach().cpu() + for expert_idx in range(gate.shape[0]): + bridge_state[f"{ckpt_prefix}.ffn.experts.{expert_idx}.w1.weight"] = gate[expert_idx].contiguous() + bridge_state[f"{ckpt_prefix}.ffn.experts.{expert_idx}.w3.weight"] = up[expert_idx].contiguous() + bridge_state[f"{ckpt_prefix}.ffn.experts.{expert_idx}.w2.weight"] = down[expert_idx].contiguous() + + copy(f"{hf_prefix}.mlp.shared_experts.gate_proj.weight", f"{ckpt_prefix}.ffn.shared_experts.w1.weight") + copy(f"{hf_prefix}.mlp.shared_experts.up_proj.weight", f"{ckpt_prefix}.ffn.shared_experts.w3.weight") + copy(f"{hf_prefix}.mlp.shared_experts.down_proj.weight", f"{ckpt_prefix}.ffn.shared_experts.w2.weight") + copy(f"{hf_prefix}.attn_hc.fn", f"{ckpt_prefix}.hc_attn_fn") + copy(f"{hf_prefix}.attn_hc.base", f"{ckpt_prefix}.hc_attn_base") + copy(f"{hf_prefix}.attn_hc.scale", f"{ckpt_prefix}.hc_attn_scale") + copy(f"{hf_prefix}.ffn_hc.fn", f"{ckpt_prefix}.hc_ffn_fn") + copy(f"{hf_prefix}.ffn_hc.base", f"{ckpt_prefix}.hc_ffn_base") + copy(f"{hf_prefix}.ffn_hc.scale", f"{ckpt_prefix}.hc_ffn_scale") + + return bridge_state + + +class TestDeepSeekV4Conversion: + """Toy HF-to-Megatron roundtrip coverage for DeepSeek V4.""" + + @pytest.fixture(scope="class") + def deepseek_v4_toy_model_path(self, tmp_path_factory): + import torch + from safetensors.torch import save_file + from tokenizers import Tokenizer + from tokenizers.models import WordLevel + from tokenizers.pre_tokenizers import Whitespace + from transformers import DeepseekV4Config, DeepseekV4ForCausalLM, PreTrainedTokenizerFast + + temp_dir = tmp_path_factory.mktemp("deepseek_v4_toy_model") + model_dir = temp_dir / "deepseek_v4_toy" + model_dir.mkdir() + + torch.manual_seed(1234) + config = DeepseekV4Config(**HF_DEEPSEEK_V4_TOY_MODEL_CONFIG) + config.torch_dtype = torch.bfloat16 + model = DeepseekV4ForCausalLM(config).bfloat16() + bridge_state = _hf_to_bridge_state_dict(model.state_dict(), config.num_hidden_layers) + for key in list(bridge_state): + if key.endswith(".tid2eid"): + bridge_state[key] = bridge_state[key].to(torch.int32) + + vocab = {"": 0, "": 1, "": 2, "": 3} + vocab.update({f"tok_{idx}": idx for idx in range(4, 128)}) + tokenizer_model = Tokenizer(WordLevel(vocab=vocab, unk_token="")) + tokenizer_model.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer_model, + bos_token="", + eos_token="", + pad_token="", + unk_token="", + ) + tokenizer.save_pretrained(model_dir) + + config.save_pretrained(model_dir) + with open(model_dir / "config.json", "w") as f: + json.dump(config.to_dict(), f, indent=2) + save_file(bridge_state, model_dir / "model.safetensors") + return str(model_dir) + + @pytest.mark.run_only_on("GPU") + def test_deepseek_v4_roundtrip_ep(self, deepseek_v4_toy_model_path, tmp_path): + test_output_dir = tmp_path / "deepseek_v4_ep" + test_output_dir.mkdir(exist_ok=True) + + cmd = [ + "python", + "-m", + "torch.distributed.run", + "--nproc_per_node=2", + "--nnodes=1", + "-m", + "coverage", + "run", + "--data-file=/opt/Megatron-Bridge/.coverage", + "--source=/opt/Megatron-Bridge/", + "--parallel-mode", + "examples/conversion/hf_megatron_roundtrip_multi_gpu.py", + "--hf-model-id", + deepseek_v4_toy_model_path, + "--output-dir", + str(test_output_dir), + "--tp", + "1", + "--pp", + "1", + "--ep", + "2", + ] + + result = subprocess.run( + cmd, capture_output=True, text=True, cwd=Path(__file__).parent.parent.parent.parent.parent.parent + ) + + if result.returncode != 0: + print(f"STDOUT: {result.stdout}") + print(f"STDERR: {result.stderr}") + assert result.returncode == 0, f"DeepSeek V4 conversion failed with {result.returncode}" + + converted_dir = test_output_dir / Path(deepseek_v4_toy_model_path).name + assert (converted_dir / "config.json").exists() + assert list(converted_dir.glob("*.safetensors")) diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py new file mode 100644 index 0000000000..b3055bf161 --- /dev/null +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the DeepSeek-V4 bridge mapping registry. + +Locks in the MTP mapping layout: per-MTP-layer HC head, separate ``e_proj`` +and ``h_proj`` mappings, and no deprecated concatenated ``eh_proj`` path. +""" + +from types import SimpleNamespace + +import pytest + +from megatron.bridge.models.conversion.param_mapping import AutoMapping, ReplicatedMapping +from megatron.bridge.models.deepseek.deepseek_v4_bridge import ( + DeepSeekV4Bridge, + _dsv4_compress_ratios, + _dsv4_num_hash_layers, +) + + +@pytest.fixture +def bridge_with_mtp(): + """A DSv4 bridge with hf_config stubbed for a single MTP layer.""" + bridge = DeepSeekV4Bridge() + # mapping_registry only reads num_nextn_predict_layers from hf_config. + bridge.hf_config = SimpleNamespace(num_nextn_predict_layers=1) + return bridge + + +@pytest.fixture +def bridge_without_mtp(): + """A DSv4 bridge with hf_config that has zero MTP layers.""" + bridge = DeepSeekV4Bridge() + bridge.hf_config = SimpleNamespace(num_nextn_predict_layers=0) + return bridge + + +def _by_megatron(registry): + """Index mappings by megatron_param for quick lookup in assertions.""" + return {m.megatron_param: m for m in registry.mappings} + + +class TestNativeDeepSeekV4ConfigTranslation: + """Native Transformers DSv4 config fields must map back to MCore fields.""" + + def test_compress_ratios_from_native_layer_types(self): + hf_config = SimpleNamespace( + num_hidden_layers=4, + num_nextn_predict_layers=1, + layer_types=[ + "sliding_attention", + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + ], + compress_rates={ + "compressed_sparse_attention": 4, + "heavily_compressed_attention": 128, + }, + ) + + assert _dsv4_compress_ratios(hf_config) == [0, 0, 4, 128, 0] + + def test_legacy_compress_ratios_still_work(self): + hf_config = SimpleNamespace( + num_hidden_layers=4, + num_nextn_predict_layers=1, + compress_ratios=[0, 0, 4, 128, 0], + ) + + assert _dsv4_compress_ratios(hf_config) == [0, 0, 4, 128, 0] + + def test_hash_layers_from_native_mlp_layer_types(self): + hf_config = SimpleNamespace( + mlp_layer_types=["hash_moe", "hash_moe", "hash_moe", "moe", "moe"], + ) + + assert _dsv4_num_hash_layers(hf_config) == 3 + + def test_hash_layers_must_be_prefix(self): + hf_config = SimpleNamespace(mlp_layer_types=["hash_moe", "moe", "hash_moe"]) + + with pytest.raises(ValueError, match="contiguous prefix"): + _dsv4_num_hash_layers(hf_config) + + +class TestDecoderHCHeadMappings: + """The global decoder HC-head triplet must be replicated mappings.""" + + @pytest.mark.parametrize( + "name", + ["decoder.hc_head_fn", "decoder.hc_head_base", "decoder.hc_head_scale"], + ) + def test_decoder_hc_head_replicated(self, bridge_with_mtp, name): + registry = bridge_with_mtp.mapping_registry() + mapping = _by_megatron(registry).get(name) + assert mapping is not None, f"missing decoder HC-head mapping: {name}" + assert isinstance(mapping, ReplicatedMapping) + # HF side drops the 'decoder.' prefix. + assert mapping.hf_param == name.removeprefix("decoder.") + + +class TestMTPHCHeadMappings: + """Per-MTP-layer HC head must mirror the decoder pattern.""" + + @pytest.mark.parametrize( + "suffix", + ["hc_head_fn", "hc_head_base", "hc_head_scale"], + ) + def test_mtp_hc_head_replicated(self, bridge_with_mtp, suffix): + registry = bridge_with_mtp.mapping_registry() + mapping = _by_megatron(registry).get(f"mtp.layers.0.{suffix}") + assert mapping is not None, f"missing MTP HC-head mapping: mtp.layers.0.{suffix}" + assert isinstance(mapping, ReplicatedMapping) + assert mapping.hf_param == f"mtp.0.{suffix}" + + def test_mtp_hc_head_absent_when_no_mtp(self, bridge_without_mtp): + registry = bridge_without_mtp.mapping_registry() + names = _by_megatron(registry) + for suffix in ("hc_head_fn", "hc_head_base", "hc_head_scale"): + assert f"mtp.layers.0.{suffix}" not in names + + +class TestMTPEHProjSplit: + """MTP e_proj and h_proj are separate ColumnParallelLinear modules. + + The bridge must use two AutoMappings (which auto-detect column parallelism), + not the deprecated concatenated eh_proj path. + """ + + @pytest.mark.parametrize("name", ["e_proj", "h_proj"]) + def test_split_proj_automapping(self, bridge_with_mtp, name): + registry = bridge_with_mtp.mapping_registry() + mapping = _by_megatron(registry).get(f"mtp.layers.0.{name}.weight") + assert mapping is not None, f"missing MTP projection: {name}" + assert isinstance(mapping, AutoMapping) + assert mapping.hf_param == f"mtp.0.{name}.weight" + + def test_eh_proj_not_in_registry(self, bridge_with_mtp): + registry = bridge_with_mtp.mapping_registry() + for mapping in registry.mappings: + assert "eh_proj" not in mapping.megatron_param, ( + f"deprecated eh_proj reference found in megatron_param: {mapping.megatron_param}" + ) + hf_param = mapping.hf_param + if isinstance(hf_param, str): + assert "eh_proj" not in hf_param, f"deprecated eh_proj reference found in hf_param: {hf_param}" + elif isinstance(hf_param, dict): + for v in hf_param.values(): + assert "eh_proj" not in v, f"deprecated eh_proj reference found in hf_param dict value: {v}" From 30ab9156f6c67f247c1c2a090c937ad80141d3ba Mon Sep 17 00:00:00 2001 From: Lingrui Mei Date: Thu, 11 Jun 2026 01:37:37 +0800 Subject: [PATCH 3/8] fix(model): keep full rotary_percent for DeepSeek-V4 MLA rope (#4271) Signed-off-by: Lingrui Mei Co-authored-by: Claude Fable 5 --- .../models/deepseek/deepseek_v4_bridge.py | 6 +++ .../deepseek/test_deepseek_v4_bridge.py | 54 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 768aa30295..7a94eca698 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -348,6 +348,12 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider # head_dim = 512 (nope_dim + rope_dim = 448 + 64) provider.v_head_dim = hf_config.head_dim # 512 provider.qk_pos_emb_head_dim = hf_config.qk_rope_head_dim # 64 + # HF's partial_rotary_factor (0.125) is relative to head_dim (512); the rope split is + # already fully encoded by qk_pos_emb_head_dim (64). The generic partial_rotary_factor + # -> rotary_percent mapping would shrink the rope cache to 64*0.125 = 8 dims: the + # unfused path then silently rotates only 8 of 64 rope dims, and the fused MLA rope + # kernel reads cos/sin out of bounds (garbage values -> the SFT loss NaN). + provider.rotary_percent = 1.0 # qk_head_dim and kv_lora_rank derived automatically in DSv4HybridConfig provider.q_lora_rank = hf_config.q_lora_rank # 1024 provider.o_groups = hf_config.o_groups # 8 diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index b3055bf161..1f3e0b166b 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -160,3 +160,57 @@ def test_eh_proj_not_in_registry(self, bridge_with_mtp): elif isinstance(hf_param, dict): for v in hf_param.values(): assert "eh_proj" not in v, f"deprecated eh_proj reference found in hf_param dict value: {v}" + + +class TestDeepSeekV4RotaryPercent: + """Regression: HF partial_rotary_factor (relative to head_dim=512) must not shrink + the Megatron rope cache — qk_pos_emb_head_dim (64) already encodes the rope split. + rotary_percent=0.125 yields an 8-dim cos/sin cache: the unfused path silently + rotates 8/64 dims and the fused MLA rope kernel reads cos/sin out of bounds (SFT NaN).""" + + def test_provider_bridge_forces_full_rotary_percent(self): + from unittest.mock import MagicMock, patch + + from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge + from megatron.bridge.models.deepseek.deepseek_v4_bridge import DeepSeekV4Bridge + + hf_config = SimpleNamespace( + head_dim=512, + qk_rope_head_dim=64, + q_lora_rank=1024, + o_groups=8, + o_lora_rank=1024, + rope_theta=10000, + compress_rope_theta=160000, + rope_scaling={"factor": 16, "original_max_position_embeddings": 65536}, + num_hidden_layers=4, + num_nextn_predict_layers=1, + num_hash_layers=3, + compress_ratios=[0, 4, 128, 4, 0], + sliding_window=128, + index_n_heads=64, + index_head_dim=128, + index_topk=512, + hc_mult=4, + hc_sinkhorn_iters=20, + scoring_func="sqrtsoftplus", + num_experts_per_tok=6, + norm_topk_prob=True, + routed_scaling_factor=1.5, + vocab_size=129280, + swiglu_limit=10.0, + moe_intermediate_size=1024, + n_shared_experts=1, + tie_word_embeddings=False, + ) + hf_pretrained = MagicMock() + hf_pretrained.config = hf_config + provider = MagicMock() + # what the generic partial_rotary_factor -> rotary_percent mapping produces + provider.rotary_percent = 0.125 + + bridge = DeepSeekV4Bridge.__new__(DeepSeekV4Bridge) + with patch.object(MegatronModelBridge, "provider_bridge", return_value=provider): + out = bridge.provider_bridge(hf_pretrained) + + assert out.rotary_percent == 1.0 From 3be1707040df58111a8448448fa577829b3c5828 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=84=8D=F0=9D=95=A0=F0=9D=95=9D=F0=9D=95=9D=F0=9D=95=A0?= =?UTF-8?q?=F0=9D=95=A8=20=F0=9D=95=84=F0=9D=95=92=F0=9D=95=9F?= Date: Fri, 10 Jul 2026 11:57:33 -0700 Subject: [PATCH 4/8] DSV4: Making compress_rope_theta priotitzed than rope_scaling->rope_theta (#4802) Signed-off-by: Hollow Man --- src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py | 2 +- tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 7a94eca698..0a918ac32a 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -375,7 +375,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider compress_rope_params = rope_params provider.rotary_base = float(main_rope_params.get("rope_theta", hf_config.rope_theta)) # 10000 provider.csa_compress_rotary_base = float( - compress_rope_params.get("rope_theta", getattr(hf_config, "compress_rope_theta", provider.rotary_base)) + getattr(hf_config, "compress_rope_theta", compress_rope_params.get("rope_theta", provider.rotary_base)) ) # 160000 provider.rotary_scaling_factor = float(compress_rope_params["factor"]) # 16 provider.original_max_position_embeddings = int( diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index 1f3e0b166b..accafd4536 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -214,3 +214,4 @@ def test_provider_bridge_forces_full_rotary_percent(self): out = bridge.provider_bridge(hf_pretrained) assert out.rotary_percent == 1.0 + assert out.csa_compress_rotary_base == 160000 From 2cd18453dc621835b001d814296c89d2068a5085 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:01:04 +0000 Subject: [PATCH 5/8] fix(deepseek-v4): support nested rope theta config Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py | 5 ++++- .../unit_tests/models/deepseek/test_deepseek_v4_bridge.py | 8 +++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 0a918ac32a..41463cfd52 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -373,7 +373,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider else: main_rope_params = rope_params compress_rope_params = rope_params - provider.rotary_base = float(main_rope_params.get("rope_theta", hf_config.rope_theta)) # 10000 + main_rope_theta = main_rope_params.get("rope_theta") + if main_rope_theta is None: + main_rope_theta = hf_config.rope_theta + provider.rotary_base = float(main_rope_theta) # 10000 provider.csa_compress_rotary_base = float( getattr(hf_config, "compress_rope_theta", compress_rope_params.get("rope_theta", provider.rotary_base)) ) # 160000 diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index accafd4536..d8383cc4d2 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -180,9 +180,10 @@ def test_provider_bridge_forces_full_rotary_percent(self): q_lora_rank=1024, o_groups=8, o_lora_rank=1024, - rope_theta=10000, - compress_rope_theta=160000, - rope_scaling={"factor": 16, "original_max_position_embeddings": 65536}, + rope_scaling={ + "main": {"rope_theta": 10000}, + "compress": {"rope_theta": 160000, "factor": 16, "original_max_position_embeddings": 65536}, + }, num_hidden_layers=4, num_nextn_predict_layers=1, num_hash_layers=3, @@ -214,4 +215,5 @@ def test_provider_bridge_forces_full_rotary_percent(self): out = bridge.provider_bridge(hf_pretrained) assert out.rotary_percent == 1.0 + assert out.rotary_base == 10000 assert out.csa_compress_rotary_base == 160000 From ca1a03de528c4b2605904dd517467dbc167c22f7 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:42:03 +0000 Subject: [PATCH 6/8] [conversion] fix: skip absent auto mappings across PP Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../bridge/models/conversion/param_mapping.py | 9 +++++++-- tests/unit_tests/models/test_param_mapping.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index 63b4cc371d..39df84b82b 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -1366,9 +1366,14 @@ def megatron_to_hf( self._detected_type = self.broadcast_obj_from_pp_rank(self._detected_type, "detected_type") else: # Receive from owning rank - self._detected_type = self.broadcast_obj_from_pp_rank(None, "detected_type") + try: + self._detected_type = self.broadcast_obj_from_pp_rank(None, "detected_type") + except ValueError as error: + if str(error) != "Object must exist on at least one PP rank": + raise + self._detected_type = None if self._detected_type is None: - # PP group likely has 1 member - skipping. + # No PP stage owns this optional parameter. return {} self._mapping = self._get_or_create_mapping(self._detected_type) diff --git a/tests/unit_tests/models/test_param_mapping.py b/tests/unit_tests/models/test_param_mapping.py index 6050ea28ff..332da158a4 100644 --- a/tests/unit_tests/models/test_param_mapping.py +++ b/tests/unit_tests/models/test_param_mapping.py @@ -273,6 +273,17 @@ class MyCustomRow(torch.nn.Module): with pytest.raises(ValueError): mapping._detect_parallelism_type(torch.nn.Linear(5, 5)) + def test_megatron_to_hf_skips_parameter_missing_from_all_pp_stages(self, mock_distributed_env): + _, mock_dist = mock_distributed_env(pp_size=2, pp_rank=0) + mapping = AutoMapping(megatron_param="optional.weight", hf_param="hf.optional.weight") + + mock_dist.all_gather_object.side_effect = lambda output, obj, group: output.__setitem__( + slice(None), [False, False] + ) + + assert mapping.megatron_to_hf(None, None) == {} + mock_dist.broadcast_object_list.assert_not_called() + def test_detect_parallelism_type_dynamic_module(self): mtq = pytest.importorskip("modelopt.torch.quantization") DynamicModule = pytest.importorskip("modelopt.torch.opt.dynamic").DynamicModule From 6f48f43b5d59713edfeb1514aae23f7cd0532110 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:08:56 +0000 Subject: [PATCH 7/8] [model] fix: Map DeepSeek V4 hash layer config Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../models/deepseek/deepseek_v4_bridge.py | 6 ++--- .../deepseek/test_deepseek_v4_bridge.py | 27 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 41463cfd52..7ece0fa782 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -48,7 +48,7 @@ F32 via ``.to(torch.float32)`` and selects the tile expansion automatically. All weights are dequantised to bfloat16 during import. -MoE router note: Hash-routing layers (layer_number <= moe_n_hash_layers) +MoE router note: Hash-routing layers (layer_number <= dsv4_n_hash_layers) contain a `tid2eid` buffer (int32 vocab→expert lookup table). Buffers are not parameters, so Megatron does not expose them via `named_parameters()`. The bridge handles `tid2eid` via `maybe_modify_loaded_hf_weight()` and @@ -435,7 +435,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.moe_router_topk_scaling_factor = hf_config.routed_scaling_factor # 1.5 # Hash routing - provider.moe_n_hash_layers = _dsv4_num_hash_layers(hf_config) # 3 for DSv4 Flash + provider.dsv4_n_hash_layers = _dsv4_num_hash_layers(hf_config) # 3 for DSv4 Flash provider.actual_vocab_size = hf_config.vocab_size # 129280 # SwiGLU activation clamp @@ -473,7 +473,7 @@ def megatron_to_hf_config(cls, provider: MLAModelProvider) -> dict: hf_cfg["num_nextn_predict_layers"] = getattr(provider, "mtp_num_layers", None) or 0 num_hidden_layers = hf_cfg.get("num_hidden_layers", getattr(provider, "num_layers", 0)) - num_hash_layers = getattr(provider, "moe_n_hash_layers", 0) + num_hash_layers = getattr(provider, "dsv4_n_hash_layers", 0) hf_cfg["num_hash_layers"] = num_hash_layers hf_cfg["mlp_layer_types"] = ["hash_moe"] * min(num_hidden_layers, num_hash_layers) + ["moe"] * max( 0, num_hidden_layers - num_hash_layers diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index d8383cc4d2..f11e55c207 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -95,6 +95,32 @@ def test_hash_layers_must_be_prefix(self): with pytest.raises(ValueError, match="contiguous prefix"): _dsv4_num_hash_layers(hf_config) + def test_export_hash_layers_from_mcore_field(self): + from unittest.mock import patch + + from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge + + provider = SimpleNamespace( + dsv4_n_hash_layers=3, + num_layers=5, + mtp_num_layers=None, + activation_func_clamp_value=10.0, + csa_compress_ratios=None, + csa_window_size=128, + num_residual_streams=4, + mhc_sinkhorn_iterations=20, + moe_shared_expert_intermediate_size=2048, + ) + with patch.object( + MegatronModelBridge, + "megatron_to_hf_config", + return_value={"num_hidden_layers": 5, "moe_intermediate_size": 2048}, + ): + hf_config = DeepSeekV4Bridge.megatron_to_hf_config(provider) + + assert hf_config["num_hash_layers"] == 3 + assert hf_config["mlp_layer_types"] == ["hash_moe", "hash_moe", "hash_moe", "moe", "moe"] + class TestDecoderHCHeadMappings: """The global decoder HC-head triplet must be replicated mappings.""" @@ -217,3 +243,4 @@ def test_provider_bridge_forces_full_rotary_percent(self): assert out.rotary_percent == 1.0 assert out.rotary_base == 10000 assert out.csa_compress_rotary_base == 160000 + assert out.dsv4_n_hash_layers == 3 From 8eaf21125096b14e242508754196deadc9843abc Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:50:33 +0000 Subject: [PATCH 8/8] test: clarify isolated DSV4 export coverage Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py index f11e55c207..303608764a 100644 --- a/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py +++ b/tests/unit_tests/models/deepseek/test_deepseek_v4_bridge.py @@ -111,6 +111,7 @@ def test_export_hash_layers_from_mcore_field(self): mhc_sinkhorn_iterations=20, moe_shared_expert_intermediate_size=2048, ) + # Stub the generic export to isolate the DSV4 fields without constructing a full Megatron provider. with patch.object( MegatronModelBridge, "megatron_to_hf_config",