Preserve convolution shapes and scaling in Muon updates - #8552
Open
vineethsaivs wants to merge 1 commit into
Open
vineethsaivs wants to merge 1 commit into
vineethsaivs wants to merge 1 commit into
Conversation
Restore the original shape before the overflow guard and derive the scale from the matrix being orthogonalized. Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
vineethsaivs
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 16, 2026 18:59
Contributor
Author
|
Could you approve the CI run? The six CPU tests and changed-file pre-commit checks pass. |
pengdurice
reviewed
Sep 16, 2026
Contributor
There was a problem hiding this comment.
is this a good place for the new test?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Muon crashes when updating convolution weights with four-dimensional gradients.
The overflow guard combines a flattened update with the original four-dimensional gradient. The aspect-ratio correction also reads the kernel dimensions instead of the flattened matrix. Restore the original shape before the overflow guard and derive the scale from the matrix being orthogonalized.
Test:
TORCHDYNAMO_DISABLE=1 DS_ACCELERATOR=cpu PYTHONPATH=. python -m pytest tests/unit/v1/ops/muon/test_muon_overflow.py -k 'convolution or overflowed_gradient or finite_gradient' -q. Four convolution regressions fail before. Six CPU tests pass after, including two Conv2d training steps with both Newton-Schulz methods and overflow checks. Changed-file pre-commit passes. Executed on Apple M2 Pro CPU; no CUDA or distributed training run.