DeepSpeed
Add Gram Newton-Schulz orthogonalization for Muon optimizer
#7953

Open

Add Gram Newton-Schulz orthogonalization for Muon optimizer #7953

delock wants to merge 7 commits into master from gma/gram_muon

delock requested a review from

tjruwase 3 days ago

delock requested a review from

tohtana 3 days ago

delock requested a review from

loadams 3 days ago

chatgpt-codex-connector commented on 2026-04-03

delock force pushed from 4fe6a0b4 to d17212ef 3 days ago

Add Gram Newton-Schulz iteration for Muon optimizer

381d8b7a

docs: add ns_method parameter to Muon optimizer documentation

e9beb2de

fix: correct Gram Newton-Schulz reference URL

d17212ef

Use accelerator API for dtype selection in Newton-Schulz iterations

54930203

Fix non-contiguous tensor output from Gram NS for tall matrices

e5de42ce

Fold transpose into matmul in Gram NS for tall matrices

a6cf6b69

Use fused addmm and eliminate eye allocation in Gram NS

61095611

PKUWZP requested a review from

PKUWZP 2 days ago

Reviewers

tjruwase

tohtana

loadams

PKUWZP

Assignees

No one assigned

Labels

None yet

Milestone

No milestone