File size: 2,244 Bytes
8535e80 15b5d41 8535e80 80187d6 8535e80 b0f46c7 8535e80 b0f46c7 8535e80 b0f46c7 8535e80 80187d6 3261444 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
---
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)
|