feat: model.init = scratch (init from config with per-layer init)

docs/training.md

@@ -256,6 +256,7 @@ The shared script passes the trainer rank-0 hostname directly to the online-eval
 
 | Knob | What it controls |
 |---|---|
+| `model.init` | `pretrained` (default) loads the weights at `model.name`; `scratch` builds the model from the HF config at `model.name` only, with untied embeddings and fresh weights drawn with the config's `initializer_range` (set `tokenizer.name` if `model.name` has no tokenizer) |
 | `data.name` | HF dataset name or local path |
 | `data.batch_size` | Tokens per trainer step (packed) |
 | `data.seq_len` | Per-sample sequence length |

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -135,9 +135,6 @@ class DebugModelConfig(BaseConfig):
     num_layers: int | None = None
     """Override the number of transformer layers (truncates the model)."""
 
-    random_init: bool = False
-    """Randomly initialize the model instead of loading weights."""
-
     force_balanced_routing: bool = False
     """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent."""
 
@@ -267,6 +264,9 @@ class MoERuntimeConfig(BaseConfig):
 
 
 class ModelConfig(BaseModelConfig):
+    init: Literal["pretrained", "scratch"] = "pretrained"
+    """How to initialize the weights. ``pretrained`` loads the checkpoint at ``name``. ``scratch`` reads only the HF config at ``name`` (hub id or local directory with a ``config.json``), unties the embeddings and initializes every weight with the architecture's init: linears, embeddings and experts from ``N(0, initializer_range)`` of that config, norms to identity, biases to zero. Set ``tokenizer.name`` when ``name`` has no tokenizer."""
+
     conversion_dir: Path | None = None
     """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory."""
 
@@ -341,6 +341,14 @@ class ModelConfig(BaseModelConfig):
     fused_lm_head_token_chunk_size: int | Literal["disabled"] = 8192
     """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head. In SFT the fused head computes the summed cross-entropy and its gradients chunk by chunk, holding one chunk's full-vocab logits at a time."""
 
+    @model_validator(mode="after")
+    def validate_scratch_init(self):
+        if self.init == "scratch" and (self.vlm is not None or self.lora is not None):
+            raise ValueError(
+                "model.init='scratch' supports full training of text models only; unset model.vlm and model.lora"
+            )
+        return self
+
     @model_validator(mode="after")
     def vlm_cp_requires_ulysses(self):
         if self.vlm is not None and self.cp > 1 and self.cp_style != "ulysses":

src/prime_rl/entrypoints/rl.py

@@ -532,7 +532,7 @@ def rl(config: RLConfig):
     if not config.dry_run:
         from prime_rl.trainer.model import pre_download_model
 
-        pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.random_init)
+        pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.init == "scratch")
 
     if config.slurm is not None:
         rl_slurm(config)

src/prime_rl/entrypoints/sft.py

@@ -485,7 +485,7 @@ def sft(config: SFTConfig):
         from prime_rl.trainer.model import pre_download_model
         from prime_rl.trainer.sft.data import pre_download_data
 
-        pre_download_model(config.model.name, skip_weights=config.model.debug.random_init)
+        pre_download_model(config.model.name, skip_weights=config.model.init == "scratch")
         pre_download_data(config.data, config.env_vars)
         if config.val is not None:
             pre_download_data(config.val.data, config.env_vars)

src/prime_rl/trainer/model.py

@@ -76,14 +76,16 @@
 def pre_download_model(model_name: str, *, skip_weights: bool = False) -> None:
     """Pre-download model from HuggingFace Hub so all nodes have cached weights before training.
 
-    With ``skip_weights`` (random-init debug runs), only config and tokenizer files are fetched.
+    With ``skip_weights`` (``model.init = "scratch"``), only config and tokenizer files are fetched.
     """
     if Path(model_name).exists():
         get_logger().info(f"Model {model_name} found at local path, skipping download")
         return
     t0 = time.perf_counter()
     if skip_weights:
-        get_logger().info(f"Pre-downloading config and tokenizer for {model_name} (random init, skipping weights)")
+        get_logger().info(
+            f"Pre-downloading config and tokenizer for {model_name} (init from scratch, skipping weights)"
+        )
         path = snapshot_download(
             repo_id=model_name, repo_type="model", allow_patterns=["*.json", "*.txt", "tokenizer*", "*.jinja"]
         )
@@ -356,6 +358,8 @@ def get_model(
         ),
     )
     model_config.use_cache = False
+    if config.init == "scratch":
+        model_config.tie_word_embeddings = False
     is_vlm_arch = is_vlm_architecture(model_config)
 
     if is_vlm_training:
@@ -715,8 +719,9 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral
     model.init_buffers_post_meta()
 
     logger = get_logger()
-    if config.debug.random_init:
-        logger.warning("Randomly initializing model. Skipping loading weights from HF.")
+    if config.init == "scratch":
+        logger.info(f"Initializing model weights from scratch (config: {config.name})")
+        model.initialize_weights()
         _move_buffers_to_cuda(model, config)
         return
 

src/prime_rl/trainer/models/base.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass
 from typing import Literal
 
-from torch import Tensor
+from torch import Tensor, nn
 from transformers.modeling_utils import PreTrainedModel
 
+from prime_rl.trainer.models.layers.moe import GroupedExperts, TokenChoiceTopKRouter
+from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm
 from prime_rl.utils.cp import CPContext
 
 CPStyle = Literal["ring", "ulysses"]
@@ -143,6 +145,24 @@ def convert_adapter_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tenso
         """
         return state_dict
 
+    def _init_weights(self, module: nn.Module) -> None:
+        """Initialize the parameters and persistent buffers that ``module`` owns directly.
+
+        ``initialize_weights()`` applies this to every module. Architectures extend it for their own layer types.
+        """
+        std = self.config.get_text_config().initializer_range
+        # The qkv fusion leaves the q/k/v projections without parameters
+        if isinstance(module, (nn.Linear, nn.Embedding)) and module.weight is not None:
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            if getattr(module, "bias", None) is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv1d):
+            module.reset_parameters()
+        elif isinstance(module, GroupedExperts):
+            module.init_weights(std)
+        elif isinstance(module, (RMSNorm, LayerNorm, TokenChoiceTopKRouter)):
+            module.init_weights()
+
     def init_buffers_post_meta(self) -> None:
         """
         Initialize buffers that are not in the state dict after loading with meta device.

src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py

@@ -121,7 +121,7 @@ def _init_weights(self, module: nn.Module) -> None:
         if isinstance(module, (DeepseekV4Attention, DeepseekV4HyperConnection, DeepseekV4HyperHead)):
             module.init_weights(init_std)
         elif isinstance(module, DeepseekV4MoE):
-            module.init_weights(init_std, module.tokens_per_expert.device)
+            module.init_weights()
         elif isinstance(module, DeepseekV4RotaryEmbedding):
             module.init_buffers_post_meta()
 

src/prime_rl/trainer/models/deepseek_v4/moe.py

@@ -7,7 +7,6 @@
 
 import torch
 import torch.nn.functional as F
-from torch import nn
 
 from prime_rl.trainer.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config
 from prime_rl.trainer.models.layers.mlp import FeedForward
@@ -99,7 +98,7 @@ class DeepseekV4HashRouter(DeepseekV4Router):
     """A bootstrap layer's router: selection is a frozen token-id lookup, gating still learned.
 
     `tid2eid` replaces the top-k of the scores with `tid2eid[token_id]`, read from the checkpoint
-    and zeros until one fills it. A frozen selection cannot be steered, so a hash layer has no use
+    or drawn by `DeepseekV4MoE.init_weights`. A frozen selection cannot be steered, so a hash layer has no use
     for the aux-loss-free load-balancing bias, and HF's `DeepseekV4HashRouter` carries no
     `e_score_correction_bias` to load into it either; `selection_bias=False` leaves that buffer
     unbuilt, keeping the state dict aligned with HF's.
@@ -132,22 +131,13 @@ class DeepseekV4Experts(GroupedExperts):
 
     The shared `GroupedExperts` already stores the stacked `gate_proj`/`up_proj`/`down_proj`
     that the on-disk per-expert `w1`/`w2`/`w3` convert into, and its `forward` reaches the
-    activation through `self.activation`, so only that attribute and the initialization
-    spread differ from the base class.
+    activation through `self.activation`, so only that attribute differs from the base class.
     """
 
     def __init__(self, dim: int, hidden_dim: int, num_experts: int, swiglu_limit: float) -> None:
         super().__init__(dim, hidden_dim, num_experts, expert_type="gated")
         self.activation = ClampedSwiglu(swiglu_limit)
 
-    def init_weights(self, init_std: float) -> None:
-        # Both halves of HF's fused gate_up_proj are drawn from the same std=0.02
-        # distribution, so gate and up match that here despite being separate tensors.
-        # The base class scales up_proj by init_std instead.
-        nn.init.trunc_normal_(self.gate_proj, mean=0.0, std=0.02)
-        nn.init.trunc_normal_(self.up_proj, mean=0.0, std=0.02)
-        nn.init.trunc_normal_(self.down_proj, mean=0.0, std=init_std)
-
 
 class DeepseekV4MLP(FeedForward):
     """Dense SwiGLU MLP with V4's clamp, used as the MoE layer's shared expert.
@@ -235,6 +225,15 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int):
         self.layer_idx = layer_idx
         self.is_hash = is_hash
 
+    def init_weights(self) -> None:
+        if self.is_hash:
+            # Seeded by layer so every rank builds the same token-to-expert table.
+            generator = torch.Generator(device="cpu").manual_seed(self.layer_idx)
+            scores = torch.rand(
+                self.router.tid2eid.shape[0], self.router.num_experts, generator=generator, device="cpu"
+            )
+            self.router.tid2eid.copy_(scores.topk(self.router.top_k, dim=-1).indices)
+
     def forward(
         self,
         x: torch.Tensor,

src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py

@@ -93,6 +93,13 @@ class GptOssPreTrainedModel(PreTrainedModelPrimeRL):
     _supports_attention_backend = True
     _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"]
 
+    def _init_weights(self, module: nn.Module) -> None:
+        super()._init_weights(module)
+        if isinstance(module, GptOssRMSNorm):
+            nn.init.ones_(module.weight)
+        elif isinstance(module, GptOssAttention):
+            nn.init.normal_(module.sinks, mean=0.0, std=self.config.initializer_range)
+
     @classmethod
     def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool:
         return is_hf_state_dict(state_dict)

src/prime_rl/trainer/models/layers/mlp.py

@@ -30,13 +30,3 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -
         gate = self.gate_proj(x) if self.gate_proj is not None else None
         up = self.up_proj(x)
         return self.down_proj(self.activation.apply(gate, up))
-
-    def init_weights(self, init_std: float = 0.02) -> None:
-        first_projection = self.gate_proj if self.gate_proj is not None else self.up_proj
-        nn.init.trunc_normal_(first_projection.weight, mean=0.0, std=0.02)
-        remaining = (self.up_proj, self.down_proj) if self.gate_proj is not None else (self.down_proj,)
-        for linear in remaining:
-            nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
-        for linear in (self.gate_proj, self.up_proj, self.down_proj):
-            if linear is not None and linear.bias is not None:
-                nn.init.zeros_(linear.bias)

src/prime_rl/trainer/models/layers/moe.py

@@ -113,20 +113,12 @@ def forward(
     ) -> torch.Tensor:
         return self.compute(self, x, num_tokens_per_expert)
 
-    def init_weights(self, init_std: float):
-        if self.gate_up_proj is None:
-            first_projection = self.gate_proj if self.gate_proj is not None else self.up_proj
-            nn.init.trunc_normal_(first_projection, mean=0.0, std=0.02)
-            remaining = (self.up_proj, self.down_proj) if self.gate_proj is not None else (self.down_proj,)
-        else:
-            gate_proj, up_proj = self.gate_up_proj.chunk(2, dim=1)
-            nn.init.trunc_normal_(gate_proj, mean=0.0, std=0.02)
-            remaining = (up_proj, self.down_proj)
-        for weight in remaining:
-            nn.init.trunc_normal_(weight, mean=0.0, std=init_std)
-        for bias in (self.gate_proj_bias, self.up_proj_bias, self.down_proj_bias):
-            if bias is not None:
-                nn.init.zeros_(bias)
+    def init_weights(self, init_std: float) -> None:
+        for name, param in self.named_parameters(recurse=False):
+            if name.endswith("_bias"):
+                nn.init.zeros_(param)
+            else:
+                nn.init.normal_(param, mean=0.0, std=init_std)
 
 
 class TokenChoiceTopKRouter(nn.Module):
@@ -266,10 +258,9 @@ def forward(
 
         return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum
 
-    def init_weights(self, init_std: float):
-        nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std)
-        if self.gate.bias is not None:
-            nn.init.zeros_(self.gate.bias)
+    def init_weights(self) -> None:
+        if self.selection_bias is not None:
+            nn.init.zeros_(self.selection_bias)
 
 
 class MoE(nn.Module):
@@ -412,19 +403,3 @@ def forward(
             routed_output = routed_output + shared_output
 
         return routed_output.reshape(bs, slen, dim)
-
-    def init_weights(
-        self,
-        init_std: float,
-        buffer_device: torch.device,
-    ):
-        self.experts.init_weights(init_std)
-        self.router.init_weights(init_std)
-        if self.shared_expert is not None:
-            self.shared_expert.init_weights(init_std)
-
-        with torch.device(buffer_device):
-            self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32)
-            self.routing_confidence_sum = torch.tensor(0.0, dtype=torch.float32)
-            if self.router.selection_bias is not None:
-                self.router.selection_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)

src/prime_rl/trainer/models/layers/norms.py

@@ -43,6 +43,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight * hidden_states.to(input_dtype)
 
+    def init_weights(self) -> None:
+        nn.init.ones_(self.weight)
+
     def extra_repr(self):
         return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
 
@@ -55,5 +58,9 @@ def __init__(self, dim: int, eps: float = 1e-6):
         self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
         self.bias = nn.Parameter(torch.zeros(dim, dtype=torch.float32))
 
+    def init_weights(self) -> None:
+        nn.init.ones_(self.weight)
+        nn.init.zeros_(self.bias)
+
     def forward(self, x: torch.Tensor):
         return F.layer_norm(x.float(), (self.dim,), self.weight.float(), self.bias.float(), self.eps).type_as(x)

src/prime_rl/trainer/models/nemotron_h/mamba.py

@@ -92,6 +92,9 @@ def __init__(self, config: NemotronHConfig) -> None:
         self.conv_size = self.intermediate_size + 2 * self.group_state_size
         self.chunk_size = config.chunk_size
         self.time_step_limit = config.time_step_limit
+        self.time_step_min = config.time_step_min
+        self.time_step_max = config.time_step_max
+        self.time_step_floor = config.time_step_floor
         self.activation = config.mamba_hidden_act
 
         projection_size = self.intermediate_size + self.conv_size + self.num_heads
@@ -105,8 +108,8 @@ def __init__(self, config: NemotronHConfig) -> None:
             bias=config.use_conv_bias,
         )
         self.dt_bias = nn.Parameter(torch.empty(self.num_heads))
-        self.A_log = nn.Parameter(torch.arange(1, self.num_heads + 1, dtype=torch.get_default_dtype()).log())
-        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.A_log = nn.Parameter(torch.empty(self.num_heads))
+        self.D = nn.Parameter(torch.empty(self.num_heads))
         self.norm = GatedRMSNorm(
             self.intermediate_size,
             group_size=self.intermediate_size // self.num_groups,
@@ -116,17 +119,19 @@ def __init__(self, config: NemotronHConfig) -> None:
         self.causal_conv1d = causal_conv1d
         self.prepare_sequence_ids = prepare_sequence_ids
         self.scan = mamba_chunk_scan_combined
-
-        time_steps = torch.exp(
-            torch.rand(self.num_heads, dtype=torch.float32)
-            * (math.log(config.time_step_max) - math.log(config.time_step_min))
-            + math.log(config.time_step_min)
-        ).clamp(min=config.time_step_floor)
-        with torch.no_grad():
-            self.dt_bias.copy_((time_steps + torch.log(-torch.expm1(-time_steps))).to(self.dt_bias.dtype))
+        self.init_weights()
 
         self.cp_context = CPContext()
 
+    @torch.no_grad()
+    def init_weights(self) -> None:
+        # Mamba-2 defaults: A ~ U(1, 16), D = 1, and dt_bias = softplus^-1(dt) for a log-uniform dt.
+        self.A_log.uniform_(1, 16).log_()
+        self.D.fill_(1.0)
+        dt = self.dt_bias.uniform_(math.log(self.time_step_min), math.log(self.time_step_max))
+        dt.exp_().clamp_(min=self.time_step_floor)
+        dt.add_(torch.log(-torch.expm1(-dt)))
+
     def forward(self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
         projected_states = self.in_proj(hidden_states)

src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py

@@ -14,7 +14,7 @@
     is_hf_state_dict,
     is_prime_state_dict,
 )
-from prime_rl.trainer.models.nemotron_h.mamba import NemotronHMamba2
+from prime_rl.trainer.models.nemotron_h.mamba import GatedRMSNorm, NemotronHMamba2
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
 
 
@@ -138,6 +138,13 @@ def cp_support(cls, config) -> CPSupport:
             "Mamba layers require Ulysses to reconstruct full sequences while sharding Mamba heads",
         )
 
+    def _init_weights(self, module: nn.Module) -> None:
+        super()._init_weights(module)
+        if isinstance(module, GatedRMSNorm):
+            nn.init.ones_(module.weight)
+        elif isinstance(module, NemotronHMamba2):
+            module.init_weights()
+
     @classmethod
     def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool:
         return name.endswith(("mamba.A_log", "mamba.D", "mlp.router.selection_bias"))