Skip to content

feat(muon): split fused QKV projections - #2030

Open
RangiLyu wants to merge 2 commits into
InternLM:mainfrom
RangiLyu:lcq/muon-qkv-split
Open

feat(muon): split fused QKV projections#2030
RangiLyu wants to merge 2 commits into
InternLM:mainfrom
RangiLyu:lcq/muon-qkv-split

Conversation

@RangiLyu

Copy link
Copy Markdown
Collaborator

Motivation

Muon should orthogonalize the logical projections inside a fused QKV weight independently.

Treating a fused QKV weight as a single matrix couples unrelated Q, K, and V updates during Newton-Schulz orthogonalization and applies learning-rate scaling based on the combined matrix shape.

This PR uses the get_muon_split_sizes() interface introduced in #2001 to expose the logical row blocks of fused QKV parameters.

Changes

  • Add Muon split metadata for GatedDeltaNet in_proj_qkv.weight:
    • Q: key_dim
    • K: key_dim
    • V: value_dim
  • Add Muon split metadata for Qwen3-VL Vision qkv.weight:
    • Q: dim
    • K: dim
    • V: dim
  • Keep the forward computation and parameter layout unchanged.

The Muon optimizer configuration discovers these mappings through the module-level get_muon_split_sizes() interface and processes each logical projection independently.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant