Skip to content

QK logits clipping (non-split version) - #1929

Merged
BoxiangW merged 40 commits into
NVIDIA:mainfrom
BoxiangW:boxiangw/muon-clip
Dec 10, 2025
Merged

QK logits clipping (non-split version)#1929
BoxiangW merged 40 commits into
NVIDIA:mainfrom
BoxiangW:boxiangw/muon-clip

Conversation

@BoxiangW

@BoxiangW BoxiangW commented Oct 24, 2025

Copy link
Copy Markdown
Contributor

Added MLA and MHA(GQA) clipping support

@BoxiangW
Add qk-clip
0ef342c 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
w/o EP working
a5d2ea8 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Lint
978b99b 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix MLA usage and add logging for max_attention_score
b77a6d6 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Removed Attention MuonClip since it is not correct. Added MLA MuonClip
a83b68c 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Add TE bug WAR and make sure MLA works
eae6a5a 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix qkclip issue
c370cdd 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix bug
6f8f77d 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix TP usage
527075f 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix Lint
3cebc1f 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Fix import and error log
e3c468f 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
lint
ef995ea 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Address comments
7f5f1e2 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Lint
65fc829 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Added GQA QK Clipping
75ee38c 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Lint
a1c0ff6 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
BoxiangW requested review from a team as code owners October 24, 2025 22:00
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2025

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@BoxiangW BoxiangW self-assigned this Oct 24, 2025
@BoxiangW

BoxiangW commented Oct 24, 2025

Copy link
Copy Markdown
Contributor Author

TE's NVIDIA/TransformerEngine#2195 (2.9.0) is needed for this PR

@BoxiangW
Remove comment
762bdb5 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Rename max_score to max_logit
ba5c073 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Change name
7917e68 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@skyw

skyw commented Oct 28, 2025

Copy link
Copy Markdown
Contributor

TE's NVIDIA/TransformerEngine#2195 (2.9.0) is needed for this PR

It has been merged.

@BoxiangW BoxiangW added Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. Run tests labels Oct 31, 2025
@BoxiangW BoxiangW added this to the Core 0.15 milestone Oct 31, 2025
@BoxiangW

BoxiangW commented Nov 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 7917e68

@BoxiangW

BoxiangW commented Nov 10, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test b63c573

@BoxiangW BoxiangW mentioned this pull request Nov 10, 2025
Merged
6 tasks
@BoxiangW
Add te version checks into PR
bdda5e9 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
yaoyu-33
yaoyu-33 reviewed Nov 11, 2025
View reviewed changes
Comment thread megatron/core/extensions/transformer_engine.py
if config.qk_clip and is_te_min_version("2.9.0"):
# TE 2.9.0 introduces return_max_logit for qk-clip getting the max attention logits
extra_kwargs["return_max_logit"] = True
self.current_max_attn_logits = None

@yaoyu-33 yaoyu-33 Nov 11, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what if not is_te_min_version but config.qk_clip? might raise an error

@BoxiangW BoxiangW Nov 24, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right now it will raise an error if te version is wrong

yaoyu-33
yaoyu-33 reviewed Nov 11, 2025
View reviewed changes
Comment thread megatron/core/optimizer/qk_clip.py Outdated
"""

log_max_attention_logit = 0
for model_chunk in model:

@yaoyu-33 yaoyu-33 Nov 11, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i dont remember, is model always a list?

@BoxiangW BoxiangW Nov 12, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

According to https://github.com/NVIDIA/Megatron-LM/pull/1929/files#diff-7fc2a91343e49be3c1705b8237bbe437029132d8826e5c963e1b0b84d457bc38L1250. I think it is true for MLM, is it true for MBridge as well?

@BoxiangW
Address comments
bd0436b 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Address comments
646aea5 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
DP all reduce and switch to mul_ inplace op
4d16029 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW
Update both main_params and non
2b890d0
@BoxiangW
Address comments
95fdba3 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW

BoxiangW commented Nov 20, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 95fdba3

@deepakn94

deepakn94 commented Nov 30, 2025

Copy link
Copy Markdown
Contributor

Can we re-name this PR? It should just be "QK logits clipping" or something similar?

jaredcasper
jaredcasper reviewed Dec 2, 2025
View reviewed changes
Comment thread megatron/core/optimizer/qk_clip.py Outdated

import torch

from megatron.core import mpu

@jaredcasper jaredcasper Dec 2, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should use parrallel_state, mpu is a legacy name.

@BoxiangW BoxiangW Dec 3, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching this, changed to parrallel_state instead

Comment thread megatron/core/optimizer/qk_clip.py Outdated
torch.distributed.all_reduce(
transformer_layer.self_attention.core_attention.current_max_attn_logits,
op=torch.distributed.ReduceOp.MAX,
group=mpu.get_data_parallel_group(with_context_parallel=True),

@jaredcasper jaredcasper Dec 2, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any way to get this group in via args? We are trying to move away from global process groups.

@BoxiangW
Change to parrallel_state
6562a52 
Signed-off-by: Boxiang Wang <boxiangw@nvidia.com>
@BoxiangW

BoxiangW commented Dec 3, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 6562a52

@BoxiangW

BoxiangW commented Dec 8, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test a2e7407

@BoxiangW

BoxiangW commented Dec 9, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test d83c38e

@BoxiangW

BoxiangW commented Dec 9, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test dd82679

@BoxiangW

BoxiangW commented Dec 9, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test eed3c2d

@BoxiangW

BoxiangW commented Dec 9, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 99c7631

@BoxiangW

BoxiangW commented Dec 10, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 24e6c88

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

Reviewers

@jaredcasper jaredcasper jaredcasper approved these changes
@ko3n1g ko3n1g ko3n1g approved these changes
@yaoyu-33 yaoyu-33 yaoyu-33 approved these changes

Assignees

@BoxiangW BoxiangW

Labels

complexity: medium Final Review PR is in the "final review" stage Run tests

Projects

None yet

Milestone

Core 0.15

Development

Successfully merging this pull request may close these issues.

7 participants

@BoxiangW @skyw @deepakn94 @jaredcasper @ko3n1g @yaoyu-33 @ericharper