Add MLA support to Muon - #5015
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 3cb3990 |
|
@Phlip79 , this code is not being covered by an expert group. Can we cover it? |
wujingyue
left a comment
wujingyue
left a comment
There was a problem hiding this comment.
Getting some expert reviews can help too.
|
/ok to test 8df9baa |
|
Reviews (1): Last reviewed commit: "Improve documentation" | Re-trigger Greptile |
|
/ok to test 1e02dc5 |
FDecaYed
left a comment
FDecaYed
left a comment
There was a problem hiding this comment.
the code change itself looks ok, but at high level I feel splitting fused should be a general thing we support with a refined API. we cannot sustain keep adding new cases as the workload coming. Ideally we should generalize this.
Also kinda related, we had the change in Emerging-optimizers to support batched muon processing. instead of supporting arbitrary batch of different shaped weights, it only support batch of weight all with same shape, and it could be useful for cases where same shape weights are fused in the first place
Specific for that, the plan is for a version of grouped linear has 3d batch shaped weights in TE or other places. |
Signed-off-by: janEbert <janpabloe@nvidia.com>
In the future, we could create an actual enum for this. Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Avoid duplicate meanings of `None`: 1. `None` implies that `split_qkv = False`. 2. `None` implies that the parameter should not be split. The new version separates these cases clearly, which requires introducing additional nesting. Signed-off-by: janEbert <janpabloe@nvidia.com>
|
/ok to test 08c7b4c |
|
@janEbert Hi, any updates on this PR? Are we going to merge it? And what's the plan for the KV down proj split? Thanks! |
|
Hey @buptzyb, there are sadly still missing expert reviews on this PR. As for the down-proj split, I didn't want to work on it before this PR is actually merged. |
|
Sounds good, but what do you mean by "merge"? It feels like both our PRs don't really address the design point, if I see correctly. |
Sorry by merge, I meant unify the PRs into a common interface that is more general than either of the 2 PRs |
|
I see, yeah, sounds good! Do you (or some of the others already in the discussion) already have something in mind? If not, I can give it some thought. |
No one else is working on it, though @philipcmonk might have thoughts. There's also an interaction with gated attention, where we split into |
11 participants
What does this PR do ?
Add MLA support to Muon by extending the existing metadata tagging scheme for splitting the MLA up-projections (and the sole Q-projection when
q_lora_rank=None).The down-projs are deliberately left out of the scope of this PR and will be included in a follow-up PR with the down-proj splitting being off-by-default behind an additional feature flag.
Issue tracking
Fix #4091.
Rebased on top of #4728.