# 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:

```bash
./run_test.sh
```

- To check the other available options, you can use:

```bash
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.
...
```