pip package for optimizers

[thib-s]

This pull request intend to make the dion optimizer pip installable in a standalone way. The objective is to allow the user to install the optimizer by running:

pip install git+https://github.com/microsoft/dion.git

This would install the package named dion, in which optimizers can be imported with:

from dion import Dion, DionParamConfig, Muon

The requirement to run the training script can be installed by cloning the repo and running:

pip install -e .[train] # can also install [dev] for developer tools
torchrun --standalone --nproc_per_node=8 train.py --config configs/dion_160m.yaml

During the process, the optimizers folder was renamed to dion to avoid users installing a package named optimizers, which is too generic and could cause conflicts. I would understand if you reject the PR if you think that this change is too important. I'm also open to modifications about names or global repository organisation. Feel free to ask for modifications!

Finally, I moved the testing functions in the tests folder to be compliant with pytest; these tests can be run in the following way:

pyest tests/

Similarly, the benchmarking can be done using the standalone app in the benchmark folder:

python benchmark/benchmark_newton_shultz.py --help
usage: benchmark_newton_shultz.py [-h] (--grid | --plot) [--m M] [--n N] [--batch_size BATCH_SIZE] [--expansion EXPANSION] [--dtype {float16,bfloat16}]

Benchmarks for Newton–Schulz Triton kernels

optional arguments:
  -h, --help            show this help message and exit
  --grid                sweep a list of sizes
  --plot                generate TFLOPS curves and write plots
  --m M                 rows
  --n N                 cols (defaults to m)
  --batch_size BATCH_SIZE
  --expansion EXPANSION
                        n = m * expansion (grid mode)
  --dtype {float16,bfloat16}
                        input dtype

Package Installation:

• setup.py: Created a setup.py file to enable installation of the dion package with optional extras for development and training. This enhances usability and simplifies dependency management for different use cases.

Dependency Management:

• Requirements files: Added separate requirements files (requirements_dion.txt, requirements_dev.txt, requirements_train.txt) to manage dependencies for standalone usage, development, and training. [1] [2] [3]

Modularization and File Restructuring:

• File renaming and imports: The optimizers module was renamed to dion, and corresponding imports were updated across files, such as in train.py and dion/__init__.py. This improves clarity and aligns with the package's naming convention. [1] [2] [3]

Benchmarking:

• Benchmark script: Added benchmark_newton_shultz.py to provide detailed benchmarks for Newton-Schulz kernel performance, including one-off timing, grid sweeps, and TFLOPS plots. This script helps evaluate and optimize computational efficiency.

Testing:

• Unit tests: Added tests/test_newton_shultz.py with comprehensive tests for functions like ns_line_1, ns_line_2, and newton_schulz_triton. These tests ensure correctness and reliability of the Triton-based implementations.

[thib-s]

@microsoft-github-policy-service agree

[byronxu99]

Thank you for your contribution! A pip package release was part of our future plans, and it's very helpful for you to get this done so quickly.

I'll take a look over the PR soon and leave some feedback.

[thib-s]

Thanks for the review, I'll update the code according to these comments.
Also, could it be relevant to remove the test folder and rebase this code on #8 (which seems to contain more tests)?

[byronxu99]

Looks good, I'm merging it

[byronxu99]

Thank you very much for your contribution!