feat(models): add native Qwen3.5 family implementation by S1ro1 · Pull Request #3429 · PrimeIntellect-ai/prime-rl

S1ro1 · 2026-08-28T20:14:03Z

Summary

Unify Qwen3.5 dense, MoE, and VLM training under a native implementation, including the Qwen3.8 checkpoint that uses this architecture. Remove the separate MoE package and global Transformers model patches.

Keep Qwen3_5PreTrainedModel(PreTrainedModelPrimeRL) shared by the backbone and causal LM; remove unused Transformers TP/PP plans, split-module metadata, cache arguments, and redundant initialization. Modules own their initialization, including zero-centered RMSNorm.
Use Prime's shared attention and MoE runtime. Gated full attention consumes the fused QKV projection with the correct doubled query width; configured MoE load balancing allocates its router selection bias.
Own packed DeltaNet, zero-centered RMSNorm, MRoPE, and the vision encoder. Hybrid layers support Ulysses CP through FLA boundary-state exchange; full attention uses common attention dispatch.
Convert official fused MoE expert weights to Prime's split projections and back, excluding MTP weights.
Keep the narrow compiler boundary for FLA CP convolution, whose collective layout fails Dynamo tracing. Vision FlashAttention has no compiler disable and uses each backend's declared return format directly.

Qwen tests live under tests/unit/train/models/qwen3_5/: test_model.py parametrizes initialization, context-parallel setup, and forward/backward/packing checks over dense and MoE configs; test_vlm.py groups vision tests with CPU MRoPE checks. MoE router replay remains a dedicated check. Tests use native configs and models without Transformers reference models, HF weight comparisons, or Hub downloads. Public model helpers use AutoModelForCausalLMPrimeRL.from_config with the trainer's shared auto-attention resolver, without hardcoded backends or private attention config writes. Remove the class/source-introspection test and duplicate shared-attention checks. Remove the blanket BF16 optimizer/reduction restriction for RL VLMs; FP32 defaults remain unchanged and FSDP still computes in BF16.

Text and vision consume Prime's resolved model.attn configuration directly, including under Ulysses CP. Remove Qwen constructor fallback/propagation logic, the forced vision SDPA override, and its unused vision branch.

MTP and video training are outside this runtime's scope. scripts/mini_moe.py is unchanged from main and is outside this PR's validation.

Math KL validation

Both checkpoints completed 20 steps on math-env (PrimeIntellect/Hendrycks-Math, train split), with batch size 64. Every step-mean KL is below 0.015.

Checkpoint	Steps	Mean KL over 20 steps	Highest step-mean KL
`Qwen/Qwen3.8-27B`	20	0.000359685	0.000953884
`Qwen/Qwen3.5-35B-A3B`	20	0.000208713	0.000462048

Runs executed concurrently on separate H200 nodes using uv run rl @ <config>, with four trainer and four inference GPUs per model. Both Slurm jobs completed with exit code 0: dense job 387, MoE job 386. Tested commit: 9983e7bed.

Settings: main's debug algorithm with constant per-token advantage 1.0; original task sampler; math-verify scoring without the optional LLM judge; Prime VM runtime enforcing the task network policy; sequence length 4096, maximum completion length 2048, group size 8, learning rate 1e-6, compilation, selective activation checkpointing, FlashAttention 3, and NCCL weight updates. MoE additionally uses EP=4 and router replay. Inference uses max_num_seqs=64. Optimization and reduction dtypes retain their FP32 defaults.

All 20 batches per model contained 64 trainable episodes, and all gradient norms were finite. These are asynchronous E2E measurements of mismatch_kl/all/mean; the summary averages the 20 step means and includes normal policy staleness. The debug algorithm exercises training with synthetic advantages, so this is a correctness check, not a learning-quality benchmark. These 20-step runs use text inputs; packed image-conditioned validation is recorded below.

All 20 step-mean KL values

Step	Dense mean KL	MoE mean KL
1	0.000087817	0.000103068
2	0.000133938	0.000132430
3	0.000112734	0.000115124
4	0.000193951	0.000099179
5	0.000205668	0.000214333
6	0.000359477	0.000348763
7	0.000838096	0.000161281
8	0.000150568	0.000207813
9	0.000190916	0.000158351
10	0.000246897	0.000258552
11	0.000492465	0.000462048
12	0.000236148	0.000177551
13	0.000311170	0.000316216
14	0.000953884	0.000286797
15	0.000817523	0.000233271
16	0.000320892	0.000184788
17	0.000567697	0.000175972
18	0.000280278	0.000164684
19	0.000264224	0.000196673
20	0.000429352	0.000177372

Packed VLM validation

Both models completed five packed image RL updates on MMMU Math dev, with batch size 32 and an 8192-token trainer packing limit. All 160 samples per model included images.

Model	Steps	Mean KL	Maximum step-mean KL	Images per packed microbatch
Qwen/Qwen3.8-27B	5	0.000501153	0.000729994	5–6
Qwen/Qwen3.5-35B-A3B	5	0.000346154	0.000458347	5–6

Step	Dense KL	MoE KL
1	0.000231250	0.000167199
2	0.000387219	0.000343682
3	0.000553381	0.000360291
4	0.000729994	0.000401253
5	0.000603919	0.000458347

Native models; trainable vision encoder; FA3; compiled decoder blocks and selective activation checkpointing; FP32 optimization and gradient reduction; CP1; MoE EP4 and router replay; debug advantages; group size 8; maximum completion length 1024; inference context 4096; NCCL weight broadcasts. These are correctness smokes on four repeated diagram questions, not learning-quality benchmarks.

Packed MRoPE exactly matched individual-document MRoPE for all 320 training samples, and image-grid/pixel-patch counts matched each document’s image-token spans. Both Slurm jobs exited 0 (dense 400, MoE 399). MoE needed allocator retries at the H200 memory limit.

Image+CP+compile remains separately unresolved; these E2E runs use CP1.

Tested runtime contents: dcdf770df (subsequent commits only change tests). Run configs and detailed batch/position evidence are under /home/matej/prime-rl-verify/pr3429-20260910/.

Other validation

Attention cleanup at a2173601f: H200 Qwen suite 24 passed, one skipped; official dense/MoE model loading preserves configured attention for text and vision at CP sizes 1 and 2. An additional compiled image/CP backward check exposed Autograd not support dtype: Int with selective activation checkpointing; the same error reproduces using the previous SDPA vision path. Compiled image/CP validation remains unresolved. The completed math runs above use CP=1.
Native-only H200 Qwen suite: 20 passed, one expected single-rank CP skip, with HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 (Slurm 407; consolidated dense/MoE suite). Existing config suite: 175 passed. Ruff check/format pass. Both dense and MoE run the shared packed-versus-unpacked comparison. MoE replays identical expert selections to isolate document boundaries; ordinary routing forward/backward remains covered.
CI at 9983e7bed: CPU unit tests, Ruff, slim install, and CodeQL pass. GPU CI passes the Qwen tests; its five failures are in unchanged DeepSeek V4 tests (three packed-batch assertions reproduced on unmodified main bad80b1a3, two FP8 architecture failures). GPU CI: 66 passed, 82 skipped, 599 deselected, five failed.
Two-rank H200 compiled packed DeltaNet output and input-gradient checks passed with the CP convolution compiler boundary.
Exact official checkpoint key/shape conversion and reverse conversion passed: dense 1,199 source keys to 1,184 runtime keys; MoE 1,811 source keys to 1,066 runtime keys after excluding MTP.
Full H200 unit suite: 737 passed, 12 skipped, three DeepSeek V4 failures also reproduced on unmodified main. The weight-transfer test imports the unified model.

Note

High Risk
Large change to core training models (hybrid attention, MoE weight conversion, packed multimodal RL) with removed Transformers workarounds; regressions would affect Qwen3.5/3.8 RL and VLM correctness rather than peripheral code.

Overview
Replaces the split Qwen3.5 dense + MoE packages and Transformers runtime monkey-patches with one native PrimeRL implementation under models/qwen3_5/. Dense text, MoE text, and composite VLMs all dispatch to Qwen3_5ForCausalLM, driven by local configs (Qwen3_5TextConfig, Qwen3_5MoeTextConfig, vision/composite configs) and updated AutoConfig / VLM registration.

The rewrite drops the HF-backed text stack in favor of Prime FlashAttention, packed varlen GatedDeltaNet (FLA + Ulysses CP), zero-centered RMSNorm, MRoPE, a custom vision encoder, and MoE layers with HF↔Prime expert weight conversion (MTP dropped). VLM prep keeps dummy vision runs for FSDP symmetry and shards routed_experts under CP. model.py no longer applies Qwen3.5 checkpoint/MRoPE/varlen patches or forces vision SDPA under CP.

Trainer/config cleanup: removes the VLM must use bfloat16 validator and the duplicate check at model load; parallel_dims treats MoE on nested text_config. Ring/Ulysses patching no longer lists Qwen-specific attention classes because Qwen3_5Attention subclasses FlashAttention.

Tests/docs: deletes HF-comparison tests and the old qwen3_5_moe tree; adds tests/unit/train/models/qwen3_5/ (forward/pack/CP/VLM/MRoPE). development.md points VLM authors at the unified package.

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

cursor

Stale Bugbot comment from a previous run.

cursor

Stale Bugbot comment from a previous run.

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c48531b. Configure here.}

garrett361

LGTM, various small comments/questions

garrett361 · 2026-09-11T13:30:00Z

+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.attn_output_gate = config.attn_output_gate
+        q_output_size = self.num_heads * self.head_dim * (2 if self.attn_output_gate else 1)


Little awkward with the overrides and redefinitions of sub modules

garrett361 · 2026-09-11T13:31:04Z

+        self.q_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps)
+
+    def forward(


Similar nit: we are overriding so much of this that it’s barely a subclass any more

For me, the smell here is that I'd expect the FlashAttention class to be less responsible.I’d expect it to consume qkv with rope already applied as necessary as inputs, and not be responsible for creating these.

Wdyt? Not a change for this PR.

agree ye, i think we might actually want to move attn to not even be subclassable at some point

garrett361 · 2026-09-11T13:35:19Z

+
+from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
+
+# FLA's CP convolution uses an all-gather layout that Dynamo cannot trace.


Can you explain why we need this? Though we'd just graph break. Would like to understand why we need to disable

this was here before the PR, it's on my TODO to explore this more

garrett361 · 2026-09-11T13:37:40Z

+
+        context = None
+        if self.context_parallel_group is not None:
+            context = build_cp_context(


Future optimization: can probably just build once per model, not once per GDN layer

garrett361 · 2026-09-11T13:43:55Z


-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
+    def set_input_embeddings(self, embeddings: nn.Embedding) -> None:


Can we delete all of the {get,set}_input_embeddings methods everywhere? They seem to be a transformers holdover that we don't use, IIUC.

I think we use those somewhere, noted though to research it

garrett361 · 2026-09-11T14:07:06Z

    )


-def _patch_qwen3_5_moe_conversion_mapping():


Great. Love deleting patches

garrett361 · 2026-09-11T14:37:25Z

 _CUSTOM_VLM_MAPPING: dict[str, type] = {
    "qwen3_5": Qwen3_5ForCausalLM,
-    "qwen3_5_moe": Qwen3_5MoeForCausalLM,
+    "qwen3_5_moe": Qwen3_5ForCausalLM,


Just verifying the class change here was intended

Yeah now both MoE and dense are a single class, which is imo a bit cleaner

garrett361 · 2026-09-11T14:47:55Z

    from transformers import AutoConfig

    model_config = AutoConfig.from_pretrained(config.name, trust_remote_code=config.trust_remote_code)
+    model_config = getattr(model_config, "text_config", model_config)


Looks like we could use a similar fix in is_tt_moe_model

garrett361 · 2026-09-11T14:48:51Z

+
+
+@pytest.mark.gpu
+def test_norms_remain_zero_centered_after_model_init(text_config):


can probably drop this test?

garrett361 · 2026-09-11T14:52:17Z

+    if int(os.environ.get("WORLD_SIZE", 1)) != 2:
+        pytest.skip("run with torchrun --nproc-per-node=2")


I think this always skips in CI, then? Maybe can delete if so, good enough it currently runs and passes, and can keep it in this branch's commit history

Follow up on Garrett's review of #3429. Composite MoE models store expert counts in `text_config`, so `is_tt_moe_model` missed them and the trainers skipped their load-balancing statistics. Resolve the text config before checking expert counts, with a CPU check covering dense/MoE text and VLM configs. Remove the unused-argument `del` statements while preserving the public signatures. Drop the zero-initialization assertion and the two-rank GDN test that ordinary CI always skipped, as requested in review. Keep the native forward/backward, packing, vision, MRoPE, and router replay coverage. Retain the existing compiler disable for CP convolution and link its exact upstream bug. Two-rank H200 testing on PyTorch 2.13.0+cu130 / FLA a954753 reproduces a Dynamo `copy_` shape error without the disable: the functional collective returns `[6, 8192]`, while FLA supplies a stacked `[2, 3, 8192]` output buffer. The same operation works eagerly; an isolated compiled collective using a concatenated buffer followed by reshape also passes. Upstream status checked September 12: FLA main (`516143e`) still has the identical communication helper. PyTorch issues [#138795](pytorch/pytorch#138795) and [#155632](pytorch/pytorch#155632) remain open. The matching [fix #177205](pytorch/pytorch#177205) closed without merging, and current PyTorch main still lacks the reshape. FLA's [graph-capture RFC #1155](fla-org/flash-linear-attention#1155 (comment)) defers CP; no exact FLA fix was found. Validation: - H200 Slurm 486: Qwen3.5 and shared MoE suite, **33 passed**. CPU MoE suite: **15 passed**. Ruff check/format and whitespace checks pass. - Standalone GDN probes use two H200 ranks and `torch.compile` with default settings (graph breaks allowed), two packed layouts, native unsharded output/input-gradient comparisons, and finite parameter gradients. These isolate the convolution boundary; they do not establish full trainer or image+CP compilation support. | CP convolution disable | Selective AC | Result | Slurm job | |---|---|---|---| | Present | Off | Both layouts pass on both ranks | 484 | | Removed in probe | Off | Dynamo collective shape error | 485 | | Present | On | Both layouts pass on both ranks | 487 | | Removed in probe | On | Both layouts pass on both ranks | 488 | The checkpointed probe passing does not make removal safe: the supported non-checkpointed configuration still fails. The isolated stacked/concatenated collective comparison passed its expected failure and parity checks (Slurm 489). Probe scripts and logs are retained outside the repository under `/home/matej/prime-rl-verify/qwen35-review-20260912/`.  --- > [!NOTE] > **Low Risk** > Targeted config lookup and test/doc cleanup; MoE detection fix reduces risk of missing load-balance metrics on composite MoE models. > > **Overview** > **MoE detection** now reads expert counts from `text_config` when present, so Qwen3.5 MoE VLMs are recognized and trainers collect load-balancing stats instead of treating them as dense models. A CPU parametrized test covers dense/MoE text and VLM config classes. > > Review cleanup removes unused `del` stubs on public APIs (`from_config`, RoPE helpers), drops CI-skipped Qwen3.5 tests (zero-init norm assertion and two-rank gated-delta-net CP parity), and tightens the comment on why context-parallel `causal_conv1d` stays outside `torch.compile` (Dynamo all-gather vs FLA buffer layout, [pytorch#155632](pytorch/pytorch#155632)). > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 0520d09. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>

S1ro1 marked this pull request as ready for review August 28, 2026 20:30

S1ro1 force-pushed the feat/qwen3-5-native branch 2 times, most recently from 8f26b48 to 47852a8 Compare August 30, 2026 17:53

S1ro1 changed the base branch from feat/nemotron-h-native to feat/selective-activation-checkpointing August 30, 2026 17:56

S1ro1 mentioned this pull request Aug 30, 2026

feat(models): add native Qwen3.8 Flash Next #3434

Open

S1ro1 force-pushed the feat/qwen3-5-native branch from 47852a8 to 767d8c4 Compare August 31, 2026 03:10

cursor Bot reviewed Aug 31, 2026

View reviewed changes

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

Base automatically changed from feat/selective-activation-checkpointing to main August 31, 2026 20:52

feat(models): add native Qwen3.5 family implementation

7178a5a

S1ro1 force-pushed the feat/qwen3-5-native branch from 767d8c4 to 7178a5a Compare August 31, 2026 20:52

S1ro1 added 2 commits September 10, 2026 12:08

Merge main into native Qwen3.5 implementation

71c3311

fix(qwen3.5): clean native model and honor shared runtime fusions

0e48de1

cursor Bot reviewed Sep 10, 2026

View reviewed changes

Comment thread docs/development.md

S1ro1 added 8 commits September 10, 2026 12:30

fix(qwen3.5): retain existing vision compiler boundary

a3eeb04

test(qwen3.5): fix expert choices when checking packed boundaries

1902f60

fix(qwen3.5): remove vision attention compiler disable

9983e7b

refactor(qwen3.5): use Prime attention configuration directly

a217360

chore: restore mini_moe script to main

133107b

fix: allow VLM RL training with FP32 optimizer state

dcdf770

test(qwen3.5): remove Transformers reference models

70d1445

test(qwen3.5): use public model factory and automatic attention

c48531b

cursor Bot reviewed Sep 10, 2026

View reviewed changes

Comment thread tests/unit/train/models/test_qwen3_5_moe.py

S1ro1 added 2 commits September 10, 2026 17:19

test(qwen3.5): use public helper names consistently

5ba7db5

test: consolidate native Qwen3.5 model coverage

d9162ce

S1ro1 mentioned this pull request Sep 10, 2026

test: remove HF reference suites and organize native model coverage #3527

Draft

garrett361 approved these changes Sep 11, 2026

View reviewed changes

S1ro1 merged commit 534888b into main Sep 12, 2026
17 of 19 checks passed

S1ro1 deleted the feat/qwen3-5-native branch September 12, 2026 15:08

S1ro1 mentioned this pull request Sep 12, 2026

fix(qwen3.5): address follow-up review cleanup #3534

Merged


		from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig

		# FLA's CP convolution uses an all-gather layout that Dynamo cannot trace.



		@pytest.mark.gpu
		def test_norms_remain_zero_centered_after_model_init(text_config):

		if int(os.environ.get("WORLD_SIZE", 1)) != 2:
		pytest.skip("run with torchrun --nproc-per-node=2")

Conversation

S1ro1 commented Aug 28, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Math KL validation

Packed VLM validation

Other validation

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

garrett361 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!

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!

2 participants

S1ro1 commented Aug 28, 2026 •

edited by cursor Bot

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading