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