[faresobeid]

Stacked on #3833 (chore/remove-transformers-modeling). The PR base is that branch, so the diff shows only this change. Retarget to main once #3833 merges.

Summary

PrimeModel.init_buffers_post_meta is now one shared method. It walks self.modules() and calls reset_buffers() on each module that defines it. This replaces the 13 per-arch overrides, which differed from each other.

• Rotary modules: the shared RotaryEmbedding, Qwen3.5 vision rotary, Qwen3.8 rotary, and NGramEmbedding rename reset_parameters to reset_buffers. DeepseekV4RotaryEmbedding.init_buffers_post_meta becomes reset_buffers. The old name was unsafe for a generic walk because the norms use reset_parameters to re-init their weights.
• MoE.reset_buffers zeroes tokens_per_expert, routing_confidence_sum and selection_bias.
• _reset_runtime_moe_buffers in trainer/model.py is deleted. It existed only because half the archs did not zero their MoE stats in init_buffers_post_meta.

Diff: 21 files, +19 / −107.

Behavior change

None for training.

• Correction to the review. The v2 review (rev-models Grpo loss #2) said GLM-5.3 / glm4_moe / qwen3_moe / minimax_m2 report first-step max_vio from uninitialized memory. That is wrong. setup_model already calls _reset_runtime_moe_buffers after both load paths, so the stats were zero before any forward. This PR keeps that result in one place.
• selection_bias. It is now zeroed before weight loading on every arch. Before, only laguna, qwen3_5 and qwen3_8 zeroed it; the others left it uninitialized until the checkpoint load overwrote it. With a checkpoint that ships the bias, the result is the same. Without one, it is now zero instead of garbage.
• tid2eid (DeepSeek-V4 hash routing) is still left to the checkpoint.

Checks

• Buffer check, every arch (CPU tiny models; qwen3_8 cannot be built on CPU because prime_kernels has no indexed_attention). Steps: build on meta, to_empty, fill every buffer with NaN, init_buffers_post_meta, then the old _reset_runtime_moe_buffers where it exists. Result:
  • every non-persistent buffer and every rotary attention_scaling equals a freshly built model, on Chore/remove transformers modeling #3833 and on this branch;
  • on this branch no buffer is still NaN after init_buffers_post_meta alone.
• MoE metrics (glm4_moe, qwen3_moe, minimax_m2, afmoe, laguna, gpt_oss). Steps: meta init, garbage buffers, init, load weights, one forward, get_load_balance_stats. max_vio, routing_confidence and tokens_per_expert are identical between Chore/remove transformers modeling #3833 and this branch.
• Forward/backward for llama, qwen3, qwen3_moe, glm4_moe, minimax_m2, afmoe, laguna, gpt_oss on CPU (attention replaced by an SDPA varlen reference). Logits, logprobs, entropy, all grads and buffers are bitwise identical to Chore/remove transformers modeling #3833.
• pytest tests/unit/train tests/unit/test_configs.py -m "not gpu": 351 passed, 1 failed. The failure is test_cpu_adam native, which fails the same way on Chore/remove transformers modeling #3833 (environment issue).
• ruff check and ruff format: clean.

GPU checks (1×8 H200, SLURM jobs 2386/2393)

• GPU buffer check, all archs except qwen3_8: to_empty(cuda) with NaN buffers, then init_buffers_post_meta. Every non-persistent buffer and rotary scale matches a fresh model, on Chore/remove transformers modeling #3833 and on this branch. On this branch no MoE stat or selection_bias is still NaN after init.
• test_deepseek_v4.py -k "init_buffers_post_meta or hash_table_survives": 2 passed (Chore/remove transformers modeling #3833: 2 passed).
• test_qwen3_8_flash_next.py + qwen3_5/test_vlm.py: 15 passed (Chore/remove transformers modeling #3833: 15 passed).
• qwen3_8 cannot be built in the equivalence script on this node (prime_kernels has no indexed_attention). It is covered by its test file above.

Conflicts

• Depends on Chore/remove transformers modeling #3833.
• No overlap with feat: Total Router Recall (replay sampler MoE routing weights) #3793, feat(trainer): Composer 2-style plausibility filter for router replay #3796 or feat(trainer): freeze MoE router by default in RL #3797. Those touch MoE.forward, the routers and model.py freeze logic, not reset_buffers.

---

Note

Low Risk
Refactor of model materialization before weight load; normal checkpoint training should be unchanged, with only a minor fix for uninitialized selection_bias when biases are absent from the checkpoint.

Overview
Unifies post–meta-device buffer initialization across all PrimeRL models instead of thirteen architecture-specific init_buffers_post_meta overrides.

PrimeModel.init_buffers_post_meta now walks self.modules() and calls reset_buffers() on any submodule that defines it. Rotary layers, DeepSeek-V4 RoPE, Qwen3.5 vision RoPE, Qwen3.8 RoPE/NGramEmbedding, and MoE implement reset_buffers (renamed from reset_parameters where that name would collide with weight re-init). MoE.reset_buffers zeroes routing stats and optional selection_bias.

trainer/model.py drops _reset_runtime_moe_buffers; MoE stat clearing is covered by the shared init path before checkpoint load.

^{Reviewed by Cursor Bugbot for commit 6f8fa80. Bugbot is set up for automated code reviews on this repo. Configure here.}