[Naeemkh]

Summary

• init_weights skipped every parameter with fewer than two dimensions, on the
  assumption they keep their constructor defaults. That holds for a direct build
  but not after to_empty(), which replaces every parameter and buffer with
  uninitialised memory. Norm weights therefore held whatever the allocator
  returned.
• Reset the modules that own construction-time state before the parameter loop,
  via an explicit list (RMSNorm, nn.LayerNorm, SigmoidTopKRouter) and
  deterministic fills. Not a blanket reset_parameters() over every module:
  nn.Linear's would overwrite the initialisation below and consume RNG, which
  can leave replicas disagreeing across ranks.
• Two further pieces shared the root cause and are fixed here: the sigmoid
  router's expert_bias is one-dimensional, and its expert_ema is a buffer
  — which the parameter loop could never have reached, since it only ever walked
  parameters.
• The trigger is the meta-device build, not tensor parallelism; tp > 1 is
  simply the only thing that currently gates that path. Pure PP, plain FSDP,
  EP-only and single-GPU runs build directly on device and were never affected.
  Nine shipped configs set tp.

Testing

• uv run ruff check kempnerforge/ tests/ passes
• uv run ruff format --check kempnerforge/ tests/ scripts/ passes
• uv run pyright kempnerforge/ passes (0 errors)
• uv run pytest tests/unit/ -v --timeout=60 passes
• If distributed code changed: uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -v
• If training loop / parallelism / optimizers changed: uv run pytest tests/e2e/ --e2e -v

Closes #214

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
kempnerforge/model/init.py	91.66% <100.00%> (+1.66%)	⬆️
kempnerforge/model/router.py	97.84% <100.00%> (+0.07%)	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[timothyngo]

Thanks for addressing all of the comments, LGTM.