Kernels
File size: 1,298 Bytes
e2b41e5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# 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.
...
```