feat(models): dense Qwen3.5 custom attention#2984
Conversation
Implement dense Qwen3.5 (Qwen3_5ForCausalLM / Qwen3_5Model) on the PrimeRL attention modules, reusing the Qwen3.5-MoE building blocks (GatedDeltaNet, gated SDPA/flash attention, RMSNorm, rotary embedding) so ring/ulysses CP substitution patches the dense and MoE models consistently: - New models/qwen3_5 package: decoder layer dispatching on layer_types (linear_attention / full_attention), flash varlen cu_seqlens from position_ids, identity state-dict conversion (dense HF layout matches PrimeRL layout). - Shared normalize_qwen3_5_attn_implementation on the MoE module (eager→sdpa, kernels-community/vllm-flash-attn3→flash_attention_3), consumed by both models' attention factories and _check_and_adjust_attn_implementation. - Register Qwen3_5TextConfig / Qwen3_5ForCausalLM in AutoConfig and the custom causal-LM mapping; hook dense Qwen3_5GatedFlashAttention into substitute_ring_attn / substitute_ulysses_attn. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 0637368. Configure here.
| return BaseModelOutputWithPast(last_hidden_state=hidden_states) | ||
|
|
||
|
|
||
| class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): |
There was a problem hiding this comment.
Missing KL validation table
Medium Severity
This PR adds a new custom Qwen3_5ForCausalLM, but the description does not include the required mean KL mismatch table (20 steps, math environment, batch_size=64, all entries below 0.015). Narrative KL figures are not a substitute for that table.
Triggered by project rule: BugBot Instructions
Reviewed by Cursor Bugbot for commit 0637368. Configure here.
| torch._dynamo.mark_dynamic(cu_seqlens, 0) | ||
| else: | ||
| cu_seqlens = None | ||
| max_seqlen = None |
There was a problem hiding this comment.
3D positions break cu_seqlens
Medium Severity
When flash attention is enabled, Qwen3_5Model builds cu_seqlens via get_cu_seqlens_from_position_ids on raw position_ids. For 3D MRoPE inputs, flattening yields wrong segment boundaries, unlike the MoE model’s explicit 3D handling.
Reviewed by Cursor Bugbot for commit 0637368. Configure here.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0637368614
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _CUSTOM_CAUSAL_LM_MAPPING.register(MiniMaxM2Config, MiniMaxM2ForCausalLM, exist_ok=True) | ||
| _CUSTOM_CAUSAL_LM_MAPPING.register(NemotronHConfig, NemotronHForCausalLM, exist_ok=True) | ||
| _CUSTOM_CAUSAL_LM_MAPPING.register(Qwen3MoeConfig, Qwen3MoeForCausalLM, exist_ok=True) | ||
| _CUSTOM_CAUSAL_LM_MAPPING.register(Qwen3_5TextConfig, Qwen3_5ForCausalLM, exist_ok=True) |
There was a problem hiding this comment.
Registering only Qwen3_5TextConfig leaves the shipped Qwen/Qwen3.5-* configs, whose top-level config is model_type='qwen3_5' with a nested text_config, outside the custom path. get_model() classifies those configs as VLMs and consults _CUSTOM_VLM_MAPPING, which still has no qwen3_5 entry, so impl='auto' selects the HF implementation and even impl='custom' goes through AutoModelForImageTextToText instead of this new PrimeRL attention implementation. This means the new dense Qwen3.5 custom model is not used for the primary Qwen3.5 checkpoints unless the checkpoint has been rewritten as a qwen3_5_text config.
Useful? React with 👍 / 👎.
| else: | ||
| cu_seqlens = None | ||
| max_seqlen = None |
There was a problem hiding this comment.
When dense Qwen3.5 is run with attn='sdpa'/'eager' on normal packed RL batches, this branch drops the boundary information encoded by reset position_ids. The model still has linear_attention layers, and Qwen3_5MoeGatedDeltaNet only resets its causal conv/DeltaNet state when cu_seqlens is supplied, so packed examples can leak state into each other; the SDPA full-attention layers also get no mask to prevent cross-example attention. Please either compute/use packed boundaries for this path or reject packed non-flash Qwen3.5 runs.
Useful? React with 👍 / 👎.
|
Superseded by #2943; keeping the existing stacked PR and rebasing it onto main after its prerequisite merges. |
2 participants
Summary
Qwen3_5ForCausalLMimplementation using PrimeRL attention modulesValidation
uv run --no-sync ruff check ...— passeduv run --no-sync pytest tests/unit/train/models/test_qwen3_5.py -q --tb=short— blocked during collection becauseflash_attnis not installed in this environmentThe custom implementation previously matched the HF path to 1e-4 loss at 0.8B and 9B in the linked stacked-PR runs, with mismatch KL 0.0003 across 20 RL steps.
Made with Cursor