Update torch-ext/optimizer/muon.py

Update torch-ext/optimizer/muon.py

Update muon.py

Merge pull request #10 from MotifTechnologies/fix_a2a_gs_assert

Add built binary [skip-build]

fix assert in a2a gather scatter

fix: use [skip-build] keyword to prevent build [skip-build]

Merge pull request #9 from MotifTechnologies/all2all_gather_scatter

Add built binary [ci skip]

delete state in split_func

Add built binary [ci skip]

change owner_params to owned_params

modify pre step (overlap step) can get from arsgs

add doc strings + init self rank on init_assign_params

Add built binary [ci skip]

Merge branch 'all2all_gather_scatter' of https://github.com/MotifTechnologies/optimizer into all2all_gather_scatter

license added for flash_muon

Add built binary [ci skip]

apply pre-commit hook

consider multi node

misc

use inpalce op in update_g

use COMM_DTYPE instead of hardcoded dtype

apply all2all scatter gather

feat(muon_clip) : add muon clip (#6)

chore: add nix build workflow (#8)

feat(muon) : add tuned-abc-values & blfoat16 communication

feat: update muon to receive paramgroups, not model (#4)

fix(muon): add update_p stage and dealloc tensors properly

feat(muon): add test for muon

chore: add .gitignore

chore: add push-to-hf workflow

add more exclude files

Merge pull request #1 from MotifTechnologies/pre-commit_test_and_apply_lint

applied lint

add pre commit action w/o pre-commit applied

Update README.md

fix(optimizer): resolve bug where weight decay was multiplied by wrong lr value (#5)

Support HSDP (#4)

support torch 2.8 (#3)

fix(muon): handle un-distributed env

refactor(muon): change argument adam_wd to weight_decay and handle params' type

fix(muon): free tensors that are no longer needed

chore(muon): update comment

chore(muon): clean build and update doc

fix(muon): delete intermediate tensors immediately to lower peak mem usage

chore: initial commit

initial commit