|
|
--- |
|
|
tags: |
|
|
- kernel |
|
|
--- |
|
|
|
|
|
# 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, |
|
|
) |
|
|
``` |
|
|
|
