Metadata-Version: 2.4
Name: muon_fsdp2
Version: 0.2.1
Summary: Muon FSDP2 is a implementation of Muon optimizer for FSDP2
Author-email: Sami Jaghouar <sami.jaghouar@gmail.com>
Project-URL: Homepage, https://github.com/samsja/muon_fsdp_2
Project-URL: Repository, https://github.com/samsja/muon_fsdp_2
Project-URL: Documentation, https://github.com/samsja/muon_fsdp_2#muon-fsdp2-package
Project-URL: Issues, https://github.com/samsja/muon_fsdp_2/issues
Keywords: muon,optimizer,fsdp2,pytorch,deep-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Dynamic: license-file

# Muon FSDP2

muon implementation for FSDP2. Mainly copied [this gist](https://gist.github.com/main-horse/7314170780e36f7443d1926418d75823) credits to [@main-horse](https://github.com/main-horse)

## Usage

```python
from muon_fsdp2 import Muon


optimizer = Muon([
    dict(
        params=model.square_params(),
        lr=2e-2,
        use_muon=True
    ),
    dict(
        params=model.non_square_params(),
        lr=3e-4,
        use_muon=False
    )
])

```
