[faresobeid]

Part of the pretraining plan (pretraining = a data mode of the SFT trainer). This PR only adds init from scratch.

What

• New model.init: "pretrained" | "scratch" (default pretrained). It replaces model.debug.random_init, which is removed. random_init skipped loading and left the to_empty() memory in the weights.

• scratch reads only the HF config at model.name. That can be a hub id or a local dir with a config.json. No weights are downloaded. The tokenizer comes from tokenizer.name (existing field), or from model.name if that is unset.

• Scratch models always have untied embeddings: tie_word_embeddings is forced to False. The exported config.json says so too.

• After FSDP/EP sharding, the scratch path calls model.initialize_weights(). This walks every module through the HF _init_weights hook, which DSv4 already used:

  • PreTrainedModelPrimeRL._init_weights handles the shared layer types:
    • nn.Linear / nn.Embedding (incl. lm_head and router gate): normal(0, initializer_range), bias 0
    • GroupedExperts: normal(0, initializer_range), biases 0; works on the packed gate_up_proj too
    • RMSNorm / LayerNorm: ones (LayerNorm bias zeros)
    • TokenChoiceTopKRouter: zeros the persistent selection_bias, for every MoE arch
    • nn.Conv1d: torch default
  • The std is the config's initializer_range.
  • Each arch extends the hook for its own layers:
    • gpt-oss: norm, sinks normal(std)
    • nemotron_h: gated norm, and Mamba-2 A ~ U(1,16), D = 1, log-uniform dt → dt_bias
    • qwen3.5 / qwen3.8-flash-next: zero-centred norms, FusedRMSNormGated, GDN dt_bias = 1, A ~ U(0,16), PLE conv zeros
    • DSv4: existing hyperconnection / sink / compressor inits, plus a random tid2eid hash table, seeded per layer so every rank builds the same table

• Cleanup made possible by the shared hook. Removed:

  • FeedForward.init_weights
  • MoE.init_weights
  • SigmoidOutputGatedFeedForward.init_weights
  • DeepseekV4Experts.init_weights
  • flash-next's post-post_init reset loop

  The odd hard-coded trunc_normal(0.02) first projection is gone. Models built directly on a real device (tests) now get the same init through post_init.

• model.init = "scratch" together with model.vlm or model.lora is rejected, because vision towers and adapters are not initialized here.

Config for a rung:

[model]
name = "/path/to/rung-config-dir"   # config.json only; initializer_range sets the std
init = "scratch"

[tokenizer]
name = "Qwen/Qwen3-0.6B"

Tests

tests/unit/train/models/test_init.py covers 14 configs: llama, qwen3, qwen3_moe, glm4_moe, glm_moe_dsa, minimax_m2, laguna, afmoe, gpt_oss, nemotron_h, qwen3_5 (dense and MoE), qwen3_8_flash_next, deepseek_v4.

• CPU: build on meta → to_empty → NaN-fill → init_buffers_post_meta + initialize_weights. Every state-dict tensor must be finite. Linear, embedding and expert weights must have std ≈ initializer_range and mean ≈ 0. Biases must be 0 and RMSNorm weights 1.
• GPU (-m gpu): the same build on CUDA, then a forward pass. The first-token loss must be within 0.25 of ln(vocab). flash-next is CPU-only here: its indexed attention needs a hub kernel.

Rebase / ordering

This touches configs/trainer.py (ModelConfig), trainer/model.py (get_model, load_dcp_from_hf) and models/layers/{moe,mlp,norms}.py plus each arch's PreTrainedModel. It is independent of #3851/#3849/#3855. It overlaps #3833 (models refactor) and #3843 (moe_runtime) only in small hunks (GroupedExperts/router init_weights, _init_weights overrides). It is easiest to land this first, or to rebase it on top of #3833; the conflicts are mechanical.

Test plan (run)

• pytest tests/unit/orchestrator tests/unit/test_configs.py -m "not gpu": 291 passed.
• pytest tests/unit/train -m "not gpu": all pass except test_cpu_adam, which fails because ninja is missing in the local env. That failure is unrelated to this PR.
• ruff check / ruff format: clean.
• One H200 node:
  • tests/unit/train/models/test_init.py: 28 passed (CPU + GPU).
  • -m gpu tests for the touched archs: all pass. Covered: nemotron_h, cp_mamba, qwen3_5, gpt_oss, moe, qwen3_moe, glm4_moe, deepseek_v4 (after the CPU-generator fix).
• End-to-end sft runs with model.init = "scratch" on 2 GPUs (FSDP), with fusions, AC and compile at their defaults:
  • qwen3, ~0.1B config dir, real data (willcb/R1-reverse-wikipedia-paragraphs-v1-1000, lr 1e-3): loss 12.07 → 5.12 in 40 steps.
  • qwen3_moe with EP=2, fake data: 30 steps, finite loss.
  • nemotron_h with AC offloading and compile off: 5 steps, finite loss.
• Export: tools/convert_dcp_to_bf16.py, then AutoModelForCausalLM.from_pretrained loads with no missing or unexpected keys for qwen3 and qwen3_moe. The exported tie_word_embeddings is false. HF loss on a probe sequence ≈ ln V, and embed std = 0.02 (initializer_range).
• Known, already on main: a qwen3_5_moe_text SFT with EP=2 fails in the backward recompute with Pointer argument cannot be accessed from Triton (cpu tensor?) in generate_permute_indices. main with debug.random_init fails the same way, so this PR does not cause it.