fix(deepseek-v4): initialize random training state

[akoumpa]

Summary

• initialize DeepSeek-V4 HyperConnection and HyperHead parameters for random-init training, including MTP blocks
• initialize HashGate weights and build deterministic, unique, balanced fallback expert routes before checkpoint loading
• add focused regression coverage for HyperConnection state and HashGate route validity

Root cause

With load_base_model=false, HyperConnection parameters allocated with torch.empty were never initialized. The HashGate route table also remained all zeros, so every token sent every top-k route to expert 0. DeepEP requires distinct expert IDs within a token's top-k routes; duplicate routes poisoned the active routed-expert backward pass and produced a NaN gradient norm at step 0.

User impact

DeepSeek-V4 random-init training with pipeline parallelism, expert parallelism, and DeepEP now starts with finite gradients instead of failing on the first backward pass. Checkpoint-loaded routes are unchanged because checkpoint loading still overwrites the initialized fallback table.

Validation

• uv run --no-sync pytest -q tests/unit_tests/models/deepseek_v4: 174 passed, 17 skipped
• ruff format --check on all five changed files
• ruff check on the three changed production files
• git diff --check
• cw-dfw, 2 nodes x 8 H100, PP2/EP8, DeepEP, bf16, random init:
  • baseline job 13633831: step-0 loss 12.1582, grad_norm=nan, validation loss nan
  • instrumented job 13634773: non-finite gradients localized to rank 0, layer 1 routed-expert gate_and_up_projs and down_projs
  • fixed job 13634923: all 16 ranks reported zero non-finite gradients; step-0 loss 12.1606, grad_norm=2.4152, validation loss 12.1979; exit code 0

Fixes NVBug 6317402.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test 9357554

[akoumpa]

/ok to test 9357554

[yuhezhang-ai]

Approving — well-diagnosed fix. The all-zeros-route → duplicate-top-k → DeepEP NaN chain is precisely identified, and the route formula (token_id * topk + offset) % n_experts is a clean choice: it enumerates contiguous integers so routes are distinct per token (for topk <= n_experts, always true) and load-balanced to within 1 by construction — exactly what the two new tests assert. I traced the full init path (initialize_weights → model.init_weights → per-layer + mtp.layers) and HC/HashGate coverage is complete for both the main stack and MTP blocks. _init_weights correctly skips HashGate (not a Gate subclass), so the explicit gate.init_weights() call is necessary, and the if self.hc_head is not None guard is safe under PP nulling.

One suggestion worth acting on, plus a minor note:

1. (robustness) Future-proof MTP hash-route init. MTP.init_weights does not mirror the Blocks if isinstance(self.mlp.gate, DeepseekV4HashGate): self.mlp.gate.init_weights(...) guard. That is correct today only because MTP builds a plain MoE(...) and never swaps in a HashGate — an implicit invariant. If MTP ever gains hash routing, its tid2eid would silently regress to all-zeros and reintroduce this exact NaN. Adding the same isinstance guard to MTP.init_weights is zero-cost insurance against re-breaking the very bug being fixed.

2. (minor) Signature-change note. HashGate.init_weights changed from (buffer_device=None) to (init_std=0.02), diverging from the buffer_device-first convention of its sibling init_weights methods. The only caller uses the init_std= keyword and nothing passes a device positionally (the generic path is .apply(_init_weights), not .init_weights()), so it is safe — a one-line comment noting the intentional divergence would save the next reader a double-take.