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

docs/development.md

@@ -141,10 +141,10 @@ In the PR that adds the new model, you also need to provide a table covering the
 
 ## Adding a Custom VLM Implementation
 
-VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementation-only: `get_model` rejects models without a custom PrimeRL VLM class at load time. To make a new VLM family trainable, extend a custom text model with a composite VLM body — the Qwen3.5 dense (`models/qwen3_5/`) and MoE (`models/qwen3_5_moe/`) implementations are the reference. The pieces, in dependency order:
+VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementation-only: `get_model` rejects models without a custom PrimeRL VLM class at load time. To make a new VLM family trainable, extend a custom text model with a composite VLM body. The unified dense/MoE Qwen3.5 implementation in `models/qwen3_5/` is the reference. The pieces, in dependency order:
 
 1. **Custom text model first.** The VLM body wraps a custom `*ForCausalLM` (see [Adding a New Model](#adding-a-new-model)), so the text side — including its state-dict conversion and KL-mismatch table — comes first.
-2. **Composite VLM body.** A `*VLMModel` that holds the HF vision encoder and the custom text model, with a `prepare_inputs_embeds_and_position_ids` step: embed tokens, run the vision encoder, scatter image embeddings over placeholder tokens, and build MRoPE 3D positions from `mm_token_type_ids` (the renderer owns the token→modality mapping). The unified `*ForCausalLM` dispatches on the config: composite config → VLM path, text config → text path.
+2. **Composite VLM body.** A `*VLMModel` that holds the vision encoder and custom text model, embeds tokens, runs vision, scatters image embeddings over placeholder tokens, and builds MRoPE 3D positions from `mm_token_type_ids` (the renderer owns the token→modality mapping). The unified `*ForCausalLM` dispatches on the config: composite config → VLM path, text config → text path.
 3. **Always run the vision encoder.** Text-only micro-batches must feed the encoder dummy pixels and graft the result into the graph with zero contribution (`inputs_embeds + image_embeds.sum() * 0.0`) so FSDP/EP collectives stay symmetric across ranks when the encoder is trainable.
 4. **Packed-boundary consumption.** Samples pack into shared rows with per-document boundaries in `seq_lens`; every custom model's `forward()` declares the typed `seq_lens`/`seq_lens_are_pre_shard` parameters (the trainer passes them unconditionally) and must honor the boundaries — varlen flash `cu_seqlens`, linear-attention state resets per document, and a loud rejection on attention paths that can't (see the packed-batch guard in any modeling file). Set `supports_packed_multimodal_training` on the VLM model once packed rows are handled — RL fails loudly at startup for VLM models without it.
 5. **Registration.** Register the composite `model_type` in `_CUSTOM_VLM_MAPPING` (`models/__init__.py`) so `get_model` dispatches to the custom class, and describe the family in `VLM_REGISTRY` (`utils/vlm.py`).

scripts/mini_moe.py

@@ -28,7 +28,7 @@
 from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head
 from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config
 from prime_rl.trainer.models.minimax_m2 import MiniMaxM2ForCausalLM as PrimeRLMiniMaxM2ForCausalLM
-from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM as PrimeRLQwen3_5MoeVLM
+from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM as PrimeRLQwen3_5MoeVLM
 from prime_rl.utils.logger import setup_logger
 from prime_rl.utils.utils import default_dtype
 
@@ -264,6 +264,9 @@ def verify(arch: str, model_dir: Path) -> None:
     vocab_size = text_config.vocab_size
 
     hf_model = _load_hf_model(preset, model_dir, config).to(device="cuda", dtype=torch.float32)
+    config._attn_implementation = "flash_attention_2"
+    if hasattr(config, "text_config"):
+        config.text_config._attn_implementation = "flash_attention_2"
     with torch.device("cuda"), default_dtype(torch.float32):
         prime_model = preset["prime_model_class"]._from_config(config)
 
@@ -281,7 +284,11 @@ def verify(arch: str, model_dir: Path) -> None:
         position_ids = torch.arange(1, 65).unsqueeze(0)
 
     hf_output = hf_model(input_ids=input_ids, position_ids=position_ids)
-    prime_output = prime_model(input_ids, position_ids)
+    prime_output = prime_model(
+        input_ids,
+        position_ids,
+        seq_lens=torch.tensor([input_ids.shape[1]], device=input_ids.device),
+    )
 
     if is_vlm:
         # HF GatedDeltaNet has a dtype bug in float32 mode; just verify non-NaN output

src/prime_rl/trainer/model.py

@@ -56,7 +56,6 @@
 from prime_rl.trainer.parallel_dims import ParallelDims
 from prime_rl.trainer.world import get_world
 from prime_rl.utils.logger import get_logger
-from prime_rl.utils.sequence import get_cu_seqlens_from_position_ids
 from prime_rl.utils.utils import format_time
 from prime_rl.utils.vlm import get_language_model, get_vision_encoder, is_vlm_architecture
 from prime_rl.utils.weights import (
@@ -88,277 +87,6 @@ def pre_download_model(model_name: str, *, skip_weights: bool = False) -> None:
     )
 
 
-def _patch_qwen3_5_moe_conversion_mapping():
-    """Fix Qwen3.5 MoE conversion mapping incorrectly applying qwen2_moe expert weight splitting.
-
-    Qwen3.5 MoE stores expert weights as fused 3D tensors natively in the checkpoint
-    (e.g. experts.gate_up_proj [num_experts, 2*intermediate, hidden]). The upstream mapping
-    incorrectly maps qwen3_5_moe → qwen2_moe, which assumes per-expert 2D checkpoint weights,
-    causing revert_weight_conversion to produce wrong shapes during weight broadcasting.
-
-    Remove once an official Transformers release fixes this.
-    """
-    from transformers.conversion_mapping import (
-        get_checkpoint_conversion_mapping,
-        register_checkpoint_conversion_mapping,
-    )
-
-    # qwen3_5_moe_text: keep only the qwen3_5_text renaming, remove qwen2_moe expert conversion
-    qwen3_5_text_mapping = get_checkpoint_conversion_mapping("qwen3_5_text")
-    if qwen3_5_text_mapping is not None:
-        register_checkpoint_conversion_mapping("qwen3_5_moe_text", qwen3_5_text_mapping, overwrite=True)
-
-    # qwen3_5_moe: remove the qwen2_moe fallback entirely
-    register_checkpoint_conversion_mapping("qwen3_5_moe", [], overwrite=True)
-
-
-def _patch_qwen3_5_text_position_ids():
-    """Fix Qwen3.5 passing 3D MRoPE position_ids to decoder layers instead of 2D text_position_ids.
-
-    Upstream fix: https://github.com/huggingface/transformers/pull/44399
-    Remove once an official Transformers release includes this fix.
-    """
-    import inspect
-
-    from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5TextModel
-    from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeDecoderLayer, Qwen3_5MoeTextModel
-
-    for text_model_cls, decoder_layer_cls in [
-        (Qwen3_5TextModel, Qwen3_5DecoderLayer),
-        (Qwen3_5MoeTextModel, Qwen3_5MoeDecoderLayer),
-    ]:
-        source = inspect.getsource(text_model_cls.forward)
-        if "decoder_layer" in source and "position_ids=text_position_ids" in source.split("decoder_layer")[-1]:
-            continue  # already fixed upstream
-
-        _original_forward = decoder_layer_cls.forward
-
-        def _make_patched_forward(original):
-            def _patched_forward(self, hidden_states, position_ids=None, **kwargs):
-                if position_ids is not None and position_ids.ndim == 3:
-                    position_ids = position_ids[0]
-                return original(self, hidden_states, position_ids=position_ids, **kwargs)
-
-            return _patched_forward
-
-        decoder_layer_cls.forward = _make_patched_forward(_original_forward)
-
-
-def _patch_qwen3_5_linear_attn_varlen():
-    """Thread cu_seqlens through Qwen3.5 GatedDeltaNet so packed batches don't
-    leak conv/SSM state across sequences.
-
-    HF's forward hardcodes seq_idx=None for causal_conv1d and omits cu_seqlens
-    for chunk_gated_delta_rule, so packed RL training sees ~0.23 Mismatch KL vs
-    vLLM (target <0.01).
-    """
-    import torch.nn.functional as F
-    from transformers.models.qwen3_5.modeling_qwen3_5 import (
-        Qwen3_5DecoderLayer,
-        Qwen3_5GatedDeltaNet,
-        Qwen3_5TextModel,
-        apply_mask_to_padding_states,
-    )
-
-    try:
-        from fla.modules.convolution import causal_conv1d as fla_causal_conv1d
-        from fla.ops.cp import build_cp_context
-    except ImportError:
-        build_cp_context = None
-        fla_causal_conv1d = None
-
-    if fla_causal_conv1d is not None:
-        # The CP boundary-state exchange inside fla's conv is not dynamo-traceable;
-        # graph-break deliberately so the rest of the layer still compiles.
-        fla_causal_conv1d = torch.compiler.disable(fla_causal_conv1d)
-
-    if getattr(Qwen3_5GatedDeltaNet.forward, "_prl_varlen_patched", False):
-        return
-
-    _gdn_orig = Qwen3_5GatedDeltaNet.forward
-
-    def _build_cp_context(self, local_seq_len: int, device: torch.device, cu_seqlens=None):
-        cp_group = getattr(self, "cp_group", None)
-        if cp_group is None or build_cp_context is None:
-            return None
-        global_seq_len = local_seq_len * self.cp_world_size
-        if cu_seqlens is not None and int(cu_seqlens[-1].item()) == global_seq_len:
-            global_cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32)
-        else:
-            global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device)
-        return build_cp_context(
-            cu_seqlens=global_cu_seqlens,
-            group=cp_group,
-            conv1d_kernel_size=self.conv_kernel_size,
-        )
-
-    def _gdn_forward(self, hidden_states, cache_params=None, attention_mask=None, cu_seqlens=None):
-        if cu_seqlens is None or cache_params is not None:
-            return _gdn_orig(self, hidden_states, cache_params=cache_params, attention_mask=attention_mask)
-
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
-        batch_size, seq_len, _ = hidden_states.shape
-
-        mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2)
-        z = self.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, self.head_v_dim)
-        b = self.in_proj_b(hidden_states)
-        a = self.in_proj_a(hidden_states)
-
-        cp_context = _build_cp_context(self, seq_len, hidden_states.device, cu_seqlens)
-
-        if cp_context is not None and fla_causal_conv1d is not None:
-            mixed_qkv, _ = fla_causal_conv1d(
-                x=mixed_qkv.transpose(1, 2),
-                weight=self.conv1d.weight.squeeze(1),
-                bias=self.conv1d.bias,
-                activation=self.activation,
-                cp_context=cp_context,
-            )
-            mixed_qkv = mixed_qkv.transpose(1, 2)
-        elif self.causal_conv1d_fn is not None:
-            seg_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-            seq_idx = torch.repeat_interleave(
-                torch.arange(seg_lens.numel(), dtype=torch.int32, device=hidden_states.device),
-                seg_lens,
-            ).unsqueeze(0)
-            mixed_qkv = self.causal_conv1d_fn(
-                x=mixed_qkv,
-                weight=self.conv1d.weight.squeeze(1),
-                bias=self.conv1d.bias,
-                activation=self.activation,
-                seq_idx=seq_idx,
-            )
-        else:
-            # Per-segment conv1d so the kernel-1 left pad only draws from within each sequence.
-            cu = cu_seqlens.tolist()
-            conv_outs = []
-            for i in range(len(cu) - 1):
-                s, e = cu[i], cu[i + 1]
-                if s == e:
-                    continue
-                conv_outs.append(self.conv1d(mixed_qkv[:, :, s:e])[:, :, : e - s])
-            mixed_qkv = F.silu(torch.cat(conv_outs, dim=-1))
-
-        mixed_qkv = mixed_qkv.transpose(1, 2)
-        query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
-        query = query.reshape(batch_size, seq_len, -1, self.head_k_dim)
-        key = key.reshape(batch_size, seq_len, -1, self.head_k_dim)
-        value = value.reshape(batch_size, seq_len, -1, self.head_v_dim)
-
-        beta = b.sigmoid()
-        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
-        if self.num_v_heads // self.num_k_heads > 1:
-            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
-            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
-
-        if cp_context is not None:
-            core_attn_out, _ = self.chunk_gated_delta_rule(
-                query,
-                key,
-                value,
-                g=g,
-                beta=beta,
-                use_qk_l2norm_in_kernel=True,
-                cu_seqlens=cp_context.cu_seqlens,
-                cp_context=cp_context,
-            )
-        else:
-            core_attn_out, _ = self.chunk_gated_delta_rule(
-                query,
-                key,
-                value,
-                g=g,
-                beta=beta,
-                initial_state=None,
-                output_final_state=False,
-                use_qk_l2norm_in_kernel=True,
-                cu_seqlens=cu_seqlens,
-            )
-
-        core_attn_out = core_attn_out.reshape(-1, self.head_v_dim)
-        z = z.reshape(-1, self.head_v_dim)
-        core_attn_out = self.norm(core_attn_out, z)
-        core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1)
-        return self.out_proj(core_attn_out)
-
-    _gdn_forward._prl_varlen_patched = True
-    Qwen3_5GatedDeltaNet.forward = _gdn_forward
-
-    _dec_orig = Qwen3_5DecoderLayer.forward
-
-    def _dec_forward(
-        self,
-        hidden_states,
-        position_embeddings,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        cu_seqlens=None,
-        **kwargs,
-    ):
-        if position_ids is not None and position_ids.ndim == 3:
-            position_ids = position_ids[0]
-        if self.layer_type != "linear_attention":
-            return _dec_orig(
-                self,
-                hidden_states,
-                position_embeddings=position_embeddings,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                **kwargs,
-            )
-
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.linear_attn(
-            hidden_states=hidden_states,
-            cache_params=past_key_values,
-            attention_mask=attention_mask,
-            cu_seqlens=cu_seqlens,
-        )
-        hidden_states = residual + hidden_states
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        return residual + hidden_states
-
-    Qwen3_5DecoderLayer.forward = _dec_forward
-
-    _text_orig = Qwen3_5TextModel.forward
-
-    def _text_forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        use_cache=None,
-        **kwargs,
-    ):
-        attn_impl = getattr(self.config, "_attn_implementation", None)
-        cu_seqlens = None
-        if attn_impl in ("flash_attention_2", "flash_attention_3", "flash_attention_4") and position_ids is not None:
-            pids = position_ids
-            if pids.ndim == 3:
-                pids = pids[0]
-            cu_seqlens, _ = get_cu_seqlens_from_position_ids(pids)
-        kwargs["cu_seqlens"] = cu_seqlens
-        return _text_orig(
-            self,
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            **kwargs,
-        )
-
-    Qwen3_5TextModel.forward = _text_forward
-
-
 # Add filter to the standard logging module for transformers.modeling_utils to supress the
 # flash attention dtype warnings since FSDP is used to handle mixed precision.
 transformers_modeling_utils_logger = logging.getLogger("transformers.modeling_utils")
@@ -569,11 +297,6 @@ def get_model(
                 "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference."
             )
 
-    # Qwen3.6 and Qwen3.8 reuse the Qwen3.5 architecture, so match on model_type, not repo name.
-    if getattr(model_config, "model_type", "").startswith("qwen3_5"):
-        _patch_qwen3_5_text_position_ids()
-        _patch_qwen3_5_moe_conversion_mapping()
-        _patch_qwen3_5_linear_attn_varlen()
     if is_vlm_arch and config.cp > 1 and config.cp_style == "ulysses":
         vision_config = getattr(model_config, "vision_config", None)
         if vision_config is not None: