Muon Optimizer Test

This directory contains a test script for the Muon optimizer.

Prerequisites

GPU Requirement
- All tests require 8 GPUs by default.
- If you have fewer GPUs available:
  - Modify the parallelism configurations in test_muon.py.
Model Access
- The tests require access to the private model repository:
  - Motif-Technologies/Motif-2.6B-4layer-random on Hugging Face.
- Set your Hugging Face token via the environment variable HF_TOKEN.
- If you don’t have access, please contact the maintainer.
Using a Different Model (Optional)
- You may modify the test to use a different model by:
  - Updating the model name in conftest.py::inputs.
  - Adjusting the tensor parallel rules in utils.py::_apply_tp.

Usage

To execute the test with 8 GPUs, simply run:

./run_test.sh

To check the other available options, you can use:

pytest --help
...
Custom options:
  --measure-perf        Measure execution time and peak memory usage during optimizer step.
  --do-profile          Enable profiling during tests.
  --skip-verify         Skip verification of optimizer step correctness with sequential implementation.
                        This can be useful when GPU memory is limited.
...