feat(trainer): push post-meta buffer reinit down to owning layers by garrett361 · Pull Request #3378 · PrimeIntellect-ai/prime-rl

garrett361 · 2026-08-24T23:25:42Z

Pushes down init_buffers_post_meta onto whatever module actually owns the buffer. The top-level PreTrainedModelPrimeRL.init_buffers_post_meta method on model classes then recursively calls this method on all buffer-owning submodules, raising an error if any such modules do not implement this method. Effectively introduces a contract that all buffer-owning modules in prime-rl should at least attempt to properly initialize their own buffers, making it harder to accidentally leave buffers improperly initialized, which was a historical issue.

There's an exemption escape hatch (_init_buffers_post_meta_exempt) for buffer-owning modules the codebase doesn't control and can't add the hook to. Currently only needed by qwen3_5/qwen3_5_moe, whose vision-tower rope is an upstream transformers class.
Added a completeness test per architecture, in each model's own test file, that constructs on meta, to_empty()s, and asserts every buffer comes out finite — previously only 2 of 12 architectures were ever exercised through this path in CI.

Verification

uv run pytest tests/unit/train/ -q green (excluding the pre-existing, unrelated test_nemotron_h* failures, confirmed via git stash comparison against main).
Manually confirmed the completeness check fails loudly (TypeError) when a layer's hook is removed, and that the new random_init regression test fails before the fix and passes after.

model.to_empty() gives no guarantee about buffer contents, most visible under debug.random_init=True which skips checkpoint loading entirely. Each model previously hand-implemented init_buffers_post_meta with duplicated rotary-embedding logic and inconsistent MoE coverage: four models (afmoe, glm4_moe, glm_moe_dsa, minimax_m2) never zeroed expert_bias, and nemotron_h's NemotronHRouter never zeroed its e_score_correction_bias buffer, both of which feed routing decisions directly. Push init_buffers_post_meta down to the layer classes that actually own the buffers (RotaryEmbedding, MoE, LatentMoE, NemotronHRouter, and per-model rotary subclasses for gpt_oss/laguna/qwen3_5(_moe)), and make PreTrainedModelPrimeRL.init_buffers_post_meta a generic dispatcher that walks the module tree via a runtime_checkable Protocol. A module that owns buffers but doesn't implement the hook now raises TypeError at meta-device load time instead of silently shipping undefined buffer contents. Qwen3.5's upstream vision-tower rope is the one documented exemption, since we don't own that class. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Only laguna and glm4_moe were ever exercised through meta-device construction + to_empty() + init_buffers_post_meta() in the test suite; the other architectures' test files build directly on cuda for forward/backward parity, so the completeness check added in b9ae380 never actually ran against them in CI. A regression in a shared layer's hook, or a new layer that forgets to implement it, would only surface in a real training run. Add one test per architecture, colocated in that architecture's own test file: construct on meta, to_empty(), call the model's public init_buffers_post_meta(), and assert every buffer is finite. Six existing files (test_llama.py, test_glm4_moe.py, test_qwen3_moe.py, test_nemotron_h.py, test_qwen3_5.py, test_qwen3_5_moe.py) get a new test with a small dedicated meta-only config; five architectures had no model-constructing test file at all (qwen3, gpt_oss, minimax_m2, laguna, glm_moe_dsa) and get a new one. afmoe is deliberately left uncovered for now: its test file's module-wide skip (HF afmoe doesn't support flash_attention_2) would also skip a completeness test that doesn't need forward/flash-attn at all, and untangling that skip is a separate concern from this change. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

The toy nn.Module completeness-check tests in test_base.py were useful during development but are redundant now that every architecture has its own init_buffers_post_meta test exercised against real models, plus the runtime TypeError enforcement itself in base.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

garrett361 · 2026-08-25T00:23:16Z

+    for name, buffer in model.named_buffers():
+        assert torch.isfinite(buffer).all(), f"buffer {name} is not finite after init_buffers_post_meta"


These torch.isfinite tests are super weak since in practice meta-init usually populates zeros into the tensor, but they're cheap

samsja · 2026-08-26T02:30:40Z

    supports_gradient_checkpointing = True
+    # Vision-tower rope belongs to upstream transformers' Qwen3_5VisionModel; we can't add
+    # init_buffers_post_meta to it, so it's exempted and reinitialized explicitly below instead.
+    _init_buffers_post_meta_exempt = (Qwen3_5VisionRotaryEmbedding,)


wondering if we shouldn't patch the rotary embedding class instead of rather inherit from it somewhere to override it ?

or rather can't we just patch the end model that inherit from this one ?

patched, but a better solution is to vendor ourselves, which I'll leave to a separate PR.

Qwen3_5VisionRotaryEmbedding/Qwen3_5MoeVisionRotaryEmbedding were the only buffer-owning modules needing the _init_buffers_post_meta_exempt escape hatch, since prime-rl doesn't own the upstream Qwen3_5(Moe)VisionModel.__init__ call site that constructs them. Instead of exempting them, attach init_buffers_post_meta directly onto the upstream class objects at import time; runtime_checkable Protocol isinstance checks are duck-typed regardless of where the method was defined, so the generic dispatch picks it up with no subclassing or forking needed. The replacement formula reads self.dim/self.theta (stored by upstream's own __init__) instead of reverse-engineering dim from the buffer's shape, closing a latent fragility in the old per-model override. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

garrett361 · 2026-08-26T07:17:22Z


+
+def _init_vision_rope_buffers_post_meta(self: Qwen3_5MoeVisionRotaryEmbedding) -> 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_5MoeVisionRotaryEmbedding is upstream transformers code; Qwen3_5MoeVisionModel.__init__
+# hardcodes its construction, so we can't subclass-and-inject like elsewhere. Attach the
+# buffer-init hook to the class directly instead.
+Qwen3_5MoeVisionRotaryEmbedding.init_buffers_post_meta = _init_vision_rope_buffers_post_meta
+


CC @samsja ended up just patching in this method on the HF class, as here

S1ro1

In general I like the idea, some minor nits on cleanup - would prefer to spend some time on this to future proof it, we're in no rush to get this in, so instead make it clean without patches etc

S1ro1 · 2026-08-26T11:27:39Z

            sin = emb.sin() * attention_scaling
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

+    def init_buffers_post_meta(self) -> None:


I'm honestly not a big fan of having this model specific code here - would try to think of a way to make this owned by the rope layer itself as well

@S1ro1 this is a method on LagunaRotaryEmbedding so this is owned by the rope layer, no?

Oh i see, i thought it's the parent model, then all good

S1ro1 · 2026-08-26T11:28:28Z

+# 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


We can probably vendor the rotary embeddings ourselves, would do it anyway given the aim to remove transformers so we can do it now instead of this ugly patch

Agreed, but didn't want to mix concerns here

S1ro1 · 2026-08-26T11:28:40Z

+# Qwen3_5MoeVisionRotaryEmbedding is upstream transformers code; Qwen3_5MoeVisionModel.__init__
+# hardcodes its construction, so we can't subclass-and-inject like elsewhere. Attach the
+# buffer-init hook to the class directly instead.
+Qwen3_5MoeVisionRotaryEmbedding.init_buffers_post_meta = _init_vision_rope_buffers_post_meta


dito here, maybe let's vendor

S1ro1 · 2026-08-26T11:29:22Z



+@runtime_checkable
+class _PostMetaBufferInitModule(Protocol):


let's make this PostMetaBufferInitModule, I hate this smell of ai code where it does private classes

Lol, I'll change it, but pretty sure I asked for private classes and fns in these spots since they're not things I'd expect users to need to use or know about

makes sense, just my aversion, can keep if you think so

already changed, no strong feelings either way

garrett361 and others added 2 commits August 24, 2026 21:38

garrett361 marked this pull request as ready for review August 25, 2026 00:19

garrett361 requested review from S1ro1 and samsja August 25, 2026 00:19

garrett361 commented Aug 25, 2026

View reviewed changes

samsja reviewed Aug 26, 2026

View reviewed changes

garrett361 commented Aug 26, 2026

View reviewed changes

S1ro1 reviewed Aug 26, 2026

View reviewed changes

refactor(trainer): make post-meta buffer-init helpers public

d5b1d4d

		for name, buffer in model.named_buffers():
		assert torch.isfinite(buffer).all(), f"buffer {name} is not finite after init_buffers_post_meta"



		@runtime_checkable
		class _PostMetaBufferInitModule(Protocol):

Uh oh!

Conversation

garrett361 commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Verification

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

S1ro1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

garrett361 commented Aug 24, 2026 •

edited

Loading