From 9357554df4ab26bf71489aaf8d1ec1ce8507ccfb Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Thu, 9 Jul 2026 08:48:41 -0700 Subject: [PATCH] fix(deepseek-v4): initialize random training state Signed-off-by: Alexandros Koumparoulis --- .../components/models/deepseek_v4/layers.py | 22 ++++++ .../components/models/deepseek_v4/model.py | 31 +++++--- .../components/models/deepseek_v4/mtp.py | 7 +- .../models/deepseek_v4/test_dsv4_layers.py | 20 +++-- .../deepseek_v4/test_dsv4_model_smoke.py | 75 +++++++++++++++++++ 5 files changed, 138 insertions(+), 17 deletions(-) diff --git a/nemo_automodel/components/models/deepseek_v4/layers.py b/nemo_automodel/components/models/deepseek_v4/layers.py index 8262e37c5a..5d7525f5d9 100644 --- a/nemo_automodel/components/models/deepseek_v4/layers.py +++ b/nemo_automodel/components/models/deepseek_v4/layers.py @@ -952,6 +952,17 @@ def __init__( self.base = nn.Parameter(torch.empty(mix, dtype=torch.float32)) self.scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + @torch.no_grad() + def init_weights(self, init_std: float) -> None: + """Initialize HyperConnection parameters using the DeepSeek-V4 reference scheme. + + Args: + init_std: Standard deviation for the ``fn`` weight initialization. + """ + nn.init.normal_(self.fn, mean=0.0, std=init_std) + nn.init.zeros_(self.base) + nn.init.ones_(self.scale) + def compute_weights(self, hidden_streams: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: flat = hidden_streams.flatten(start_dim=2).float() # [B, S, H*D] # HC mixer params are kept in fp32 for Sinkhorn stability — cast defensively. @@ -1010,6 +1021,17 @@ def __init__(self, hc_mult: int, hidden_size: int, hc_eps: float, rms_norm_eps: self.hc_base = nn.Parameter(torch.empty(self.hc_mult, dtype=torch.float32)) self.hc_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) + @torch.no_grad() + def init_weights(self, init_std: float) -> None: + """Initialize HyperHead parameters using the DeepSeek-V4 reference scheme. + + Args: + init_std: Standard deviation for the ``hc_fn`` weight initialization. + """ + nn.init.normal_(self.hc_fn, mean=0.0, std=init_std) + nn.init.zeros_(self.hc_base) + nn.init.ones_(self.hc_scale) + def forward(self, x: torch.Tensor) -> torch.Tensor: flat = x.flatten(2).float() mixes = torch.nn.functional.linear(_rms_norm_last_dim(flat, self.norm_eps), self.hc_fn.float()) diff --git a/nemo_automodel/components/models/deepseek_v4/model.py b/nemo_automodel/components/models/deepseek_v4/model.py index faeb985317..a8ada8c09f 100644 --- a/nemo_automodel/components/models/deepseek_v4/model.py +++ b/nemo_automodel/components/models/deepseek_v4/model.py @@ -219,14 +219,15 @@ def ffn_prepare(hidden_streams: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso dtype = x.dtype return post.to(dtype).unsqueeze(-1) * mlp_out.unsqueeze(-2) + torch.matmul(comb.transpose(-1, -2).to(dtype), x) - def init_weights(self, buffer_device: torch.device) -> None: + def init_weights(self, buffer_device: torch.device, init_std: float = 0.02) -> None: self.input_layernorm.reset_parameters() self.post_attention_layernorm.reset_parameters() - self.self_attn.init_weights(buffer_device) - self.mlp.init_weights(buffer_device) - # HC mixer params stay at whatever the checkpoint provides (init.normal_ - # on ``fn``, init.zeros_ on ``base``, init.ones_ on ``scale`` for random - # init — matches HF's _init_weights at modular_deepseek_v4.py:923-926). + self.self_attn.init_weights(buffer_device, init_std=init_std) + self.mlp.init_weights(buffer_device, init_std=init_std) + if isinstance(self.mlp.gate, DeepseekV4HashGate): + self.mlp.gate.init_weights(init_std=init_std) + self.attn_hc.init_weights(init_std) + self.ffn_hc.init_weights(init_std) class DeepseekV4HashGate(nn.Module): @@ -279,10 +280,17 @@ def set_input_ids(self, input_ids: torch.Tensor | None) -> None: def update_bias(self) -> None: """No-op for compat with callers that walk MoE gates and call update_bias.""" - def init_weights(self, buffer_device: torch.device | None = None) -> None: - nn.init.zeros_(self.weight) + def init_weights(self, init_std: float = 0.02) -> None: + """Initialize the trainable gate and a valid deterministic hash table. + + Args: + init_std: Standard deviation for the routing weight initialization. + """ + nn.init.normal_(self.weight, mean=0.0, std=init_std) with torch.no_grad(): - self.tid2eid.zero_() + token_ids = torch.arange(self.tid2eid.shape[0], device=self.tid2eid.device).unsqueeze(1) + expert_offsets = torch.arange(self.topk, device=self.tid2eid.device).unsqueeze(0) + self.tid2eid.copy_((token_ids * self.topk + expert_offsets) % self.n_experts) def forward( self, @@ -569,13 +577,16 @@ def update_moe_gate_bias(self) -> None: @torch.no_grad() def init_weights(self, buffer_device: torch.device | None = None) -> None: buffer_device = buffer_device or torch.device(f"cuda:{torch.cuda.current_device()}") + init_std = float(getattr(self.config, "initializer_range", 0.02)) with buffer_device: if self.embed_tokens is not None: nn.init.normal_(self.embed_tokens.weight) if self.norm is not None: self.norm.reset_parameters() + if self.hc_head is not None: + self.hc_head.init_weights(init_std) for layer in self.layers.values(): - layer.init_weights(buffer_device=buffer_device) + layer.init_weights(buffer_device=buffer_device, init_std=init_std) class DeepseekV4ForCausalLM(HFCheckpointingMixin, nn.Module, MoEFSDPSyncMixin): diff --git a/nemo_automodel/components/models/deepseek_v4/mtp.py b/nemo_automodel/components/models/deepseek_v4/mtp.py index b4cea6384c..7854ddea0a 100644 --- a/nemo_automodel/components/models/deepseek_v4/mtp.py +++ b/nemo_automodel/components/models/deepseek_v4/mtp.py @@ -185,8 +185,11 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None: with target_device: nn.init.trunc_normal_(self.e_proj.weight, mean=0.0, std=init_std) nn.init.trunc_normal_(self.h_proj.weight, mean=0.0, std=init_std) - self.self_attn.init_weights(target_device) - self.mlp.init_weights(target_device) + self.self_attn.init_weights(target_device, init_std=init_std) + self.mlp.init_weights(target_device, init_std=init_std) + self.attn_hc.init_weights(init_std) + self.ffn_hc.init_weights(init_std) + self.hc_head.init_weights(init_std) class DeepseekV4MTPModule(nn.Module): diff --git a/tests/unit_tests/models/deepseek_v4/test_dsv4_layers.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_layers.py index 20f521a3a0..4356d6617f 100644 --- a/tests/unit_tests/models/deepseek_v4/test_dsv4_layers.py +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_layers.py @@ -693,11 +693,8 @@ class TestDeepseekV4HyperConnection: @pytest.fixture def hc(self): - # ``DeepseekV4HyperConnection`` allocates ``fn``/``base``/``scale`` - # via ``torch.empty(...)``; those are uninitialized memory and may - # contain NaN bit patterns. Zero them so the Sinkhorn-row test has - # a well-defined starting point (real model loads init from the - # checkpoint via the state-dict adapter, not via ``empty``). + # Use explicit zeros so formula-specific tests have a deterministic + # starting point independent of the production random-init scheme. m = DeepseekV4HyperConnection( hc_mult=4, hidden_size=16, @@ -711,6 +708,19 @@ def hc(self): m.scale.zero_() return m + def test_init_weights_matches_reference_scheme(self, hc): + with torch.no_grad(): + hc.fn.fill_(float("nan")) + hc.base.fill_(float("nan")) + hc.scale.fill_(float("nan")) + + hc.init_weights(0.02) + + assert torch.isfinite(hc.fn).all() + assert torch.count_nonzero(hc.fn) > 0 + torch.testing.assert_close(hc.base, torch.zeros_like(hc.base)) + torch.testing.assert_close(hc.scale, torch.ones_like(hc.scale)) + def test_parameter_dtypes_are_fp32(self, hc): # HC params must stay fp32 even when the surrounding model is bf16. assert hc.fn.dtype == torch.float32 diff --git a/tests/unit_tests/models/deepseek_v4/test_dsv4_model_smoke.py b/tests/unit_tests/models/deepseek_v4/test_dsv4_model_smoke.py index a18e4b94da..f706e20abd 100644 --- a/tests/unit_tests/models/deepseek_v4/test_dsv4_model_smoke.py +++ b/tests/unit_tests/models/deepseek_v4/test_dsv4_model_smoke.py @@ -414,6 +414,53 @@ def fail_cast(*args, **kwargs): model.initialize_weights(buffer_device=torch.device("cpu"), dtype=torch.bfloat16) + def test_initialize_weights_initializes_all_hyper_connection_parameters(self): + cfg = _tiny_config( + num_hidden_layers=1, + num_hash_layers=0, + compress_ratios=[0], + num_nextn_predict_layers=1, + ) + model = DeepseekV4ForCausalLM( + cfg, + backend=BackendConfig( + attn="sdpa", + linear="torch", + rms_norm="torch", + rope_fusion=False, + enable_hf_state_dict_adapter=False, + dispatcher="torch", + experts="torch_mm", + ), + ) + block = model.model.layers["0"] + assert model.mtp is not None + mtp_block = model.mtp.layers[0] + hyper_modules = (block.attn_hc, block.ffn_hc, mtp_block.attn_hc, mtp_block.ffn_hc) + hyper_heads = (model.model.hc_head, mtp_block.hc_head) + with torch.no_grad(): + for module in hyper_modules: + module.fn.fill_(float("nan")) + module.base.fill_(float("nan")) + module.scale.fill_(float("nan")) + for head in hyper_heads: + head.hc_fn.fill_(float("nan")) + head.hc_base.fill_(float("nan")) + head.hc_scale.fill_(float("nan")) + + model.initialize_weights(buffer_device=torch.device("cpu"), dtype=torch.float32) + + for module in hyper_modules: + assert torch.isfinite(module.fn).all() + assert torch.count_nonzero(module.fn) > 0 + torch.testing.assert_close(module.base, torch.zeros_like(module.base)) + torch.testing.assert_close(module.scale, torch.ones_like(module.scale)) + for head in hyper_heads: + assert torch.isfinite(head.hc_fn).all() + assert torch.count_nonzero(head.hc_fn) > 0 + torch.testing.assert_close(head.hc_base, torch.zeros_like(head.hc_base)) + torch.testing.assert_close(head.hc_scale, torch.ones_like(head.hc_scale)) + def test_hash_gate_tid2eid_uses_deepep_runtime_int64_dtype(self): cfg = _tiny_config(num_hidden_layers=1, num_hash_layers=1, compress_ratios=[0]) model = DeepseekV4ForCausalLM( @@ -438,6 +485,34 @@ def test_hash_gate_tid2eid_uses_deepep_runtime_int64_dtype(self): _, indices, _ = gate(torch.zeros(3, cfg.hidden_size), torch.ones(3, dtype=torch.bool)) assert indices.dtype == torch.long + def test_initialize_weights_builds_valid_hash_routes(self): + cfg = _tiny_config(num_hidden_layers=1, num_hash_layers=1, compress_ratios=[0]) + model = DeepseekV4ForCausalLM( + cfg, + backend=BackendConfig( + attn="sdpa", + linear="torch", + rms_norm="torch", + rope_fusion=False, + enable_hf_state_dict_adapter=False, + dispatcher="torch", + experts="torch_mm", + ), + ) + gate = model.model.layers["0"].mlp.gate + with torch.no_grad(): + gate.weight.fill_(float("nan")) + gate.tid2eid.zero_() + + model.initialize_weights(buffer_device=torch.device("cpu"), dtype=torch.float32) + + assert torch.isfinite(gate.weight).all() + assert torch.count_nonzero(gate.weight) > 0 + sorted_routes = gate.tid2eid.sort(dim=-1).values + assert torch.all(sorted_routes[:, 1:] != sorted_routes[:, :-1]) + expert_load = torch.bincount(gate.tid2eid.flatten(), minlength=gate.n_experts) + assert expert_load.max() - expert_load.min() <= 1 + def test_hc_comb_transpose_used_at_attn_and_mlp_sites(self): """Both HC expand sites mix residual streams as ``comb.T @ x``.