Port Multi-Latent Attention to HybridModel - #4452
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. |
3f9af18 to
bded2ae
Compare
|
/ok to test bded2ae |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ok to test a46c197 |
a46c197 to
6ae9873
Compare
|
/ok to test 6ae9873 |
Phlip79
left a comment
Phlip79
left a comment
There was a problem hiding this comment.
This MLA spec ignores --qk-layernorm and --qk-l2-norm and doesn't expose --mla-down-proj-fusion like GPTModel does.
6ae9873 to
aa68627
Compare
Addressed; the logic for QK layernorm is quite complex, so this blew the PR up. I could put this into another PR to make this one easier to review. |
aa68627 to
432d3d6
Compare
|
/ok to test 5336c52 |
5336c52 to
0f9d8df
Compare
We obviously want to return an instance here, not a type. Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
It's more likely that we encounter non-fused custom layers, so check in this direction instead. Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Previously, we rejected anything unusual. Not, we allow for unconsidered specs to be used, granting more flexibility in exchange for less certainty that the implementation isn't broken. Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
This reverts commit 7697f16. Signed-off-by: janEbert <janpabloe@nvidia.com>
Also clean up tests a bit. (Especially redundant imports.) Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
It's a lot of boilerplate code, so we try to keep `multi_latent_attention.py` clean by having the resolver functionality in the new module. We try to mark it "private" by not explicitly exporting anything from it. Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
Signed-off-by: janEbert <janpabloe@nvidia.com>
|
/ok to test 19af620 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30034259054 |
5 participants
Add Multi-Latent Attention (MLA) support to
HybridModel(similar to DSA support in #3553).Also includes a minor rename for Mamba→Hybrid that wasn't caught in the earlier large PRs.
Note that the QK norm spec resolution functionality also catches errors that the existing implementation doesn't catch when the spec is misconfigured, therefore we need to make some additional changes.