Skip to content

Fix Polar Express Coefficient Iterator for >8 step polar express - #303

Merged
skyw merged 6 commits into
NVIDIA-NeMo:mainfrom
mkhona-nvidia:mkhona/muon-norm-1p01
Oct 12, 2026
Merged

skyw merged 6 commits into
NVIDIA-NeMo:mainfrom
mkhona-nvidia:mkhona/muon-norm-1p01

Conversation

@mkhona-nvidia

Copy link
Copy Markdown
Contributor

Some fixes for polar express NS

Signed-off-by: mkhona <mkhona@nvidia.com>
Signed-off-by: mkhona <mkhona@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 23, 2026

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.

@greptile-apps

greptile-apps Bot commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues remain.

Summary

The PR adds a guarded coefficient schedule for Polar Express runs longer than eight steps.

  • Polar Express repeats a guarded coefficient before its final step.
  • Distributed normalization tests cover tiny inputs and both modes.

Reviews (2) · Last reviewed commit: "Round guarded Polar Express coefficients..." · Reviewed by Greptile

Comment thread tests/test_muon_utils.py Outdated
Comment thread emerging_optimizers/orthogonalized_optimizers/muon_utils.py Outdated
Comment thread emerging_optimizers/orthogonalized_optimizers/muon_utils.py Outdated
@mkhona-nvidia mkhona-nvidia added the wontfix This will not be worked on label Sep 24, 2026
@mkhona-nvidia
mkhona-nvidia marked this pull request as draft September 24, 2026 02:57
@skyw skyw removed the wontfix This will not be worked on label Sep 24, 2026
@skyw skyw closed this Sep 24, 2026
Signed-off-by: mkhona <mkhona@nvidia.com>
@mkhona-nvidia mkhona-nvidia reopened this Oct 12, 2026
Comment thread emerging_optimizers/orthogonalized_optimizers/muon_utils.py Outdated
Comment thread emerging_optimizers/orthogonalized_optimizers/muon_utils.py Outdated
Comment thread tests/test_muon_utils.py Outdated
Signed-off-by: mkhona <mkhona@nvidia.com>
skyw
skyw previously approved these changes Oct 12, 2026

@skyw skyw left a comment

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.

This is a bit out of original design and doesn't look clean.

Any future change we should consider redesign the whole thing.

@skyw

skyw commented Oct 12, 2026

Copy link
Copy Markdown
Contributor

@greptile

@mkhona-nvidia mkhona-nvidia changed the title Mkhona/muon norm 1p01 Fix Polar Express Coefficient Iterator for >8 step polar express Oct 12, 2026
Signed-off-by: mkhona <mkhona@nvidia.com>
Signed-off-by: mkhona <mkhona@nvidia.com>
@mkhona-nvidia
mkhona-nvidia marked this pull request as ready for review October 12, 2026 22:21
@mkhona-nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 0c484d0

@skyw
skyw enabled auto-merge (squash) October 12, 2026 22:26
@skyw
skyw merged commit 4ebce60 into NVIDIA-NeMo:main Oct 12, 2026
31 checks passed

This branch was successfully deployed

2 active deployments
public — 0c484d0b Deployed Oct 12, 2026 by copy-pr-bot[bot] via release / finalize / notify #276
test — 0c484d0b Deployed Oct 12, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #1079
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.

2 participants