[lora] Support GDN in_proj_ba adapters for Qwen3.5

[yushengsu-thu]

Motivation

Qwen3.5's GatedDeltaNet (GDN) has two packed input projections: in_proj_qkvz (already LoRA-supported) and in_proj_ba — the delta-rule β/α gate projections. LoRA adapter weights covering in_proj_ba are currently silently dropped at load time (not in _KNOWN_LORA_TARGET_MODULES, warning only). An adapter trained against the fused megatron-side in_proj (Megatron-Bridge exports it split 4-way as in_proj_qkv/z/b/a) therefore serves base weights for the β/α slices — a quiet train/serve mismatch for RL.

Modifications

• models/qwen3_5.py: add in_proj_ba to supported_lora_modules + get_hidden_dim (2 * linear_num_value_heads)
• lora/utils.py: register in_proj_ba in _KNOWN_LORA_TARGET_MODULES, get_stacked_multiply (=2), and in_proj_b/in_proj_a → in_proj_ba name normalization
• lora/lora.py: new _normalize_in_proj_ba (split b+a concat, merged repeat×2); extend _normalize_in_proj_qkvz to accept the 2-way in_proj_qkv + in_proj_z layout emitted by Megatron-Bridge adapter export (shared lora_A repeated ×3 for the q|k|v slices) — previously only the 4-way PEFT layout and pre-merged form were handled

MergedColumnParallelLinearWithLoRA wraps the base layer generically (2-equal-slice path, same as gate_up) — no layer/mem-pool/manager changes needed.

Validation

• Normalization unit checks: MB 2-way qkvz, b+a split, merged ba, and 4-way PEFT regression — all pass
• End-to-end RL (miles) on Qwen/Qwen3.5-35B-A3B, 8×H200 TP2/EP8 colocate, LoRA on attention + MoE experts + shared expert + GDN in_proj/out_proj: rollout-vs-train logprob_abs_diff = 0.0100, zero adapter weights skipped

---

CI States

Latest PR Test (Base): ❌ Run #28927579933
Latest PR Test (Extra): ❌ Run #28978966210

Model coverage

Also covers Qwen3.6 with no additional changes: Qwen3.6 checkpoints declare the same Qwen3_5(Moe)ForConditionalGeneration architectures / qwen3_5* model_type, so they dispatch to qwen3_5.py and pick up this LoRA support directly. Verified with a 10-step LoRA RL run on Qwen/Qwen3.6-35B-A3B (full target set incl. GDN in_proj_qkvz/in_proj_ba/out_proj): rollout↔train logprob_abs_diff 0.010 at step 0, KL ≤ 0.004, reward 0.19 → 0.47, zero adapter weights skipped.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[gemini-code-assist]

Code Review

This pull request adds support for normalizing and stacking in_proj_ba weights for GatedDeltaNet (GDN) layers, such as those in Qwen 3.5, and adds support for 2-way split in_proj_qkv and in_proj_z weights. The changes span LoRA weight normalization, utility functions, and the Qwen 3.5 model definition. Feedback suggests updating get_hidden_dim in qwen3_5.py to use config.linear_num_value_heads_cpu when running on CPU to avoid shape mismatches when loading LoRA adapters.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[yushengsu-thu]

Additional end-to-end evidence: 10-step LoRA RL run (miles, GRPO/gsm8k) on Qwen/Qwen3.6-35B-A3B (same Qwen3_5* architecture/classes) with the full target set including GDN in_proj_qkvz + in_proj_ba + out_proj: step-0 rollout↔train logprob_abs_diff 0.010, KL ≤ 0.004 across 10 steps, reward 0.19 → 0.47, zero adapter weights skipped.