muon-fsdp2 0.2.1

Muon FSDP2 is a implementation of Muon optimizer for FSDP2

Muon FSDP2

muon implementation for FSDP2. Mainly copied this gist credits to @main-horse

Usage

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