From 6f8fa803cfc75897449cd6b58d4f8a00c8f991b6 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sun, 4 Oct 2026 12:14:44 +0000 Subject: [PATCH] refactor(models): one init_buffers_post_meta walking reset_buffers hooks --- src/prime_rl/trainer/model.py | 8 -------- .../trainer/models/afmoe/modeling_afmoe.py | 3 --- src/prime_rl/trainer/models/base.py | 14 ++++++-------- .../models/deepseek_v4/modeling_deepseek_v4.py | 14 -------------- src/prime_rl/trainer/models/deepseek_v4/rotary.py | 2 +- .../trainer/models/glm4_moe/modeling_glm4_moe.py | 3 --- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 3 --- .../trainer/models/gpt_oss/modeling_gpt_oss.py | 7 ------- .../trainer/models/laguna/modeling_laguna.py | 10 ---------- src/prime_rl/trainer/models/layers/moe.py | 6 ++++++ src/prime_rl/trainer/models/layers/rotary_emb.py | 2 +- .../trainer/models/llama/modeling_llama.py | 3 --- .../models/minimax_m2/modeling_minimax_m2.py | 3 --- .../models/nemotron_h/modeling_nemotron_h.py | 6 ------ .../trainer/models/qwen3/modeling_qwen3.py | 3 --- .../trainer/models/qwen3_5/modeling_qwen3_5.py | 12 ------------ src/prime_rl/trainer/models/qwen3_5/vision.py | 4 ++-- .../modeling_qwen3_8_flash_next.py | 14 -------------- .../models/qwen3_8_flash_next/ngram_embedding.py | 2 +- .../models/qwen3_8_flash_next/rotary_embedding.py | 4 ++-- .../trainer/models/qwen3_moe/modeling_qwen3_moe.py | 3 --- 21 files changed, 19 insertions(+), 107 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 98d71b36b4d..9c8790b304b 100644 --- a/src/prime_rl/trainer/model.py +++ b/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 diff --git a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py b/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py index 9935d8faaea..5a755a8dd2f 100644 --- a/src/prime_rl/trainer/models/afmoe/modeling_afmoe.py +++ b/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() diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py index 58efd0e8c6a..caa2d251d8c 100644 --- a/src/prime_rl/trainer/models/base.py +++ b/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"] diff --git a/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py b/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py index 4105bb98bf3..b278fe20db2 100644 --- a/src/prime_rl/trainer/models/deepseek_v4/modeling_deepseek_v4.py +++ b/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", diff --git a/src/prime_rl/trainer/models/deepseek_v4/rotary.py b/src/prime_rl/trainer/models/deepseek_v4/rotary.py index c377f766069..129c1a718b4 100644 --- a/src/prime_rl/trainer/models/deepseek_v4/rotary.py +++ b/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 diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 298161846a8..4d52492d2eb 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/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() diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py index a11c0c3d6e6..11e9a417e88 100644 --- a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/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() diff --git a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py index 0caa0bc942b..923fb856c3a 100644 --- a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py +++ b/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", diff --git a/src/prime_rl/trainer/models/laguna/modeling_laguna.py b/src/prime_rl/trainer/models/laguna/modeling_laguna.py index 6048d1f33e5..1f9db46b619 100644 --- a/src/prime_rl/trainer/models/laguna/modeling_laguna.py +++ b/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"] diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 99a25cb4e4a..732c1e76300 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/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 diff --git a/src/prime_rl/trainer/models/layers/rotary_emb.py b/src/prime_rl/trainer/models/layers/rotary_emb.py index 874a32305d9..150a2cad8aa 100644 --- a/src/prime_rl/trainer/models/layers/rotary_emb.py +++ b/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 diff --git a/src/prime_rl/trainer/models/llama/modeling_llama.py b/src/prime_rl/trainer/models/llama/modeling_llama.py index fe9a9e250eb..66b0091df96 100644 --- a/src/prime_rl/trainer/models/llama/modeling_llama.py +++ b/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() diff --git a/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py b/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py index e21d799c64e..43fdb4e565d 100644 --- a/src/prime_rl/trainer/models/minimax_m2/modeling_minimax_m2.py +++ b/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() diff --git a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py index f4f6bf9e425..6907da8cf52 100644 --- a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py +++ b/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", diff --git a/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py b/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py index 9b7d91a6c9a..e6ae606e00f 100644 --- a/src/prime_rl/trainer/models/qwen3/modeling_qwen3.py +++ b/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() diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index 3d0004a845c..425c1300cf9 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/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", diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 4fc3bcaa4aa..7307a0e5243 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/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))) diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py index 3ea78b7fdd3..d6dc67e31a3 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py +++ b/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", diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/ngram_embedding.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/ngram_embedding.py index 221ceea5f29..39b2530e602 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/ngram_embedding.py +++ b/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, diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/rotary_embedding.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/rotary_embedding.py index 35846fab6ed..5cf97b75268 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/rotary_embedding.py +++ b/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))) diff --git a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py b/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py index fb58c38ad30..96450cd6344 100644 --- a/src/prime_rl/trainer/models/qwen3_moe/modeling_qwen3_moe.py +++ b/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()