[kadirnar]

#39537

cc @muellerzr @SunMarc

[SabaPivot]

Well coded. We need muon optimizer for sure.

[SunMarc]

Thanks for the PR @kadirnar ! Definitely intestered in having this in transformers. Can we directly use muon library for this instead of re-coding them ? I saw that there were 4 classes that we might want to have : SingleDeviceMuon, Muon, MuonWithAuxAdam and SingleDeviceMuonWithAuxAdam. Right now in this PR, we deal with the case SingleDeviceMuon. Also, if you can add a couple of tests that we have for all optimizers, that would be great !

[SunMarc]

Thanks for iterating, left a couple of comments

[kadirnar]

@SunMarc Is there any obstacle to adding a Muon optimizer?

[kadirnar]

@ArthurZucker Could you provide feedback on this PR? Are you planning to merge it? Is there any obstacle?

[ArthurZucker]

Ooo exciting! 🔥 sorry did not see until now!
@3outeille will help review he know more about me in these regards!

[3outeille]

Hello @kadirnar, was running the 4 tests and it seems to fail !

self = <[AttributeError("'Muon' object has no attribute 'param_groups'") raised in repr()] Muon object at 0x7f9cf00cf890>
params = [{'params': [Parameter containing:
tensor([[ 1.1534e-03,  1.1191e-02, -2.1830e-02,  ..., -3.9372e-03,
         -4.8418..., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       device='cuda:0', requires_grad=True), ...], 'weight_decay': 0.0}], lr = 0.02, weight_decay = 0
momentum = 0.95

    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95):
        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum)
>       assert isinstance(params, list) and len(params) >= 1 and isinstance(params[0], torch.nn.Parameter)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       AssertionError

../../env_transformers_core/lib/python3.11/site-packages/muon.py:67: AssertionError

[kadirnar]

Hello @kadirnar, was running the 4 tests and it seems to fail !

self = <[AttributeError("'Muon' object has no attribute 'param_groups'") raised in repr()] Muon object at 0x7f9cf00cf890>
params = [{'params': [Parameter containing:
tensor([[ 1.1534e-03,  1.1191e-02, -2.1830e-02,  ..., -3.9372e-03,
         -4.8418..., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       device='cuda:0', requires_grad=True), ...], 'weight_decay': 0.0}], lr = 0.02, weight_decay = 0
momentum = 0.95

    def __init__(self, params, lr=0.02, weight_decay=0, momentum=0.95):
        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum)
>       assert isinstance(params, list) and len(params) >= 1 and isinstance(params[0], torch.nn.Parameter)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       AssertionError

../../env_transformers_core/lib/python3.11/site-packages/muon.py:67: AssertionError

Could you check it again?

[S1ro1]

The implementation in the original repo of Muon is not distributed friendly, can we do something about that? If the parameters of the linear layers are sharded this will either fail or produce incorrect results (been a while since I tried that). For an initial implementation it's fine IMO, though we should warn/raise correctly

[kadirnar]

The implementation in the original repo of Muon is not distributed friendly, can we do something about that? If the parameters of the linear layers are sharded this will either fail or produce incorrect results (been a while since I tried that). For an initial implementation it's fine IMO, though we should warn/raise correctly

Actually, after this optimization, I also want to add the flash-muon and Dion repositories. How would Dion be?

https://github.com/microsoft/dion/

[S1ro1]

The implementation in the original repo of Muon is not distributed friendly, can we do something about that? If the parameters of the linear layers are sharded this will either fail or produce incorrect results (been a while since I tried that). For an initial implementation it's fine IMO, though we should warn/raise correctly

Actually, after this optimization, I also want to add the flash-muon and Dion repositories. How would Dion be?

https://github.com/microsoft/dion/

I don't think adding multiple repositories just for one optimizer (not familiar with flash-muon) is a good approach. I think we should make the optimizer work with distributed settings too. For Dion you'd be facing the same approach, as the low rank orthogonalization needs to be distributed aware as well.

I propose either rolling our own implementation (that's distributed aware) or finding a repo that does it. Afaik there's a repo from microsoft research with both Muon and Dion which are DTensor compatible for starters, we'd need to hack something for our TP integration later.

[kadirnar]

The implementation in the original repo of Muon is not distributed friendly, can we do something about that? If the parameters of the linear layers are sharded this will either fail or produce incorrect results (been a while since I tried that). For an initial implementation it's fine IMO, though we should warn/raise correctly

Actually, after this optimization, I also want to add the flash-muon and Dion repositories. How would Dion be?
https://github.com/microsoft/dion/

I don't think adding multiple repositories just for one optimizer (not familiar with flash-muon) is a good approach. I think we should make the optimizer work with distributed settings too. For Dion you'd be facing the same approach, as the low rank orthogonalization needs to be distributed aware as well.

I propose either rolling our own implementation (that's distributed aware) or finding a repo that does it. Afaik there's a repo from microsoft research with both Muon and Dion which are DTensor compatible for starters, we'd need to hack something for our TP integration later.

https://github.com/samsja/muon_fsdp_2 ? They've tested it with multi-GPU. Do you think this is suitable?

[kadirnar]

@S1ro1 https://github.com/pytorch/pytorch/blob/main/torch/optim/_muon.py PyTorch support has been added.

[VarunGumma]

@kadirnar is there any update on this (using the pytorch Muon itself)?

I would really like to use Muon to finetune my model. Is there any fork/branch I can use?
I can implement it, but will it break too many files in my codebase

[sglucas]

@kadirnar is there any update on this (using the pytorch Muon itself)?

I would really like to use Muon to finetune my model. Is there any fork/branch I can use? I can implement it, but will it break too many files in my codebase

I have the same question. Is there a great repo to reproduce Muon and use Muon to fine-tune a pre-trained model.

[Rocketknight1]

Closing this since Muon is now available in PyTorch.