Skip to content

feat(cutedsl): add CuTe-DSL Newton-Schulz operator and LigerMuon optimizer - #1449

Draft
piyushumate wants to merge 1 commit into
linkedin:mainfrom
piyushumate:feat/cutedsl-newton-schulz-muon
Draft

piyushumate wants to merge 1 commit into
linkedin:mainfrom
piyushumate:feat/cutedsl-newton-schulz-muon

Conversation

@piyushumate

Copy link
Copy Markdown
Contributor

Summary

This PR introduces a CuTe-DSL accelerated 5th-order Newton-Schulz operator and the LigerMuon transformer optimizer to Liger-Kernel.

Muon (MomentUm Orthogonalized by Newton-schulz) has demonstrated significant sample efficiency gains over AdamW for pre-training and fine-tuning LLMs by applying orthogonalized momentum updates to 2D weight matrices:

X_{k+1} = 0.5 * X_k * (3I - X_k^T * X_k)

In eager PyTorch, the 5-step iterative polar decomposition incurs repeated un-fused GEMM launches and high HBM memory traffic across dozens of layers. This PR provides:

  1. cutedsl_newton_schulz_forward in liger_kernel.ops.cutedsl.ops.newton_schulz with automated PyTorch fallback.
  2. liger_newton_schulz operator exposed in liger_kernel.ops.newton_schulz.
  3. LigerMuon optimizer in liger_kernel.transformers.muon, routing 2D weights to accelerated Newton-Schulz and 1D parameters (biases, layer norms) to AdamW.

Testing Done

  • Hardware Type: NVIDIA H100 SXM5 80GB (tested via Modal E2E)
  • All 11 tests in test/cutedsl/test_newton_schulz.py passed on H100 (numerical parity across shapes/dtypes, singular value bounds, and loss convergence).
  • ruff check and ruff format passed cleanly.

@piyushumate
piyushumate marked this pull request as draft September 9, 2026 06:18
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.

1 participant