From 7178a5ac82a96c45353ea1cfa494463d2bccb595 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Fri, 28 Aug 2026 20:13:18 +0000 Subject: [PATCH 01/12] feat(models): add native Qwen3.5 family implementation --- docs/development.md | 4 +- scripts/mini_moe.py | 11 +- src/prime_rl/trainer/model.py | 277 ----- src/prime_rl/trainer/models/__init__.py | 22 +- src/prime_rl/trainer/models/base.py | 3 +- src/prime_rl/trainer/models/layers/attn.py | 8 - .../trainer/models/layers/ulysses_attn.py | 8 - .../trainer/models/qwen3_5/__init__.py | 18 +- .../trainer/models/qwen3_5/attention.py | 89 ++ .../models/qwen3_5/configuration_qwen3_5.py | 244 +++++ .../models/qwen3_5/converting_qwen3_5.py | 48 + .../trainer/models/qwen3_5/gated_delta_net.py | 127 +++ .../models/qwen3_5/modeling_qwen3_5.py | 523 +++++---- src/prime_rl/trainer/models/qwen3_5/norm.py | 25 + .../models/qwen3_5/rotary_embedding.py | 185 ++++ src/prime_rl/trainer/models/qwen3_5/vision.py | 286 +++++ .../trainer/models/qwen3_5_moe/__init__.py | 13 - .../qwen3_5_moe/configuration_qwen3_5_moe.py | 106 -- .../qwen3_5_moe/converting_qwen3_5_moe.py | 25 - .../qwen3_5_moe/modeling_qwen3_5_moe.py | 988 ------------------ .../trainer/models/qwen3_5_moe/mrope.py | 130 --- src/prime_rl/trainer/parallel_dims.py | 1 + tests/unit/train/models/test_qwen3_5.py | 124 ++- tests/unit/train/models/test_qwen3_5_moe.py | 63 +- .../train/models/test_qwen3_5_moe_mrope.py | 14 +- .../unit/train/models/test_qwen3_5_moe_vlm.py | 27 +- 26 files changed, 1448 insertions(+), 1921 deletions(-) create mode 100644 src/prime_rl/trainer/models/qwen3_5/attention.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/converting_qwen3_5.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/norm.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py create mode 100644 src/prime_rl/trainer/models/qwen3_5/vision.py delete mode 100644 src/prime_rl/trainer/models/qwen3_5_moe/__init__.py delete mode 100644 src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py delete mode 100644 src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py delete mode 100644 src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py delete mode 100644 src/prime_rl/trainer/models/qwen3_5_moe/mrope.py diff --git a/docs/development.md b/docs/development.md index ed3bcfe0211..f45e645ecf0 100644 --- a/docs/development.md +++ b/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`). diff --git a/scripts/mini_moe.py b/scripts/mini_moe.py index c151d0a7555..9d884dcb4e7 100644 --- a/scripts/mini_moe.py +++ b/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 diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 095795b8438..ddc1917ecfb 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -50,7 +50,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 ( @@ -82,277 +81,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). Mirrors the NemotronH mamba fix. - """ - 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") @@ -575,11 +303,6 @@ def get_model( _hub_kernels._kernels_enabled = True - # 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: diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index c9b84c19e7c..2fbade4f9f2 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -8,7 +8,6 @@ from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.qwen3.configuration_qwen3 import Qwen3Config -from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig from prime_rl.trainer.models.afmoe import AfmoeConfig, AfmoeForCausalLM from prime_rl.trainer.models.base import PreTrainedModelPrimeRL @@ -21,8 +20,14 @@ from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config, MiniMaxM2ForCausalLM from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM from prime_rl.trainer.models.qwen3 import Qwen3ForCausalLM -from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeForCausalLM +from prime_rl.trainer.models.qwen3_5 import ( + Qwen3_5Config, + Qwen3_5ForCausalLM, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, + Qwen3_5VisionConfig, +) from prime_rl.trainer.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM # Make custom config discoverable by AutoConfig @@ -33,8 +38,11 @@ AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True) AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True) AutoConfig.register("qwen3_moe", Qwen3MoeConfig, exist_ok=True) +AutoConfig.register("qwen3_5", Qwen3_5Config, exist_ok=True) AutoConfig.register("qwen3_5_text", Qwen3_5TextConfig, exist_ok=True) -AutoConfig.register("qwen3_5_moe_text", Qwen3_5MoeConfig, exist_ok=True) +AutoConfig.register("qwen3_5_moe", Qwen3_5MoeConfig, exist_ok=True) +AutoConfig.register("qwen3_5_moe_text", Qwen3_5MoeTextConfig, exist_ok=True) +AutoConfig.register("qwen3_5_vision", Qwen3_5VisionConfig, exist_ok=True) # GptOssConfig is just HF's class - already registered by transformers, no override needed. _CUSTOM_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, OrderedDict()) @@ -52,7 +60,9 @@ (NemotronHConfig, NemotronHForCausalLM), (Qwen3MoeConfig, Qwen3MoeForCausalLM), (Qwen3_5TextConfig, Qwen3_5ForCausalLM), - (Qwen3_5MoeConfig, Qwen3_5MoeForCausalLM), + (Qwen3_5MoeTextConfig, Qwen3_5ForCausalLM), + (Qwen3_5Config, Qwen3_5ForCausalLM), + (Qwen3_5MoeConfig, Qwen3_5ForCausalLM), (GptOssConfig, GptOssForCausalLM), ) for config_cls, model_cls in _CUSTOM_CAUSAL_LM_MODELS: @@ -92,7 +102,7 @@ def supports_custom_impl(model_config: PretrainedConfig) -> bool: # Points to the same unified class — the config drives text-only vs VLM behavior. _CUSTOM_VLM_MAPPING: dict[str, type] = { "qwen3_5": Qwen3_5ForCausalLM, - "qwen3_5_moe": Qwen3_5MoeForCausalLM, + "qwen3_5_moe": Qwen3_5ForCausalLM, } diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py index 363702b1f22..03701476cc2 100644 --- a/src/prime_rl/trainer/models/base.py +++ b/src/prime_rl/trainer/models/base.py @@ -20,8 +20,9 @@ def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool: return False @classmethod - def from_config(cls, config, **kwargs): + def from_config(cls, config, trust_remote_code: bool = False, **kwargs): """Public from_config that mirrors the Auto class API.""" + del trust_remote_code return cls._from_config(config, **kwargs) @classmethod diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index 58dfb98ac61..23828d18492 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -199,11 +199,3 @@ def _ring_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention AfmoeFlashAttention._compute_attention = _ring_compute_attention - - from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention - - Qwen3_5MoeGatedFlashAttention._compute_attention = _ring_compute_attention - - from prime_rl.trainer.models.qwen3_5.modeling_qwen3_5 import Qwen3_5GatedFlashAttention - - Qwen3_5GatedFlashAttention._compute_attention = _ring_compute_attention diff --git a/src/prime_rl/trainer/models/layers/ulysses_attn.py b/src/prime_rl/trainer/models/layers/ulysses_attn.py index 92c390a5b0f..fa6e5d22968 100644 --- a/src/prime_rl/trainer/models/layers/ulysses_attn.py +++ b/src/prime_rl/trainer/models/layers/ulysses_attn.py @@ -221,14 +221,6 @@ def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): AfmoeFlashAttention._compute_attention = _ulysses_compute_attention - from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention - - Qwen3_5MoeGatedFlashAttention._compute_attention = _ulysses_compute_attention - - from prime_rl.trainer.models.qwen3_5.modeling_qwen3_5 import Qwen3_5GatedFlashAttention - - Qwen3_5GatedFlashAttention._compute_attention = _ulysses_compute_attention - def substitute_hf_ulysses_attn(process_group: dist.ProcessGroup) -> None: """Patch HF's `_flash_attention_forward` to use Ulysses all-to-all + local FA2. diff --git a/src/prime_rl/trainer/models/qwen3_5/__init__.py b/src/prime_rl/trainer/models/qwen3_5/__init__.py index c5fba8b7d45..baa0b00c409 100644 --- a/src/prime_rl/trainer/models/qwen3_5/__init__.py +++ b/src/prime_rl/trainer/models/qwen3_5/__init__.py @@ -1,3 +1,19 @@ +from .configuration_qwen3_5 import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, + Qwen3_5VisionConfig, +) from .modeling_qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5Model, Qwen3_5PreTrainedModel -__all__ = ["Qwen3_5ForCausalLM", "Qwen3_5Model", "Qwen3_5PreTrainedModel"] +__all__ = [ + "Qwen3_5Config", + "Qwen3_5ForCausalLM", + "Qwen3_5Model", + "Qwen3_5MoeConfig", + "Qwen3_5MoeTextConfig", + "Qwen3_5PreTrainedModel", + "Qwen3_5TextConfig", + "Qwen3_5VisionConfig", +] diff --git a/src/prime_rl/trainer/models/qwen3_5/attention.py b/src/prime_rl/trainer/models/qwen3_5/attention.py new file mode 100644 index 00000000000..040f3118199 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/attention.py @@ -0,0 +1,89 @@ +import torch +from torch import nn + +from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention +from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig +from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm + + +class Qwen3_5Attention(FlashAttention): + def __init__(self, config: Qwen3_5TextConfig, attention_implementation: str) -> None: + flash_attn_version = { + "flash_attention_2": 2, + "flash_attention_3": 3, + "flash_attention_4": 4, + }[attention_implementation] + super().__init__( + AttentionConfig( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + is_causal=True, + attention_bias=config.attention_bias, + output_bias=config.attention_bias, + use_qk_norm=False, + rms_norm_eps=config.rms_norm_eps, + ), + flash_attn_version=flash_attn_version, + ) + 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) + self.q_proj = nn.Linear(config.hidden_size, q_output_size, bias=config.attention_bias) + 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( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor, + max_seqlen: int, + ) -> tuple[torch.Tensor, None]: + batch_size, sequence_length, _ = hidden_states.shape + + query_states = self.q_proj(hidden_states) + if self.attn_output_gate: + query_states, output_gate = query_states.view( + batch_size, sequence_length, self.num_heads, self.head_dim * 2 + ).chunk(2, dim=-1) + output_gate = output_gate.reshape(batch_size, sequence_length, -1) + else: + query_states = query_states.view(batch_size, sequence_length, self.num_heads, self.head_dim) + output_gate = None + + key_states = self.k_proj(hidden_states).view( + batch_size, sequence_length, self.num_key_value_heads, self.head_dim + ) + value_states = self.v_proj(hidden_states).view( + batch_size, sequence_length, self.num_key_value_heads, self.head_dim + ) + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states.transpose(1, 2), + key_states.transpose(1, 2), + cos, + sin, + ) + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + + attention_output = self._compute_attention( + query_states[0], + key_states[0], + value_states[0], + cu_seqlens, + max_seqlen, + ).reshape(batch_size, sequence_length, -1) + if output_gate is not None: + attention_output = attention_output * output_gate.sigmoid() + return self.o_proj(attention_output), None + + +__all__ = ["Qwen3_5Attention"] diff --git a/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py new file mode 100644 index 00000000000..b643d8bf2fe --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py @@ -0,0 +1,244 @@ +from transformers.configuration_utils import PretrainedConfig + + +class Qwen3_5TextConfig(PretrainedConfig): + model_type = "qwen3_5_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", + "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size: int = 248320, + hidden_size: int = 4096, + intermediate_size: int = 12288, + num_hidden_layers: int = 32, + num_attention_heads: int = 16, + num_key_value_heads: int = 4, + head_dim: int = 256, + hidden_act: str = "silu", + max_position_embeddings: int = 32768, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10_000.0, + rope_scaling: dict | None = None, + rope_parameters: dict | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + attn_output_gate: bool = True, + output_gate_type: str = "silu", + linear_conv_kernel_dim: int = 4, + linear_key_head_dim: int = 128, + linear_value_head_dim: int = 128, + linear_num_key_heads: int = 16, + linear_num_value_heads: int = 32, + layer_types: list[str] | None = None, + full_attention_interval: int = 4, + mamba_ssm_dtype: str = "float32", + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int | None = None, + eos_token_id: int | list[int] | None = None, + tie_word_embeddings: bool = False, + **kwargs, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.attn_output_gate = attn_output_gate + self.output_gate_type = output_gate_type + + if rope_parameters is None: + rope_parameters = rope_scaling + if rope_parameters is None: + rope_parameters = { + "rope_type": "default", + "rope_theta": rope_theta, + "partial_rotary_factor": kwargs.pop("partial_rotary_factor", 0.25), + "mrope_interleaved": True, + } + self.rope_parameters = dict(rope_parameters) + self.rope_parameters.setdefault("rope_type", "default") + self.rope_parameters.setdefault("rope_theta", rope_theta) + self.rope_parameters.setdefault("partial_rotary_factor", kwargs.pop("partial_rotary_factor", 0.25)) + self.rope_parameters.setdefault("mrope_interleaved", True) + + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + self.layer_types = layer_types or [ + "linear_attention" if (layer_idx + 1) % full_attention_interval else "full_attention" + for layer_idx in range(num_hidden_layers) + ] + self.mamba_ssm_dtype = mamba_ssm_dtype + self.use_cache = use_cache + + kwargs.pop("mtp_num_hidden_layers", None) + kwargs.pop("mtp_use_dedicated_embeddings", None) + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class Qwen3_5MoeTextConfig(Qwen3_5TextConfig): + model_type = "qwen3_5_moe_text" + + base_model_tp_plan = { + **Qwen3_5TextConfig.base_model_tp_plan, + "layers.*.mlp.experts.gate_proj": "grouped_gemm", + "layers.*.mlp.experts.up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_expert.gate_proj": "colwise", + "layers.*.mlp.shared_expert.up_proj": "colwise", + "layers.*.mlp.shared_expert.down_proj": "rowwise", + } + base_model_ep_plan = { + "layers.*.mlp.router": "ep_router", + "layers.*.mlp.experts.gate_proj": "grouped_gemm", + "layers.*.mlp.experts.up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + } + + def __init__( + self, + moe_intermediate_size: int = 512, + shared_expert_intermediate_size: int = 512, + num_experts_per_tok: int = 8, + num_experts: int = 256, + output_router_logits: bool = False, + router_aux_loss_coef: float = 0.001, + load_balance_coeff: float | None = None, + **kwargs, + ) -> None: + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + self.load_balance_coeff = load_balance_coeff + kwargs.pop("mlp_only_layers", None) + super().__init__(**kwargs) + + +class Qwen3_5VisionConfig(PretrainedConfig): + model_type = "qwen3_5_vision" + base_config_key = "vision_config" + + def __init__( + self, + depth: int = 27, + hidden_size: int = 1152, + hidden_act: str = "gelu_pytorch_tanh", + intermediate_size: int = 4304, + num_heads: int = 16, + in_channels: int = 3, + patch_size: int = 16, + spatial_merge_size: int = 2, + temporal_patch_size: int = 2, + out_hidden_size: int = 3584, + num_position_embeddings: int = 2304, + initializer_range: float = 0.02, + **kwargs, + ) -> None: + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + kwargs.pop("deepstack_visual_indexes", None) + super().__init__(**kwargs) + + +class Qwen3_5Config(PretrainedConfig): + model_type = "qwen3_5" + sub_configs = {"vision_config": Qwen3_5VisionConfig, "text_config": Qwen3_5TextConfig} + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: dict | PretrainedConfig | None = None, + vision_config: dict | PretrainedConfig | None = None, + image_token_id: int = 248056, + video_token_id: int = 248057, + vision_start_token_id: int = 248053, + vision_end_token_id: int = 248054, + tie_word_embeddings: bool = False, + **kwargs, + ) -> None: + self.text_config = ( + self.sub_configs["text_config"](**text_config) if isinstance(text_config, dict) else text_config + ) or self.sub_configs["text_config"]() + self.vision_config = ( + self.sub_configs["vision_config"](**vision_config) if isinstance(vision_config, dict) else vision_config + ) or self.sub_configs["vision_config"]() + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + kwargs.pop("language_model_only", None) + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + +class Qwen3_5MoeConfig(Qwen3_5Config): + model_type = "qwen3_5_moe" + sub_configs = {"vision_config": Qwen3_5VisionConfig, "text_config": Qwen3_5MoeTextConfig} + + def __init__( + self, + text_config: dict | PretrainedConfig | None = None, + vision_config: dict | PretrainedConfig | None = None, + **kwargs, + ) -> None: + super().__init__(text_config=text_config, vision_config=vision_config, **kwargs) + + +__all__ = [ + "Qwen3_5Config", + "Qwen3_5MoeConfig", + "Qwen3_5MoeTextConfig", + "Qwen3_5TextConfig", + "Qwen3_5VisionConfig", +] diff --git a/src/prime_rl/trainer/models/qwen3_5/converting_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/converting_qwen3_5.py new file mode 100644 index 00000000000..f19280a3b59 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/converting_qwen3_5.py @@ -0,0 +1,48 @@ +from torch import Tensor + +from prime_rl.trainer.models.conversion_ops import ConvOp, Drop, Rename, SplitConcat +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5MoeTextConfig + + +def is_hf_state_dict(state_dict: dict[str, Tensor]) -> bool: + return any(name == "mtp.fc.weight" or name.endswith("mlp.experts.gate_up_proj") for name in state_dict) + + +def is_prime_state_dict(state_dict: dict[str, Tensor]) -> bool: + if any(name.endswith("mlp.experts.gate_proj") for name in state_dict): + return True + has_qwen_layers = any(name.endswith("linear_attn.A_log") for name in state_dict) + has_fused_experts = any(name.endswith("mlp.experts.gate_up_proj") for name in state_dict) + return has_qwen_layers and not has_fused_experts + + +def conversion_chain(config) -> list[ConvOp]: + operations: list[ConvOp] = [Drop("mtp.", is_prefix=True)] + text_config = getattr(config, "text_config", config) + if not isinstance(text_config, Qwen3_5MoeTextConfig): + return operations + + model_prefix = "model.language_model" if hasattr(config, "vision_config") else "model" + for layer_index in range(text_config.num_hidden_layers): + prefix = f"{model_prefix}.layers.{layer_index}.mlp" + operations.extend( + [ + Rename(f"{prefix}.gate.weight", f"{prefix}.router.gate.weight"), + Rename( + f"{prefix}.shared_expert_gate.weight", + f"{prefix}.shared_expert.output_gate.weight", + ), + SplitConcat( + combined=f"{prefix}.experts.gate_up_proj", + parts=[ + (f"{prefix}.experts.gate_proj", None), + (f"{prefix}.experts.up_proj", None), + ], + dim=1, + ), + ] + ) + return operations + + +__all__ = ["conversion_chain", "is_hf_state_dict", "is_prime_state_dict"] diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py new file mode 100644 index 00000000000..84ccccc1594 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py @@ -0,0 +1,127 @@ +import torch +import torch.nn.functional as F +from fla.modules import FusedRMSNormGated +from fla.modules.conv import causal_conv1d +from fla.ops.cp import build_cp_context +from fla.ops.gated_delta_rule import chunk_gated_delta_rule +from torch import nn + +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + +# FLA's CP convolution exchanges boundary state inside the kernel call, which +# Dynamo cannot trace; keep that call eager while compiling the surrounding layer. +causal_conv1d_with_context_parallelism = torch.compiler.disable(causal_conv1d) + + +class Qwen3_5GatedDeltaNet(nn.Module): + def __init__(self, config: Qwen3_5TextConfig) -> None: + super().__init__() + self.num_value_heads = config.linear_num_value_heads + self.num_key_heads = config.linear_num_key_heads + self.key_head_dim = config.linear_key_head_dim + self.value_head_dim = config.linear_value_head_dim + self.key_dim = self.key_head_dim * self.num_key_heads + self.value_dim = self.value_head_dim * self.num_value_heads + self.conv_kernel_size = config.linear_conv_kernel_dim + self.activation = config.hidden_act + + conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + conv_dim, + conv_dim, + kernel_size=self.conv_kernel_size, + groups=conv_dim, + padding=self.conv_kernel_size - 1, + bias=False, + ) + self.dt_bias = nn.Parameter(torch.ones(self.num_value_heads)) + self.A_log = nn.Parameter(torch.empty(self.num_value_heads).uniform_(0, 16).log_()) + self.norm = FusedRMSNormGated( + self.value_head_dim, + eps=config.rms_norm_eps, + activation=config.output_gate_type, + ) + self.out_proj = nn.Linear(self.value_dim, config.hidden_size, bias=False) + self.in_proj_qkv = nn.Linear(config.hidden_size, conv_dim, bias=False) + self.in_proj_z = nn.Linear(config.hidden_size, self.value_dim, bias=False) + self.in_proj_b = nn.Linear(config.hidden_size, self.num_value_heads, bias=False) + self.in_proj_a = nn.Linear(config.hidden_size, self.num_value_heads, bias=False) + self.context_parallel_group = None + self.context_parallel_world_size = 1 + + def set_context_parallel_attributes(self, process_group, world_size: int) -> None: + self.context_parallel_group = process_group + self.context_parallel_world_size = world_size + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.LongTensor, + *, + cu_seqlens_are_pre_shard: bool, + ) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + mixed_qkv = self.in_proj_qkv(hidden_states) + output_gate = self.in_proj_z(hidden_states).reshape( + batch_size, sequence_length, self.num_value_heads, self.value_head_dim + ) + beta = self.in_proj_b(hidden_states).sigmoid() + decay = -self.A_log.float().exp() * F.softplus(self.in_proj_a(hidden_states).float() + self.dt_bias) + + context = None + if self.context_parallel_group is not None: + if not cu_seqlens_are_pre_shard: + raise ValueError("Qwen3.5 context parallelism requires pre-shard sequence boundaries") + context = build_cp_context( + cu_seqlens=cu_seqlens.to(device=hidden_states.device, dtype=torch.int32), + group=self.context_parallel_group, + conv1d_kernel_size=self.conv_kernel_size, + ) + + conv_kwargs = { + "x": mixed_qkv, + "weight": self.conv1d.weight.squeeze(1), + "bias": self.conv1d.bias, + "activation": self.activation, + } + if context is None: + mixed_qkv, _ = causal_conv1d(**conv_kwargs, cu_seqlens=cu_seqlens) + else: + mixed_qkv, _ = causal_conv1d_with_context_parallelism(**conv_kwargs, cp_context=context) + + query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + query = query.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim) + key = key.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim) + value = value.reshape(batch_size, sequence_length, self.num_value_heads, self.value_head_dim) + + heads_per_key = self.num_value_heads // self.num_key_heads + if heads_per_key > 1: + query = query.repeat_interleave(heads_per_key, dim=2) + key = key.repeat_interleave(heads_per_key, dim=2) + + delta_kwargs = { + "q": query, + "k": key, + "v": value, + "g": decay, + "beta": beta, + "use_qk_l2norm_in_kernel": True, + "cu_seqlens": context.cu_seqlens if context is not None else cu_seqlens, + } + if context is None: + core_output, _ = chunk_gated_delta_rule( + **delta_kwargs, + initial_state=None, + output_final_state=False, + ) + else: + core_output, _ = chunk_gated_delta_rule(**delta_kwargs, cp_context=context) + + core_output = self.norm( + core_output.reshape(-1, self.value_head_dim), + output_gate.reshape(-1, self.value_head_dim), + ).reshape(batch_size, sequence_length, self.value_dim) + return self.out_proj(core_output) + + +__all__ = ["Qwen3_5GatedDeltaNet"] 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 69618fe1f53..88a2f32301d 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 @@ -1,317 +1,285 @@ -import functools -from typing import Optional, Union - import torch from torch import Tensor, nn -from transformers.cache_utils import Cache -from transformers.configuration_utils import PretrainedConfig -from transformers.generation import GenerationMixin -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_outputs import BaseModelOutputWithPast -from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig -from transformers.models.qwen3_5.modeling_qwen3_5 import ( - Qwen3_5PreTrainedModel as HFQwen3_5PreTrainedModel, -) -from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel -from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs +from transformers.modeling_outputs import BaseModelOutput from prime_rl.trainer.models.base import PreTrainedModelPrimeRL -from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear from prime_rl.trainer.models.layers.mlp import FeedForward -from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeGatedAttentionConfig, - Qwen3_5MoeGatedDeltaNet, - Qwen3_5MoeGatedFlashAttention, - Qwen3_5MoeRMSNorm, - Qwen3_5MoeRotaryEmbedding, - normalize_qwen3_5_attn_implementation, +from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter +from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import ( + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, +) +from prime_rl.trainer.models.qwen3_5.converting_qwen3_5 import ( + conversion_chain, + is_hf_state_dict, + is_prime_state_dict, ) -from prime_rl.trainer.models.qwen3_5_moe.mrope import build_qwen3_5_mrope_position_ids +from prime_rl.trainer.models.qwen3_5.gated_delta_net import Qwen3_5GatedDeltaNet +from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm +from prime_rl.trainer.models.qwen3_5.rotary_embedding import ( + Qwen3_5RotaryEmbedding, + build_qwen3_5_mrope_position_ids, +) +from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel from prime_rl.utils.cp import setup_cp_attention_params, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens -class Qwen3_5GatedFlashAttention(Qwen3_5MoeGatedFlashAttention): - pass - - -QWEN35_ATTN_IMPL2CLASS = { - "flash_attention_2": functools.partial(Qwen3_5GatedFlashAttention, flash_attn_version=2), - "flash_attention_3": functools.partial(Qwen3_5GatedFlashAttention, flash_attn_version=3), - "flash_attention_4": functools.partial(Qwen3_5GatedFlashAttention, flash_attn_version=4), -} - - -def _get_gated_attention(config: Qwen3_5TextConfig) -> nn.Module: - attn_config = Qwen3_5MoeGatedAttentionConfig( - hidden_size=config.hidden_size, - head_dim=config.head_dim, - num_attention_heads=config.num_attention_heads, - num_key_value_heads=config.num_key_value_heads, - rms_norm_eps=config.rms_norm_eps, - attention_bias=config.attention_bias, - attention_dropout=config.attention_dropout, - ) - - attn_impl = normalize_qwen3_5_attn_implementation(config._attn_implementation) - config._attn_implementation = attn_impl - - if attn_impl not in QWEN35_ATTN_IMPL2CLASS: - supported = list(QWEN35_ATTN_IMPL2CLASS.keys()) - raise ValueError( - f"Qwen3.5 attention does not support '{config._attn_implementation}'. " - f"Supported implementations: {supported}." +class Qwen3_5SharedExpert(FeedForward): + def __init__(self, config: Qwen3_5MoeTextConfig) -> None: + super().__init__( + dim=config.hidden_size, + hidden_dim=config.shared_expert_intermediate_size, + expert_type="gated", + activation=config.hidden_act, ) + self.output_gate = nn.Linear(config.hidden_size, 1, bias=False) - return QWEN35_ATTN_IMPL2CLASS[attn_impl](attn_config) - + def forward(self, hidden_states: torch.Tensor, routed_experts: torch.Tensor | None = None) -> torch.Tensor: + output = super().forward(hidden_states, routed_experts) + return output * self.output_gate(hidden_states).sigmoid() -def _create_rotary_emb(config: Qwen3_5TextConfig) -> Qwen3_5MoeRotaryEmbedding: - return Qwen3_5MoeRotaryEmbedding(config) - -class Qwen3_5DecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: Qwen3_5TextConfig, layer_idx: int): +class Qwen3_5DecoderLayer(nn.Module): + def __init__(self, config: Qwen3_5TextConfig, layer_index: int) -> None: super().__init__() - self.hidden_size = config.hidden_size - self.layer_type = config.layer_types[layer_idx] - + self.layer_type = config.layer_types[layer_index] if self.layer_type == "linear_attention": - self.linear_attn = Qwen3_5MoeGatedDeltaNet(config) + self.linear_attn = Qwen3_5GatedDeltaNet(config) elif self.layer_type == "full_attention": - self.self_attn = _get_gated_attention(config) + self.self_attn = Qwen3_5Attention(config, config._attn_implementation) else: raise ValueError(f"Unsupported Qwen3.5 layer type: {self.layer_type}") - self.mlp = FeedForward( - dim=config.hidden_size, - hidden_dim=config.intermediate_size, - activation=config.hidden_act, - ) - self.input_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if isinstance(config, Qwen3_5MoeTextConfig): + router = TokenChoiceTopKRouter( + dim=config.hidden_size, + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + score_func="softmax", + route_norm=True, + route_scale=1.0, + ) + experts = GroupedExperts( + dim=config.hidden_size, + hidden_dim=config.moe_intermediate_size, + num_experts=config.num_experts, + expert_type="gated", + activation=config.hidden_act, + ) + experts.init_weights(config.initializer_range) + self.mlp = MoE( + router=router, + experts=experts, + shared_expert=Qwen3_5SharedExpert(config), + score_before_experts=False, + load_balance_coeff=config.load_balance_coeff, + ) + else: + self.mlp = FeedForward( + dim=config.hidden_size, + hidden_dim=config.intermediate_size, + expert_type="gated", + activation=config.hidden_act, + ) + + self.input_layernorm = Qwen3_5RMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = Qwen3_5RMSNorm(config.hidden_size, config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - cu_seqlens_are_pre_shard: bool = False, - ) -> torch.FloatTensor: + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor, + max_seqlen: int, + routed_experts: torch.LongTensor | None = None, + *, + cu_seqlens_are_pre_shard: bool, + ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - if self.layer_type == "linear_attention": hidden_states = self.linear_attn( - hidden_states, cu_seqlens=cu_seqlens, cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard + hidden_states, + cu_seqlens, + cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard, ) else: hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, + hidden_states, + position_embeddings, + cu_seqlens, + max_seqlen, ) 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 + return residual + self.mlp(hidden_states, routed_experts=routed_experts) -class Qwen3_5PreTrainedModel(PreTrainedModelPrimeRL, HFQwen3_5PreTrainedModel): +class Qwen3_5PreTrainedModel(PreTrainedModelPrimeRL): config_class = Qwen3_5TextConfig base_model_prefix = "model" supports_gradient_checkpointing = True - _no_split_modules = ["Qwen3_5DecoderLayer"] - _skip_keys_device_placement = ["past_key_values"] + _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock"] _supports_flash_attn = True _supports_sdpa = False _supports_flex_attn = False _supports_attention_backend = True _can_compile_fullgraph = False - _can_record_outputs = { - "hidden_states": Qwen3_5DecoderLayer, - } @classmethod def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool: return name.endswith(("linear_attn.A_log", "linear_attn.norm.weight")) - def _check_and_adjust_attn_implementation( - self, attn_implementation: str | None, is_init_check: bool = False, allow_all_kernels: bool = False - ) -> str: - attn_impl = normalize_qwen3_5_attn_implementation(attn_implementation or "flash_attention_3") - if attn_impl not in QWEN35_ATTN_IMPL2CLASS: - supported = list(QWEN35_ATTN_IMPL2CLASS.keys()) - raise ValueError( - f"Qwen3.5 attention does not support '{attn_implementation}'. Supported implementations: {supported}." - ) - return attn_impl - @classmethod def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return True + return is_hf_state_dict(state_dict) @classmethod def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - # Dense models use identical key names in HF and PrimeRL format, so we - # never claim to be in a separate PrimeRL format. This disables the - # auto-conversion path in load_dcp_from_hf and lets DCP load directly. - return False - - @classmethod - def convert_to_hf(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: - return state_dict - - @classmethod - def convert_to_prime(cls, state_dict: dict[str, Tensor]) -> dict[str, Tensor]: - return state_dict + return is_prime_state_dict(state_dict) @classmethod - def convert_layer_to_hf(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: - return state_dict - - @classmethod - def convert_layer_to_prime(cls, state_dict: dict[str, Tensor], layer_idx: int) -> dict[str, Tensor]: - return state_dict + def conversion_chain(cls, config): + return conversion_chain(config) class Qwen3_5Model(Qwen3_5PreTrainedModel): - def __init__(self, config: Qwen3_5TextConfig): - config._attn_implementation = normalize_qwen3_5_attn_implementation(config._attn_implementation) + def __init__(self, config: Qwen3_5TextConfig) -> None: super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( - [Qwen3_5DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + Qwen3_5DecoderLayer(config, layer_index) for layer_index in range(config.num_hidden_layers) ) - self.norm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = _create_rotary_emb(config) + self.norm = Qwen3_5RMSNorm(config.hidden_size, config.rms_norm_eps) + self.rotary_emb = Qwen3_5RotaryEmbedding(config) self.gradient_checkpointing = False - self.post_init() + # Qwen stores RMSNorm weights as offsets from one, while post_init treats them as direct scales. + for module in self.modules(): + if isinstance(module, Qwen3_5RMSNorm): + module.reset_parameters() - def get_input_embeddings(self): + def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens - def set_input_embeddings(self, value): - self.embed_tokens = value + def set_input_embeddings(self, embeddings: nn.Embedding) -> None: + self.embed_tokens = embeddings - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self._cp_group = cp_group - self._cp_rank = cp_rank - self._cp_world_size = cp_world_size - for layer in self.layers.modules(): - if getattr(layer, "layer_type", None) == "linear_attention": - layer.linear_attn.cp_group = cp_group - layer.linear_attn.cp_rank = cp_rank - layer.linear_attn.cp_world_size = cp_world_size + def set_context_parallel_attributes(self, process_group, rank: int, world_size: int) -> None: + self.context_parallel_group = process_group + self.context_parallel_rank = rank + self.context_parallel_world_size = world_size + for module in self.modules(): + if isinstance(module, Qwen3_5GatedDeltaNet): + module.set_context_parallel_attributes(process_group, world_size) def forward( self, - input_ids: Optional[torch.LongTensor] = None, - position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + routed_experts: torch.LongTensor | None = None, *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, - ) -> BaseModelOutputWithPast: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - + ) -> BaseModelOutput: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - + if inputs_embeds.shape[0] != 1: + raise ValueError(f"Qwen3.5 expects one packed row, got batch size {inputs_embeds.shape[0]}") if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( - seq_lens.to(device=inputs_embeds.device), + seq_lens.to(inputs_embeds.device), total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1], ) torch._dynamo.mark_dynamic(cu_seqlens, 0) - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) - cu_seqlens_are_pre_shard = seq_lens_are_pre_shard - - for decoder_layer in self.layers: + hidden_states = inputs_embeds + for layer_index, decoder_layer in enumerate(self.layers): + layer_routed_experts = routed_experts[:, :, layer_index] if routed_experts is not None else None hidden_states = decoder_layer( hidden_states, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard, + position_embeddings, + cu_seqlens, + max_seqlen, + routed_experts=layer_routed_experts, + cu_seqlens_are_pre_shard=seq_lens_are_pre_shard, ) - - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast(last_hidden_state=hidden_states) + return BaseModelOutput(last_hidden_state=self.norm(hidden_states)) class Qwen3_5VLMModel(nn.Module): - """Composite VLM body: HF vision encoder + custom PrimeRL dense text model.""" - - def __init__(self, config: PretrainedConfig): + def __init__(self, config) -> None: super().__init__() self.config = config - self.visual = Qwen3_5VisionModel._from_config(config.vision_config) + self.visual = Qwen3_5VisionModel(config.vision_config) self.language_model = Qwen3_5Model(config.text_config) - def get_input_embeddings(self): + def get_input_embeddings(self) -> nn.Embedding: return self.language_model.get_input_embeddings() - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) + def set_input_embeddings(self, embeddings: nn.Embedding) -> None: + self.language_model.set_input_embeddings(embeddings) - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self.language_model.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) + def set_context_parallel_attributes(self, process_group, rank: int, world_size: int) -> None: + self.language_model.set_context_parallel_attributes(process_group, rank, world_size) - def _dummy_vision_inputs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: - vcfg = self.config.vision_config - m = vcfg.spatial_merge_size - num_patches = m * m - patch_dim = vcfg.in_channels * vcfg.temporal_patch_size * vcfg.patch_size * vcfg.patch_size - pixel_values = torch.zeros(num_patches, patch_dim, device=device, dtype=self.visual.dtype) - grid_thw = torch.tensor([[1, m, m]], dtype=torch.long, device=device) - return pixel_values, grid_thw - - def prepare_inputs_embeds_and_position_ids( + def prepare_inputs( self, input_ids: torch.LongTensor, - position_ids: torch.LongTensor | None = None, - pixel_values: torch.Tensor | None = None, - image_grid_thw: torch.LongTensor | None = None, - mm_token_type_ids: torch.LongTensor | None = None, - *, + position_ids: torch.LongTensor | None, + pixel_values: torch.Tensor | None, + image_grid_thw: torch.LongTensor | None, + mm_token_type_ids: torch.LongTensor | None, seq_lens: torch.LongTensor, - ) -> tuple[torch.FloatTensor, torch.LongTensor]: + ) -> tuple[torch.Tensor, torch.LongTensor]: inputs_embeds = self.language_model.embed_tokens(input_ids) - has_images = pixel_values is not None - vision_grid_thw = image_grid_thw + vision_grid = image_grid_thw if has_images: - pixel_values = pixel_values.type(self.visual.dtype) + pixel_values = pixel_values.to(self.visual.dtype) else: - pixel_values, vision_grid_thw = self._dummy_vision_inputs(inputs_embeds.device) - - vision_output = self.visual(pixel_values, grid_thw=vision_grid_thw, return_dict=True) - image_embeds = vision_output.pooler_output.to(inputs_embeds.device, inputs_embeds.dtype) + merge_size = self.config.vision_config.spatial_merge_size + num_patches = merge_size**2 + patch_dim = ( + self.config.vision_config.in_channels + * self.config.vision_config.temporal_patch_size + * self.config.vision_config.patch_size**2 + ) + pixel_values = torch.zeros( + num_patches, + patch_dim, + device=inputs_embeds.device, + dtype=self.visual.dtype, + ) + vision_grid = torch.tensor([[1, merge_size, merge_size]], device=inputs_embeds.device) + image_embeds = self.visual(pixel_values, vision_grid).pooler_output.to(inputs_embeds.dtype) if has_images: - image_mask = input_ids == self.config.image_token_id - image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + image_mask = (input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) else: + # Every rank must retain the vision graph so FSDP collectives stay symmetric. inputs_embeds = inputs_embeds + image_embeds.sum() * 0.0 if position_ids is None: - if image_grid_thw is not None: + if image_grid_thw is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) + else: + if mm_token_type_ids is None: + raise ValueError("mm_token_type_ids are required with Qwen3.5 image inputs") position_ids = build_qwen3_5_mrope_position_ids( input_ids=input_ids, mm_token_type_ids=mm_token_type_ids, @@ -319,9 +287,6 @@ def prepare_inputs_embeds_and_position_ids( spatial_merge_size=self.config.vision_config.spatial_merge_size, seq_lens=seq_lens, ) - else: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - return inputs_embeds, position_ids def forward( @@ -331,105 +296,107 @@ def forward( pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, mm_token_type_ids: torch.LongTensor | None = None, + routed_experts: torch.LongTensor | None = None, *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, - ) -> BaseModelOutputWithPast: - inputs_embeds, position_ids = self.prepare_inputs_embeds_and_position_ids( - input_ids=input_ids, - position_ids=position_ids, - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - mm_token_type_ids=mm_token_type_ids, - seq_lens=seq_lens, + ) -> BaseModelOutput: + inputs_embeds, position_ids = self.prepare_inputs( + input_ids, + position_ids, + pixel_values, + image_grid_thw, + mm_token_type_ids, + seq_lens, ) - - cp_group = getattr(self.language_model, "_cp_group", None) - if image_grid_thw is not None and cp_group is not None: - cp_rank = self.language_model._cp_rank - cp_world_size = self.language_model._cp_world_size - setup_cp_attention_params(position_ids, cp_group=cp_group, cp_style="ulysses", seq_lens=seq_lens) - inputs_embeds = shard_for_cp(inputs_embeds, cp_rank=cp_rank, cp_world_size=cp_world_size) - position_ids = shard_position_ids_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size) + process_group = getattr(self.language_model, "context_parallel_group", None) + if image_grid_thw is not None and process_group is not None: + rank = self.language_model.context_parallel_rank + world_size = self.language_model.context_parallel_world_size + setup_cp_attention_params(position_ids, cp_group=process_group, cp_style="ulysses", seq_lens=seq_lens) + inputs_embeds = shard_for_cp(inputs_embeds, cp_rank=rank, cp_world_size=world_size) + position_ids = shard_position_ids_for_cp(position_ids, cp_rank=rank, cp_world_size=world_size) + if routed_experts is not None: + routed_experts = shard_for_cp(routed_experts, cp_rank=rank, cp_world_size=world_size) seq_lens_are_pre_shard = True return self.language_model( inputs_embeds=inputs_embeds, position_ids=position_ids, + routed_experts=routed_experts, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, ) -class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): - """Unified dense Qwen3.5 model for both text-only and VLM configs.""" - +class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _checkpoint_conversion_mapping = {} _tp_plan = {"lm_head": "colwise_gather_output"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - def __init__(self, config, **kwargs): - super().__init__(config, **kwargs) - self._is_vlm = hasattr(config, "vision_config") - self.supports_packed_multimodal_training = self._is_vlm + def __init__(self, config) -> None: + super().__init__(config) + self.is_vlm = hasattr(config, "vision_config") + text_config = config.text_config if self.is_vlm else config + attention_implementation = ( + getattr(config, "_attn_implementation", None) + or getattr(text_config, "_attn_implementation", None) + or "flash_attention_3" + ) + text_config._attn_implementation = attention_implementation - if self._is_vlm: + if self.is_vlm: + if getattr(config.vision_config, "_attn_implementation_internal", None) is None: + config.vision_config._attn_implementation = attention_implementation self.model = Qwen3_5VLMModel(config) - text_config = config.text_config self._tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} else: self.model = Qwen3_5Model(config) - text_config = config + self.supports_packed_multimodal_training = self.is_vlm self.vocab_size = text_config.vocab_size - self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) + self.lm_head = VanillaOutputLinear(text_config.hidden_size, text_config.vocab_size) + if isinstance(text_config, Qwen3_5MoeTextConfig): + self.num_experts = text_config.num_experts + self.num_experts_per_tok = text_config.num_experts_per_tok self.post_init() - def get_input_embeddings(self): + def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() - def set_input_embeddings(self, value): - self.model.set_input_embeddings(value) - - def set_decoder(self, decoder): - self.model = decoder + def set_input_embeddings(self, embeddings: nn.Embedding) -> None: + self.model.set_input_embeddings(embeddings) - def get_decoder(self): - return self.model - - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self.model.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) + def set_context_parallel_attributes(self, process_group, rank: int, world_size: int) -> None: + self.model.set_context_parallel_attributes(process_group, rank, world_size) def forward( self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - logits_to_keep: Union[int, torch.Tensor] = 0, - temperature: Union[torch.Tensor, None] = None, - pixel_values: Optional[torch.Tensor] = None, - image_grid_thw: Optional[torch.LongTensor] = None, - mm_token_type_ids: Optional[torch.LongTensor] = None, + input_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + temperature: torch.Tensor | None = None, + routed_experts: torch.LongTensor | None = None, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + mm_token_type_ids: torch.LongTensor | None = None, *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, - **kwargs: Unpack[TransformersKwargs], + **kwargs, ) -> PrimeLmOutput: - assert use_cache is None, "use_cache is not supported for custom qwen3_5 for now" - assert past_key_values is None, "past_key_values is not supported for custom qwen3_5 for now" - - if self._is_vlm: + if kwargs.get("use_cache") is not None or kwargs.get("past_key_values") is not None: + raise ValueError("Qwen3.5 custom training does not support KV caching") + if self.is_vlm: outputs = self.model( input_ids=input_ids, position_ids=position_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, + routed_experts=routed_experts, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, ) @@ -438,42 +405,38 @@ def forward( input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, + routed_experts=routed_experts, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, ) - hidden_states = outputs.last_hidden_state - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + if isinstance(logits_to_keep, int): + slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None) + else: + slice_indices = logits_to_keep return self.lm_head( - hidden_states[:, slice_indices, :], + outputs.last_hidden_state[:, slice_indices], labels[:, slice_indices] if labels is not None else None, - temperature=temperature, + temperature=temperature[:, slice_indices] if temperature is not None else None, ) - def init_buffers_post_meta(self): - if self._is_vlm: - lm_rope = self.model.language_model.rotary_emb - else: - lm_rope = self.model.rotary_emb - - if hasattr(lm_rope, "rope_init_fn"): - inv_freq, lm_rope.attention_scaling = lm_rope.rope_init_fn(lm_rope.config, lm_rope.inv_freq.device) - lm_rope.inv_freq.copy_(inv_freq) - - if self._is_vlm: - vis_rope = self.model.visual.rotary_pos_emb - if hasattr(vis_rope, "inv_freq"): - dim = vis_rope.inv_freq.shape[0] - inv_freq = 1.0 / ( - 10000.0 - ** (torch.arange(0, dim * 2, 2, dtype=torch.float32, device=vis_rope.inv_freq.device) / (dim * 2)) - ) - vis_rope.inv_freq.copy_(inv_freq) + 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", + "Qwen3_5DecoderLayer", "Qwen3_5ForCausalLM", - "Qwen3_5GatedFlashAttention", "Qwen3_5Model", "Qwen3_5PreTrainedModel", ] diff --git a/src/prime_rl/trainer/models/qwen3_5/norm.py b/src/prime_rl/trainer/models/qwen3_5/norm.py new file mode 100644 index 00000000000..721ab2f605a --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/norm.py @@ -0,0 +1,25 @@ +import torch +from torch import nn + + +class Qwen3_5RMSNorm(nn.Module): + """Qwen's zero-centered RMSNorm, whose effective scale is ``1 + weight``.""" + + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(hidden_size)) + self.variance_epsilon = eps + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.zeros_(self.weight) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + variance = hidden_states.square().mean(dim=-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return ((1.0 + self.weight.float()) * hidden_states).to(input_dtype) + + +__all__ = ["Qwen3_5RMSNorm"] diff --git a/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py b/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py new file mode 100644 index 00000000000..769eaf4cad6 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py @@ -0,0 +1,185 @@ +import itertools + +import torch +from torch import nn +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update + +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + + +class Qwen3_5RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor + + def __init__(self, config: Qwen3_5TextConfig, device=None) -> None: + super().__init__() + self.config = config + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + self.rope_type = config.rope_parameters["rope_type"] + self.rope_init_fn = self.compute_default_rope_parameters + if self.rope_type != "default": + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + self.mrope_section = config.rope_parameters.get("mrope_section") + if self.mrope_section is None: + self.mrope_section = self.scaled_mrope_section(inv_freq.numel()) + if not config.rope_parameters["mrope_interleaved"]: + raise ValueError("Qwen3.5 requires interleaved MRoPE") + if sum(self.mrope_section) != inv_freq.numel(): + raise ValueError(f"mrope_section must contain {inv_freq.numel()} rotary pairs, got {self.mrope_section}") + + @staticmethod + def scaled_mrope_section(num_rotary_pairs: int) -> list[int]: + default_section = [11, 11, 10] + total = sum(default_section) + section = [num_rotary_pairs * size // total for size in default_section] + remainder_order = sorted( + range(len(section)), + key=lambda index: num_rotary_pairs * default_section[index] % total, + reverse=True, + ) + for index in remainder_order[: num_rotary_pairs - sum(section)]: + section[index] += 1 + return section + + @staticmethod + def compute_default_rope_parameters( + config: Qwen3_5TextConfig, + device: torch.device | None = None, + seq_len: int | None = None, + ) -> tuple[torch.Tensor, float]: + del seq_len + rope_parameters = config.rope_parameters + rotary_dim = int(config.head_dim * rope_parameters.get("partial_rotary_factor", 1.0)) + positions = torch.arange(0, rotary_dim, 2, dtype=torch.int64, device=device).float() + return 1.0 / (rope_parameters["rope_theta"] ** (positions / rotary_dim)), 1.0 + + def reset_parameters(self) -> None: + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, self.inv_freq.device) + self.inv_freq.copy_(inv_freq) + self.original_inv_freq.copy_(inv_freq) + + @torch.no_grad() + @dynamic_rope_update + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.ndim == 2: + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + elif position_ids.ndim != 3: + raise ValueError(f"Qwen3.5 position_ids must be 2D or 3D, got {tuple(position_ids.shape)}") + + position_ids = position_ids.to(hidden_states.device) + inv_freq = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + positions = position_ids[:, :, None, :].float() + device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + frequencies = (inv_freq @ positions).transpose(2, 3) + interleaved = frequencies[0].clone() + for dimension, offset in enumerate((1, 2), start=1): + index = slice(offset, self.mrope_section[dimension] * 3, 3) + interleaved[..., index] = frequencies[dimension, ..., index] + embeddings = torch.cat((interleaved, interleaved), dim=-1) + cos = embeddings.cos() * self.attention_scaling + sin = embeddings.sin() * self.attention_scaling + return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype) + + +def get_qwen3_5_vision_position_ids( + *, + start_position: int, + grid_thw: torch.Tensor, + spatial_merge_size: int, + temporal_merge_size: int = 1, + time_interval: int = 1, + device: torch.device | None = None, +) -> torch.LongTensor: + grid_t = int(grid_thw[0].item()) // temporal_merge_size + grid_h = int(grid_thw[1].item()) // spatial_merge_size + grid_w = int(grid_thw[2].item()) // spatial_merge_size + + temporal = torch.arange(grid_t, device=device) * time_interval + start_position + height = torch.arange(grid_h, device=device) + start_position + width = torch.arange(grid_w, device=device) + start_position + return torch.stack( + [ + temporal.repeat_interleave(grid_h * grid_w), + height.repeat_interleave(grid_w).repeat(grid_t), + width.repeat(grid_h * grid_t), + ], + dim=0, + ) + + +def build_qwen3_5_mrope_position_ids( + *, + input_ids: torch.LongTensor, + mm_token_type_ids: torch.LongTensor, + image_grid_thw: torch.LongTensor | None, + spatial_merge_size: int, + seq_lens: torch.Tensor, +) -> torch.LongTensor: + if input_ids.ndim != 2 or input_ids.shape[0] != 1: + raise ValueError(f"Packed Qwen3.5 input_ids must have shape (1, tokens), got {tuple(input_ids.shape)}") + if mm_token_type_ids.shape != input_ids.shape: + raise ValueError("mm_token_type_ids must have the same shape as input_ids") + + seq_lens = seq_lens.to(device=input_ids.device, dtype=torch.long) + if seq_lens.ndim != 1 or bool((seq_lens <= 0).any()) or int(seq_lens.sum()) != input_ids.shape[1]: + raise ValueError("seq_lens must contain positive lengths summing to the packed sequence length") + + image_grids = iter(image_grid_thw) if image_grid_thw is not None else None + position_ids = torch.empty(3, 1, input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device) + offset = 0 + for sequence_length_tensor in seq_lens: + sequence_length = int(sequence_length_tensor) + token_types = mm_token_type_ids[0, offset : offset + sequence_length] + current_position = 0 + sequence_positions = [] + + for modality, indexed_group in itertools.groupby(enumerate(token_types.tolist()), lambda item: item[1]): + group = list(indexed_group) + group_length = group[-1][0] - group[0][0] + 1 + if modality == 0: + text_positions = torch.arange(group_length, device=input_ids.device) + current_position + sequence_positions.append(text_positions.unsqueeze(0).expand(3, -1)) + current_position += group_length + continue + if modality == 2: + raise ValueError("Qwen3.5 video MRoPE is not supported") + if modality != 1: + raise ValueError(f"Unsupported Qwen3.5 multimodal token type: {modality}") + if image_grids is None: + raise ValueError("image_grid_thw is required for image tokens") + + remaining = group_length + while remaining: + grid = next(image_grids, None) + if grid is None: + raise ValueError("Not enough image_grid_thw rows for the image tokens") + image_positions = get_qwen3_5_vision_position_ids( + start_position=current_position, + grid_thw=grid, + spatial_merge_size=spatial_merge_size, + device=input_ids.device, + ) + if image_positions.shape[1] > remaining: + raise ValueError("Image token group length does not match image_grid_thw") + sequence_positions.append(image_positions) + current_position += max(int(grid[1]), int(grid[2])) // spatial_merge_size + remaining -= image_positions.shape[1] + + positions = torch.cat(sequence_positions, dim=1) + if positions.shape[1] != sequence_length: + raise ValueError("Built MRoPE positions do not match the packed sequence length") + position_ids[:, 0, offset : offset + sequence_length] = positions + offset += sequence_length + + if image_grids is not None and next(image_grids, None) is not None: + raise ValueError("image_grid_thw contains unused rows") + return position_ids + + +__all__ = ["Qwen3_5RotaryEmbedding", "build_qwen3_5_mrope_position_ids", "get_qwen3_5_vision_position_ids"] diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py new file mode 100644 index 00000000000..2b4f3a43bd3 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -0,0 +1,286 @@ +import math + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.modeling_outputs import BaseModelOutputWithPooling + +from prime_rl.trainer.models.layers.attn import ( + flash_attn_3_varlen_func, + flash_attn_4_varlen_func, + flash_attn_varlen_func, +) +from prime_rl.trainer.models.layers.rotary_emb import rotate_half +from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5VisionConfig + + +class Qwen3_5VisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10_000.0) -> None: + super().__init__() + self.dim = dim + self.theta = theta + self.register_buffer("inv_freq", torch.empty(dim // 2), persistent=False) + self.reset_parameters() + + def reset_parameters(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))) + + def forward(self, sequence_length: int) -> torch.Tensor: + positions = torch.arange(sequence_length, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + return torch.outer(positions, self.inv_freq) + + +class Qwen3_5VisionPatchEmbed(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + self.in_channels = config.in_channels + self.temporal_patch_size = config.temporal_patch_size + self.patch_size = config.patch_size + kernel_size = (self.temporal_patch_size, self.patch_size, self.patch_size) + self.proj = nn.Conv3d( + self.in_channels, + config.hidden_size, + kernel_size=kernel_size, + stride=kernel_size, + bias=True, + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + pixel_values = pixel_values.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + return self.proj(pixel_values.to(self.proj.weight.dtype)).view(pixel_values.shape[0], -1) + + +class Qwen3_5VisionMLP(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + self.linear_fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states), approximate="tanh")) + + +class Qwen3_5VisionAttention(nn.Module): + FLASH_ATTENTION_FUNCTIONS = { + "flash_attention_2": flash_attn_varlen_func, + "flash_attention_3": flash_attn_3_varlen_func, + "flash_attention_4": flash_attn_4_varlen_func, + } + + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + self.num_heads = config.num_heads + self.head_dim = config.hidden_size // config.num_heads + self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=True) + self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + attention_implementation = getattr(config, "_attn_implementation", None) or "flash_attention_3" + self.attention_implementation = attention_implementation + self.flash_attention = ( + None if attention_implementation == "sdpa" else self.FLASH_ATTENTION_FUNCTIONS[attention_implementation] + ) + if attention_implementation == "flash_attention_4": + self.flash_attention = torch.compiler.disable(self.flash_attention) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.IntTensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + sequence_length = hidden_states.shape[0] + query, key, value = ( + self.qkv(hidden_states).view(sequence_length, 3, self.num_heads, self.head_dim).unbind(dim=1) + ) + cos, sin = position_embeddings + cos = cos.unsqueeze(1).float() + sin = sin.unsqueeze(1).float() + query_dtype = query.dtype + query = ((query.float() * cos) + (rotate_half(query.float()) * sin)).to(query_dtype) + key = ((key.float() * cos) + (rotate_half(key.float()) * sin)).to(query_dtype) + + if self.attention_implementation == "sdpa": + lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + outputs = [] + for query_chunk, key_chunk, value_chunk in zip( + query.split(lengths), + key.split(lengths), + value.split(lengths), + ): + output = F.scaled_dot_product_attention( + query_chunk.transpose(0, 1).unsqueeze(0), + key_chunk.transpose(0, 1).unsqueeze(0), + value_chunk.transpose(0, 1).unsqueeze(0), + is_causal=False, + ) + outputs.append(output.squeeze(0).transpose(0, 1)) + attention_output = torch.cat(outputs, dim=0) + elif self.attention_implementation == "flash_attention_4": + attention_output = self.flash_attention( + query, + key, + value, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + causal=False, + ) + else: + max_sequence_length = int((cu_seqlens[1:] - cu_seqlens[:-1]).max()) + attention_output = self.flash_attention( + query, + key, + value, + cu_seqlens, + cu_seqlens, + max_sequence_length, + max_sequence_length, + causal=False, + ) + if isinstance(attention_output, tuple): + attention_output = attention_output[0] + return self.proj(attention_output.reshape(sequence_length, -1)) + + +class Qwen3_5VisionBlock(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3_5VisionAttention(config) + self.mlp = Qwen3_5VisionMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.IntTensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), cu_seqlens, position_embeddings) + return hidden_states + self.mlp(self.norm2(hidden_states)) + + +class Qwen3_5VisionPatchMerger(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + merged_hidden_size = config.hidden_size * config.spatial_merge_size**2 + self.merged_hidden_size = merged_hidden_size + self.norm = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(merged_hidden_size, merged_hidden_size) + self.linear_fc2 = nn.Linear(merged_hidden_size, config.out_hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(hidden_states).view(-1, self.merged_hidden_size) + return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states))) + + +class Qwen3_5VisionModel(nn.Module): + def __init__(self, config: Qwen3_5VisionConfig) -> None: + super().__init__() + self.config = config + self.spatial_merge_size = config.spatial_merge_size + self.patch_embed = Qwen3_5VisionPatchEmbed(config) + self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size) + self.num_grid_per_side = math.isqrt(config.num_position_embeddings) + self.rotary_pos_emb = Qwen3_5VisionRotaryEmbedding(config.hidden_size // config.num_heads // 2) + self.blocks = nn.ModuleList(Qwen3_5VisionBlock(config) for _ in range(config.depth)) + self.merger = Qwen3_5VisionPatchMerger(config) + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + def rotary_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor: + grids = grid_thw.tolist() + frequency_table = self.rotary_pos_emb(max(max(height, width) for _, height, width in grids)) + coordinates = [] + merge_size = self.spatial_merge_size + for frames, height, width in grids: + block_rows = torch.arange(height // merge_size, device=frequency_table.device) + block_columns = torch.arange(width // merge_size, device=frequency_table.device) + intra_rows = torch.arange(merge_size, device=frequency_table.device) + intra_columns = torch.arange(merge_size, device=frequency_table.device) + rows = block_rows[:, None, None, None] * merge_size + intra_rows[None, None, :, None] + columns = block_columns[None, :, None, None] * merge_size + intra_columns[None, None, None, :] + rows = rows.expand(height // merge_size, width // merge_size, merge_size, merge_size).reshape(-1) + columns = columns.expand(height // merge_size, width // merge_size, merge_size, merge_size).reshape(-1) + coordinates.append(torch.stack((rows, columns), dim=-1).repeat(frames, 1)) + return frequency_table[torch.cat(coordinates)].flatten(1) + + def interpolated_position_embeddings(self, grid_thw: torch.Tensor) -> torch.Tensor: + embeddings = [] + merge_size = self.spatial_merge_size + for frames, height, width in grid_thw.tolist(): + row_positions = torch.linspace( + 0, + self.num_grid_per_side - 1, + height, + device=self.pos_embed.weight.device, + ) + column_positions = torch.linspace( + 0, + self.num_grid_per_side - 1, + width, + device=self.pos_embed.weight.device, + ) + row_floor = row_positions.floor().long() + column_floor = column_positions.floor().long() + row_ceil = (row_floor + 1).clamp_max(self.num_grid_per_side - 1) + column_ceil = (column_floor + 1).clamp_max(self.num_grid_per_side - 1) + row_fraction = row_positions - row_floor + column_fraction = column_positions - column_floor + + indices = torch.stack( + [ + (row_floor[:, None] * self.num_grid_per_side + column_floor[None, :]).flatten(), + (row_floor[:, None] * self.num_grid_per_side + column_ceil[None, :]).flatten(), + (row_ceil[:, None] * self.num_grid_per_side + column_floor[None, :]).flatten(), + (row_ceil[:, None] * self.num_grid_per_side + column_ceil[None, :]).flatten(), + ] + ) + weights = torch.stack( + [ + ((1 - row_fraction)[:, None] * (1 - column_fraction)[None, :]).flatten(), + ((1 - row_fraction)[:, None] * column_fraction[None, :]).flatten(), + (row_fraction[:, None] * (1 - column_fraction)[None, :]).flatten(), + (row_fraction[:, None] * column_fraction[None, :]).flatten(), + ] + ).to(self.pos_embed.weight.dtype) + position_embedding = (self.pos_embed(indices) * weights[..., None]).sum(dim=0).repeat(frames, 1) + position_embedding = ( + position_embedding.view( + frames, + height // merge_size, + merge_size, + width // merge_size, + merge_size, + -1, + ) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + embeddings.append(position_embedding) + return torch.cat(embeddings) + + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> BaseModelOutputWithPooling: + hidden_states = self.patch_embed(pixel_values) + hidden_states = hidden_states + self.interpolated_position_embeddings(grid_thw) + rotary = self.rotary_embeddings(grid_thw) + rotary = torch.cat((rotary, rotary), dim=-1) + position_embeddings = (rotary.cos(), rotary.sin()) + sequence_lengths = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + cu_seqlens = F.pad(sequence_lengths.cumsum(dim=0, dtype=torch.int32), (1, 0)) + for block in self.blocks: + hidden_states = block(hidden_states, cu_seqlens, position_embeddings) + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=self.merger(hidden_states), + ) + + +__all__ = ["Qwen3_5VisionModel"] diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py b/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py deleted file mode 100644 index 17573101dd3..00000000000 --- a/src/prime_rl/trainer/models/qwen3_5_moe/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from prime_rl.trainer.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig -from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeForCausalLM, - Qwen3_5MoeModel, - Qwen3_5MoePreTrainedModel, -) - -__all__ = [ - "Qwen3_5MoeConfig", - "Qwen3_5MoeForCausalLM", - "Qwen3_5MoeModel", - "Qwen3_5MoePreTrainedModel", -] diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py deleted file mode 100644 index 09aff4568ee..00000000000 --- a/src/prime_rl/trainer/models/qwen3_5_moe/configuration_qwen3_5_moe.py +++ /dev/null @@ -1,106 +0,0 @@ -from transformers.configuration_utils import PretrainedConfig - - -class Qwen3_5MoeConfig(PretrainedConfig): - r""" - Configuration class for the custom PrimeRL Qwen3.5-MoE model. - - This matches the HuggingFace `Qwen3_5MoeTextConfig` architecture: - hybrid attention (GatedDeltaNet linear + gated softmax), MoE with gated shared expert, - and (1+weight) RMSNorm parameterization. - - Defaults match Qwen3.5-35B-A3B. - """ - - model_type = "qwen3_5_moe_text" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=248320, - hidden_size=2048, - num_hidden_layers=40, - num_attention_heads=16, - num_key_value_heads=2, - hidden_act="silu", - max_position_embeddings=32768, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - attention_dropout=0.0, - head_dim=256, - # Linear attention (GatedDeltaNet) params - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - # MoE params - moe_intermediate_size=512, - shared_expert_intermediate_size=512, - num_experts_per_tok=8, - num_experts=256, - output_router_logits=False, - router_aux_loss_coef=0.001, - # Layer type configuration - layer_types=None, - # PrimeRL additions - load_balance_coeff=None, - pad_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.hidden_act = hidden_act - self.initializer_range = initializer_range - self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.head_dim = head_dim - self.pad_token_id = pad_token_id - - # Partial RoPE: only rotate 25% of head_dim - kwargs.setdefault("partial_rotary_factor", 0.25) - - # Layer types: default every 4th layer is full_attention, rest linear_attention - self.layer_types = layer_types - if self.layer_types is None: - interval_pattern = kwargs.pop("full_attention_interval", 4) - self.layer_types = [ - "linear_attention" if bool((i + 1) % interval_pattern) else "full_attention" - for i in range(self.num_hidden_layers) - ] - - # Linear attention (GatedDeltaNet) params - self.linear_conv_kernel_dim = linear_conv_kernel_dim - self.linear_key_head_dim = linear_key_head_dim - self.linear_value_head_dim = linear_value_head_dim - self.linear_num_key_heads = linear_num_key_heads - self.linear_num_value_heads = linear_num_value_heads - - # MoE params - self.moe_intermediate_size = moe_intermediate_size - self.shared_expert_intermediate_size = shared_expert_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.num_experts = num_experts - self.output_router_logits = output_router_logits - self.router_aux_loss_coef = router_aux_loss_coef - - # PrimeRL additions - self.load_balance_coeff = load_balance_coeff - - super().__init__( - tie_word_embeddings=tie_word_embeddings, - **kwargs, - ) diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py deleted file mode 100644 index ed1f2aba749..00000000000 --- a/src/prime_rl/trainer/models/qwen3_5_moe/converting_qwen3_5_moe.py +++ /dev/null @@ -1,25 +0,0 @@ -"""HF<->prime weight conversion for Qwen3.5-MoE, as a declarative op chain. - -Per layer: the router and shared-expert output gate are nested under their -PrimeRL modules, and routed experts use stacked canonical projections. -""" - -from __future__ import annotations - -from prime_rl.trainer.models.conversion_ops import ConvOp, Rename, routed_experts_op - - -def _conversion_chain(config, model_prefix: str) -> list[ConvOp]: - ops: list[ConvOp] = [] - for i in range(config.num_hidden_layers): - p = f"{model_prefix}.layers.{i}" - # Router: mlp.gate.weight -> mlp.router.gate.weight - ops.append(Rename(f"{p}.mlp.gate.weight", f"{p}.mlp.router.gate.weight")) - ops.append(Rename(f"{p}.mlp.shared_expert_gate.weight", f"{p}.mlp.shared_expert.output_gate.weight")) - ops.append(routed_experts_op(p, hf_experts="mlp.experts", prime_experts="mlp.experts", fused=True)) - return ops - - -def conversion_chain(config) -> list[ConvOp]: - text_config = getattr(config, "text_config", config) - return _conversion_chain(text_config, "model") + _conversion_chain(text_config, "model.language_model") diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py deleted file mode 100644 index 81b2faab856..00000000000 --- a/src/prime_rl/trainer/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ /dev/null @@ -1,988 +0,0 @@ -import functools -from dataclasses import dataclass -from typing import Optional, Union - -import torch -import torch.nn.functional as F -from fla.modules import FusedRMSNormGated -from fla.modules.conv import causal_conv1d as fla_causal_conv1d -from fla.ops.cp import FLACPContext, build_cp_context -from fla.ops.gated_delta_rule import chunk_gated_delta_rule -from torch import Tensor, nn -from transformers.cache_utils import Cache -from transformers.configuration_utils import PretrainedConfig -from transformers.generation import GenerationMixin -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_outputs import MoeModelOutputWithPast -from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeVisionModel -from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs, logging - -from prime_rl.trainer.models.base import PreTrainedModelPrimeRL -from prime_rl.trainer.models.layers.attn import ( - flash_attn_3_varlen_func, - flash_attn_4_varlen_func, - flash_attn_varlen_func, -) -from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput -from prime_rl.trainer.models.layers.mlp import FeedForward -from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb -from prime_rl.utils.cp import setup_cp_attention_params, shard_for_cp, shard_position_ids_for_cp -from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens - -from .configuration_qwen3_5_moe import Qwen3_5MoeConfig -from .converting_qwen3_5_moe import conversion_chain -from .mrope import build_qwen3_5_mrope_position_ids - -logger = logging.get_logger(__name__) - - -@torch.compiler.disable -def _fla_causal_conv1d_cp( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - activation: str, - cp_context: FLACPContext, -) -> torch.Tensor: - output, _ = fla_causal_conv1d( - x=x, - weight=weight, - bias=bias, - activation=activation, - cp_context=cp_context, - ) - return output - - -# --------------------------------------------------------------------------- -# RMSNorm variants -# --------------------------------------------------------------------------- - - -class Qwen3_5MoeRMSNorm(nn.Module): - """RMSNorm with (1+weight) parameterization. Weight initialized to zeros.""" - - def __init__(self, dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.zeros(dim)) - - def _norm(self, x): - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) - - def forward(self, x): - output = self._norm(x.float()) - output = output * (1.0 + self.weight.float()) - return output.type_as(x) - - -# --------------------------------------------------------------------------- -# GatedDeltaNet linear attention -# --------------------------------------------------------------------------- - - -class Qwen3_5MoeGatedDeltaNet(nn.Module): - """GatedDeltaNet linear attention with Conv1d, beta/gamma gates, and chunk delta rule.""" - - def __init__(self, config: Qwen3_5MoeConfig): - super().__init__() - self.hidden_size = config.hidden_size - self.num_v_heads = config.linear_num_value_heads - self.num_k_heads = config.linear_num_key_heads - self.head_k_dim = config.linear_key_head_dim - self.head_v_dim = config.linear_value_head_dim - self.key_dim = self.head_k_dim * self.num_k_heads - self.value_dim = self.head_v_dim * self.num_v_heads - - self.conv_kernel_size = config.linear_conv_kernel_dim - self.activation = config.hidden_act - self.layer_norm_epsilon = config.rms_norm_eps - - # QKV convolution (depthwise) - self.conv_dim = self.key_dim * 2 + self.value_dim - self.conv1d = nn.Conv1d( - in_channels=self.conv_dim, - out_channels=self.conv_dim, - bias=False, - kernel_size=self.conv_kernel_size, - groups=self.conv_dim, - padding=self.conv_kernel_size - 1, - ) - - # Time step projection - self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) - A = torch.empty(self.num_v_heads).uniform_(0, 16) - self.A_log = nn.Parameter(torch.log(A)) - - self.norm = FusedRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) - - self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) - - # Input projections - self.in_proj_qkv = nn.Linear(self.hidden_size, self.key_dim * 2 + self.value_dim, bias=False) - self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False) - self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) - self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) - - def _build_cp_context( - self, - device: torch.device, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_are_pre_shard: bool = False, - ) -> "FLACPContext | None": - """Build the FLA CP context from full pre-shard sequence boundaries.""" - cp_group = getattr(self, "cp_group", None) - if cp_group is None: - return None - if cu_seqlens is None or not cu_seqlens_are_pre_shard: - raise ValueError("Qwen3.5 context parallelism requires full pre-shard sequence boundaries") - global_cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32) - return build_cp_context( - cu_seqlens=global_cu_seqlens, - group=cp_group, - conv1d_kernel_size=self.conv_kernel_size, - ) - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_are_pre_shard: bool = False, - ) -> torch.Tensor: - 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 = self._build_cp_context(hidden_states.device, cu_seqlens, cu_seqlens_are_pre_shard) - - # Causal conv1d — must reset at sequence boundaries for packed batches, - # otherwise the kernel-1 left pad leaks state across sequences. - conv_input = mixed_qkv.transpose(1, 2) - conv_weight = self.conv1d.weight.squeeze(1) - if cp_context is None: - mixed_qkv, _ = fla_causal_conv1d( - x=conv_input, - weight=conv_weight, - bias=self.conv1d.bias, - activation=self.activation, - cu_seqlens=cu_seqlens, - ) - else: - mixed_qkv = _fla_causal_conv1d_cp( - x=conv_input, - weight=conv_weight, - bias=self.conv1d.bias, - activation=self.activation, - cp_context=cp_context, - ) - 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: - cu_seqlens = cp_context.cu_seqlens - core_attn_out, _ = chunk_gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - cp_context=cp_context, - ) - else: - core_attn_out, _ = 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, - ) - - # Gated RMSNorm - 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) - - -# --------------------------------------------------------------------------- -# Gated softmax attention (for full_attention layers) -# --------------------------------------------------------------------------- - - -@dataclass -class Qwen3_5MoeGatedAttentionConfig: - hidden_size: int - head_dim: int - num_attention_heads: int - num_key_value_heads: int - rms_norm_eps: float - attention_bias: bool = False - attention_dropout: float = 0.0 - - -def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - if n_rep == 1: - return hidden_states - batch, num_kv_heads, slen, head_dim = hidden_states.shape - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim) - - -class Qwen3_5MoeGatedAttentionBase(nn.Module): - """Base class for gated softmax attention (Q projects 2x: query + gate).""" - - def __init__(self, config: Qwen3_5MoeGatedAttentionConfig): - super().__init__() - self.head_dim = config.head_dim - self.num_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.scaling = self.head_dim**-0.5 - self.attention_dropout = config.attention_dropout - - # Q projects 2x: query + gate - self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim * 2, bias=config.attention_bias) - self.k_proj = nn.Linear( - config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.v_proj = nn.Linear( - config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias - ) - self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias) - - # QK normalization with (1+weight) parameterization - self.q_norm = Qwen3_5MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = Qwen3_5MoeRMSNorm(self.head_dim, eps=config.rms_norm_eps) - - -class Qwen3_5MoeGatedFlashAttention(Qwen3_5MoeGatedAttentionBase): - """Gated softmax attention using Flash Attention varlen functions.""" - - _funcs = { - 2: flash_attn_varlen_func, - 3: flash_attn_3_varlen_func, - 4: flash_attn_4_varlen_func, - } - - def __init__(self, config: Qwen3_5MoeGatedAttentionConfig, flash_attn_version: int = 4): - super().__init__(config) - self._flash_attn_version = flash_attn_version - self.func = self._funcs[flash_attn_version] - self._flash_attn_call = self.func - if self._flash_attn_version == 4: - self._flash_attn_call = torch._dynamo.disable(self.func) - - def _compute_attention(self, q, k, v, cu_seqlens, max_seqlen): - """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" - kwargs: dict = {"causal": True} - sliding_window = getattr(self, "sliding_window", None) - if sliding_window is not None: - kwargs["window_size"] = (sliding_window - 1, 0) - if self._flash_attn_version == 4: - # FA4's flash_attn_varlen_func has qv as the 4th positional arg, - # so cu_seqlens must be passed as keyword args to avoid misalignment. - kwargs["cu_seqlens_q"] = cu_seqlens - kwargs["cu_seqlens_k"] = cu_seqlens - out = self._flash_attn_call(q, k, v, **kwargs) - else: - out = self._flash_attn_call(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, **kwargs) - if isinstance(out, tuple): - out = out[0] - return out - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states, gate = torch.chunk( - self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1 - ) - gate = gate.reshape(*input_shape, -1) - - query_states = self.q_norm(query_states.view(hidden_shape)) - key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)) - value_states = self.v_proj(hidden_states).view(hidden_shape) - - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - - attn_output = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) - attn_output = attn_output.contiguous().view(*input_shape, -1) - attn_output = attn_output * torch.sigmoid(gate) - return self.o_proj(attn_output), None - - -QWEN35MOE_ATTN_IMPL2CLASS = { - "flash_attention_2": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=2), - "flash_attention_3": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=3), - "flash_attention_4": functools.partial(Qwen3_5MoeGatedFlashAttention, flash_attn_version=4), -} - - -def normalize_qwen3_5_attn_implementation(attn_impl: str) -> str: - if attn_impl == "kernels-community/vllm-flash-attn3": - return "flash_attention_3" - return attn_impl - - -# --------------------------------------------------------------------------- -# Decoder layer -# --------------------------------------------------------------------------- - - -def _get_gated_attention(config: Qwen3_5MoeConfig) -> nn.Module: - attn_config = Qwen3_5MoeGatedAttentionConfig( - hidden_size=config.hidden_size, - head_dim=config.head_dim, - num_attention_heads=config.num_attention_heads, - num_key_value_heads=config.num_key_value_heads, - rms_norm_eps=config.rms_norm_eps, - attention_bias=config.attention_bias, - attention_dropout=config.attention_dropout, - ) - - attn_impl = normalize_qwen3_5_attn_implementation(config._attn_implementation) - config._attn_implementation = attn_impl - - if attn_impl not in QWEN35MOE_ATTN_IMPL2CLASS: - supported = list(QWEN35MOE_ATTN_IMPL2CLASS.keys()) - raise ValueError( - f"Qwen3.5-MoE attention does not support '{config._attn_implementation}'. " - f"Supported implementations: {supported}." - ) - - return QWEN35MOE_ATTN_IMPL2CLASS[attn_impl](attn_config) - - -class Qwen3_5MoeSharedExpert(FeedForward): - def __init__(self, config: Qwen3_5MoeConfig) -> None: - super().__init__( - dim=config.hidden_size, - hidden_dim=config.shared_expert_intermediate_size, - expert_type="gated", - activation=config.hidden_act, - ) - self.output_gate = nn.Linear(config.hidden_size, 1, bias=False) - - def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -> torch.Tensor: - return torch.sigmoid(self.output_gate(x)) * super().forward(x, routed_experts) - - def init_weights(self, init_std: float = 0.02) -> None: - super().init_weights(init_std) - nn.init.trunc_normal_(self.output_gate.weight, mean=0.0, std=init_std) - - -class Qwen3_5MoeDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: Qwen3_5MoeConfig, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - self.layer_type = config.layer_types[layer_idx] - - # Token mixer: either GatedDeltaNet or gated softmax attention - if self.layer_type == "linear_attention": - self.linear_attn = Qwen3_5MoeGatedDeltaNet(config) - elif self.layer_type == "full_attention": - self.self_attn = _get_gated_attention(config) - - moe_args = MoEArgs( - num_experts=config.num_experts, - expert_type="gated", - activation=config.hidden_act, - score_func="softmax", - route_norm=True, - route_scale=1.0, - score_before_experts=False, - top_k=config.num_experts_per_tok, - load_balance_coeff=config.load_balance_coeff, - ) - self.mlp = MoE.from_args( - moe_args, - dim=config.hidden_size, - hidden_dim=config.moe_intermediate_size, - shared_expert=Qwen3_5MoeSharedExpert(config), - ) - - # Layer norms with (1+weight) parameterization - self.input_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - routed_experts: Optional[torch.LongTensor] = None, - cu_seqlens_are_pre_shard: bool = False, - ) -> torch.FloatTensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - # Token mixer - if self.layer_type == "linear_attention": - hidden_states = self.linear_attn( - hidden_states, cu_seqlens=cu_seqlens, cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard - ) - elif self.layer_type == "full_attention": - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - ) - - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = residual + self.mlp(hidden_states, routed_experts=routed_experts) - return hidden_states - - -# --------------------------------------------------------------------------- -# Rotary embedding -# --------------------------------------------------------------------------- - - -class Qwen3_5MoeRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor - - def __init__(self, config: Qwen3_5MoeConfig, device=None): - super().__init__() - self.config = config - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - - rope_parameters = getattr(config, "rope_parameters", None) - if rope_parameters is None: - config.standardize_rope_params() - rope_parameters = config.rope_parameters - - self.rope_type = rope_parameters.get("rope_type", rope_parameters.get("type", "default")) - self.rope_init_fn = self.compute_default_rope_parameters - if self.rope_type != "default": - self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] - - inv_freq, self.attention_scaling = self.rope_init_fn(config, device) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.original_inv_freq = self.inv_freq.clone() - - self.mrope_section = rope_parameters.get("mrope_section") - if self.mrope_section is None: - self.mrope_section = self._scaled_default_mrope_section(self.inv_freq.numel()) - if not rope_parameters.get("mrope_interleaved", True): - raise ValueError("Qwen3.5 MoE custom model expects interleaved MRoPE") - if sum(self.mrope_section) != self.inv_freq.numel(): - raise ValueError( - "Qwen3.5 mrope_section must sum to rotary_dim // 2: " - f"mrope_section={self.mrope_section}, rotary_dim={self.inv_freq.numel() * 2}" - ) - - @staticmethod - def _scaled_default_mrope_section(num_rotary_pairs: int) -> list[int]: - default_section = [11, 11, 10] - default_total = sum(default_section) - scaled_section = [num_rotary_pairs * section // default_total for section in default_section] - remainder = num_rotary_pairs - sum(scaled_section) - remainder_order = sorted( - range(len(default_section)), - key=lambda idx: num_rotary_pairs * default_section[idx] % default_total, - reverse=True, - ) - for idx in remainder_order[:remainder]: - scaled_section[idx] += 1 - return scaled_section - - @staticmethod - def compute_default_rope_parameters( - config: Qwen3_5MoeConfig | None = None, - device: Optional[torch.device] = None, - seq_len: int | None = None, - ) -> tuple[torch.Tensor, float]: - rope_parameters = config.rope_parameters - base = rope_parameters["rope_theta"] - partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0) - head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads - dim = int(head_dim * partial_rotary_factor) - inv_freq = 1.0 / ( - base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) - ) - return inv_freq, 1.0 - - @torch.no_grad() - @dynamic_rope_update - def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - if position_ids.ndim == 2: - position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) - elif position_ids.ndim != 3: - raise ValueError(f"Qwen3.5 position_ids must be 2D or 3D, got shape={tuple(position_ids.shape)}") - - position_ids = position_ids.to(device=x.device) - inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) - position_ids_expanded = position_ids[:, :, None, :].float() - - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) - freqs = self._apply_interleaved_mrope(freqs) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: - freqs_t = freqs[0].clone() - for dim, offset in enumerate((1, 2), start=1): - length = self.mrope_section[dim] * 3 - idx = slice(offset, length, 3) - freqs_t[..., idx] = freqs[dim, ..., idx] - return freqs_t - - -def _create_rotary_emb(config: Qwen3_5MoeConfig) -> Qwen3_5MoeRotaryEmbedding: - return Qwen3_5MoeRotaryEmbedding(config) - - -# --------------------------------------------------------------------------- -# Model classes -# --------------------------------------------------------------------------- - - -class Qwen3_5MoePreTrainedModel(PreTrainedModelPrimeRL): - config_class = Qwen3_5MoeConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["Qwen3_5MoeDecoderLayer"] - _skip_keys_device_placement = ["past_key_values"] - _supports_flash_attn = True - _supports_sdpa = False - _supports_flex_attn = False - _supports_attention_backend = True - _can_compile_fullgraph = False - _can_record_outputs = { - "hidden_states": Qwen3_5MoeDecoderLayer, - } - - @classmethod - def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool: - return name.endswith(("linear_attn.A_log", "linear_attn.norm.weight")) - - def _check_and_adjust_attn_implementation( - self, attn_implementation: str | None, is_init_check: bool = False, allow_all_kernels: bool = False - ) -> str: - attn_impl = normalize_qwen3_5_attn_implementation(attn_implementation or "flash_attention_3") - if attn_impl not in QWEN35MOE_ATTN_IMPL2CLASS: - supported = list(QWEN35MOE_ATTN_IMPL2CLASS.keys()) - raise ValueError( - f"Qwen3.5-MoE attention does not support '{attn_implementation}'. Supported implementations: {supported}." - ) - return attn_impl - - @classmethod - def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict.keys()) - - @classmethod - def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.gate_proj" in name for name in state_dict.keys()) - - @classmethod - def conversion_chain(cls, config): - return conversion_chain(config) - - -class Qwen3_5MoeModel(Qwen3_5MoePreTrainedModel): - def __init__(self, config: Qwen3_5MoeConfig): - config._attn_implementation = normalize_qwen3_5_attn_implementation(config._attn_implementation) - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - [Qwen3_5MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - self.norm = Qwen3_5MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = _create_rotary_emb(config) - self.gradient_checkpointing = False - - self.post_init() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, value): - self.embed_tokens = value - - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self._cp_group = cp_group - self._cp_rank = cp_rank - self._cp_world_size = cp_world_size - for layer in self.layers.modules(): - if getattr(layer, "layer_type", None) == "linear_attention": - layer.linear_attn.cp_group = cp_group - layer.linear_attn.cp_rank = cp_rank - layer.linear_attn.cp_world_size = cp_world_size - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - routed_experts: Optional[torch.LongTensor] = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - ) -> MoeModelOutputWithPast: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if position_ids is None: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - - cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( - seq_lens.to(device=inputs_embeds.device), - total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1], - ) - torch._dynamo.mark_dynamic(cu_seqlens, 0) - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - cu_seqlens_are_pre_shard = seq_lens_are_pre_shard - - for layer_idx, decoder_layer in enumerate(self.layers): - routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None - hidden_states = decoder_layer( - hidden_states, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - routed_experts=routed_experts_layer, - cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard, - ) - - hidden_states = self.norm(hidden_states) - return MoeModelOutputWithPast(last_hidden_state=hidden_states) - - -# --------------------------------------------------------------------------- -# VLM composite model body -# --------------------------------------------------------------------------- - - -def _build_text_config(composite_config: PretrainedConfig) -> Qwen3_5MoeConfig: - """Build custom PrimeRL text config from HF's composite VLM config.""" - text_dict = composite_config.text_config.to_dict() - text_config = Qwen3_5MoeConfig(**text_dict) - attn_impl = getattr( - composite_config.text_config, - "_attn_implementation", - getattr(composite_config, "_attn_implementation", None), - ) - if attn_impl is not None: - text_config._attn_implementation = attn_impl - return text_config - - -class Qwen3_5MoeVLMModel(nn.Module): - """Composite VLM body: HF vision encoder + custom PrimeRL text model.""" - - def __init__(self, config: PretrainedConfig): - super().__init__() - self.config = config - self.visual = Qwen3_5MoeVisionModel._from_config(config.vision_config) - self.language_model = Qwen3_5MoeModel(_build_text_config(config)) - - def get_input_embeddings(self): - return self.language_model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.language_model.set_input_embeddings(value) - - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self.language_model.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) - - def _dummy_vision_inputs(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: - """Smallest valid vision input: a single merged token (grid [1, m, m]).""" - vcfg = self.config.vision_config - m = vcfg.spatial_merge_size - num_patches = m * m - patch_dim = vcfg.in_channels * vcfg.temporal_patch_size * vcfg.patch_size * vcfg.patch_size - pixel_values = torch.zeros(num_patches, patch_dim, device=device, dtype=self.visual.dtype) - grid_thw = torch.tensor([[1, m, m]], dtype=torch.long, device=device) - return pixel_values, grid_thw - - def prepare_inputs_embeds_and_position_ids( - self, - input_ids: torch.LongTensor, - position_ids: torch.LongTensor | None = None, - pixel_values: torch.Tensor | None = None, - image_grid_thw: torch.LongTensor | None = None, - mm_token_type_ids: torch.LongTensor | None = None, - *, - seq_lens: torch.LongTensor, - ) -> tuple[torch.FloatTensor, torch.LongTensor]: - inputs_embeds = self.language_model.embed_tokens(input_ids) - - # Always run the vision encoder for collective symmetry: under FSDP + EP - # all ranks share one process group, so every rank must issue the vision - # encoder's collectives in the same order each step. Text-only microbatches - # keep the dummy embeds in the graph with zero contribution so trainable - # vision modules also participate in backward collectives. - has_images = pixel_values is not None - vision_grid_thw = image_grid_thw - if has_images: - pixel_values = pixel_values.type(self.visual.dtype) - else: - pixel_values, vision_grid_thw = self._dummy_vision_inputs(inputs_embeds.device) - - vision_output = self.visual(pixel_values, grid_thw=vision_grid_thw, return_dict=True) - image_embeds = vision_output.pooler_output.to(inputs_embeds.device, inputs_embeds.dtype) - - if has_images: - image_mask = input_ids == self.config.image_token_id - image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) - else: - inputs_embeds = inputs_embeds + image_embeds.sum() * 0.0 - - if position_ids is None: - if image_grid_thw is not None: - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=self.config.vision_config.spatial_merge_size, - seq_lens=seq_lens, - ) - else: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - - return inputs_embeds, position_ids - - def forward( - self, - input_ids: torch.LongTensor, - position_ids: torch.LongTensor | None = None, - pixel_values: torch.Tensor | None = None, - image_grid_thw: torch.LongTensor | None = None, - mm_token_type_ids: torch.LongTensor | None = None, - routed_experts: torch.LongTensor | None = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - ) -> MoeModelOutputWithPast: - inputs_embeds, position_ids = self.prepare_inputs_embeds_and_position_ids( - input_ids=input_ids, - position_ids=position_ids, - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - mm_token_type_ids=mm_token_type_ids, - seq_lens=seq_lens, - ) - - cp_group = getattr(self.language_model, "_cp_group", None) - if image_grid_thw is not None and cp_group is not None: - cp_rank = self.language_model._cp_rank - cp_world_size = self.language_model._cp_world_size - setup_cp_attention_params(position_ids, cp_group=cp_group, cp_style="ulysses", seq_lens=seq_lens) - inputs_embeds = shard_for_cp(inputs_embeds, cp_rank=cp_rank, cp_world_size=cp_world_size) - position_ids = shard_position_ids_for_cp(position_ids, cp_rank=cp_rank, cp_world_size=cp_world_size) - if routed_experts is not None: - routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_world_size) - seq_lens_are_pre_shard = True - - return self.language_model( - inputs_embeds=inputs_embeds, - position_ids=position_ids, - routed_experts=routed_experts, - seq_lens=seq_lens, - seq_lens_are_pre_shard=seq_lens_are_pre_shard, - ) - - -# --------------------------------------------------------------------------- -# Unified CausalLM / VLM class -# --------------------------------------------------------------------------- - - -class Qwen3_5MoeForCausalLM(Qwen3_5MoePreTrainedModel, GenerationMixin): - """Unified Qwen3.5 MoE model for both text-only and VLM configs. - - When config has a vision_config, creates a composite model with HF's frozen - vision encoder + custom text model. Otherwise creates a text-only model. - """ - - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _checkpoint_conversion_mapping = {} - _tp_plan = {"lm_head": "colwise_rep"} - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - - def __init__(self, config, **kwargs): - super().__init__(config, **kwargs) - self._is_vlm = hasattr(config, "vision_config") - self.supports_packed_multimodal_training = self._is_vlm - - if self._is_vlm: - self.model = Qwen3_5MoeVLMModel(config) - text_config = config.text_config - self._tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} - else: - self.model = Qwen3_5MoeModel(config) - text_config = config - - self.vocab_size = text_config.vocab_size - self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) - self.post_init() - - def get_input_embeddings(self): - return self.model.get_input_embeddings() - - def set_input_embeddings(self, value): - self.model.set_input_embeddings(value) - - def set_decoder(self, decoder): - self.model = decoder - - def get_decoder(self): - return self.model - - def set_context_parallel_attributes(self, cp_group, cp_rank: int, cp_world_size: int) -> None: - self.model.set_context_parallel_attributes(cp_group, cp_rank, cp_world_size) - - # ------------------------------------------------------------------ - # State dict detection & conversion (handles both text-only and VLM) - # ------------------------------------------------------------------ - - @classmethod - def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.gate_up_proj" in name or "mlp.experts.1.up_proj" in name for name in state_dict) - - @classmethod - def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.gate_proj" in name for name in state_dict) - - # ------------------------------------------------------------------ - # Forward - # ------------------------------------------------------------------ - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - logits_to_keep: Union[int, torch.Tensor] = 0, - temperature: Union[torch.Tensor, None] = None, - routed_experts: Optional[torch.LongTensor] = None, - pixel_values: Optional[torch.Tensor] = None, - image_grid_thw: Optional[torch.LongTensor] = None, - mm_token_type_ids: Optional[torch.LongTensor] = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - **kwargs: Unpack[TransformersKwargs], - ) -> PrimeLmOutput: - assert use_cache is None, "use_cache is not supported for custom qwen3_5_moe for now" - assert past_key_values is None, "past_key_values is not supported for custom qwen3_5_moe for now" - - if self._is_vlm: - outputs: MoeModelOutputWithPast = self.model( - input_ids=input_ids, - position_ids=position_ids, - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - mm_token_type_ids=mm_token_type_ids, - routed_experts=routed_experts, - seq_lens=seq_lens, - seq_lens_are_pre_shard=seq_lens_are_pre_shard, - ) - else: - outputs = self.model( - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - routed_experts=routed_experts, - seq_lens=seq_lens, - seq_lens_are_pre_shard=seq_lens_are_pre_shard, - ) - - hidden_states = outputs.last_hidden_state - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - return self.lm_head( - hidden_states[:, slice_indices, :], - labels[:, slice_indices] if labels is not None else None, - temperature=temperature, - ) - - # ------------------------------------------------------------------ - # Buffer init after meta-device loading - # ------------------------------------------------------------------ - - def init_buffers_post_meta(self): - if self._is_vlm: - lm_rope = self.model.language_model.rotary_emb - else: - lm_rope = self.model.rotary_emb - - if hasattr(lm_rope, "rope_init_fn"): - inv_freq, lm_rope.attention_scaling = lm_rope.rope_init_fn(lm_rope.config, lm_rope.inv_freq.device) - lm_rope.inv_freq.copy_(inv_freq) - - if self._is_vlm: - vis_rope = self.model.visual.rotary_pos_emb - if hasattr(vis_rope, "inv_freq"): - dim = vis_rope.inv_freq.shape[0] - inv_freq = 1.0 / ( - 10000.0 - ** (torch.arange(0, dim * 2, 2, dtype=torch.float32, device=vis_rope.inv_freq.device) / (dim * 2)) - ) - vis_rope.inv_freq.copy_(inv_freq) - - -__all__ = [ - "Qwen3_5MoeForCausalLM", - "Qwen3_5MoeModel", - "Qwen3_5MoePreTrainedModel", -] diff --git a/src/prime_rl/trainer/models/qwen3_5_moe/mrope.py b/src/prime_rl/trainer/models/qwen3_5_moe/mrope.py deleted file mode 100644 index a2d1c208e43..00000000000 --- a/src/prime_rl/trainer/models/qwen3_5_moe/mrope.py +++ /dev/null @@ -1,130 +0,0 @@ -from __future__ import annotations - -import itertools - -import torch - - -def get_qwen3_5_vision_position_ids( - *, - start_position: int, - grid_thw: torch.Tensor, - spatial_merge_size: int, - temp_merge_size: int = 1, - time_interval: int = 1, - device: torch.device | None = None, -) -> torch.LongTensor: - llm_grid_t = int(grid_thw[0].item()) // temp_merge_size - llm_grid_h = int(grid_thw[1].item()) // spatial_merge_size - llm_grid_w = int(grid_thw[2].item()) // spatial_merge_size - - position_temporal = torch.arange(llm_grid_t, device=device) * time_interval - position_height = torch.arange(llm_grid_h, device=device) + start_position - position_width = torch.arange(llm_grid_w, device=device) + start_position - - position_width = position_width.repeat(llm_grid_h * llm_grid_t) - position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t) - position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position - return torch.stack([position_temporal, position_height, position_width], dim=0) - - -def build_qwen3_5_mrope_position_ids( - *, - input_ids: torch.LongTensor, - mm_token_type_ids: torch.LongTensor, - image_grid_thw: torch.LongTensor | None, - spatial_merge_size: int, - seq_lens: torch.Tensor, -) -> torch.LongTensor: - if input_ids.ndim != 2: - raise ValueError(f"input_ids must be 2D, got shape={tuple(input_ids.shape)}") - if mm_token_type_ids.shape != input_ids.shape: - raise ValueError( - "mm_token_type_ids must match input_ids shape: " - f"mm_token_type_ids={tuple(mm_token_type_ids.shape)}, input_ids={tuple(input_ids.shape)}" - ) - if input_ids.shape[0] != 1: - raise ValueError("Packed Qwen3.5 MRoPE builder currently supports batch size 1") - - total_seq_len = input_ids.shape[1] - seq_lens = seq_lens.to(device=input_ids.device, dtype=torch.long) - - if seq_lens.ndim != 1: - raise ValueError(f"seq_lens must be 1D, got shape={tuple(seq_lens.shape)}") - if bool((seq_lens <= 0).any().item()): - raise ValueError(f"seq_lens must be positive, got {seq_lens.tolist()}") - if int(seq_lens.sum().item()) != total_seq_len: - raise ValueError(f"seq_lens sum must equal sequence length: {seq_lens.tolist()} vs {total_seq_len}") - - image_grids = iter(image_grid_thw) if image_grid_thw is not None else None - all_positions = torch.empty(3, 1, total_seq_len, dtype=input_ids.dtype, device=input_ids.device) - - offset = 0 - for seq_len_tensor in seq_lens: - seq_len = int(seq_len_tensor.item()) - token_types = mm_token_type_ids[0, offset : offset + seq_len] - current_pos = 0 - segment_positions = [] - - for modality_type, group in itertools.groupby(enumerate(token_types.tolist()), lambda item: item[1]): - group = list(group) - group_len = group[-1][0] - group[0][0] + 1 - - if modality_type == 0: - segment_positions.append( - torch.arange(group_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos - ) - current_pos += group_len - continue - - if modality_type == 2: - raise ValueError("Qwen3.5 video MRoPE is not supported by the custom trainer path yet") - if modality_type != 1: - raise ValueError(f"Unsupported Qwen3.5 multimodal token type: {modality_type}") - if image_grids is None: - raise ValueError("image_grid_thw is required when mm_token_type_ids contains image tokens") - - remaining = group_len - while remaining > 0: - try: - grid_thw = next(image_grids) - except StopIteration as exc: - raise ValueError("Not enough image_grid_thw rows for image token groups") from exc - - vision_positions = get_qwen3_5_vision_position_ids( - start_position=current_pos, - grid_thw=grid_thw, - spatial_merge_size=spatial_merge_size, - device=input_ids.device, - ) - image_len = vision_positions.shape[1] - if image_len > remaining: - raise ValueError( - "Image token group length does not match image_grid_thw-derived token count: " - f"group_len={group_len}, grid_thw={grid_thw.tolist()}, " - f"derived={image_len}, spatial_merge_size={spatial_merge_size}" - ) - segment_positions.append(vision_positions) - current_pos += max(int(grid_thw[1].item()), int(grid_thw[2].item())) // spatial_merge_size - remaining -= image_len - - if not segment_positions: - raise ValueError("Cannot build MRoPE positions for an empty sequence segment") - - segment_position_ids = torch.cat(segment_positions, dim=1) - if segment_position_ids.shape[1] != seq_len: - raise ValueError( - f"Built MRoPE length does not match segment length: {segment_position_ids.shape[1]} vs {seq_len}" - ) - all_positions[:, 0, offset : offset + seq_len] = segment_position_ids - offset += seq_len - - if image_grids is not None: - try: - extra_grid = next(image_grids) - except StopIteration: - pass - else: - raise ValueError(f"Unused image_grid_thw row while building MRoPE positions: {extra_grid.tolist()}") - - return all_positions diff --git a/src/prime_rl/trainer/parallel_dims.py b/src/prime_rl/trainer/parallel_dims.py index 0cd5568d395..86a4a255faa 100644 --- a/src/prime_rl/trainer/parallel_dims.py +++ b/src/prime_rl/trainer/parallel_dims.py @@ -280,6 +280,7 @@ def _is_moe_model(config: ModelConfig) -> bool: 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) return hasattr(model_config, "num_experts") or hasattr(model_config, "n_routed_experts") diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py index 1fd28ec4a3a..28b7f0d47eb 100644 --- a/tests/unit/train/models/test_qwen3_5.py +++ b/tests/unit/train/models/test_qwen3_5.py @@ -1,15 +1,24 @@ import inspect +import os from unittest.mock import MagicMock import pytest import torch -from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig, Qwen3_5VisionConfig +import torch.distributed as dist from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM as HFQwen3_5ForCausalLM from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn -from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5Model -from prime_rl.trainer.models.qwen3_5.modeling_qwen3_5 import Qwen3_5GatedFlashAttention -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig +from prime_rl.trainer.models.qwen3_5 import ( + Qwen3_5Config, + Qwen3_5ForCausalLM, + Qwen3_5Model, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, + Qwen3_5VisionConfig, +) +from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention +from prime_rl.trainer.models.qwen3_5.gated_delta_net import Qwen3_5GatedDeltaNet +from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm from prime_rl.utils.cp import setup_model_cp @@ -56,8 +65,8 @@ def _tiny_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: return config -def _tiny_moe_config(attn_impl: str = "flash_attention_2") -> Qwen3_5MoeConfig: - config = Qwen3_5MoeConfig( +def _tiny_moe_config(attn_impl: str = "flash_attention_2") -> Qwen3_5MoeTextConfig: + config = Qwen3_5MoeTextConfig( vocab_size=128, hidden_size=64, num_hidden_layers=2, @@ -84,7 +93,9 @@ def _tiny_moe_config(attn_impl: str = "flash_attention_2") -> Qwen3_5MoeConfig: def test_qwen3_5_dense_matches_hf_state_keys_on_meta(): config = _tiny_text_config() with torch.device("meta"): + config._attn_implementation = "eager" hf_model = HFQwen3_5ForCausalLM(config) + config._attn_implementation = "flash_attention_2" prime_model = Qwen3_5ForCausalLM(config) assert set(prime_model.state_dict()) == set(hf_model.state_dict()) @@ -92,29 +103,22 @@ def test_qwen3_5_dense_matches_hf_state_keys_on_meta(): assert tensor.shape == hf_model.state_dict()[name].shape, name -@pytest.mark.parametrize("attn_impl", ["flash_attention_3", "kernels-community/vllm-flash-attn3"]) -def test_qwen3_5_full_attention_uses_custom_class(attn_impl: str): - config = _tiny_text_config(attn_impl=attn_impl) +def test_qwen3_5_full_attention_uses_custom_class(): + config = _tiny_text_config(attn_impl="flash_attention_3") with torch.device("meta"): model = Qwen3_5Model(config) - assert isinstance(model.layers[1].self_attn, Qwen3_5GatedFlashAttention) + assert isinstance(model.layers[1].self_attn, Qwen3_5Attention) assert model.config._attn_implementation == "flash_attention_3" assert "ALL_ATTENTION_FUNCTIONS" not in inspect.getsource(type(model.layers[1].self_attn).forward) -def test_qwen3_5_moe_full_attention_normalizes_fa3_hub_alias(): - from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeGatedFlashAttention, - Qwen3_5MoeModel, - ) - - config = _tiny_moe_config(attn_impl="kernels-community/vllm-flash-attn3") - with torch.device("meta"): - model = Qwen3_5MoeModel(config) +def test_qwen3_5_norms_remain_zero_centered_after_model_init(): + model = Qwen3_5ForCausalLM(_tiny_text_config()) - assert isinstance(model.layers[1].self_attn, Qwen3_5MoeGatedFlashAttention) - assert model.config._attn_implementation == "flash_attention_3" + norms = [module for module in model.modules() if isinstance(module, Qwen3_5RMSNorm)] + assert norms + assert all(torch.count_nonzero(norm.weight) == 0 for norm in norms) def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): @@ -124,10 +128,10 @@ def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): linear_layer = text_model.model.layers[0] text_model.model.layers[0] = torch.nn.Sequential(linear_layer) setup_model_cp(text_model, cp_group, cp_rank=1, cp_world_size=2) - assert text_model.model._cp_group is cp_group - assert text_model.model._cp_rank == 1 - assert text_model.model._cp_world_size == 2 - assert linear_layer.linear_attn.cp_group is cp_group + assert text_model.model.context_parallel_group is cp_group + assert text_model.model.context_parallel_rank == 1 + assert text_model.model.context_parallel_world_size == 2 + assert linear_layer.linear_attn.context_parallel_group is cp_group vlm_config = _tiny_vlm_config() vlm_config.vision_config._attn_implementation = "sdpa" @@ -135,8 +139,65 @@ def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): with torch.device("meta"): vlm_model = Qwen3_5ForCausalLM(vlm_config) setup_model_cp(vlm_model, cp_group, cp_rank=0, cp_world_size=2) - assert vlm_model.model.language_model._cp_group is cp_group - assert vlm_model.model.language_model.layers[0].linear_attn.cp_world_size == 2 + assert vlm_model.model.language_model.context_parallel_group is cp_group + assert vlm_model.model.language_model.layers[0].linear_attn.context_parallel_world_size == 2 + + +@pytest.mark.gpu +def test_qwen3_5_gated_delta_net_context_parallel(): + if int(os.environ.get("WORLD_SIZE", 1)) != 2: + pytest.skip("run with torchrun --nproc-per-node=2") + + dist.init_process_group("nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + + try: + torch.manual_seed(0) + config = _tiny_text_config() + config.linear_key_head_dim = 128 + config.linear_value_head_dim = 128 + config.linear_num_key_heads = 16 + config.linear_num_value_heads = 32 + reference = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) + context_parallel = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) + context_parallel.load_state_dict(reference.state_dict()) + context_parallel.set_context_parallel_attributes(dist.group.WORLD, world_size=2) + + hidden_states = torch.randn(1, 16, config.hidden_size, device="cuda", dtype=torch.bfloat16) + dist.broadcast(hidden_states, src=0) + cu_seqlens = torch.tensor([0, 5, 11, 16], device="cuda", dtype=torch.int32) + + reference_input = hidden_states.detach().clone().requires_grad_() + expected = reference( + reference_input, + cu_seqlens, + cu_seqlens_are_pre_shard=False, + ) + output_gradient = torch.randn_like(expected) + dist.broadcast(output_gradient, src=0) + expected.backward(output_gradient) + + local_slice = slice(local_rank * 8, (local_rank + 1) * 8) + local_input = hidden_states[:, local_slice].detach().clone().requires_grad_() + actual = context_parallel( + local_input, + cu_seqlens, + cu_seqlens_are_pre_shard=True, + ) + actual.backward(output_gradient[:, local_slice]) + + gathered = [torch.empty_like(actual) for _ in range(2)] + dist.all_gather(gathered, actual) + torch.testing.assert_close(torch.cat(gathered, dim=1), expected, rtol=3e-2, atol=3e-2) + torch.testing.assert_close( + local_input.grad, + reference_input.grad[:, local_slice], + rtol=5e-2, + atol=5e-2, + ) + finally: + dist.destroy_process_group() def test_setup_model_cp_requires_hook_only_for_hybrid_models(): @@ -160,15 +221,12 @@ class SoftmaxOnlyModel: def test_qwen3_5_ring_patches_dense_flash_attention(): from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention - from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention - originals = { - cls: cls._compute_attention - for cls in (FlashAttention, AfmoeFlashAttention, Qwen3_5MoeGatedFlashAttention, Qwen3_5GatedFlashAttention) - } + originals = {cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention)} try: substitute_ring_attn(process_group=MagicMock(), heads_k_stride=1) - assert Qwen3_5GatedFlashAttention._compute_attention is not originals[Qwen3_5GatedFlashAttention] + assert Qwen3_5Attention._compute_attention is FlashAttention._compute_attention + assert Qwen3_5Attention._compute_attention is not originals[FlashAttention] finally: for cls, method in originals.items(): cls._compute_attention = method diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py index 1400962b02b..f465adb5bb2 100644 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -3,8 +3,7 @@ from transformers import Qwen3_5MoeForCausalLM as HFQwen3_5MoeForCausalLM from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM as PrimeRLQwen3_5MoeForCausalLM +from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeTextConfig from prime_rl.utils.cp import setup_model_cp from prime_rl.utils.utils import default_dtype @@ -12,7 +11,7 @@ def get_model_pairs(): - config = Qwen3_5MoeConfig( + config = Qwen3_5MoeTextConfig( vocab_size=256, hidden_size=256, num_hidden_layers=4, @@ -34,7 +33,7 @@ def get_model_pairs(): config._attn_implementation = "flash_attention_2" with torch.device("cuda"), default_dtype(torch.bfloat16): hf_model = HFQwen3_5MoeForCausalLM._from_config(config) - prime_model = PrimeRLQwen3_5MoeForCausalLM._from_config(config) + prime_model = Qwen3_5ForCausalLM._from_config(config) with torch.no_grad(): state_dict = hf_model.state_dict() prime_state_keys = prime_model.state_dict().keys() @@ -68,9 +67,28 @@ def test_qwen3_5_moe(): grad_diff = hf_model.model.embed_tokens.weight.grad - prime_model.model.embed_tokens.weight.grad assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=1000), f"Max grad diff: {grad_diff.abs().max()}" + packed_position_ids = torch.arange(1, 51, device="cuda").repeat(2).unsqueeze(0) + with torch.no_grad(): + packed = prime_model( + input_ids, + position_ids=packed_position_ids, + seq_lens=torch.tensor([50, 50], device="cuda"), + )["logits"] + unpacked = torch.cat( + [ + prime_model( + input_ids[:, start : start + 50], + position_ids=packed_position_ids[:, :50], + seq_lens=torch.tensor([50], device="cuda"), + )["logits"] + for start in (0, 50) + ], + dim=1, + ) + torch.testing.assert_close(packed, unpacked, atol=0.03, rtol=0.01) + def test_qwen3_5_moe_roundtrip(): - """Verify HF → PrimeRL → HF weight conversion is lossless at the state_dict level.""" hf_model, prime_model = get_model_pairs() # Get original HF state_dict and the PrimeRL-converted version @@ -81,18 +99,12 @@ def test_qwen3_5_moe_roundtrip(): assert prime_model.is_prime_state_dict(prime_sd) assert not prime_model.is_hf_state_dict(prime_sd) - # Convert PrimeRL → per-expert HF format converted_hf_sd = prime_model.convert_to_hf(dict(prime_sd)) - - # Also convert original HF (fused) to per-expert format for comparison - - # First convert original HF → PrimeRL, then back to per-expert HF orig_prime_sd = dict(original_hf_sd) prime_model.convert_to_prime(orig_prime_sd) orig_roundtripped = dict(orig_prime_sd) prime_model.convert_to_hf(orig_roundtripped) - # All non-expert keys should match exactly, expert keys should match after roundtrip for key in orig_roundtripped: assert key in converted_hf_sd, f"Missing key: {key}" assert torch.equal(orig_roundtripped[key], converted_hf_sd[key]), f"Mismatch at {key}" @@ -128,23 +140,18 @@ def test_qwen3_5_moe_router_replay(): def test_qwen3_5_moe_cp_patching(): - """Verify substitute_ring_attn patches Qwen3_5MoeGatedFlashAttention._compute_attention.""" from unittest.mock import MagicMock from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn - from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedFlashAttention - - # substitute_ring_attn rewrites _compute_attention on all three classes; - # snapshot every one so the patch can't leak into later tests via the - # untouched siblings. - originals = { - cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention, Qwen3_5MoeGatedFlashAttention) - } + from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention + + originals = {cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention)} try: mock_group = MagicMock() substitute_ring_attn(process_group=mock_group, heads_k_stride=1) - assert Qwen3_5MoeGatedFlashAttention._compute_attention is not originals[Qwen3_5MoeGatedFlashAttention] + assert Qwen3_5Attention._compute_attention is FlashAttention._compute_attention + assert Qwen3_5Attention._compute_attention is not originals[FlashAttention] finally: for cls, method in originals.items(): cls._compute_attention = method @@ -153,7 +160,7 @@ def test_qwen3_5_moe_cp_patching(): def test_qwen3_5_moe_context_parallel_setup_hook(): from unittest.mock import MagicMock - config = Qwen3_5MoeConfig( + config = Qwen3_5MoeTextConfig( vocab_size=128, hidden_size=64, num_hidden_layers=2, @@ -174,18 +181,18 @@ def test_qwen3_5_moe_context_parallel_setup_hook(): ) config._attn_implementation = "flash_attention_2" with torch.device("meta"): - model = PrimeRLQwen3_5MoeForCausalLM(config) + model = Qwen3_5ForCausalLM(config) linear_layer = model.model.layers[0] model.model.layers[0] = torch.nn.Sequential(linear_layer) cp_group = MagicMock() setup_model_cp(model, cp_group, cp_rank=1, cp_world_size=2) - assert model.model._cp_group is cp_group - assert model.model._cp_rank == 1 - assert model.model._cp_world_size == 2 - assert linear_layer.linear_attn.cp_group is cp_group - assert linear_layer.linear_attn.cp_world_size == 2 + assert model.model.context_parallel_group is cp_group + assert model.model.context_parallel_rank == 1 + assert model.model.context_parallel_world_size == 2 + assert linear_layer.linear_attn.context_parallel_group is cp_group + assert linear_layer.linear_attn.context_parallel_world_size == 2 if __name__ == "__main__": diff --git a/tests/unit/train/models/test_qwen3_5_moe_mrope.py b/tests/unit/train/models/test_qwen3_5_moe_mrope.py index d562daa021a..2814f5f560d 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_mrope.py +++ b/tests/unit/train/models/test_qwen3_5_moe_mrope.py @@ -4,13 +4,15 @@ Qwen3_5MoeTextRotaryEmbedding as HFQwen3_5MoeRotaryEmbedding, ) -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeConfig -from prime_rl.trainer.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeRotaryEmbedding -from prime_rl.trainer.models.qwen3_5_moe.mrope import build_qwen3_5_mrope_position_ids +from prime_rl.trainer.models.qwen3_5 import Qwen3_5MoeTextConfig +from prime_rl.trainer.models.qwen3_5.rotary_embedding import ( + Qwen3_5RotaryEmbedding, + build_qwen3_5_mrope_position_ids, +) -def _tiny_config() -> Qwen3_5MoeConfig: - return Qwen3_5MoeConfig( +def _tiny_config() -> Qwen3_5MoeTextConfig: + return Qwen3_5MoeTextConfig( vocab_size=128, hidden_size=128, num_attention_heads=2, @@ -169,7 +171,7 @@ def test_qwen35_mrope_rejects_video_tokens(): def test_qwen35_rotary_matches_hf_for_2d_and_3d_positions(): config = _tiny_config() - prime_rotary = Qwen3_5MoeRotaryEmbedding(config) + prime_rotary = Qwen3_5RotaryEmbedding(config) hf_rotary = HFQwen3_5MoeRotaryEmbedding(config) hidden_states = torch.randn(1, 6, config.hidden_size) diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/test_qwen3_5_moe_vlm.py index 14ef465afaa..09e748116bc 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_vlm.py +++ b/tests/unit/train/models/test_qwen3_5_moe_vlm.py @@ -7,7 +7,7 @@ from prime_rl.trainer.model import can_reinit_empty_buffers from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM +from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM from prime_rl.utils.utils import default_dtype pytestmark = [pytest.mark.gpu] @@ -79,7 +79,7 @@ def test_vlm_forward(): """Custom VLM produces logits for both text-only and multimodal inputs.""" config = _tiny_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5MoeForCausalLM(config) + model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) vocab = config.text_config.vocab_size @@ -111,7 +111,7 @@ def test_vlm_backward(): """Gradients flow through both vision scatter and text model.""" config = _tiny_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5MoeForCausalLM(config) + model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) @@ -142,7 +142,7 @@ def test_vlm_weight_load_from_hf(): config = _tiny_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): hf_model = HFQwen3_5MoeVLM._from_config(config) - prime_model = Qwen3_5MoeForCausalLM(config) + prime_model = Qwen3_5ForCausalLM(config) # Copy weights: HF -> PrimeRL (with MoE conversion) with torch.no_grad(): @@ -156,6 +156,16 @@ def test_vlm_weight_load_from_hf(): prime_param = dict(prime_model.model.visual.named_parameters())[name] assert torch.equal(param, prime_param), f"Vision weight mismatch: {name}" + pixel_values, image_grid_thw, _ = _make_image_inputs(config) + with torch.no_grad(): + hf_image_embeds = hf_model.model.visual( + pixel_values, + grid_thw=image_grid_thw, + return_dict=True, + ).pooler_output + prime_image_embeds = prime_model.model.visual(pixel_values, image_grid_thw).pooler_output + torch.testing.assert_close(prime_image_embeds, hf_image_embeds, atol=0.03, rtol=0.01) + # Verify model produces output after weight loading input_ids = torch.randint(0, 200, (1, 20), device="cuda") position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) @@ -169,7 +179,7 @@ def test_vlm_weight_roundtrip(): config = _tiny_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): hf_model = HFQwen3_5MoeVLM._from_config(config) - prime_model = Qwen3_5MoeForCausalLM(config) + prime_model = Qwen3_5ForCausalLM(config) hf_sd = hf_model.state_dict() original_vision_key = "model.visual.blocks.0.mlp.linear_fc1.weight" @@ -202,7 +212,7 @@ def test_vlm_router_replay(): """routed_experts bypasses router computation in VLM multimodal forward.""" config = _tiny_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5MoeForCausalLM(config) + model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) vocab = config.text_config.vocab_size @@ -235,7 +245,7 @@ def test_vlm_meta_device_and_buffer_reinit(): """Model can be created on meta device and buffers reinitialized.""" config = _tiny_vlm_config() with torch.device("meta"): - model = Qwen3_5MoeForCausalLM.from_config(config) + model = Qwen3_5ForCausalLM.from_config(config, trust_remote_code=False) assert can_reinit_empty_buffers(model) @@ -243,8 +253,11 @@ def test_vlm_meta_device_and_buffer_reinit(): model.init_buffers_post_meta() lm_inv = model.model.language_model.rotary_emb.inv_freq + lm_original_inv = model.model.language_model.rotary_emb.original_inv_freq vis_inv = model.model.visual.rotary_pos_emb.inv_freq assert lm_inv.device.type == "cuda" + assert lm_original_inv.device.type == "cuda" assert vis_inv.device.type == "cuda" assert lm_inv.abs().sum() > 0 + assert lm_original_inv.abs().sum() > 0 assert vis_inv.abs().sum() > 0 From 0e48de13f7360b7226a8331afd3d064d4fe6ae37 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 12:23:57 +0000 Subject: [PATCH 02/12] fix(qwen3.5): clean native model and honor shared runtime fusions --- .../trainer/models/qwen3_5/attention.py | 11 ++-- .../models/qwen3_5/configuration_qwen3_5.py | 37 ------------ .../trainer/models/qwen3_5/gated_delta_net.py | 53 +++++++---------- .../models/qwen3_5/modeling_qwen3_5.py | 59 +++---------------- .../models/qwen3_5/rotary_embedding.py | 13 ---- src/prime_rl/trainer/models/qwen3_5/vision.py | 6 +- tests/unit/train/models/test_qwen3_5.py | 2 - tests/unit/train/models/test_qwen3_5_moe.py | 22 +++---- .../train/models/test_qwen3_5_moe_mrope.py | 29 --------- .../unit/train/models/test_qwen3_5_moe_vlm.py | 11 ---- tests/unit/utils/test_weights.py | 3 +- 11 files changed, 41 insertions(+), 205 deletions(-) diff --git a/src/prime_rl/trainer/models/qwen3_5/attention.py b/src/prime_rl/trainer/models/qwen3_5/attention.py index 040f3118199..c8c8acb24c2 100644 --- a/src/prime_rl/trainer/models/qwen3_5/attention.py +++ b/src/prime_rl/trainer/models/qwen3_5/attention.py @@ -32,6 +32,7 @@ def __init__(self, config: Qwen3_5TextConfig, attention_implementation: str) -> 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) + self.qkv_sizes = (q_output_size, self.qkv_sizes[1], self.qkv_sizes[2]) self.q_proj = nn.Linear(config.hidden_size, q_output_size, bias=config.attention_bias) self.q_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = Qwen3_5RMSNorm(self.head_dim, config.rms_norm_eps) @@ -45,7 +46,7 @@ def forward( ) -> tuple[torch.Tensor, None]: batch_size, sequence_length, _ = hidden_states.shape - query_states = self.q_proj(hidden_states) + query_states, key_states, value_states = self.project_qkv(hidden_states) if self.attn_output_gate: query_states, output_gate = query_states.view( batch_size, sequence_length, self.num_heads, self.head_dim * 2 @@ -55,12 +56,8 @@ def forward( query_states = query_states.view(batch_size, sequence_length, self.num_heads, self.head_dim) output_gate = None - key_states = self.k_proj(hidden_states).view( - batch_size, sequence_length, self.num_key_value_heads, self.head_dim - ) - value_states = self.v_proj(hidden_states).view( - batch_size, sequence_length, self.num_key_value_heads, self.head_dim - ) + key_states = key_states.view(batch_size, sequence_length, self.num_key_value_heads, self.head_dim) + value_states = value_states.view(batch_size, sequence_length, self.num_key_value_heads, self.head_dim) query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) diff --git a/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py index b643d8bf2fe..10363c46549 100644 --- a/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/configuration_qwen3_5.py @@ -4,26 +4,8 @@ class Qwen3_5TextConfig(PretrainedConfig): model_type = "qwen3_5_text" base_config_key = "text_config" - keys_to_ignore_at_inference = ["past_key_values"] ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce", - "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - def __init__( self, vocab_size: int = 248320, @@ -116,24 +98,6 @@ def __init__( class Qwen3_5MoeTextConfig(Qwen3_5TextConfig): model_type = "qwen3_5_moe_text" - base_model_tp_plan = { - **Qwen3_5TextConfig.base_model_tp_plan, - "layers.*.mlp.experts.gate_proj": "grouped_gemm", - "layers.*.mlp.experts.up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - "layers.*.mlp.shared_expert.gate_proj": "colwise", - "layers.*.mlp.shared_expert.up_proj": "colwise", - "layers.*.mlp.shared_expert.down_proj": "rowwise", - } - base_model_ep_plan = { - "layers.*.mlp.router": "ep_router", - "layers.*.mlp.experts.gate_proj": "grouped_gemm", - "layers.*.mlp.experts.up_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - } - def __init__( self, moe_intermediate_size: int = 512, @@ -195,7 +159,6 @@ def __init__( class Qwen3_5Config(PretrainedConfig): model_type = "qwen3_5" sub_configs = {"vision_config": Qwen3_5VisionConfig, "text_config": Qwen3_5TextConfig} - keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, diff --git a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py index 84ccccc1594..0f8c9381417 100644 --- a/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py +++ b/src/prime_rl/trainer/models/qwen3_5/gated_delta_net.py @@ -8,8 +8,7 @@ from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig -# FLA's CP convolution exchanges boundary state inside the kernel call, which -# Dynamo cannot trace; keep that call eager while compiling the surrounding layer. +# FLA's CP convolution uses an all-gather layout that Dynamo cannot trace. causal_conv1d_with_context_parallelism = torch.compiler.disable(causal_conv1d) @@ -57,8 +56,6 @@ def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.LongTensor, - *, - cu_seqlens_are_pre_shard: bool, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape mixed_qkv = self.in_proj_qkv(hidden_states) @@ -70,24 +67,21 @@ def forward( context = None if self.context_parallel_group is not None: - if not cu_seqlens_are_pre_shard: - raise ValueError("Qwen3.5 context parallelism requires pre-shard sequence boundaries") context = build_cp_context( cu_seqlens=cu_seqlens.to(device=hidden_states.device, dtype=torch.int32), group=self.context_parallel_group, conv1d_kernel_size=self.conv_kernel_size, ) - conv_kwargs = { - "x": mixed_qkv, - "weight": self.conv1d.weight.squeeze(1), - "bias": self.conv1d.bias, - "activation": self.activation, - } - if context is None: - mixed_qkv, _ = causal_conv1d(**conv_kwargs, cu_seqlens=cu_seqlens) - else: - mixed_qkv, _ = causal_conv1d_with_context_parallelism(**conv_kwargs, cp_context=context) + convolution = causal_conv1d_with_context_parallelism if context is not None else causal_conv1d + mixed_qkv, _ = convolution( + x=mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + cu_seqlens=cu_seqlens, + cp_context=context, + ) query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) query = query.reshape(batch_size, sequence_length, self.num_key_heads, self.key_head_dim) @@ -99,23 +93,16 @@ def forward( query = query.repeat_interleave(heads_per_key, dim=2) key = key.repeat_interleave(heads_per_key, dim=2) - delta_kwargs = { - "q": query, - "k": key, - "v": value, - "g": decay, - "beta": beta, - "use_qk_l2norm_in_kernel": True, - "cu_seqlens": context.cu_seqlens if context is not None else cu_seqlens, - } - if context is None: - core_output, _ = chunk_gated_delta_rule( - **delta_kwargs, - initial_state=None, - output_final_state=False, - ) - else: - core_output, _ = chunk_gated_delta_rule(**delta_kwargs, cp_context=context) + core_output, _ = chunk_gated_delta_rule( + q=query, + k=key, + v=value, + g=decay, + beta=beta, + use_qk_l2norm_in_kernel=True, + cu_seqlens=context.cu_seqlens if context is not None else cu_seqlens, + cp_context=context, + ) core_output = self.norm( core_output.reshape(-1, self.value_head_dim), 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 ed44040b54d..172b018a531 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 @@ -50,8 +50,6 @@ def __init__(self, config: Qwen3_5TextConfig, layer_index: int) -> None: self.linear_attn = Qwen3_5GatedDeltaNet(config) elif self.layer_type == "full_attention": self.self_attn = Qwen3_5Attention(config, config._attn_implementation) - else: - raise ValueError(f"Unsupported Qwen3.5 layer type: {self.layer_type}") if isinstance(config, Qwen3_5MoeTextConfig): router = TokenChoiceTopKRouter( @@ -61,6 +59,7 @@ def __init__(self, config: Qwen3_5TextConfig, layer_index: int) -> None: score_func="softmax", route_norm=True, route_scale=1.0, + selection_bias=config.load_balance_coeff is not None, ) experts = GroupedExperts( dim=config.hidden_size, @@ -95,8 +94,6 @@ def forward( cu_seqlens: torch.LongTensor, max_seqlen: int, routed_experts: torch.LongTensor | None = None, - *, - cu_seqlens_are_pre_shard: bool, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -104,7 +101,6 @@ def forward( hidden_states = self.linear_attn( hidden_states, cu_seqlens, - cu_seqlens_are_pre_shard=cu_seqlens_are_pre_shard, ) else: hidden_states, _ = self.self_attn( @@ -122,14 +118,6 @@ def forward( class Qwen3_5PreTrainedModel(PreTrainedModelPrimeRL): config_class = Qwen3_5TextConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock"] - _supports_flash_attn = True - _supports_sdpa = False - _supports_flex_attn = False - _supports_attention_backend = True - _can_compile_fullgraph = False @classmethod def cp_support(cls, config) -> CPSupport: @@ -138,9 +126,7 @@ def cp_support(cls, config) -> CPSupport: if "linear_attention" in (getattr(text_config, "layer_types", None) or ()): return CPSupport( frozenset({"ulysses"}), - "ring CP is a softmax-attention algorithm and cannot run this model's DeltaNet " - "layers, whereas ulysses' all-to-all on Q/K/V leaves the linear-attention kernel " - "unchanged", + "DeltaNet layers require Ulysses sequence sharding with FLA boundary-state exchange", ) return CPSupport(ALL_CP_STYLES) @@ -164,20 +150,12 @@ def conversion_chain(cls, config): class Qwen3_5Model(Qwen3_5PreTrainedModel): def __init__(self, config: Qwen3_5TextConfig) -> None: super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id) self.layers = nn.ModuleList( Qwen3_5DecoderLayer(config, layer_index) for layer_index in range(config.num_hidden_layers) ) self.norm = Qwen3_5RMSNorm(config.hidden_size, config.rms_norm_eps) self.rotary_emb = Qwen3_5RotaryEmbedding(config) - self.gradient_checkpointing = False - self.post_init() - # Qwen stores RMSNorm weights as offsets from one, while post_init treats them as direct scales. - for module in self.modules(): - if isinstance(module, Qwen3_5RMSNorm): - module.reset_parameters() def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens @@ -203,12 +181,8 @@ def forward( seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, ) -> BaseModelOutput: - if (input_ids is None) == (inputs_embeds is None): - raise ValueError("Specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - if inputs_embeds.shape[0] != 1: - raise ValueError(f"Qwen3.5 expects one packed row, got batch size {inputs_embeds.shape[0]}") if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) @@ -228,7 +202,6 @@ def forward( cu_seqlens, max_seqlen, routed_experts=layer_routed_experts, - cu_seqlens_are_pre_shard=seq_lens_are_pre_shard, ) return BaseModelOutput(last_hidden_state=self.norm(hidden_states)) @@ -343,10 +316,6 @@ def forward( class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel): - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - def __init__(self, config) -> None: super().__init__(config) self.is_vlm = hasattr(config, "vision_config") @@ -362,17 +331,11 @@ def __init__(self, config) -> None: if getattr(config.vision_config, "_attn_implementation_internal", None) is None: config.vision_config._attn_implementation = attention_implementation self.model = Qwen3_5VLMModel(config) - self._tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} else: self.model = Qwen3_5Model(config) self.supports_packed_multimodal_training = self.is_vlm - self.vocab_size = text_config.vocab_size self.lm_head = VanillaOutputLinear(text_config.hidden_size, text_config.vocab_size) - if isinstance(text_config, Qwen3_5MoeTextConfig): - self.num_experts = text_config.num_experts - self.num_experts_per_tok = text_config.num_experts_per_tok - self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() @@ -389,8 +352,8 @@ def forward( position_ids: torch.LongTensor | None = None, inputs_embeds: torch.Tensor | None = None, labels: torch.LongTensor | None = None, - logits_to_keep: int | torch.Tensor = 0, temperature: torch.Tensor | None = None, + sampling_mask: torch.Tensor | None = None, routed_experts: torch.LongTensor | None = None, pixel_values: torch.Tensor | None = None, image_grid_thw: torch.LongTensor | None = None, @@ -398,10 +361,7 @@ def forward( *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, - **kwargs, ) -> PrimeLmOutput: - if kwargs.get("use_cache") is not None or kwargs.get("past_key_values") is not None: - raise ValueError("Qwen3.5 custom training does not support KV caching") if self.is_vlm: outputs = self.model( input_ids=input_ids, @@ -423,14 +383,11 @@ def forward( seq_lens_are_pre_shard=seq_lens_are_pre_shard, ) - if isinstance(logits_to_keep, int): - slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None) - else: - slice_indices = logits_to_keep return self.lm_head( - outputs.last_hidden_state[:, slice_indices], - labels[:, slice_indices] if labels is not None else None, - temperature=temperature[:, slice_indices] if temperature is not None else None, + outputs.last_hidden_state, + labels, + temperature=temperature, + sampling_mask=sampling_mask, ) def init_buffers_post_meta(self) -> None: diff --git a/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py b/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py index 769eaf4cad6..9e78d4c8419 100644 --- a/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py +++ b/src/prime_rl/trainer/models/qwen3_5/rotary_embedding.py @@ -27,10 +27,6 @@ def __init__(self, config: Qwen3_5TextConfig, device=None) -> None: self.mrope_section = config.rope_parameters.get("mrope_section") if self.mrope_section is None: self.mrope_section = self.scaled_mrope_section(inv_freq.numel()) - if not config.rope_parameters["mrope_interleaved"]: - raise ValueError("Qwen3.5 requires interleaved MRoPE") - if sum(self.mrope_section) != inv_freq.numel(): - raise ValueError(f"mrope_section must contain {inv_freq.numel()} rotary pairs, got {self.mrope_section}") @staticmethod def scaled_mrope_section(num_rotary_pairs: int) -> list[int]: @@ -68,8 +64,6 @@ def reset_parameters(self) -> None: def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: if position_ids.ndim == 2: position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) - elif position_ids.ndim != 3: - raise ValueError(f"Qwen3.5 position_ids must be 2D or 3D, got {tuple(position_ids.shape)}") position_ids = position_ids.to(hidden_states.device) inv_freq = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) @@ -121,14 +115,7 @@ def build_qwen3_5_mrope_position_ids( spatial_merge_size: int, seq_lens: torch.Tensor, ) -> torch.LongTensor: - if input_ids.ndim != 2 or input_ids.shape[0] != 1: - raise ValueError(f"Packed Qwen3.5 input_ids must have shape (1, tokens), got {tuple(input_ids.shape)}") - if mm_token_type_ids.shape != input_ids.shape: - raise ValueError("mm_token_type_ids must have the same shape as input_ids") - seq_lens = seq_lens.to(device=input_ids.device, dtype=torch.long) - if seq_lens.ndim != 1 or bool((seq_lens <= 0).any()) or int(seq_lens.sum()) != input_ids.shape[1]: - raise ValueError("seq_lens must contain positive lengths summing to the packed sequence length") image_grids = iter(image_grid_thw) if image_grid_thw is not None else None position_ids = torch.empty(3, 1, input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device) diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 2b4f3a43bd3..58c588aff58 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -85,8 +85,6 @@ def __init__(self, config: Qwen3_5VisionConfig) -> None: self.flash_attention = ( None if attention_implementation == "sdpa" else self.FLASH_ATTENTION_FUNCTIONS[attention_implementation] ) - if attention_implementation == "flash_attention_4": - self.flash_attention = torch.compiler.disable(self.flash_attention) def forward( self, @@ -122,7 +120,7 @@ def forward( outputs.append(output.squeeze(0).transpose(0, 1)) attention_output = torch.cat(outputs, dim=0) elif self.attention_implementation == "flash_attention_4": - attention_output = self.flash_attention( + attention_output, _ = self.flash_attention( query, key, value, @@ -142,8 +140,6 @@ def forward( max_sequence_length, causal=False, ) - if isinstance(attention_output, tuple): - attention_output = attention_output[0] return self.proj(attention_output.reshape(sequence_length, -1)) diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py index 28b7f0d47eb..7e4faf17d55 100644 --- a/tests/unit/train/models/test_qwen3_5.py +++ b/tests/unit/train/models/test_qwen3_5.py @@ -172,7 +172,6 @@ def test_qwen3_5_gated_delta_net_context_parallel(): expected = reference( reference_input, cu_seqlens, - cu_seqlens_are_pre_shard=False, ) output_gradient = torch.randn_like(expected) dist.broadcast(output_gradient, src=0) @@ -183,7 +182,6 @@ def test_qwen3_5_gated_delta_net_context_parallel(): actual = context_parallel( local_input, cu_seqlens, - cu_seqlens_are_pre_shard=True, ) actual.backward(output_gradient[:, local_slice]) diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py index f465adb5bb2..f586803762f 100644 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -2,6 +2,7 @@ import torch from transformers import Qwen3_5MoeForCausalLM as HFQwen3_5MoeForCausalLM +from prime_rl.trainer.models.fusions import apply_model_fusions from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeTextConfig from prime_rl.utils.cp import setup_model_cp @@ -45,27 +46,18 @@ def get_model_pairs(): def test_qwen3_5_moe(): - hf_model, prime_model = get_model_pairs() - - with torch.device("cuda"), default_dtype(torch.bfloat16): - input_ids = torch.randint(0, hf_model.config.vocab_size, (1, 100)) - position_ids = torch.arange(1, 101).unsqueeze(0) - - hf_output = hf_model(input_ids, position_ids=position_ids) + _, prime_model = get_model_pairs() + apply_model_fusions(prime_model, ["qkv", "gate_up"]) + input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100), device="cuda") + position_ids = torch.arange(1, 101, device="cuda").unsqueeze(0) prime_output = prime_model( input_ids, position_ids=position_ids, seq_lens=torch.tensor([input_ids.shape[1]], device="cuda"), ) - hf_output.logits.sum().backward() prime_output["logits"].sum().backward() - - logits_diff = prime_output["logits"] - hf_output.logits - assert torch.allclose(logits_diff, torch.zeros_like(logits_diff), atol=1e-0), ( - f"Max logits diff: {logits_diff.abs().max()}" - ) - grad_diff = hf_model.model.embed_tokens.weight.grad - prime_model.model.embed_tokens.weight.grad - assert torch.allclose(grad_diff, torch.zeros_like(grad_diff), atol=1000), f"Max grad diff: {grad_diff.abs().max()}" + assert torch.isfinite(prime_output["logits"]).all() + assert torch.isfinite(prime_model.model.embed_tokens.weight.grad).all() packed_position_ids = torch.arange(1, 51, device="cuda").repeat(2).unsqueeze(0) with torch.no_grad(): diff --git a/tests/unit/train/models/test_qwen3_5_moe_mrope.py b/tests/unit/train/models/test_qwen3_5_moe_mrope.py index 2814f5f560d..708f084379d 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_mrope.py +++ b/tests/unit/train/models/test_qwen3_5_moe_mrope.py @@ -1,12 +1,8 @@ import pytest import torch -from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeTextRotaryEmbedding as HFQwen3_5MoeRotaryEmbedding, -) from prime_rl.trainer.models.qwen3_5 import Qwen3_5MoeTextConfig from prime_rl.trainer.models.qwen3_5.rotary_embedding import ( - Qwen3_5RotaryEmbedding, build_qwen3_5_mrope_position_ids, ) @@ -167,28 +163,3 @@ def test_qwen35_mrope_rejects_video_tokens(): spatial_merge_size=2, seq_lens=torch.tensor([input_ids.shape[1]]), ) - - -def test_qwen35_rotary_matches_hf_for_2d_and_3d_positions(): - config = _tiny_config() - prime_rotary = Qwen3_5RotaryEmbedding(config) - hf_rotary = HFQwen3_5MoeRotaryEmbedding(config) - hidden_states = torch.randn(1, 6, config.hidden_size) - - text_positions = torch.arange(6).unsqueeze(0) - prime_cos, prime_sin = prime_rotary(hidden_states, text_positions) - hf_cos, hf_sin = hf_rotary(hidden_states, text_positions) - torch.testing.assert_close(prime_cos, hf_cos) - torch.testing.assert_close(prime_sin, hf_sin) - - mrope_positions = torch.stack( - [ - torch.tensor([[0, 1, 2, 2, 2, 4]]), - torch.tensor([[0, 1, 2, 2, 3, 4]]), - torch.tensor([[0, 1, 2, 3, 2, 4]]), - ] - ) - prime_cos, prime_sin = prime_rotary(hidden_states, mrope_positions) - hf_cos, hf_sin = hf_rotary(hidden_states, mrope_positions) - torch.testing.assert_close(prime_cos, hf_cos) - torch.testing.assert_close(prime_sin, hf_sin) diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/test_qwen3_5_moe_vlm.py index 09e748116bc..fe48bcdca5f 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_vlm.py +++ b/tests/unit/train/models/test_qwen3_5_moe_vlm.py @@ -136,7 +136,6 @@ def test_vlm_backward(): def test_vlm_weight_load_from_hf(): """Weights from HF VLM checkpoint load correctly into custom VLM after conversion. - Text model numerical match is already validated by test_qwen3_5_moe.py::test_qwen3_5_moe. This test verifies that VLM weight conversion + loading produces a working model. """ config = _tiny_vlm_config() @@ -156,16 +155,6 @@ def test_vlm_weight_load_from_hf(): prime_param = dict(prime_model.model.visual.named_parameters())[name] assert torch.equal(param, prime_param), f"Vision weight mismatch: {name}" - pixel_values, image_grid_thw, _ = _make_image_inputs(config) - with torch.no_grad(): - hf_image_embeds = hf_model.model.visual( - pixel_values, - grid_thw=image_grid_thw, - return_dict=True, - ).pooler_output - prime_image_embeds = prime_model.model.visual(pixel_values, image_grid_thw).pooler_output - torch.testing.assert_close(prime_image_embeds, hf_image_embeds, atol=0.03, rtol=0.01) - # Verify model produces output after weight loading input_ids = torch.randint(0, 200, (1, 20), device="cuda") position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) diff --git a/tests/unit/utils/test_weights.py b/tests/unit/utils/test_weights.py index 8f94856cf06..0a9bdc52108 100644 --- a/tests/unit/utils/test_weights.py +++ b/tests/unit/utils/test_weights.py @@ -6,7 +6,6 @@ from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaForCausalLM from prime_rl.trainer.models.nemotron_h import NemotronHForCausalLM from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM -from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM from prime_rl.utils.weights import resolve_wire_dtype @@ -31,7 +30,7 @@ def keep_in_fp32(key: str) -> bool: [ (NemotronHForCausalLM, "model.layers.0.mamba.A_log", "model.layers.0.mamba.in_proj.weight"), (Qwen3_5ForCausalLM, "model.layers.0.linear_attn.A_log", "model.layers.0.self_attn.q_proj.weight"), - (Qwen3_5MoeForCausalLM, "model.layers.0.linear_attn.norm.weight", "model.layers.0.mlp.experts.gate_proj"), + (Qwen3_5ForCausalLM, "model.layers.0.linear_attn.norm.weight", "model.layers.0.mlp.experts.gate_proj"), (Glm4MoeForCausalLM, "model.layers.0.mlp.router.selection_bias", "model.layers.0.self_attn.q_proj.weight"), (GlmMoeDsaForCausalLM, "model.layers.0.mlp.router.selection_bias", "model.layers.0.self_attn.q_proj.weight"), (DeepseekV4ForCausalLM, "model.layers.0.attn_hc.fn", "model.layers.0.self_attn.q_b_proj.weight"), From a3eeb047634c7ea24da8dd10c40e267fb58283ae Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 12:30:31 +0000 Subject: [PATCH 03/12] fix(qwen3.5): retain existing vision compiler boundary --- src/prime_rl/trainer/models/qwen3_5/vision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 58c588aff58..48bc048717c 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -71,7 +71,7 @@ class Qwen3_5VisionAttention(nn.Module): FLASH_ATTENTION_FUNCTIONS = { "flash_attention_2": flash_attn_varlen_func, "flash_attention_3": flash_attn_3_varlen_func, - "flash_attention_4": flash_attn_4_varlen_func, + "flash_attention_4": torch.compiler.disable(flash_attn_4_varlen_func), } def __init__(self, config: Qwen3_5VisionConfig) -> None: From 1902f607c9b132b026facfbd494377a713300844 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 12:44:00 +0000 Subject: [PATCH 04/12] test(qwen3.5): fix expert choices when checking packed boundaries --- tests/unit/train/models/test_qwen3_5_moe.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py index f586803762f..509efcbbed6 100644 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -60,11 +60,19 @@ def test_qwen3_5_moe(): assert torch.isfinite(prime_model.model.embed_tokens.weight.grad).all() packed_position_ids = torch.arange(1, 51, device="cuda").repeat(2).unsqueeze(0) + # Keep expert selection fixed to isolate packed sequence boundaries. + config = prime_model.config + routed_experts = ( + torch.rand(1, 100, config.num_hidden_layers, config.num_experts, device="cuda") + .topk(config.num_experts_per_tok, dim=-1) + .indices + ) with torch.no_grad(): packed = prime_model( input_ids, position_ids=packed_position_ids, seq_lens=torch.tensor([50, 50], device="cuda"), + routed_experts=routed_experts, )["logits"] unpacked = torch.cat( [ @@ -72,6 +80,7 @@ def test_qwen3_5_moe(): input_ids[:, start : start + 50], position_ids=packed_position_ids[:, :50], seq_lens=torch.tensor([50], device="cuda"), + routed_experts=routed_experts[:, start : start + 50], )["logits"] for start in (0, 50) ], From 9983e7bedfa92b806fd2998bd24f3d6c5231e01d Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 13:34:19 +0000 Subject: [PATCH 05/12] fix(qwen3.5): remove vision attention compiler disable --- src/prime_rl/trainer/models/qwen3_5/vision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 48bc048717c..58c588aff58 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -71,7 +71,7 @@ class Qwen3_5VisionAttention(nn.Module): FLASH_ATTENTION_FUNCTIONS = { "flash_attention_2": flash_attn_varlen_func, "flash_attention_3": flash_attn_3_varlen_func, - "flash_attention_4": torch.compiler.disable(flash_attn_4_varlen_func), + "flash_attention_4": flash_attn_4_varlen_func, } def __init__(self, config: Qwen3_5VisionConfig) -> None: From a2173601f387d4cc0c44e681de46e9bc5ed3fade Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 16:20:44 +0000 Subject: [PATCH 06/12] refactor(qwen3.5): use Prime attention configuration directly --- src/prime_rl/trainer/model.py | 7 ------ .../models/qwen3_5/modeling_qwen3_5.py | 8 ------ src/prime_rl/trainer/models/qwen3_5/vision.py | 25 +++---------------- tests/unit/train/models/test_qwen3_5.py | 5 ++-- 4 files changed, 5 insertions(+), 40 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 87b832a8f7b..a70ee2b7538 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -297,13 +297,6 @@ def get_model( "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." ) - 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: - logger.info("Using SDPA for VLM vision encoder under CP") - vision_config._attn_implementation = "sdpa" - if hasattr(vision_config, "_attn_implementation_internal"): - vision_config._attn_implementation_internal = "sdpa" for subconfig_key in getattr(model_config, "sub_configs", {}): subconfig = getattr(model_config, subconfig_key, None) if subconfig is not None and hasattr(subconfig, "use_cache"): 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 172b018a531..c2df2e8e25c 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 @@ -320,16 +320,8 @@ def __init__(self, config) -> None: super().__init__(config) self.is_vlm = hasattr(config, "vision_config") text_config = config.text_config if self.is_vlm else config - attention_implementation = ( - getattr(config, "_attn_implementation", None) - or getattr(text_config, "_attn_implementation", None) - or "flash_attention_3" - ) - text_config._attn_implementation = attention_implementation if self.is_vlm: - if getattr(config.vision_config, "_attn_implementation_internal", None) is None: - config.vision_config._attn_implementation = attention_implementation self.model = Qwen3_5VLMModel(config) else: self.model = Qwen3_5Model(config) diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 58c588aff58..902c202839a 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -80,11 +80,8 @@ def __init__(self, config: Qwen3_5VisionConfig) -> None: self.head_dim = config.hidden_size // config.num_heads self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=True) self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) - attention_implementation = getattr(config, "_attn_implementation", None) or "flash_attention_3" - self.attention_implementation = attention_implementation - self.flash_attention = ( - None if attention_implementation == "sdpa" else self.FLASH_ATTENTION_FUNCTIONS[attention_implementation] - ) + self.attention_implementation = config._attn_implementation + self.flash_attention = self.FLASH_ATTENTION_FUNCTIONS[self.attention_implementation] def forward( self, @@ -103,23 +100,7 @@ def forward( query = ((query.float() * cos) + (rotate_half(query.float()) * sin)).to(query_dtype) key = ((key.float() * cos) + (rotate_half(key.float()) * sin)).to(query_dtype) - if self.attention_implementation == "sdpa": - lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() - outputs = [] - for query_chunk, key_chunk, value_chunk in zip( - query.split(lengths), - key.split(lengths), - value.split(lengths), - ): - output = F.scaled_dot_product_attention( - query_chunk.transpose(0, 1).unsqueeze(0), - key_chunk.transpose(0, 1).unsqueeze(0), - value_chunk.transpose(0, 1).unsqueeze(0), - is_causal=False, - ) - outputs.append(output.squeeze(0).transpose(0, 1)) - attention_output = torch.cat(outputs, dim=0) - elif self.attention_implementation == "flash_attention_4": + if self.attention_implementation == "flash_attention_4": attention_output, _ = self.flash_attention( query, key, diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py index 7e4faf17d55..0aa35f8dacf 100644 --- a/tests/unit/train/models/test_qwen3_5.py +++ b/tests/unit/train/models/test_qwen3_5.py @@ -61,7 +61,6 @@ def _tiny_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: vision_end_token_id=123, ) config._attn_implementation = attn_impl - config.text_config._attn_implementation = attn_impl return config @@ -134,10 +133,10 @@ def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): assert linear_layer.linear_attn.context_parallel_group is cp_group vlm_config = _tiny_vlm_config() - vlm_config.vision_config._attn_implementation = "sdpa" - vlm_config.vision_config._attn_implementation_internal = "sdpa" with torch.device("meta"): vlm_model = Qwen3_5ForCausalLM(vlm_config) + assert vlm_model.model.language_model.layers[1].self_attn._flash_attn_version == 2 + assert vlm_model.model.visual.blocks[0].attn.attention_implementation == "flash_attention_2" setup_model_cp(vlm_model, cp_group, cp_rank=0, cp_world_size=2) assert vlm_model.model.language_model.context_parallel_group is cp_group assert vlm_model.model.language_model.layers[0].linear_attn.context_parallel_world_size == 2 From 133107bcdac65c045fa896bb5657d97c6c07e090 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 16:27:17 +0000 Subject: [PATCH 07/12] chore: restore mini_moe script to main --- scripts/mini_moe.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/scripts/mini_moe.py b/scripts/mini_moe.py index 9d884dcb4e7..c151d0a7555 100644 --- a/scripts/mini_moe.py +++ b/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 import Qwen3_5ForCausalLM as PrimeRLQwen3_5MoeVLM +from prime_rl.trainer.models.qwen3_5_moe import Qwen3_5MoeForCausalLM as PrimeRLQwen3_5MoeVLM from prime_rl.utils.logger import setup_logger from prime_rl.utils.utils import default_dtype @@ -264,9 +264,6 @@ 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) @@ -284,11 +281,7 @@ 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, - seq_lens=torch.tensor([input_ids.shape[1]], device=input_ids.device), - ) + prime_output = prime_model(input_ids, position_ids) if is_vlm: # HF GatedDeltaNet has a dtype bug in float32 mode; just verify non-NaN output From dcdf770df34a5aaeb119cd7f35f8aff97a4f9137 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 17:05:19 +0000 Subject: [PATCH 08/12] fix: allow VLM RL training with FP32 optimizer state --- .../prime-rl-configs/src/prime_rl/configs/trainer.py | 10 ---------- src/prime_rl/trainer/model.py | 4 ---- 2 files changed, 14 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 5c2edba1250..44b6ac5416a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -741,16 +741,6 @@ def full_optimizer_offload_disables_grad_clipping(self): self.optim.max_norm = None return self - @model_validator(mode="after") - def vlms_require_bfloat16(self): - if self.model.vlm is not None and ( - self.model.optimization_dtype != "bfloat16" or self.model.reduce_dtype != "bfloat16" - ): - raise ValueError( - "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." - ) - return self - @model_validator(mode="after") def vlm_freeze_incompatible_with_lora(self): if self.model.vlm is not None and not self.model.vlm.freeze_vision_encoder and self.model.lora is not None: diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index a70ee2b7538..37a9429e052 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -292,10 +292,6 @@ def get_model( if is_vlm_training: logger.info(f"Detected vision-language model: {config.name}") - if config.optimization_dtype != "bfloat16" or config.reduce_dtype != "bfloat16": - raise ValueError( - "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." - ) for subconfig_key in getattr(model_config, "sub_configs", {}): subconfig = getattr(model_config, subconfig_key, None) From 70d144599d2b69e01da45eb158408f6a00e25bb4 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 17:06:34 +0000 Subject: [PATCH 09/12] test(qwen3.5): remove Transformers reference models --- tests/unit/train/models/test_qwen3_5.py | 40 ----- tests/unit/train/models/test_qwen3_5_moe.py | 40 +---- .../unit/train/models/test_qwen3_5_moe_vlm.py | 152 ++++++------------ 3 files changed, 50 insertions(+), 182 deletions(-) diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py index 0aa35f8dacf..bae8871afe6 100644 --- a/tests/unit/train/models/test_qwen3_5.py +++ b/tests/unit/train/models/test_qwen3_5.py @@ -5,14 +5,12 @@ import pytest import torch import torch.distributed as dist -from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM as HFQwen3_5ForCausalLM from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, Qwen3_5ForCausalLM, Qwen3_5Model, - Qwen3_5MoeTextConfig, Qwen3_5TextConfig, Qwen3_5VisionConfig, ) @@ -64,44 +62,6 @@ def _tiny_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: return config -def _tiny_moe_config(attn_impl: str = "flash_attention_2") -> Qwen3_5MoeTextConfig: - config = Qwen3_5MoeTextConfig( - vocab_size=128, - hidden_size=64, - num_hidden_layers=2, - layer_types=["linear_attention", "full_attention"], - num_attention_heads=4, - num_key_value_heads=2, - head_dim=16, - max_position_embeddings=128, - linear_key_head_dim=8, - linear_value_head_dim=8, - linear_num_key_heads=4, - linear_num_value_heads=8, - linear_conv_kernel_dim=4, - moe_intermediate_size=128, - shared_expert_intermediate_size=128, - num_experts=4, - num_experts_per_tok=2, - ) - config._attn_implementation = attn_impl - return config - - -@pytest.mark.gpu -def test_qwen3_5_dense_matches_hf_state_keys_on_meta(): - config = _tiny_text_config() - with torch.device("meta"): - config._attn_implementation = "eager" - hf_model = HFQwen3_5ForCausalLM(config) - config._attn_implementation = "flash_attention_2" - prime_model = Qwen3_5ForCausalLM(config) - - assert set(prime_model.state_dict()) == set(hf_model.state_dict()) - for name, tensor in prime_model.state_dict().items(): - assert tensor.shape == hf_model.state_dict()[name].shape, name - - def test_qwen3_5_full_attention_uses_custom_class(): config = _tiny_text_config(attn_impl="flash_attention_3") with torch.device("meta"): diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py index 509efcbbed6..e9f9bc15007 100644 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -1,6 +1,5 @@ import pytest import torch -from transformers import Qwen3_5MoeForCausalLM as HFQwen3_5MoeForCausalLM from prime_rl.trainer.models.fusions import apply_model_fusions from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head @@ -11,7 +10,7 @@ pytestmark = [pytest.mark.gpu] -def get_model_pairs(): +def _tiny_model(): config = Qwen3_5MoeTextConfig( vocab_size=256, hidden_size=256, @@ -33,20 +32,13 @@ def get_model_pairs(): ) config._attn_implementation = "flash_attention_2" with torch.device("cuda"), default_dtype(torch.bfloat16): - hf_model = HFQwen3_5MoeForCausalLM._from_config(config) - prime_model = Qwen3_5ForCausalLM._from_config(config) - with torch.no_grad(): - state_dict = hf_model.state_dict() - prime_state_keys = prime_model.state_dict().keys() - prime_model.convert_to_prime(state_dict) - prime_model.load_state_dict(state_dict) + prime_model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(prime_model, chunk_size=None) - assert set(prime_state_keys) - set(state_dict.keys()) == set() - return hf_model, prime_model + return prime_model def test_qwen3_5_moe(): - _, prime_model = get_model_pairs() + prime_model = _tiny_model() apply_model_fusions(prime_model, ["qkv", "gate_up"]) input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100), device="cuda") position_ids = torch.arange(1, 101, device="cuda").unsqueeze(0) @@ -89,31 +81,9 @@ def test_qwen3_5_moe(): torch.testing.assert_close(packed, unpacked, atol=0.03, rtol=0.01) -def test_qwen3_5_moe_roundtrip(): - hf_model, prime_model = get_model_pairs() - - # Get original HF state_dict and the PrimeRL-converted version - original_hf_sd = hf_model.state_dict() - prime_sd = prime_model.state_dict() - assert prime_model.is_hf_state_dict(original_hf_sd) - assert not prime_model.is_prime_state_dict(original_hf_sd) - assert prime_model.is_prime_state_dict(prime_sd) - assert not prime_model.is_hf_state_dict(prime_sd) - - converted_hf_sd = prime_model.convert_to_hf(dict(prime_sd)) - orig_prime_sd = dict(original_hf_sd) - prime_model.convert_to_prime(orig_prime_sd) - orig_roundtripped = dict(orig_prime_sd) - prime_model.convert_to_hf(orig_roundtripped) - - for key in orig_roundtripped: - assert key in converted_hf_sd, f"Missing key: {key}" - assert torch.equal(orig_roundtripped[key], converted_hf_sd[key]), f"Mismatch at {key}" - - def test_qwen3_5_moe_router_replay(): """When routed_experts are provided, the model uses them instead of computing routing.""" - _, prime_model = get_model_pairs() + prime_model = _tiny_model() with torch.device("cuda"), default_dtype(torch.bfloat16): input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/test_qwen3_5_moe_vlm.py index fe48bcdca5f..98cab80df47 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_vlm.py +++ b/tests/unit/train/models/test_qwen3_5_moe_vlm.py @@ -1,57 +1,59 @@ import pytest import torch -from transformers import AutoConfig -from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeForConditionalGeneration as HFQwen3_5MoeVLM, -) from prime_rl.trainer.model import can_reinit_empty_buffers from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head -from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM +from prime_rl.trainer.models.qwen3_5 import ( + Qwen3_5ForCausalLM, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5VisionConfig, +) from prime_rl.utils.utils import default_dtype pytestmark = [pytest.mark.gpu] def _tiny_vlm_config(): - """HF composite config shrunk for unit testing.""" - config = AutoConfig.from_pretrained( - "Qwen/Qwen3.5-35B-A3B", trust_remote_code=True, attn_implementation="flash_attention_2" + return Qwen3_5MoeConfig( + text_config=Qwen3_5MoeTextConfig( + vocab_size=256, + hidden_size=256, + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + moe_intermediate_size=128, + shared_expert_intermediate_size=128, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=512, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=8, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10_000_000.0, + "partial_rotary_factor": 0.25, + "mrope_section": [3, 3, 2], + "mrope_interleaved": True, + }, + ), + vision_config=Qwen3_5VisionConfig( + depth=2, + hidden_size=128, + intermediate_size=256, + num_heads=4, + out_hidden_size=256, + ), + image_token_id=250, + video_token_id=251, + vision_start_token_id=252, + vision_end_token_id=253, + attn_implementation="flash_attention_2", ) - config.use_cache = False - tc = config.text_config - tc.vocab_size = 256 - tc.hidden_size = 256 - tc.num_hidden_layers = 2 - tc.layer_types = ["linear_attention", "full_attention"] - tc.num_attention_heads = 4 - tc.num_key_value_heads = 2 - tc.head_dim = 64 - tc.moe_intermediate_size = 128 - tc.shared_expert_intermediate_size = 128 - tc.num_experts = 4 - tc.num_experts_per_tok = 2 - tc.max_position_embeddings = 512 - tc.linear_key_head_dim = 32 - tc.linear_value_head_dim = 32 - tc.linear_num_key_heads = 4 - tc.linear_num_value_heads = 8 - tc.use_cache = False - tc.rope_parameters["mrope_section"] = [3, 3, 2] - - vc = config.vision_config - vc.depth = 2 - vc.hidden_size = 128 - vc.intermediate_size = 256 - vc.num_heads = 4 - vc.out_hidden_size = tc.hidden_size - - # Special token IDs must fit within the tiny vocab - config.image_token_id = 250 - config.video_token_id = 251 - config.vision_start_token_id = 252 - config.vision_end_token_id = 253 - return config def _make_image_inputs(config, device="cuda", dtype=torch.bfloat16): @@ -133,70 +135,6 @@ def test_vlm_backward(): assert model.model.visual.patch_embed.proj.weight.grad is not None -def test_vlm_weight_load_from_hf(): - """Weights from HF VLM checkpoint load correctly into custom VLM after conversion. - - This test verifies that VLM weight conversion + loading produces a working model. - """ - config = _tiny_vlm_config() - with torch.device("cuda"), default_dtype(torch.bfloat16): - hf_model = HFQwen3_5MoeVLM._from_config(config) - prime_model = Qwen3_5ForCausalLM(config) - - # Copy weights: HF -> PrimeRL (with MoE conversion) - with torch.no_grad(): - hf_sd = hf_model.state_dict() - prime_model.convert_to_prime(hf_sd) - prime_model.load_state_dict(hf_sd) - inject_prime_lm_head(prime_model) - - # Verify vision encoder weights match exactly (should be untouched by conversion) - for name, param in hf_model.model.visual.named_parameters(): - prime_param = dict(prime_model.model.visual.named_parameters())[name] - assert torch.equal(param, prime_param), f"Vision weight mismatch: {name}" - - # Verify model produces output after weight loading - input_ids = torch.randint(0, 200, (1, 20), device="cuda") - position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) - out = prime_model(input_ids=input_ids, position_ids=position_ids, seq_lens=_seq_lens(input_ids)) - assert out["logits"].shape[2] == config.text_config.vocab_size - assert not torch.isnan(out["logits"]).any() - - -def test_vlm_weight_roundtrip(): - """HF -> PrimeRL -> HF weight conversion is lossless (vision keys untouched, text keys converted).""" - config = _tiny_vlm_config() - with torch.device("cuda"), default_dtype(torch.bfloat16): - hf_model = HFQwen3_5MoeVLM._from_config(config) - prime_model = Qwen3_5ForCausalLM(config) - - hf_sd = hf_model.state_dict() - original_vision_key = "model.visual.blocks.0.mlp.linear_fc1.weight" - original_vision_weight = hf_sd[original_vision_key].clone() - - # HF -> PrimeRL - prime_sd = dict(hf_sd) - prime_model.convert_to_prime(prime_sd) - assert any("language_model" in k and "mlp.experts.gate_proj" in k for k in prime_sd) - assert original_vision_key in prime_sd - - # PrimeRL -> HF - roundtripped = dict(prime_sd) - prime_model.convert_to_hf(roundtripped) - - # Original HF also needs roundtrip for expert format normalization - orig_rt = dict(hf_sd) - prime_model.convert_to_prime(orig_rt) - prime_model.convert_to_hf(orig_rt) - - for key in orig_rt: - assert key in roundtripped, f"Missing key: {key}" - assert torch.equal(orig_rt[key], roundtripped[key]), f"Mismatch at {key}" - - # Vision weights preserved through the whole roundtrip - assert torch.equal(roundtripped[original_vision_key], original_vision_weight) - - def test_vlm_router_replay(): """routed_experts bypasses router computation in VLM multimodal forward.""" config = _tiny_vlm_config() @@ -234,7 +172,7 @@ def test_vlm_meta_device_and_buffer_reinit(): """Model can be created on meta device and buffers reinitialized.""" config = _tiny_vlm_config() with torch.device("meta"): - model = Qwen3_5ForCausalLM.from_config(config, trust_remote_code=False) + model = Qwen3_5ForCausalLM(config) assert can_reinit_empty_buffers(model) From c48531bbd424d46b2303454886bc4874bb08f611 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 17:13:52 +0000 Subject: [PATCH 10/12] test(qwen3.5): use public model factory and automatic attention --- tests/unit/train/models/test_qwen3_5_moe.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py index e9f9bc15007..66885138903 100644 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ b/tests/unit/train/models/test_qwen3_5_moe.py @@ -1,6 +1,9 @@ import pytest import torch +from prime_rl.configs.trainer import ModelConfig +from prime_rl.trainer.model import resolve_auto_attn +from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL from prime_rl.trainer.models.fusions import apply_model_fusions from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeTextConfig @@ -10,7 +13,9 @@ pytestmark = [pytest.mark.gpu] -def _tiny_model(): +def get_model(): + runtime_config = ModelConfig() + resolve_auto_attn(runtime_config) config = Qwen3_5MoeTextConfig( vocab_size=256, hidden_size=256, @@ -30,15 +35,16 @@ def _tiny_model(): linear_num_key_heads=4, linear_num_value_heads=8, ) - config._attn_implementation = "flash_attention_2" - with torch.device("cuda"), default_dtype(torch.bfloat16): - prime_model = Qwen3_5ForCausalLM(config) + with torch.device("cuda"): + prime_model = AutoModelForCausalLMPrimeRL.from_config( + config, attn_implementation=runtime_config.attn, dtype=torch.bfloat16 + ) inject_prime_lm_head(prime_model, chunk_size=None) return prime_model def test_qwen3_5_moe(): - prime_model = _tiny_model() + prime_model = get_model() apply_model_fusions(prime_model, ["qkv", "gate_up"]) input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100), device="cuda") position_ids = torch.arange(1, 101, device="cuda").unsqueeze(0) @@ -83,7 +89,7 @@ def test_qwen3_5_moe(): def test_qwen3_5_moe_router_replay(): """When routed_experts are provided, the model uses them instead of computing routing.""" - prime_model = _tiny_model() + prime_model = get_model() with torch.device("cuda"), default_dtype(torch.bfloat16): input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) @@ -150,7 +156,6 @@ def test_qwen3_5_moe_context_parallel_setup_hook(): linear_num_key_heads=4, linear_num_value_heads=8, ) - config._attn_implementation = "flash_attention_2" with torch.device("meta"): model = Qwen3_5ForCausalLM(config) From 5ba7db5ac6aeaa5fc82b76623bb9bb06c8ab1eea Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 17:19:11 +0000 Subject: [PATCH 11/12] test(qwen3.5): use public helper names consistently --- tests/unit/train/models/test_qwen3_5.py | 16 ++++----- .../train/models/test_qwen3_5_moe_mrope.py | 19 ---------- .../unit/train/models/test_qwen3_5_moe_vlm.py | 36 +++++++++---------- 3 files changed, 26 insertions(+), 45 deletions(-) diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py index bae8871afe6..30af50a74cd 100644 --- a/tests/unit/train/models/test_qwen3_5.py +++ b/tests/unit/train/models/test_qwen3_5.py @@ -20,7 +20,7 @@ from prime_rl.utils.cp import setup_model_cp -def _tiny_text_config(attn_impl: str = "flash_attention_2") -> Qwen3_5TextConfig: +def get_text_config(attn_impl: str = "flash_attention_2") -> Qwen3_5TextConfig: config = Qwen3_5TextConfig( vocab_size=128, hidden_size=64, @@ -41,8 +41,8 @@ def _tiny_text_config(attn_impl: str = "flash_attention_2") -> Qwen3_5TextConfig return config -def _tiny_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: - text_config = _tiny_text_config(attn_impl) +def get_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: + text_config = get_text_config(attn_impl) vision_config = Qwen3_5VisionConfig( depth=1, hidden_size=64, @@ -63,7 +63,7 @@ def _tiny_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: def test_qwen3_5_full_attention_uses_custom_class(): - config = _tiny_text_config(attn_impl="flash_attention_3") + config = get_text_config(attn_impl="flash_attention_3") with torch.device("meta"): model = Qwen3_5Model(config) @@ -73,7 +73,7 @@ def test_qwen3_5_full_attention_uses_custom_class(): def test_qwen3_5_norms_remain_zero_centered_after_model_init(): - model = Qwen3_5ForCausalLM(_tiny_text_config()) + model = Qwen3_5ForCausalLM(get_text_config()) norms = [module for module in model.modules() if isinstance(module, Qwen3_5RMSNorm)] assert norms @@ -83,7 +83,7 @@ def test_qwen3_5_norms_remain_zero_centered_after_model_init(): def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): cp_group = MagicMock() - text_model = Qwen3_5ForCausalLM(_tiny_text_config()) + text_model = Qwen3_5ForCausalLM(get_text_config()) linear_layer = text_model.model.layers[0] text_model.model.layers[0] = torch.nn.Sequential(linear_layer) setup_model_cp(text_model, cp_group, cp_rank=1, cp_world_size=2) @@ -92,7 +92,7 @@ def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): assert text_model.model.context_parallel_world_size == 2 assert linear_layer.linear_attn.context_parallel_group is cp_group - vlm_config = _tiny_vlm_config() + vlm_config = get_vlm_config() with torch.device("meta"): vlm_model = Qwen3_5ForCausalLM(vlm_config) assert vlm_model.model.language_model.layers[1].self_attn._flash_attn_version == 2 @@ -113,7 +113,7 @@ def test_qwen3_5_gated_delta_net_context_parallel(): try: torch.manual_seed(0) - config = _tiny_text_config() + config = get_text_config() config.linear_key_head_dim = 128 config.linear_value_head_dim = 128 config.linear_num_key_heads = 16 diff --git a/tests/unit/train/models/test_qwen3_5_moe_mrope.py b/tests/unit/train/models/test_qwen3_5_moe_mrope.py index 708f084379d..eb2366c6228 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_mrope.py +++ b/tests/unit/train/models/test_qwen3_5_moe_mrope.py @@ -1,30 +1,11 @@ import pytest import torch -from prime_rl.trainer.models.qwen3_5 import Qwen3_5MoeTextConfig from prime_rl.trainer.models.qwen3_5.rotary_embedding import ( build_qwen3_5_mrope_position_ids, ) -def _tiny_config() -> Qwen3_5MoeTextConfig: - return Qwen3_5MoeTextConfig( - vocab_size=128, - hidden_size=128, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=64, - num_hidden_layers=1, - rope_parameters={ - "rope_type": "default", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25, - "mrope_section": [3, 3, 2], - "mrope_interleaved": True, - }, - ) - - def test_qwen35_mrope_text_only_positions(): input_ids = torch.tensor([[10, 11, 12, 13]]) mm_token_type_ids = torch.zeros_like(input_ids) diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/test_qwen3_5_moe_vlm.py index 98cab80df47..6715c924e16 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_vlm.py +++ b/tests/unit/train/models/test_qwen3_5_moe_vlm.py @@ -14,7 +14,7 @@ pytestmark = [pytest.mark.gpu] -def _tiny_vlm_config(): +def get_vlm_config(): return Qwen3_5MoeConfig( text_config=Qwen3_5MoeTextConfig( vocab_size=256, @@ -56,7 +56,7 @@ def _tiny_vlm_config(): ) -def _make_image_inputs(config, device="cuda", dtype=torch.bfloat16): +def get_image_inputs(config, device="cuda", dtype=torch.bfloat16): """Create minimal image inputs matching the vision config.""" vc = config.vision_config patch_dim = vc.in_channels * vc.temporal_patch_size * vc.patch_size * vc.patch_size @@ -67,19 +67,19 @@ def _make_image_inputs(config, device="cuda", dtype=torch.bfloat16): return pixel_values, image_grid_thw, num_image_tokens -def _make_mm_token_type_ids(input_ids, image_token_id): +def get_mm_token_type_ids(input_ids, image_token_id): mm_token_type_ids = torch.zeros_like(input_ids) mm_token_type_ids[input_ids == image_token_id] = 1 return mm_token_type_ids -def _seq_lens(input_ids: torch.Tensor) -> torch.Tensor: +def get_seq_lens(input_ids: torch.Tensor) -> torch.Tensor: return torch.tensor([input_ids.shape[1]], device=input_ids.device) def test_vlm_forward(): """Custom VLM produces logits for both text-only and multimodal inputs.""" - config = _tiny_vlm_config() + config = get_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) @@ -89,45 +89,45 @@ def test_vlm_forward(): # Text-only (avoid special token range 250-253) input_ids = torch.randint(0, 200, (1, 20), device="cuda") position_ids = torch.arange(1, 21, device="cuda").unsqueeze(0) - out_text = model(input_ids=input_ids, position_ids=position_ids, seq_lens=_seq_lens(input_ids)) + out_text = model(input_ids=input_ids, position_ids=position_ids, seq_lens=get_seq_lens(input_ids)) assert out_text["logits"].shape == (1, 20, vocab) # Multimodal - pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + pixel_values, image_grid_thw, n_img_tokens = get_image_inputs(config) text_part = torch.randint(0, 200, (1, 10), device="cuda") img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") input_ids_mm = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) - mm_token_type_ids = _make_mm_token_type_ids(input_ids_mm, config.image_token_id) + mm_token_type_ids = get_mm_token_type_ids(input_ids_mm, config.image_token_id) out_mm = model( input_ids=input_ids_mm, pixel_values=pixel_values, image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, - seq_lens=_seq_lens(input_ids_mm), + seq_lens=get_seq_lens(input_ids_mm), ) assert out_mm["logits"].shape == (1, input_ids_mm.shape[1], vocab) def test_vlm_backward(): """Gradients flow through both vision scatter and text model.""" - config = _tiny_vlm_config() + config = get_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) - pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + pixel_values, image_grid_thw, n_img_tokens = get_image_inputs(config) text_part = torch.randint(0, 200, (1, 10), device="cuda") img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") input_ids = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) - mm_token_type_ids = _make_mm_token_type_ids(input_ids, config.image_token_id) + mm_token_type_ids = get_mm_token_type_ids(input_ids, config.image_token_id) out = model( input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, - seq_lens=_seq_lens(input_ids), + seq_lens=get_seq_lens(input_ids), ) out["logits"].sum().backward() @@ -137,17 +137,17 @@ def test_vlm_backward(): def test_vlm_router_replay(): """routed_experts bypasses router computation in VLM multimodal forward.""" - config = _tiny_vlm_config() + config = get_vlm_config() with torch.device("cuda"), default_dtype(torch.bfloat16): model = Qwen3_5ForCausalLM(config) inject_prime_lm_head(model) vocab = config.text_config.vocab_size - pixel_values, image_grid_thw, n_img_tokens = _make_image_inputs(config) + pixel_values, image_grid_thw, n_img_tokens = get_image_inputs(config) text_part = torch.randint(0, 200, (1, 10), device="cuda") img_part = torch.full((1, n_img_tokens), config.image_token_id, device="cuda") input_ids = torch.cat([text_part[:, :5], img_part, text_part[:, 5:]], dim=1) - mm_token_type_ids = _make_mm_token_type_ids(input_ids, config.image_token_id) + mm_token_type_ids = get_mm_token_type_ids(input_ids, config.image_token_id) seq_len = input_ids.shape[1] num_layers = config.text_config.num_hidden_layers @@ -160,7 +160,7 @@ def test_vlm_router_replay(): image_grid_thw=image_grid_thw, mm_token_type_ids=mm_token_type_ids, routed_experts=routed_experts, - seq_lens=_seq_lens(input_ids), + seq_lens=get_seq_lens(input_ids), ) assert out["logits"].shape == (1, seq_len, vocab) @@ -170,7 +170,7 @@ def test_vlm_router_replay(): def test_vlm_meta_device_and_buffer_reinit(): """Model can be created on meta device and buffers reinitialized.""" - config = _tiny_vlm_config() + config = get_vlm_config() with torch.device("meta"): model = Qwen3_5ForCausalLM(config) From d9162ce07425b56e01f6b62d3b24b141a8ff4a51 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Thu, 10 Sep 2026 17:33:48 +0000 Subject: [PATCH 12/12] test: consolidate native Qwen3.5 model coverage --- tests/unit/train/models/qwen3_5/__init__.py | 0 tests/unit/train/models/qwen3_5/test_model.py | 282 ++++++++++++++++++ .../test_vlm.py} | 180 +++++++++-- tests/unit/train/models/test_qwen3_5.py | 189 ------------ tests/unit/train/models/test_qwen3_5_moe.py | 175 ----------- .../train/models/test_qwen3_5_moe_mrope.py | 146 --------- 6 files changed, 445 insertions(+), 527 deletions(-) create mode 100644 tests/unit/train/models/qwen3_5/__init__.py create mode 100644 tests/unit/train/models/qwen3_5/test_model.py rename tests/unit/train/models/{test_qwen3_5_moe_vlm.py => qwen3_5/test_vlm.py} (53%) delete mode 100644 tests/unit/train/models/test_qwen3_5.py delete mode 100644 tests/unit/train/models/test_qwen3_5_moe.py delete mode 100644 tests/unit/train/models/test_qwen3_5_moe_mrope.py diff --git a/tests/unit/train/models/qwen3_5/__init__.py b/tests/unit/train/models/qwen3_5/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/train/models/qwen3_5/test_model.py b/tests/unit/train/models/qwen3_5/test_model.py new file mode 100644 index 00000000000..ca4960a3663 --- /dev/null +++ b/tests/unit/train/models/qwen3_5/test_model.py @@ -0,0 +1,282 @@ +import os +from unittest.mock import MagicMock + +import pytest +import torch +import torch.distributed as dist + +from prime_rl.configs.trainer import ModelConfig +from prime_rl.trainer.model import resolve_auto_attn +from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL +from prime_rl.trainer.models.fusions import apply_model_fusions +from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.qwen3_5 import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, + Qwen3_5VisionConfig, +) +from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention +from prime_rl.trainer.models.qwen3_5.gated_delta_net import Qwen3_5GatedDeltaNet +from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm +from prime_rl.utils.cp import setup_model_cp + + +def get_text_config(config_cls=Qwen3_5TextConfig) -> Qwen3_5TextConfig: + moe_config = {} + if config_cls is Qwen3_5MoeTextConfig: + moe_config = dict( + moe_intermediate_size=128, + shared_expert_intermediate_size=128, + num_experts=8, + num_experts_per_tok=2, + ) + return config_cls( + vocab_size=256, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=4, + layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"], + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=512, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=8, + linear_conv_kernel_dim=4, + **moe_config, + ) + + +@pytest.fixture(params=[Qwen3_5TextConfig, Qwen3_5MoeTextConfig], ids=["dense", "moe"]) +def text_config(request): + return get_text_config(request.param) + + +def get_vlm_config(text_config) -> Qwen3_5Config: + vision_config = Qwen3_5VisionConfig( + depth=1, + hidden_size=64, + intermediate_size=128, + num_heads=4, + out_hidden_size=text_config.hidden_size, + ) + config_cls = Qwen3_5MoeConfig if isinstance(text_config, Qwen3_5MoeTextConfig) else Qwen3_5Config + return config_cls( + text_config=text_config, + vision_config=vision_config, + image_token_id=120, + video_token_id=121, + vision_start_token_id=122, + vision_end_token_id=123, + ) + + +def get_model(config, device="cuda"): + runtime_config = ModelConfig() + resolve_auto_attn(runtime_config) + with torch.device(device): + model = AutoModelForCausalLMPrimeRL.from_config( + config, attn_implementation=runtime_config.attn, dtype=torch.bfloat16 + ) + inject_prime_lm_head(model, chunk_size=None) + return model + + +@pytest.mark.gpu +def test_norms_remain_zero_centered_after_model_init(text_config): + model = get_model(text_config) + + norms = [module for module in model.modules() if isinstance(module, Qwen3_5RMSNorm)] + assert norms + assert all(torch.count_nonzero(norm.weight) == 0 for norm in norms) + + +@pytest.mark.gpu +def test_context_parallel_setup_chain_text_and_vlm(text_config): + cp_group = MagicMock() + + text_model = get_model(text_config, device="meta") + linear_layer = text_model.model.layers[0] + text_model.model.layers[0] = torch.nn.Sequential(linear_layer) + setup_model_cp(text_model, cp_group, cp_rank=1, cp_world_size=2) + assert text_model.model.context_parallel_group is cp_group + assert text_model.model.context_parallel_rank == 1 + assert text_model.model.context_parallel_world_size == 2 + assert linear_layer.linear_attn.context_parallel_group is cp_group + assert linear_layer.linear_attn.context_parallel_world_size == 2 + + vlm_model = get_model(get_vlm_config(text_config), device="meta") + setup_model_cp(vlm_model, cp_group, cp_rank=0, cp_world_size=2) + assert vlm_model.model.language_model.context_parallel_group is cp_group + assert vlm_model.model.language_model.layers[0].linear_attn.context_parallel_world_size == 2 + + +@pytest.mark.gpu +def test_qwen3_5_gated_delta_net_context_parallel(): + if int(os.environ.get("WORLD_SIZE", 1)) != 2: + pytest.skip("run with torchrun --nproc-per-node=2") + + dist.init_process_group("nccl") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + + try: + torch.manual_seed(0) + config = get_text_config() + config.hidden_size = 64 + config.linear_key_head_dim = 128 + config.linear_value_head_dim = 128 + config.linear_num_key_heads = 16 + config.linear_num_value_heads = 32 + reference = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) + context_parallel = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) + context_parallel.load_state_dict(reference.state_dict()) + context_parallel.set_context_parallel_attributes(dist.group.WORLD, world_size=2) + + hidden_states = torch.randn(1, 16, config.hidden_size, device="cuda", dtype=torch.bfloat16) + dist.broadcast(hidden_states, src=0) + cu_seqlens = torch.tensor([0, 5, 11, 16], device="cuda", dtype=torch.int32) + + reference_input = hidden_states.detach().clone().requires_grad_() + expected = reference( + reference_input, + cu_seqlens, + ) + output_gradient = torch.randn_like(expected) + dist.broadcast(output_gradient, src=0) + expected.backward(output_gradient) + + local_slice = slice(local_rank * 8, (local_rank + 1) * 8) + local_input = hidden_states[:, local_slice].detach().clone().requires_grad_() + actual = context_parallel( + local_input, + cu_seqlens, + ) + actual.backward(output_gradient[:, local_slice]) + + gathered = [torch.empty_like(actual) for _ in range(2)] + dist.all_gather(gathered, actual) + torch.testing.assert_close(torch.cat(gathered, dim=1), expected, rtol=3e-2, atol=3e-2) + torch.testing.assert_close( + local_input.grad, + reference_input.grad[:, local_slice], + rtol=5e-2, + atol=5e-2, + ) + finally: + dist.destroy_process_group() + + +def test_setup_model_cp_requires_hook_only_for_hybrid_models(): + class HybridLayer(torch.nn.Module): + layer_type = "linear_attention" + + class Inner: + layers = torch.nn.Sequential(torch.nn.Sequential(HybridLayer())) + + class HybridNoHookModel: + model = Inner() + + with pytest.raises(ValueError, match="set_context_parallel_attributes"): + setup_model_cp(HybridNoHookModel(), MagicMock(), cp_rank=0, cp_world_size=2) + + class SoftmaxOnlyModel: + pass + + setup_model_cp(SoftmaxOnlyModel(), MagicMock(), cp_rank=0, cp_world_size=2) + + +def test_ring_patches_flash_attention(): + from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention + + originals = {cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention)} + try: + substitute_ring_attn(process_group=MagicMock(), heads_k_stride=1) + assert Qwen3_5Attention._compute_attention is FlashAttention._compute_attention + assert Qwen3_5Attention._compute_attention is not originals[FlashAttention] + finally: + for cls, method in originals.items(): + cls._compute_attention = method + + +@pytest.mark.gpu +def test_forward_backward_and_packing(text_config): + prime_model = get_model(text_config) + fusions = ["qkv", "gate_up"] if isinstance(text_config, Qwen3_5MoeTextConfig) else ["qkv"] + apply_model_fusions(prime_model, fusions) + input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100), device="cuda") + position_ids = torch.arange(1, 101, device="cuda").unsqueeze(0) + prime_output = prime_model( + input_ids, + position_ids=position_ids, + seq_lens=torch.tensor([input_ids.shape[1]], device="cuda"), + ) + prime_output["logits"].sum().backward() + assert torch.isfinite(prime_output["logits"]).all() + assert torch.isfinite(prime_model.model.embed_tokens.weight.grad).all() + + packed_position_ids = torch.arange(1, 51, device="cuda").repeat(2).unsqueeze(0) + # Keep expert selection fixed to isolate packed sequence boundaries. + config = prime_model.config + routed_experts = None + if isinstance(config, Qwen3_5MoeTextConfig): + routed_experts = ( + torch.rand(1, 100, config.num_hidden_layers, config.num_experts, device="cuda") + .topk(config.num_experts_per_tok, dim=-1) + .indices + ) + with torch.no_grad(): + packed = prime_model( + input_ids, + position_ids=packed_position_ids, + seq_lens=torch.tensor([50, 50], device="cuda"), + routed_experts=routed_experts, + )["logits"] + unpacked = torch.cat( + [ + prime_model( + input_ids[:, start : start + 50], + position_ids=packed_position_ids[:, :50], + seq_lens=torch.tensor([50], device="cuda"), + routed_experts=routed_experts[:, start : start + 50] if routed_experts is not None else None, + )["logits"] + for start in (0, 50) + ], + dim=1, + ) + torch.testing.assert_close(packed, unpacked, atol=0.03, rtol=0.01) + + +@pytest.mark.gpu +def test_moe_router_replay(): + """When routed_experts are provided, the model uses them instead of computing routing.""" + prime_model = get_model(get_text_config(Qwen3_5MoeTextConfig)) + + with torch.device("cuda"): + input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) + position_ids = torch.arange(1, 101).unsqueeze(0) + + seq_lens = torch.tensor([input_ids.shape[1]], device="cuda") + out_normal = prime_model(input_ids, position_ids=position_ids, seq_lens=seq_lens) + + num_layers = prime_model.config.num_hidden_layers + topk = prime_model.config.num_experts_per_tok + routed_experts = torch.randint(0, prime_model.config.num_experts, (1, 100, num_layers, topk), device="cuda") + + prime_model.zero_grad() + out_replay = prime_model( + input_ids, + position_ids=position_ids, + routed_experts=routed_experts, + seq_lens=seq_lens, + ) + + assert out_replay["logits"].shape == out_normal["logits"].shape + + out_replay["logits"].sum().backward() + assert prime_model.model.embed_tokens.weight.grad is not None diff --git a/tests/unit/train/models/test_qwen3_5_moe_vlm.py b/tests/unit/train/models/qwen3_5/test_vlm.py similarity index 53% rename from tests/unit/train/models/test_qwen3_5_moe_vlm.py rename to tests/unit/train/models/qwen3_5/test_vlm.py index 6715c924e16..a61a4e6a713 100644 --- a/tests/unit/train/models/test_qwen3_5_moe_vlm.py +++ b/tests/unit/train/models/qwen3_5/test_vlm.py @@ -1,17 +1,16 @@ import pytest import torch -from prime_rl.trainer.model import can_reinit_empty_buffers +from prime_rl.configs.trainer import ModelConfig +from prime_rl.trainer.model import can_reinit_empty_buffers, resolve_auto_attn +from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head from prime_rl.trainer.models.qwen3_5 import ( - Qwen3_5ForCausalLM, Qwen3_5MoeConfig, Qwen3_5MoeTextConfig, Qwen3_5VisionConfig, ) -from prime_rl.utils.utils import default_dtype - -pytestmark = [pytest.mark.gpu] +from prime_rl.trainer.models.qwen3_5.rotary_embedding import build_qwen3_5_mrope_position_ids def get_vlm_config(): @@ -52,10 +51,20 @@ def get_vlm_config(): video_token_id=251, vision_start_token_id=252, vision_end_token_id=253, - attn_implementation="flash_attention_2", ) +def get_model(config, device="cuda"): + runtime_config = ModelConfig() + resolve_auto_attn(runtime_config) + with torch.device(device): + model = AutoModelForCausalLMPrimeRL.from_config( + config, attn_implementation=runtime_config.attn, dtype=torch.bfloat16 + ) + inject_prime_lm_head(model) + return model + + def get_image_inputs(config, device="cuda", dtype=torch.bfloat16): """Create minimal image inputs matching the vision config.""" vc = config.vision_config @@ -77,12 +86,11 @@ def get_seq_lens(input_ids: torch.Tensor) -> torch.Tensor: return torch.tensor([input_ids.shape[1]], device=input_ids.device) +@pytest.mark.gpu def test_vlm_forward(): """Custom VLM produces logits for both text-only and multimodal inputs.""" config = get_vlm_config() - with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5ForCausalLM(config) - inject_prime_lm_head(model) + model = get_model(config) vocab = config.text_config.vocab_size @@ -109,12 +117,11 @@ def test_vlm_forward(): assert out_mm["logits"].shape == (1, input_ids_mm.shape[1], vocab) +@pytest.mark.gpu def test_vlm_backward(): """Gradients flow through both vision scatter and text model.""" config = get_vlm_config() - with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5ForCausalLM(config) - inject_prime_lm_head(model) + model = get_model(config) pixel_values, image_grid_thw, n_img_tokens = get_image_inputs(config) text_part = torch.randint(0, 200, (1, 10), device="cuda") @@ -135,12 +142,11 @@ def test_vlm_backward(): assert model.model.visual.patch_embed.proj.weight.grad is not None +@pytest.mark.gpu def test_vlm_router_replay(): """routed_experts bypasses router computation in VLM multimodal forward.""" config = get_vlm_config() - with torch.device("cuda"), default_dtype(torch.bfloat16): - model = Qwen3_5ForCausalLM(config) - inject_prime_lm_head(model) + model = get_model(config) vocab = config.text_config.vocab_size pixel_values, image_grid_thw, n_img_tokens = get_image_inputs(config) @@ -168,11 +174,11 @@ def test_vlm_router_replay(): assert model.model.language_model.embed_tokens.weight.grad is not None +@pytest.mark.gpu def test_vlm_meta_device_and_buffer_reinit(): """Model can be created on meta device and buffers reinitialized.""" config = get_vlm_config() - with torch.device("meta"): - model = Qwen3_5ForCausalLM(config) + model = get_model(config, device="meta") assert can_reinit_empty_buffers(model) @@ -188,3 +194,143 @@ def test_vlm_meta_device_and_buffer_reinit(): assert lm_inv.abs().sum() > 0 assert lm_original_inv.abs().sum() > 0 assert vis_inv.abs().sum() > 0 + + +def test_qwen35_mrope_text_only_positions(): + input_ids = torch.tensor([[10, 11, 12, 13]]) + mm_token_type_ids = torch.zeros_like(input_ids) + + position_ids = build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=None, + spatial_merge_size=2, + seq_lens=torch.tensor([input_ids.shape[1]]), + ) + + expected = torch.arange(4).view(1, 1, -1).expand(3, 1, -1) + torch.testing.assert_close(position_ids, expected) + + +def test_qwen35_mrope_single_image_with_surrounding_text(): + input_ids = torch.tensor([[10, 11, 99, 99, 99, 99, 12]]) + mm_token_type_ids = torch.tensor([[0, 0, 1, 1, 1, 1, 0]]) + image_grid_thw = torch.tensor([[1, 4, 4]]) + + position_ids = build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + seq_lens=torch.tensor([input_ids.shape[1]]), + ) + + expected = torch.tensor( + [ + [[0, 1, 2, 2, 2, 2, 4]], + [[0, 1, 2, 2, 3, 3, 4]], + [[0, 1, 2, 3, 2, 3, 4]], + ] + ) + torch.testing.assert_close(position_ids, expected) + + +def test_qwen35_mrope_adjacent_images_consume_multiple_grids(): + input_ids = torch.tensor([[10, 99, 99, 99, 99, 99, 99, 99, 99, 11]]) + mm_token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 0]]) + image_grid_thw = torch.tensor([[1, 4, 4], [1, 4, 4]]) + + position_ids = build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + seq_lens=torch.tensor([input_ids.shape[1]]), + ) + + expected = torch.tensor( + [ + [[0, 1, 1, 1, 1, 3, 3, 3, 3, 5]], + [[0, 1, 1, 2, 2, 3, 3, 4, 4, 5]], + [[0, 1, 2, 1, 2, 3, 4, 3, 4, 5]], + ] + ) + torch.testing.assert_close(position_ids, expected) + + +def test_qwen35_mrope_packed_segments_reset_independently(): + input_ids = torch.tensor([[10, 99, 99, 99, 99, 11, 99, 99, 99, 99]]) + mm_token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 0, 1, 1, 1, 1]]) + image_grid_thw = torch.tensor([[1, 4, 4], [1, 4, 4]]) + seq_lens = torch.tensor([5, 5]) + + position_ids = build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + seq_lens=seq_lens, + ) + + expected_segment = torch.tensor( + [ + [0, 1, 1, 1, 1], + [0, 1, 1, 2, 2], + [0, 1, 2, 1, 2], + ] + ) + torch.testing.assert_close(position_ids[:, 0, :5], expected_segment) + torch.testing.assert_close(position_ids[:, 0, 5:], expected_segment) + + +def test_qwen35_mrope_packed_matches_standalone_segments(): + segment_input_ids = torch.tensor([[10, 99, 99, 99, 99, 11]]) + segment_token_types = torch.tensor([[0, 1, 1, 1, 1, 0]]) + image_grid_thw = torch.tensor([[1, 4, 4]]) + + standalone = build_qwen3_5_mrope_position_ids( + input_ids=segment_input_ids, + mm_token_type_ids=segment_token_types, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + seq_lens=torch.tensor([segment_input_ids.shape[1]]), + ) + packed = build_qwen3_5_mrope_position_ids( + input_ids=torch.cat([segment_input_ids, segment_input_ids], dim=1), + mm_token_type_ids=torch.cat([segment_token_types, segment_token_types], dim=1), + image_grid_thw=torch.cat([image_grid_thw, image_grid_thw], dim=0), + spatial_merge_size=2, + seq_lens=torch.tensor([segment_input_ids.shape[1], segment_input_ids.shape[1]]), + ) + + torch.testing.assert_close(packed[:, :, : segment_input_ids.shape[1]], standalone) + torch.testing.assert_close(packed[:, :, segment_input_ids.shape[1] :], standalone) + + +def test_qwen35_mrope_rejects_image_length_grid_mismatch(): + input_ids = torch.tensor([[10, 99, 99]]) + mm_token_type_ids = torch.tensor([[0, 1, 1]]) + image_grid_thw = torch.tensor([[1, 4, 4]]) + + with pytest.raises(ValueError, match="Image token group length"): + build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=2, + seq_lens=torch.tensor([input_ids.shape[1]]), + ) + + +def test_qwen35_mrope_rejects_video_tokens(): + input_ids = torch.tensor([[10, 99]]) + mm_token_type_ids = torch.tensor([[0, 2]]) + + with pytest.raises(ValueError, match="video MRoPE"): + build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=None, + spatial_merge_size=2, + seq_lens=torch.tensor([input_ids.shape[1]]), + ) diff --git a/tests/unit/train/models/test_qwen3_5.py b/tests/unit/train/models/test_qwen3_5.py deleted file mode 100644 index 30af50a74cd..00000000000 --- a/tests/unit/train/models/test_qwen3_5.py +++ /dev/null @@ -1,189 +0,0 @@ -import inspect -import os -from unittest.mock import MagicMock - -import pytest -import torch -import torch.distributed as dist - -from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn -from prime_rl.trainer.models.qwen3_5 import ( - Qwen3_5Config, - Qwen3_5ForCausalLM, - Qwen3_5Model, - Qwen3_5TextConfig, - Qwen3_5VisionConfig, -) -from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention -from prime_rl.trainer.models.qwen3_5.gated_delta_net import Qwen3_5GatedDeltaNet -from prime_rl.trainer.models.qwen3_5.norm import Qwen3_5RMSNorm -from prime_rl.utils.cp import setup_model_cp - - -def get_text_config(attn_impl: str = "flash_attention_2") -> Qwen3_5TextConfig: - config = Qwen3_5TextConfig( - vocab_size=128, - hidden_size=64, - intermediate_size=128, - num_hidden_layers=2, - layer_types=["linear_attention", "full_attention"], - num_attention_heads=4, - num_key_value_heads=2, - head_dim=16, - max_position_embeddings=128, - linear_key_head_dim=8, - linear_value_head_dim=8, - linear_num_key_heads=4, - linear_num_value_heads=8, - linear_conv_kernel_dim=4, - ) - config._attn_implementation = attn_impl - return config - - -def get_vlm_config(attn_impl: str = "flash_attention_2") -> Qwen3_5Config: - text_config = get_text_config(attn_impl) - vision_config = Qwen3_5VisionConfig( - depth=1, - hidden_size=64, - intermediate_size=128, - num_heads=4, - out_hidden_size=text_config.hidden_size, - ) - config = Qwen3_5Config( - text_config=text_config, - vision_config=vision_config, - image_token_id=120, - video_token_id=121, - vision_start_token_id=122, - vision_end_token_id=123, - ) - config._attn_implementation = attn_impl - return config - - -def test_qwen3_5_full_attention_uses_custom_class(): - config = get_text_config(attn_impl="flash_attention_3") - with torch.device("meta"): - model = Qwen3_5Model(config) - - assert isinstance(model.layers[1].self_attn, Qwen3_5Attention) - assert model.config._attn_implementation == "flash_attention_3" - assert "ALL_ATTENTION_FUNCTIONS" not in inspect.getsource(type(model.layers[1].self_attn).forward) - - -def test_qwen3_5_norms_remain_zero_centered_after_model_init(): - model = Qwen3_5ForCausalLM(get_text_config()) - - norms = [module for module in model.modules() if isinstance(module, Qwen3_5RMSNorm)] - assert norms - assert all(torch.count_nonzero(norm.weight) == 0 for norm in norms) - - -def test_qwen3_5_context_parallel_setup_chain_text_and_vlm(): - cp_group = MagicMock() - - text_model = Qwen3_5ForCausalLM(get_text_config()) - linear_layer = text_model.model.layers[0] - text_model.model.layers[0] = torch.nn.Sequential(linear_layer) - setup_model_cp(text_model, cp_group, cp_rank=1, cp_world_size=2) - assert text_model.model.context_parallel_group is cp_group - assert text_model.model.context_parallel_rank == 1 - assert text_model.model.context_parallel_world_size == 2 - assert linear_layer.linear_attn.context_parallel_group is cp_group - - vlm_config = get_vlm_config() - with torch.device("meta"): - vlm_model = Qwen3_5ForCausalLM(vlm_config) - assert vlm_model.model.language_model.layers[1].self_attn._flash_attn_version == 2 - assert vlm_model.model.visual.blocks[0].attn.attention_implementation == "flash_attention_2" - setup_model_cp(vlm_model, cp_group, cp_rank=0, cp_world_size=2) - assert vlm_model.model.language_model.context_parallel_group is cp_group - assert vlm_model.model.language_model.layers[0].linear_attn.context_parallel_world_size == 2 - - -@pytest.mark.gpu -def test_qwen3_5_gated_delta_net_context_parallel(): - if int(os.environ.get("WORLD_SIZE", 1)) != 2: - pytest.skip("run with torchrun --nproc-per-node=2") - - dist.init_process_group("nccl") - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - - try: - torch.manual_seed(0) - config = get_text_config() - config.linear_key_head_dim = 128 - config.linear_value_head_dim = 128 - config.linear_num_key_heads = 16 - config.linear_num_value_heads = 32 - reference = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) - context_parallel = Qwen3_5GatedDeltaNet(config).cuda().to(torch.bfloat16) - context_parallel.load_state_dict(reference.state_dict()) - context_parallel.set_context_parallel_attributes(dist.group.WORLD, world_size=2) - - hidden_states = torch.randn(1, 16, config.hidden_size, device="cuda", dtype=torch.bfloat16) - dist.broadcast(hidden_states, src=0) - cu_seqlens = torch.tensor([0, 5, 11, 16], device="cuda", dtype=torch.int32) - - reference_input = hidden_states.detach().clone().requires_grad_() - expected = reference( - reference_input, - cu_seqlens, - ) - output_gradient = torch.randn_like(expected) - dist.broadcast(output_gradient, src=0) - expected.backward(output_gradient) - - local_slice = slice(local_rank * 8, (local_rank + 1) * 8) - local_input = hidden_states[:, local_slice].detach().clone().requires_grad_() - actual = context_parallel( - local_input, - cu_seqlens, - ) - actual.backward(output_gradient[:, local_slice]) - - gathered = [torch.empty_like(actual) for _ in range(2)] - dist.all_gather(gathered, actual) - torch.testing.assert_close(torch.cat(gathered, dim=1), expected, rtol=3e-2, atol=3e-2) - torch.testing.assert_close( - local_input.grad, - reference_input.grad[:, local_slice], - rtol=5e-2, - atol=5e-2, - ) - finally: - dist.destroy_process_group() - - -def test_setup_model_cp_requires_hook_only_for_hybrid_models(): - class HybridLayer(torch.nn.Module): - layer_type = "linear_attention" - - class Inner: - layers = torch.nn.Sequential(torch.nn.Sequential(HybridLayer())) - - class HybridNoHookModel: - model = Inner() - - with pytest.raises(ValueError, match="set_context_parallel_attributes"): - setup_model_cp(HybridNoHookModel(), MagicMock(), cp_rank=0, cp_world_size=2) - - class SoftmaxOnlyModel: - pass - - setup_model_cp(SoftmaxOnlyModel(), MagicMock(), cp_rank=0, cp_world_size=2) - - -def test_qwen3_5_ring_patches_dense_flash_attention(): - from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention - - originals = {cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention)} - try: - substitute_ring_attn(process_group=MagicMock(), heads_k_stride=1) - assert Qwen3_5Attention._compute_attention is FlashAttention._compute_attention - assert Qwen3_5Attention._compute_attention is not originals[FlashAttention] - finally: - for cls, method in originals.items(): - cls._compute_attention = method diff --git a/tests/unit/train/models/test_qwen3_5_moe.py b/tests/unit/train/models/test_qwen3_5_moe.py deleted file mode 100644 index 66885138903..00000000000 --- a/tests/unit/train/models/test_qwen3_5_moe.py +++ /dev/null @@ -1,175 +0,0 @@ -import pytest -import torch - -from prime_rl.configs.trainer import ModelConfig -from prime_rl.trainer.model import resolve_auto_attn -from prime_rl.trainer.models import AutoModelForCausalLMPrimeRL -from prime_rl.trainer.models.fusions import apply_model_fusions -from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head -from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeTextConfig -from prime_rl.utils.cp import setup_model_cp -from prime_rl.utils.utils import default_dtype - -pytestmark = [pytest.mark.gpu] - - -def get_model(): - runtime_config = ModelConfig() - resolve_auto_attn(runtime_config) - config = Qwen3_5MoeTextConfig( - vocab_size=256, - hidden_size=256, - num_hidden_layers=4, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=64, - moe_intermediate_size=128, - shared_expert_intermediate_size=128, - num_experts=8, - num_experts_per_tok=2, - max_position_embeddings=512, - rms_norm_eps=1e-6, - linear_conv_kernel_dim=4, - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_num_key_heads=4, - linear_num_value_heads=8, - ) - with torch.device("cuda"): - prime_model = AutoModelForCausalLMPrimeRL.from_config( - config, attn_implementation=runtime_config.attn, dtype=torch.bfloat16 - ) - inject_prime_lm_head(prime_model, chunk_size=None) - return prime_model - - -def test_qwen3_5_moe(): - prime_model = get_model() - apply_model_fusions(prime_model, ["qkv", "gate_up"]) - input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100), device="cuda") - position_ids = torch.arange(1, 101, device="cuda").unsqueeze(0) - prime_output = prime_model( - input_ids, - position_ids=position_ids, - seq_lens=torch.tensor([input_ids.shape[1]], device="cuda"), - ) - prime_output["logits"].sum().backward() - assert torch.isfinite(prime_output["logits"]).all() - assert torch.isfinite(prime_model.model.embed_tokens.weight.grad).all() - - packed_position_ids = torch.arange(1, 51, device="cuda").repeat(2).unsqueeze(0) - # Keep expert selection fixed to isolate packed sequence boundaries. - config = prime_model.config - routed_experts = ( - torch.rand(1, 100, config.num_hidden_layers, config.num_experts, device="cuda") - .topk(config.num_experts_per_tok, dim=-1) - .indices - ) - with torch.no_grad(): - packed = prime_model( - input_ids, - position_ids=packed_position_ids, - seq_lens=torch.tensor([50, 50], device="cuda"), - routed_experts=routed_experts, - )["logits"] - unpacked = torch.cat( - [ - prime_model( - input_ids[:, start : start + 50], - position_ids=packed_position_ids[:, :50], - seq_lens=torch.tensor([50], device="cuda"), - routed_experts=routed_experts[:, start : start + 50], - )["logits"] - for start in (0, 50) - ], - dim=1, - ) - torch.testing.assert_close(packed, unpacked, atol=0.03, rtol=0.01) - - -def test_qwen3_5_moe_router_replay(): - """When routed_experts are provided, the model uses them instead of computing routing.""" - prime_model = get_model() - - with torch.device("cuda"), default_dtype(torch.bfloat16): - input_ids = torch.randint(0, prime_model.config.vocab_size, (1, 100)) - position_ids = torch.arange(1, 101).unsqueeze(0) - - seq_lens = torch.tensor([input_ids.shape[1]], device="cuda") - out_normal = prime_model(input_ids, position_ids=position_ids, seq_lens=seq_lens) - - num_layers = prime_model.config.num_hidden_layers - topk = prime_model.config.num_experts_per_tok - routed_experts = torch.randint(0, prime_model.config.num_experts, (1, 100, num_layers, topk), device="cuda") - - prime_model.zero_grad() - out_replay = prime_model( - input_ids, - position_ids=position_ids, - routed_experts=routed_experts, - seq_lens=seq_lens, - ) - - assert out_replay["logits"].shape == out_normal["logits"].shape - - out_replay["logits"].sum().backward() - assert prime_model.model.embed_tokens.weight.grad is not None - - -def test_qwen3_5_moe_cp_patching(): - from unittest.mock import MagicMock - - from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention - from prime_rl.trainer.models.layers.attn import FlashAttention, substitute_ring_attn - from prime_rl.trainer.models.qwen3_5.attention import Qwen3_5Attention - - originals = {cls: cls._compute_attention for cls in (FlashAttention, AfmoeFlashAttention)} - try: - mock_group = MagicMock() - substitute_ring_attn(process_group=mock_group, heads_k_stride=1) - assert Qwen3_5Attention._compute_attention is FlashAttention._compute_attention - assert Qwen3_5Attention._compute_attention is not originals[FlashAttention] - finally: - for cls, method in originals.items(): - cls._compute_attention = method - - -def test_qwen3_5_moe_context_parallel_setup_hook(): - from unittest.mock import MagicMock - - config = Qwen3_5MoeTextConfig( - vocab_size=128, - hidden_size=64, - num_hidden_layers=2, - layer_types=["linear_attention", "full_attention"], - num_attention_heads=4, - num_key_value_heads=2, - head_dim=16, - moe_intermediate_size=64, - shared_expert_intermediate_size=64, - num_experts=4, - num_experts_per_tok=2, - max_position_embeddings=128, - linear_conv_kernel_dim=4, - linear_key_head_dim=8, - linear_value_head_dim=8, - linear_num_key_heads=4, - linear_num_value_heads=8, - ) - with torch.device("meta"): - model = Qwen3_5ForCausalLM(config) - - linear_layer = model.model.layers[0] - model.model.layers[0] = torch.nn.Sequential(linear_layer) - cp_group = MagicMock() - setup_model_cp(model, cp_group, cp_rank=1, cp_world_size=2) - - assert model.model.context_parallel_group is cp_group - assert model.model.context_parallel_rank == 1 - assert model.model.context_parallel_world_size == 2 - assert linear_layer.linear_attn.context_parallel_group is cp_group - assert linear_layer.linear_attn.context_parallel_world_size == 2 - - -if __name__ == "__main__": - test_qwen3_5_moe() diff --git a/tests/unit/train/models/test_qwen3_5_moe_mrope.py b/tests/unit/train/models/test_qwen3_5_moe_mrope.py deleted file mode 100644 index eb2366c6228..00000000000 --- a/tests/unit/train/models/test_qwen3_5_moe_mrope.py +++ /dev/null @@ -1,146 +0,0 @@ -import pytest -import torch - -from prime_rl.trainer.models.qwen3_5.rotary_embedding import ( - build_qwen3_5_mrope_position_ids, -) - - -def test_qwen35_mrope_text_only_positions(): - input_ids = torch.tensor([[10, 11, 12, 13]]) - mm_token_type_ids = torch.zeros_like(input_ids) - - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=None, - spatial_merge_size=2, - seq_lens=torch.tensor([input_ids.shape[1]]), - ) - - expected = torch.arange(4).view(1, 1, -1).expand(3, 1, -1) - torch.testing.assert_close(position_ids, expected) - - -def test_qwen35_mrope_single_image_with_surrounding_text(): - input_ids = torch.tensor([[10, 11, 99, 99, 99, 99, 12]]) - mm_token_type_ids = torch.tensor([[0, 0, 1, 1, 1, 1, 0]]) - image_grid_thw = torch.tensor([[1, 4, 4]]) - - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - seq_lens=torch.tensor([input_ids.shape[1]]), - ) - - expected = torch.tensor( - [ - [[0, 1, 2, 2, 2, 2, 4]], - [[0, 1, 2, 2, 3, 3, 4]], - [[0, 1, 2, 3, 2, 3, 4]], - ] - ) - torch.testing.assert_close(position_ids, expected) - - -def test_qwen35_mrope_adjacent_images_consume_multiple_grids(): - input_ids = torch.tensor([[10, 99, 99, 99, 99, 99, 99, 99, 99, 11]]) - mm_token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 0]]) - image_grid_thw = torch.tensor([[1, 4, 4], [1, 4, 4]]) - - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - seq_lens=torch.tensor([input_ids.shape[1]]), - ) - - expected = torch.tensor( - [ - [[0, 1, 1, 1, 1, 3, 3, 3, 3, 5]], - [[0, 1, 1, 2, 2, 3, 3, 4, 4, 5]], - [[0, 1, 2, 1, 2, 3, 4, 3, 4, 5]], - ] - ) - torch.testing.assert_close(position_ids, expected) - - -def test_qwen35_mrope_packed_segments_reset_independently(): - input_ids = torch.tensor([[10, 99, 99, 99, 99, 11, 99, 99, 99, 99]]) - mm_token_type_ids = torch.tensor([[0, 1, 1, 1, 1, 0, 1, 1, 1, 1]]) - image_grid_thw = torch.tensor([[1, 4, 4], [1, 4, 4]]) - seq_lens = torch.tensor([5, 5]) - - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - seq_lens=seq_lens, - ) - - expected_segment = torch.tensor( - [ - [0, 1, 1, 1, 1], - [0, 1, 1, 2, 2], - [0, 1, 2, 1, 2], - ] - ) - torch.testing.assert_close(position_ids[:, 0, :5], expected_segment) - torch.testing.assert_close(position_ids[:, 0, 5:], expected_segment) - - -def test_qwen35_mrope_packed_matches_standalone_segments(): - segment_input_ids = torch.tensor([[10, 99, 99, 99, 99, 11]]) - segment_token_types = torch.tensor([[0, 1, 1, 1, 1, 0]]) - image_grid_thw = torch.tensor([[1, 4, 4]]) - - standalone = build_qwen3_5_mrope_position_ids( - input_ids=segment_input_ids, - mm_token_type_ids=segment_token_types, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - seq_lens=torch.tensor([segment_input_ids.shape[1]]), - ) - packed = build_qwen3_5_mrope_position_ids( - input_ids=torch.cat([segment_input_ids, segment_input_ids], dim=1), - mm_token_type_ids=torch.cat([segment_token_types, segment_token_types], dim=1), - image_grid_thw=torch.cat([image_grid_thw, image_grid_thw], dim=0), - spatial_merge_size=2, - seq_lens=torch.tensor([segment_input_ids.shape[1], segment_input_ids.shape[1]]), - ) - - torch.testing.assert_close(packed[:, :, : segment_input_ids.shape[1]], standalone) - torch.testing.assert_close(packed[:, :, segment_input_ids.shape[1] :], standalone) - - -def test_qwen35_mrope_rejects_image_length_grid_mismatch(): - input_ids = torch.tensor([[10, 99, 99]]) - mm_token_type_ids = torch.tensor([[0, 1, 1]]) - image_grid_thw = torch.tensor([[1, 4, 4]]) - - with pytest.raises(ValueError, match="Image token group length"): - build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=2, - seq_lens=torch.tensor([input_ids.shape[1]]), - ) - - -def test_qwen35_mrope_rejects_video_tokens(): - input_ids = torch.tensor([[10, 99]]) - mm_token_type_ids = torch.tensor([[0, 2]]) - - with pytest.raises(ValueError, match="video MRoPE"): - build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=None, - spatial_merge_size=2, - seq_lens=torch.tensor([input_ids.shape[1]]), - )