--- id: gram-newton-schulz version: "0.1.6" license: MIT license_treatment: permissive maintenance: active --- # gram-newton-schulz — Fast Newton-Schulz Algorithm with Kernels License: permissive · Maintenance: active · Downloads: 176.3K/mo ## What it is and what it does Gram Newton-Schulz is a specialized optimizer component that accelerates polar decomposition—a key operation in the Muon optimizer—by iterating on a smaller symmetric Gram matrix instead of the full rectangular matrix. This mathematical reformulation reduces floating-point operations and enables more efficient symmetric GEMM kernels on modern NVIDIA GPUs. The package provides both a standalone GramNewtonSchulz callable class and integration into a full Muon optimizer that handles mixed parameter types (2D weights for orthogonalization, scalars via an auxiliary optimizer) with autotuned restart points for numerical stability. It depends on torch for tensor operations, quack-kernels for custom GEMM implementations, and nvidia-cutlass-dsl for low-level GPU kernel generation. The algorithm is mathematically equivalent to standard Newton-Schulz with no claimed training accuracy loss, making it a direct swap for existing Muon-based training pipelines. Installation requires explicit GPU hardware (Hopper or Blackwell) and careful dependency management to avoid pip installing a CPU-only PyTorch. Use it for: - Drop-in replacement for Newton-Schulz in Muon optimizer for faster large-model training on H100 or B200/B300 clusters. - Accelerating polar decomposition in custom orthogonalization-based optimizers that need symmetric matrix operations. - Autotuning restart schedules for Newton-Schulz iterations to balance speed and numerical stability in iterative algorithms. - Training transformer models with Muon where 2D weight matrices benefit from symmetric GEMM kernels on Hopper/Blackwell hardware. ## Worth the install? AI-flagged interpretation of the facts above — verify before relying. Implements a hardware-aware Gram Newton-Schulz algorithm for polar decomposition that serves as a drop-in replacement for standard Newton-Schulz in PyTorch optimizers like Muon, achieving faster convergence through symmetric GEMM kernels on NVIDIA GPUs. Yes, if you are training on NVIDIA Hopper or Blackwell GPUs with PyTorch 2.7.1+ and use or plan to use Muon as your optimizer. The package is actively maintained, has no known vulnerabilities, carries permissive MIT licensing, and offers a mathematically equivalent but faster polar decomposition with no training accuracy tradeoff. No, if you lack the required GPU hardware or are locked to older PyTorch versions—the installation is hardware-specific and will fail gracefully without the target GPU. ## Install pip install gram-newton-schulz uv add gram-newton-schulz poetry add gram-newton-schulz ## Installing gram-newton-schulz Before you install: Low friction install via pip with a pure-Python wheel, but requires NVIDIA Hopper or Blackwell GPU hardware, PyTorch 2.7.1+, and CUDA 12.9+. Must use --no-build-isolation flag to avoid pip installing a CPU-only PyTorch. Active maintenance with recent commits. License in practice: MIT license is permissive and imposes no restrictions on use, modification, or distribution in proprietary or open-source projects. Quickstart: pip install gram-newton-schulz --no-build-isolation from gram_newton_schulz import GramNewtonSchulz, POLAR_EXPRESS_COEFFICIENTS gram_NS = GramNewtonSchulz( ns_coefficients=POLAR_EXPRESS_COEFFICIENTS, gram_newton_schulz_reset_iterations=[2] ) result = gram_NS(X) Requires NVIDIA Hopper (H100) or Blackwell (B200/B300) GPU, PyTorch 2.7.1+, and CUDA 12.9+. Must install PyTorch first before running pip install with --no-build-isolation. Verify before relying: - Actual speedup magnitude relative to standard Newton-Schulz in real training scenarios beyond the claimed 2x in the description. - Numerical stability guarantees and convergence behavior across different matrix shapes and conditioning. - Compatibility with torch.compile on Hopper GPUs and workaround effectiveness on Blackwell. ## Package facts - License: MIT (permissive) - Python support: supports_current - Install friction: low - Maintenance: active - Downloads: 176.3K/month (top 15,000 on PyPI) - Known vulnerabilities: none known ## Tags polar decomposition pytorch, newton-schulz optimizer, muon optimizer, gpu accelerated matrix algorithms, symmetric gemm kernels, hardware aware optimization, matrix factorization pytorch, gpu-optimization, matrix-algorithms, pytorch-optimizer [View on SkillFed](https://skillfed.io/packages/gram-newton-schulz) · [View on PyPI](https://pypi.org/project/gram-newton-schulz/)