Kernels
optimizer / test /README.md
wyldecat's picture
Support param group with various placements (#13)
e2b41e5 unverified
# 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.
...
```