fix(deepseek-v4): initialize random training state

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())

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):

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):

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

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``.