Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials

This repository contains the code for our paper "Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials" by Ekaterina Grishina, Matvey Smirnov and Maxim Rakhuba.

Outline

• polynomials.py contains code for finding optimal polynomials.
• notebooks/polynomials.ipynb contains examples of how to generate optimal polynomials using these functions.
• riemannian_opt.py contains code for polar retraction, Riemannian SGD and Adam optimizers on Stiefel manifold.
• /nanoGPT contains code for running NanoGPT with Muon optimizer and different polynomials. This code is based on NanoGPT speedrun repository.

NanoGPT training

Navigate to /nanoGPT directory and install requirements. To download the data for training NanoGPT, run

python data/cached_fineweb10B.py 8

To train the model

./run.sh