---
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")

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

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