Restore construction-time state after a meta-device build

kempnerforge/model/init.py

@@ -9,6 +9,16 @@
 from kempnerforge.config.schema import ModelConfig
 from kempnerforge.model.moe import MoEMLP
 from kempnerforge.model.mot import MoTAttention, MoTBlock
+from kempnerforge.model.norm import RMSNorm
+from kempnerforge.model.router import SigmoidTopKRouter
+
+# Modules holding construction-time state that ``to_empty`` discards and the
+# parameter loop below cannot reach: one-dimensional parameters and buffers.
+# Only these are reset, and only through deterministic fills. A blanket
+# ``reset_parameters()`` over every module would let ``nn.Linear`` overwrite the
+# initialisation below with its own, and would consume RNG, which can leave
+# replicas disagreeing across ranks.
+_RESETS_OWN_STATE = (RMSNorm, nn.LayerNorm, SigmoidTopKRouter)
 
 
 def init_weights(model: nn.Module, config: ModelConfig) -> None:
@@ -21,11 +31,18 @@ def init_weights(model: nn.Module, config: ModelConfig) -> None:
     - Residual output projections (o_proj, down_proj): scaled by 1/sqrt(2 * n_layers)
     - Cross-attention block residuals: zero-initialized (identity-at-init warm-start)
     - MoT per-modality residual projections: zero-initialized (identity-at-construction)
-    - Norm layers: weight=1 (already default)
+    - Norm layers and the sigmoid router's own state: reset to their
+      construction-time values, which a meta-device build does not preserve
     """
     std = config.init_std
     residual_scale = 1.0 / math.sqrt(2.0 * config.n_layers)
 
+    for module in model.modules():
+        if isinstance(module, _RESETS_OWN_STATE) and not any(
+            t.is_meta for t in (*module.parameters(recurse=False), *module.buffers(recurse=False))
+        ):
+            module.reset_parameters()
+
     for name, param in model.named_parameters():
         if param.is_meta:
             continue

kempnerforge/model/router.py

@@ -127,6 +127,16 @@ def __init__(
         self.z_loss = torch.tensor(0.0)
         self.expert_counts: torch.Tensor = torch.zeros(num_experts)
 
+    def reset_parameters(self) -> None:
+        """Restore this router's construction-time state.
+
+        ``to_empty()`` on a meta build leaves both holding whatever the
+        allocator returns, and ``init_weights`` covers neither: ``expert_bias``
+        is one-dimensional and ``expert_ema`` is a buffer.
+        """
+        nn.init.zeros_(self.expert_bias)
+        self.expert_ema.fill_(1.0 / self.num_experts)  # type: ignore[reportCallIssue]
+
     def set_step(self, step: int, max_steps: int) -> None:
         """Set current training step for adaptive bias scheduling."""
         self._step = step

tests/unit/test_model.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import math
+from dataclasses import replace
 
 import pytest
 import torch
@@ -1169,6 +1170,59 @@ def test_custom_init_std(self):
         ]
         assert all(abs(s - 0.01) < 0.01 for s in stds)
 
+    @pytest.mark.parametrize(
+        "overrides",
+        [
+            {},
+            {"qk_norm": True},
+            {"norm_type": "layernorm"},
+            {
+                "num_experts": 4,
+                "moe_top_k": 2,
+                "moe_frequency": 1,
+                "moe_router": "sigmoid_topk",
+            },
+        ],
+        ids=["dense", "qk_norm", "layernorm", "sigmoid_router"],
+    )
+    def test_materialized_state_matches_a_constructed_model(self, overrides):
+        """A meta build must end up where a direct build starts.
+
+        ``to_empty`` replaces every parameter and buffer with uninitialised
+        memory. Anything set in a constructor rather than by ``init_weights``
+        -- norm weights, the sigmoid router's bias, its utilisation buffer --
+        keeps that memory unless something restores it.
+
+        The poison step is load-bearing. Freed blocks often still hold the
+        values that were just in them, so an allocator handing back ones makes
+        the naive version of this test pass while the bug is present.
+        """
+        config = replace(_INIT_CONFIG, **overrides)
+
+        reference = Transformer(config)
+        reference.init_weights_and_freqs()
+
+        with torch.device("meta"):
+            model = Transformer(config)
+        model.to_empty(device="cpu")
+        for tensor in (*model.parameters(), *model.buffers()):
+            tensor.detach().fill_(float("nan"))
+        model.init_weights_and_freqs()
+
+        for name, param in model.named_parameters():
+            assert torch.isfinite(param).all(), f"{name} left uninitialised"
+        for name, buf in model.named_buffers():
+            assert torch.isfinite(buf).all(), f"buffer {name} left uninitialised"
+
+        # Constructor-set state is deterministic, so it must match exactly.
+        ref_params = dict(reference.named_parameters())
+        for name, param in model.named_parameters():
+            if param.dim() < 2:
+                assert torch.equal(param, ref_params[name]), name
+        ref_bufs = dict(reference.named_buffers())
+        for name, buf in model.named_buffers():
+            assert torch.equal(buf, ref_bufs[name]), f"buffer {name}"
+
     def test_skips_meta_parameters(self):
         """init_weights should skip parameters on meta device without error."""
         with torch.device("meta"):