---
tags:
- kernel
license: apache-2.0
---

# Optimizer

Optimizer is a python package that provides:
- PyTorch implementation of recent optimizer algorithms
- with support for parallelism techniques for efficient large-scale training.

## Currently implemented
- [Parallel Muon with FSDP2](./docs/muon/parallel_muon.pdf)

## Usage

```python
import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from kernels import get_kernel

optimizer = get_kernel("motif-technologies/optimizer")
get_default_muon_param_groups = optimizer.muon.get_default_muon_param_groups

model = None # your model here
fsdp_model = FSDP(model)

# muon, in nature, cannot use 1-d tensor
# we provide helper function to group such tensors
# you can use your own function, if necessary
params = get_default_muon_param_groups(model) # user can write own is_muon_func, if necessary

optim = optimizer.Muon(
    params,
    lr=0.01,
    momentum=0.9,
    weight_decay=1e-4,
)
```

## Pre-commit Hooks

This project uses [pre-commit](https://pre-commit.com/) to automatically check and format code before commits.

### Setup

1. Install pre-commit:

   ```bash
   pip install pre-commit
   ```

2. Install the git hooks:
  
```bash
   pre-commit install
   ```

Once installed, the configured hooks will run automatically on each commit.

### Included Hooks

The following tools are run via pre-commit:

- **[yapf](https://github.com/google/yapf)** – Python code formatter  
- **[typos](https://github.com/crate-ci/typos)** – Spell checker for common typos  
- **[isort](https://github.com/PyCQA/isort)** – Organizes and sorts Python imports  
- **[clang-format](https://clang.llvm.org/docs/ClangFormat.html)** – Formats C++/CUDA code (`--style=file`)  
- **[pymarkdown](https://github.com/jackdewinter/pymarkdown)** – Lints and auto-fixes Markdown files  
- **[actionlint](https://github.com/rhysd/actionlint)** – Validates GitHub Actions workflows  

### Usage

- Run all checks on the entire codebase:

   ```bash
   pre-commit run --all-files
   ```

- Run a specific hook (example: isort):
  
 ```bash
   pre-commit run isort --all-files
   ```

### Test

- There is a [simple unittest for Parallel Muon](./test/test_muon/README.md)