feat(trainer): push post-meta buffer reinit down to owning layers

docs/development.md

@@ -70,6 +70,8 @@ Bringing up a new model family is three steps: implement the modeling code, regi
 
 Drop the modeling code under `src/prime_rl/trainer/models/<arch>/` (HF-compatible config, modeling, and weight conversion). Mirror the layout of an existing family — `glm4_moe/` or `qwen3_moe/` are good starting points.
 
+**Buffer contract.** Models are constructed on the meta device, so any module that registers a buffer (`register_buffer`) must implement `init_buffers_post_meta()` giving it a reasonable value. Modules without this method will cause a runtime failure.
+
 ### Register a Mini Preset
 
 Add an entry to [`scripts/mini_moe.py`](https://github.com/PrimeIntellect-ai/prime-rl/blob/main/scripts/mini_moe.py) so the smoke-test workflow can build a ~0.5B test model in your architecture. The preset names the config class, picks small dimensions, and wires up the HF + prime-rl model classes plus a tokenizer source:
@@ -135,6 +137,7 @@ Don't expect reward to climb meaningfully in 20 steps on a random model.
 Before merging a new model, you need to ensure the following:
 
 - The model is correctly registered and defines and all the required methods - such as `convert_hf_layer_to_tt` and `convert_tt_layer_to_hf`.
+- Every buffer-owning module the new model introduces implements `init_buffers_post_meta()` (see above).
 - The small smoke test passes.
 
 In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `batch_size=64`. All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework.

src/prime_rl/trainer/models/afmoe/modeling_afmoe.py

@@ -494,15 +494,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)
-
 
 __all__ = [
     "AfmoeForCausalLM",

src/prime_rl/trainer/models/base.py

@@ -1,7 +1,37 @@
+from typing import Protocol, runtime_checkable
+
+import torch.nn as nn
 from torch import Tensor
 from transformers.modeling_utils import PreTrainedModel
 
 
+@runtime_checkable
+class _PostMetaBufferInitModule(Protocol):
+    def init_buffers_post_meta(self) -> None: ...
+
+
+def _run_init_buffers_post_meta(module: nn.Module, exempt: tuple[type[nn.Module], ...] = ()) -> None:
+    """Walk every submodule of `module` (module itself excluded) and either call its
+    `init_buffers_post_meta()` hook or, if it owns buffers but doesn't implement the hook, raise.
+
+    Standalone so it's usable both as the real post-meta-init dispatch (called from
+    `PreTrainedModelPrimeRL.init_buffers_post_meta`) and directly in tests, against either a real
+    model or a plain `nn.Module` tree with no HF/PreTrainedModel machinery involved.
+    """
+    for submodule in module.modules():
+        if submodule is module or isinstance(submodule, exempt):
+            continue
+        if isinstance(submodule, _PostMetaBufferInitModule):
+            submodule.init_buffers_post_meta()
+        elif next(submodule.buffers(recurse=False), None) is not None:
+            raise TypeError(
+                f"{type(submodule).__name__} owns buffers "
+                f"{[n for n, _ in submodule.named_buffers(recurse=False)]} but doesn't implement "
+                "init_buffers_post_meta() -- implement it (even a documented no-op) so these "
+                "buffers don't silently hold undefined values after meta-device materialization."
+            )
+
+
 class PreTrainedModelPrimeRL(PreTrainedModel):
     """
     Base class for all PrimeRL models that extends HuggingFace PreTrainedModel.
@@ -132,17 +162,21 @@ def convert_layer_to_vllm_kernel(
         """
         raise NotImplementedError(f"convert_layer_to_vllm_kernel is not implemented for {cls.__name__}")
 
+    _init_buffers_post_meta_exempt: tuple[type[nn.Module], ...] = ()
+
     def init_buffers_post_meta(self) -> None:
         """
         Initialize buffers that are not in the state dict after loading with meta device.
 
         Some models have buffers (non-trainable tensors) that are not saved in the state dict
         but need to be properly initialized after loading the model on meta device and then
-        moving to the actual device. This method should initialize such buffers.
+        moving to the actual device. Dispatches to each submodule's own `init_buffers_post_meta`
+        so every layer owns the reinitialization of the buffers it registers; submodules that own
+        buffers but don't implement the hook cause this to raise (see `_run_init_buffers_post_meta`).
 
         This is called after loading the model from a checkpoint with meta device.
         """
-        raise NotImplementedError(f"init_buffers_post_meta is not implemented for {self.__class__.__name__}")
+        _run_init_buffers_post_meta(self, exempt=self._init_buffers_post_meta_exempt)
 
 
 __all__ = ["PreTrainedModelPrimeRL"]

src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py

@@ -322,18 +322,5 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        # HF standard transformer model
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)
-
-        # TODO: Init TT MoE buffers
-        # I think .to_empty() on gpu by default fills 0 so we are ok but this might not be guaranteed behavior
-
 
 __all__ = ["Glm4MoeConfig", "Glm4MoePreTrainedModel", "Glm4MoeModel", "Glm4MoeForCausalLM"]

src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py

@@ -348,14 +348,5 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)
-
 
 __all__ = ["GlmMoeDsaConfig", "GlmMoeDsaPreTrainedModel", "GlmMoeDsaModel", "GlmMoeDsaForCausalLM"]

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

@@ -15,10 +15,13 @@
 from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import MoeModelOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 from transformers.models.gpt_oss.modeling_gpt_oss import (
     GptOssAttention,
     GptOssRMSNorm,
-    GptOssRotaryEmbedding,
+)
+from transformers.models.gpt_oss.modeling_gpt_oss import (
+    GptOssRotaryEmbedding as HFGptOssRotaryEmbedding,
 )
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
@@ -34,6 +37,18 @@
 from prime_rl.trainer.models.layers.moe import GptOssGroupedExperts
 
 
+class GptOssRotaryEmbedding(HFGptOssRotaryEmbedding):
+    """HF's GptOssRotaryEmbedding, plus the buffer-reinit hook it doesn't define upstream."""
+
+    def init_buffers_post_meta(self) -> None:
+        rope_init_fn = (
+            self.compute_default_rope_parameters if self.rope_type == "default" else ROPE_INIT_FUNCTIONS[self.rope_type]
+        )
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, self.inv_freq.device)
+        self.inv_freq.copy_(inv_freq)
+        self.original_inv_freq.copy_(inv_freq)
+
+
 class GptOssTopKRouter(nn.Module):
     """Token-choice top-k router matching HF's GptOssTopKRouter parameter naming.
 
@@ -336,22 +351,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
-
-            rope_init_fn = (
-                ROPE_INIT_FUNCTIONS[rotary_emb.rope_type]
-                if rotary_emb.rope_type != "default"
-                else rotary_emb.compute_default_rope_parameters
-            )
-            inv_freq, rotary_emb.attention_scaling = rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device)
-            rotary_emb.inv_freq.copy_(inv_freq)
-            if "model.rotary_emb.original_inv_freq" in buffer_names:
-                rotary_emb.original_inv_freq.copy_(inv_freq)
-
 
 __all__ = [
     "GptOssForCausalLM",

src/prime_rl/trainer/models/laguna/modeling_laguna.py

@@ -82,6 +82,20 @@ def forward(
             sin = emb.sin() * attention_scaling
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
+    def init_buffers_post_meta(self) -> None:
+        for layer_type in self.layer_types:
+            rope_init_fn = self.compute_default_rope_parameters
+            if self.rope_type[layer_type] != "default":
+                rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type[layer_type]]
+            inv_freq, attention_scaling = rope_init_fn(
+                self.config,
+                getattr(self, f"{layer_type}_inv_freq").device,
+                layer_type=layer_type,
+            )
+            getattr(self, f"{layer_type}_inv_freq").copy_(inv_freq)
+            getattr(self, f"{layer_type}_original_inv_freq").copy_(inv_freq)
+            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)
+
 
 def _laguna_attention_config(config: LagunaConfig, num_heads: int) -> AttentionConfig:
     return AttentionConfig(
@@ -384,27 +398,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self) -> None:
-        rotary_emb = self.model.rotary_emb
-        for layer_type in rotary_emb.layer_types:
-            rope_init_fn = rotary_emb.compute_default_rope_parameters
-            if rotary_emb.rope_type[layer_type] != "default":
-                rope_init_fn = ROPE_INIT_FUNCTIONS[rotary_emb.rope_type[layer_type]]
-            inv_freq, attention_scaling = rope_init_fn(
-                rotary_emb.config,
-                getattr(rotary_emb, f"{layer_type}_inv_freq").device,
-                layer_type=layer_type,
-            )
-            getattr(rotary_emb, f"{layer_type}_inv_freq").copy_(inv_freq)
-            getattr(rotary_emb, f"{layer_type}_original_inv_freq").copy_(inv_freq)
-            setattr(rotary_emb, f"{layer_type}_attention_scaling", attention_scaling)
-
-        for module in self.modules():
-            if isinstance(module, MoE) and module.tokens_per_expert.device.type != "meta":
-                module.tokens_per_expert.zero_()
-                if module.expert_bias is not None:
-                    module.expert_bias.zero_()
-
 
 __all__ = [
     "LagunaForCausalLM",

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

@@ -1069,6 +1069,12 @@ def init_weights(
             if self.load_balance_coeff is not None:
                 self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)
 
+    def init_buffers_post_meta(self) -> None:
+        self.tokens_per_expert.zero_()
+        self.routing_confidence_sum.zero_()
+        if self.expert_bias is not None:
+            self.expert_bias.zero_()
+
 
 @torch.compile(dynamic=True)
 def relu2(x: torch.Tensor) -> torch.Tensor:
@@ -1287,6 +1293,9 @@ def forward(
     def init_weights(self, init_std: float):
         nn.init.trunc_normal_(self.gate, mean=0.0, std=init_std)
 
+    def init_buffers_post_meta(self) -> None:
+        self.e_score_correction_bias.zero_()
+
 
 class BCNonGatedFeedForward(nn.Module):
     """Non-gated feed-forward network used as the shared expert in NemotronH.
@@ -1514,3 +1523,9 @@ def init_weights(self, init_std: float, buffer_device: torch.device):
             self.routing_confidence_sum = torch.tensor(0.0, dtype=torch.float32)
             if self.load_balance_coeff is not None:
                 self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)
+
+    def init_buffers_post_meta(self) -> None:
+        self.tokens_per_expert.zero_()
+        self.routing_confidence_sum.zero_()
+        if self.expert_bias is not None:
+            self.expert_bias.zero_()

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

@@ -57,6 +57,10 @@ def compute_default_rope_parameters(self, config=None, device=None, seq_len=None
         """Required by transformers 5.0.0 for weight initialization when rope_type is 'default'."""
         return _compute_default_rope_parameters(config or self.config, device, seq_len, layer_type)
 
+    def init_buffers_post_meta(self) -> None:
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, self.inv_freq.device)
+        self.inv_freq.copy_(inv_freq)
+
     @torch.no_grad()
     @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
     def forward(self, x, position_ids):

src/prime_rl/trainer/models/llama/modeling_llama.py

@@ -312,13 +312,3 @@ def forward(
             labels[:, slice_indices] if labels is not None else None,
             temperature=temperature,
         )
-
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        # HF standard transformer model
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)

src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py

@@ -262,15 +262,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)
-
 
 __all__ = [
     "MiniMaxM2ForCausalLM",

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

@@ -514,9 +514,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        pass
-
 
 __all__ = [
     "NemotronHForCausalLM",

src/prime_rl/trainer/models/qwen3/modeling_qwen3.py

@@ -268,14 +268,5 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        buffer_names = [name for name, _ in self.named_buffers()]
-        if "model.rotary_emb.inv_freq" in buffer_names:
-            rotary_emb = self.model.rotary_emb
-            inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(
-                rotary_emb.config, rotary_emb.inv_freq.device
-            )
-            rotary_emb.inv_freq.copy_(inv_freq)
-
 
 __all__ = ["Qwen3ForCausalLM", "Qwen3Model", "Qwen3PreTrainedModel"]

src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py

@@ -12,7 +12,7 @@
 from transformers.models.qwen3_5.modeling_qwen3_5 import (
     Qwen3_5PreTrainedModel as HFQwen3_5PreTrainedModel,
 )
-from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel
+from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel, Qwen3_5VisionRotaryEmbedding
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs
 
@@ -32,6 +32,19 @@
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
 
 
+def _init_vision_rope_buffers_post_meta(self: Qwen3_5VisionRotaryEmbedding) -> None:
+    inv_freq = 1.0 / (
+        self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=self.inv_freq.device) / self.dim)
+    )
+    self.inv_freq.copy_(inv_freq)
+
+
+# Qwen3_5VisionRotaryEmbedding is upstream transformers code; Qwen3_5VisionModel.__init__
+# hardcodes its construction, so we can't subclass-and-inject like elsewhere. Attach the
+# buffer-init hook to the class directly instead.
+Qwen3_5VisionRotaryEmbedding.init_buffers_post_meta = _init_vision_rope_buffers_post_meta
+
+
 class Qwen3_5GatedFlashAttention(Qwen3_5MoeGatedFlashAttention):
     pass
 
@@ -452,26 +465,6 @@ def forward(
             temperature=temperature,
         )
 
-    def init_buffers_post_meta(self):
-        if self._is_vlm:
-            lm_rope = self.model.language_model.rotary_emb
-        else:
-            lm_rope = self.model.rotary_emb
-
-        if hasattr(lm_rope, "rope_init_fn"):
-            inv_freq, lm_rope.attention_scaling = lm_rope.rope_init_fn(lm_rope.config, lm_rope.inv_freq.device)
-            lm_rope.inv_freq.copy_(inv_freq)
-
-        if self._is_vlm:
-            vis_rope = self.model.visual.rotary_pos_emb
-            if hasattr(vis_rope, "inv_freq"):
-                dim = vis_rope.inv_freq.shape[0]
-                inv_freq = 1.0 / (
-                    10000.0
-                    ** (torch.arange(0, dim * 2, 2, dtype=torch.float32, device=vis_rope.inv_freq.device) / (dim * 2))
-                )
-                vis_rope.inv_freq.copy_(inv_freq)
-
 
 __all__ = [
     "Qwen3_5ForCausalLM",