refactor(models): one shared init_buffers_post_meta (stacked on #3833)

src/prime_rl/trainer/model.py

@@ -787,13 +787,6 @@ def _move_buffers_to_cuda(model: nn.Module, config: ModelConfig) -> None:
             buffer.data = buffer.data.to("cuda")
 
 
-def _reset_runtime_moe_buffers(model: nn.Module) -> None:
-    for module in model.modules():
-        if isinstance(module, MoE) and module.tokens_per_expert.device.type != "meta":
-            module.tokens_per_expert.zero_()
-            module.routing_confidence_sum.zero_()
-
-
 def _validate_flash_attn_4_installed() -> None:
     """Validate that flash-attn-cute is installed and not overwritten by flash-attn.
 
@@ -915,7 +908,6 @@ def setup_model(
     else:
         load_dcp_from_hf(model, config, parallel_dims)
 
-    _reset_runtime_moe_buffers(model)
     return model
 
 

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

@@ -285,6 +285,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

src/prime_rl/trainer/models/base.py

@@ -123,16 +123,14 @@ def convert_adapter_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tenso
         return state_dict
 
     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.
+        """Rebuild the buffers ``to_empty`` leaves uninitialized: rotary tables and MoE routing stats.
 
-        This is called after loading the model from a checkpoint with meta device.
+        Runs before weight loading, so checkpoint-carried buffers such as ``selection_bias`` are
+        zeroed here and then overwritten by the load. Modules owning such buffers implement ``reset_buffers``.
         """
-        raise NotImplementedError(f"init_buffers_post_meta is not implemented for {self.__class__.__name__}")
+        for module in self.modules():
+            if hasattr(module, "reset_buffers"):
+                module.reset_buffers()
 
 
 __all__ = ["ALL_CP_STYLES", "CPStyle", "CPSupport", "PrimeModel"]

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

@@ -24,7 +24,6 @@
 from prime_rl.trainer.models.deepseek_v4.moe import DeepseekV4MoE
 from prime_rl.trainer.models.deepseek_v4.rotary import DeepseekV4RotaryEmbedding
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
-from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig
 from prime_rl.utils.cp import CPContext
 
@@ -237,19 +236,6 @@ def forward(
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        # `to_empty()` leaves every buffer uninitialized and this runs before `dcp_load`, so a
-        # buffer is restored either here or by the checkpoint. Rebuilt here are the ones no
-        # checkpoint carries: the rotary tables (non-persistent, and the one rotary every attention
-        # layer shares) and `tokens_per_expert`. The router's persistent buffers, `selection_bias` and a
-        # hash layer's `tid2eid`, are in the checkpoint that `dcp_load` applies next, so they are
-        # left to it.
-        for module in self.modules():
-            if isinstance(module, DeepseekV4RotaryEmbedding):
-                module.init_buffers_post_meta()
-            elif isinstance(module, MoE) and module.tokens_per_expert.device.type != "meta":
-                module.tokens_per_expert.zero_()
-
 
 __all__ = [
     "DeepseekV4DecoderLayer",

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

@@ -44,7 +44,7 @@ def _compute_inv_freq(self, layer_type: RopeLabel, device: torch.device | None)
             device,
         )
 
-    def init_buffers_post_meta(self) -> None:
+    def reset_buffers(self) -> None:
         """Re-derive the per-rope-type inverse frequencies and cos/sin caches in place.
 
         The tables are computed eagerly in `__init__` and registered non-persistently, so they

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

@@ -179,6 +179,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

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

@@ -212,6 +212,3 @@ def forward(
         # Sparse MLA derives document boundaries from position_ids, so seq_lens is unused.
         hidden_states = self.model(input_ids, position_ids, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

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

@@ -149,13 +149,6 @@ def forward(
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()
-        for module in self.modules():
-            if isinstance(module, MoE):
-                module.tokens_per_expert.zero_()
-                module.routing_confidence_sum.zero_()
-
 
 __all__ = [
     "GptOssDecoderLayer",

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

@@ -246,15 +246,5 @@ def forward(
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        for rotary_emb in self.model.rotary_emb.values():
-            rotary_emb.reset_parameters()
-
-        for module in self.modules():
-            if isinstance(module, MoE) and module.tokens_per_expert.device.type != "meta":
-                module.tokens_per_expert.zero_()
-                if module.router.selection_bias is not None:
-                    module.router.selection_bias.zero_()
-
 
 __all__ = ["LagunaForCausalLM", "LagunaModel"]

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

@@ -343,6 +343,12 @@ def __init__(
         )
         self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False)
 
+    def reset_buffers(self) -> None:
+        self.tokens_per_expert.zero_()
+        self.routing_confidence_sum.zero_()
+        if self.router.selection_bias is not None:
+            self.router.selection_bias.zero_()
+
     def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None:
         self.token_dispatcher = token_dispatcher
 

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

@@ -147,7 +147,7 @@ def __init__(self, rope: RopeParameters, head_dim: int, max_position_embeddings:
         inv_freq, self.attention_scaling = compute_rope_inv_freq(rope, head_dim, max_position_embeddings)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 
-    def reset_parameters(self) -> None:
+    def reset_buffers(self) -> None:
         """Recompute ``inv_freq``, which is not in the state dict, after materializing from the meta device."""
         inv_freq, self.attention_scaling = compute_rope_inv_freq(
             self.rope, self.head_dim, self.max_position_embeddings, self.inv_freq.device

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

@@ -113,6 +113,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

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

@@ -140,6 +140,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

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

@@ -218,12 +218,6 @@ def forward(
         hidden_states = self.model(input_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        for module in self.modules():
-            if isinstance(module, MoE):
-                module.tokens_per_expert.zero_()
-                module.routing_confidence_sum.zero_()
-
 
 __all__ = [
     "NemotronHDecoderLayer",

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

@@ -118,6 +118,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()

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

@@ -331,18 +331,6 @@ def forward(
             )
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        language_model = self.model.language_model if self.is_vlm else self.model
-        language_model.rotary_emb.reset_parameters()
-        if self.is_vlm:
-            self.model.visual.rotary_pos_emb.reset_parameters()
-        for module in self.modules():
-            if isinstance(module, MoE):
-                module.tokens_per_expert.zero_()
-                module.routing_confidence_sum.zero_()
-                if module.router.selection_bias is not None:
-                    module.router.selection_bias.zero_()
-
 
 __all__ = [
     "Qwen3_5Attention",

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

@@ -19,9 +19,9 @@ def __init__(self, dim: int, theta: float = 10_000.0) -> None:
         self.dim = dim
         self.theta = theta
         self.register_buffer("inv_freq", torch.empty(dim // 2), persistent=False)
-        self.reset_parameters()
+        self.reset_buffers()
 
-    def reset_parameters(self) -> None:
+    def reset_buffers(self) -> None:
         positions = torch.arange(0, self.dim, 2, dtype=torch.float32, device=self.inv_freq.device)
         self.inv_freq.copy_(1.0 / (self.theta ** (positions / self.dim)))
 

src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py

@@ -3,7 +3,6 @@
 
 from prime_rl.trainer.models.base import CPSupport, PrimeModel
 from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear
-from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.models.qwen3_8_flash_next.attention import IndexedGatedAttention
 from prime_rl.trainer.models.qwen3_8_flash_next.configuration_qwen3_8_flash_next import (
     Qwen3_8FlashNextConfig,
@@ -17,7 +16,6 @@
 from prime_rl.trainer.models.qwen3_8_flash_next.gated_delta_net import GatedDeltaNet
 from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import HyperConnection
 from prime_rl.trainer.models.qwen3_8_flash_next.moe import SigmoidOutputGatedMoE
-from prime_rl.trainer.models.qwen3_8_flash_next.ngram_embedding import NGramEmbedding
 from prime_rl.trainer.models.qwen3_8_flash_next.position_learning import PositionLearningEnhancement
 from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import RotaryEmbedding
 from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens
@@ -222,18 +220,6 @@ def forward(
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
 
-    def init_buffers_post_meta(self) -> None:
-        language_model = getattr(self.model, "language_model", self.model)
-        language_model.rotary_emb.reset_parameters()
-        for module in self.modules():
-            if isinstance(module, NGramEmbedding):
-                module.reset_parameters()
-            elif isinstance(module, MoE):
-                module.tokens_per_expert.zero_()
-                module.routing_confidence_sum.zero_()
-                if module.router.selection_bias is not None:
-                    module.router.selection_bias.zero_()
-
 
 __all__ = [
     "Qwen3_8FlashNextDecoderLayer",

src/prime_rl/trainer/models/qwen3_8_flash_next/ngram_embedding.py

@@ -132,7 +132,7 @@ def __init__(
 
         self.cp_context = CPContext()
 
-    def reset_parameters(self) -> None:
+    def reset_buffers(self) -> None:
         multipliers, sizes, offsets = build_hash_layout(
             ngram_size=self.ngram_size,
             heads_per_ngram=self.heads_per_ngram,

src/prime_rl/trainer/models/qwen3_8_flash_next/rotary_embedding.py

@@ -27,9 +27,9 @@ def __init__(self, rope: Qwen3_8FlashNextRopeParameters, head_dim: int) -> None:
         self.theta = rope.rope_theta
         self.mrope_section = rope.mrope_section
         self.register_buffer("inv_freq", torch.empty(self.rotary_dim // 2), persistent=False)
-        self.reset_parameters()
+        self.reset_buffers()
 
-    def reset_parameters(self) -> None:
+    def reset_buffers(self) -> None:
         dimensions = torch.arange(0, self.rotary_dim, 2, dtype=torch.float32, device=self.inv_freq.device)
         self.inv_freq.copy_(1.0 / (self.theta ** (dimensions / self.rotary_dim)))
 

src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py

@@ -165,6 +165,3 @@ def forward(
     ) -> PrimeLmOutput:
         hidden_states = self.model(input_ids, position_ids, seq_lens, seq_lens_are_pre_shard, routed_experts)
         return self.lm_head(hidden_states, labels, temperature=temperature, sampling_mask=sampling_mask)
-
-    def init_buffers_post_meta(self) -> None:
-        self.model.rotary_emb.reset_parameters()