[BUG] The different lr for Muon doesn't work?

[Zazexy]

I highly appreciate the improvements you've made in supporting the Muon optimizer. However, I've run into an issue when trying to set the individual muon_lr and adam_lr within the configuration JSON.
I've found that the optimizer still defaults to setting a single lr for both the Muon parameter group and the Adam parameter group. I've verified that I am on the latest commit, as the relevant processing code for this feature branch is present in my local engine.py.
All other settings in the configuration are working as expected, but the specific muon_lr and adam_lr settings are not taking effect. Has anyone else encountered this issue?
"optimizer": {
"type": "Muon",
"params": {
"lr": 1e-4,
"betas": [0.9, 0.95],
"eps": 1e-8,
"weight_decay": 0.01,
"muon_lr": 2e-2,
"adam_lr": 1e-5
}

[KeeProMise]

@Zazexy Hello! Thank you for providing the detailed information. According to your description, you have set muon_lr and adam_lr in the configuration JSON, but you found that these settings did not take effect, and the optimizer still uses a single lr value. I noticed that there is indeed logic in the code where muon_lr and adam_lr override lr. To further confirm the issue, could you please confirm whether you are using the latest version of the code? Additionally, I suggest trying to apply the Pull Request (PR) I previously submitted to see if the issue still persists. If the problem remains, please provide more debugging information so that we can better locate and resolve the issue. Looking forward to your reply!

[Zazexy]

Thanks for your patient reply! I do many checks about the version and I pull your PR as well. I found the reason is that I'm using accelerate in my project. It looks like some default settings of optimizer just override the config. What I do is create a dummy optimizer and afford it to accelerator. Result is the other parameters like eps and betas are set to optimizer, however, no matter lr or muon_lr is override by acclerate.
Is there any method to apply muon with deepspeed in accelerate?
optimizer = DummyOptim(params_to_optimize) lr_scheduler = DummyScheduler( optimizer, total_num_steps=config.trainer.max_train_steps, warmup_num_steps=config.trainer.lr_warmup_steps * accelerator.num_processes ) train_module, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( train_module, optimizer, train_dataloader, lr_scheduler )

[delock]

Hi @Zazexy do you have a small script that can run to expose this issue? I want to understand why setting a dummy optimizer would skip muon_lr overriding.

@KeeProMise thanks for submit the PR. I think the muon_lr override logic didn't change with this PR.

[Zazexy]

Hello @delock @KeeProMise ,

Thank you for looking into my issue. After more recent testing and debugging, I've found that the problem is not related to any specific PR or version (I am certain that both my initial version and my current one are correct).

To summarize my findings: The issue appears to be with the LR initialization for the scheduler within accelerate, specifically that the scheduler seems to be passed and set before the optimizer's learning rate is overwritten by the muon_lr and adam_lr parameters.

I observed that before this point, all other parameters like betas, eps, etc., were correctly passed to the optimizer, and the parameter groups were correctly set with use_muon.

✅ Temporary SolutionThe issue disappeared, and my problem was resolved to an extent, when I:

• Stopped passing any scheduler to the prepare function.
• Implemented and called the scheduler manually instead.

I apologize that I cannot provide a minimal reproducible example as I am still learning to use accelerate and deepspeed. However, during my debugging, I did confirm that:

• The kwargs in accelerate/src/accelerate/accelerator.py were indeed populated according to my provided config.
• The parameter groups were set up correctly.

Despite this, I still lack further insight into the underlying cause.

🛠️ Configuration DetailsHere is the configuration I was using when I encountered the issue, which I hope is helpful:
I used dummyscheduler and dummyotim for accelerate.prepare() and set config as followed:

{
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 16,
    "zero_allow_untested_optimizer": "true",

    "optimizer": {
            "type": "muon",
            "params": {
                "lr": 1e-4,
                "betas": [0.9, 0.95],
                "eps": 1e-8,
                "weight_decay": 0.01,
                "muon_lr": 2e-2,
                "adam_lr": 1e-5
            }
        },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 1,
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        }
    },
    "fp16": {
        "enabled": false
    },
    "bf16": {
        "enabled": true
    }
}

Please let me know if there's any other information I can provide.

[delock]

Hi @Zazexy , thanks for your information. Yet I didn't figure out how to put accelerate and deepspeed together to reproduce this bug. When you have the reproducible script available, ping me and I'll try it in my environment. Thanks!

[Zazexy]

Hi @delock,

Thank you for your patience. I've created a toy script to better illustrate my issue. In this script, I implemented a dummy scheduler with a learning rate of 1e-5 (if not set explicitly, it defaults to 1e-3), which differs from the lr or muon_lr configured in my config file.
After calling accelerator.prepare(), the model starts running. During execution, I printed the optimizer's state and observed the attributes in both parameter groups.
In my environment, the code behaves as follows: the learning rate increases from 0 to the dummy optimizer's 1e-5 (rather than muon_lr or adam_lr) during the warm-up period, and then decays to 0 over the total steps.
This behavior suggests that the scheduler might not be properly reading the optimizer's base learning rates from the DeepSpeed configuration.
If the code you can't run in your environment, I will try to provide another one. Hope for your reply.

deepspeed_muon_test.json

test_acc.log

test.py

[Zazexy]

I use the optimizer without accelerate wrapping. Result shows that adamW is not infulenced, when muon performing worse. Maybe the solution I choose doesn't solve the problem from root?

[delock]

Hi @delock,

Thank you for your patience. I've created a toy script to better illustrate my issue. In this script, I implemented a dummy scheduler with a learning rate of 1e-5 (if not set explicitly, it defaults to 1e-3), which differs from the lr or muon_lr configured in my config file. After calling accelerator.prepare(), the model starts running. During execution, I printed the optimizer's state and observed the attributes in both parameter groups. In my environment, the code behaves as follows: the learning rate increases from 0 to the dummy optimizer's 1e-5 (rather than muon_lr or adam_lr) during the warm-up period, and then decays to 0 over the total steps. This behavior suggests that the scheduler might not be properly reading the optimizer's base learning rates from the DeepSpeed configuration. If the code you can't run in your environment, I will try to provide another one. Hope for your reply.

deepspeed_muon_test.json

test_acc.log

test.py

Hi @Zazexy I tried your code. In my first run I observe no lr difference between group 0 and group 1. However, when I remove muon_lr and adam_lr in the config json file, and retain "lr" only. Then no matter what lr number I use, I always see the same lr in both group 0 and group 1 between run to run. So I suspect DummyOptim fully overrid deepspeed settings.

I see the scheduler would change the learning rate over time. In your case, if you would like different lr for muon and adam optimizer. Is it possible to create two optimizers and two schedulers in your test code, one set for muon parameters, the other set for non-muon parameters? I'm not familiar to accelerate enough to know how this could be done.

Besides that, I also don't know whether deepspeed adam and muon optimizer would take effect. From the code optimizer is created from DummyOptim and its behavior should be decided by DummyOptim. Does DummyOptim finally goes to DeepSpeed optimizer code?

[Zazexy]

Hi @delock , I understood what you said, since I face these problems as well. I tried to find out the process of accelerate. Finally, I found some weired behaviours in https://github.com/huggingface/accelerate/blob/4826b66d4d4b7c60a2722e83dfe90347fd971502/src/accelerate/accelerator.py#L2391C12-L2391C73. It looks like that accelerate fills the kwargs as super parameters table first, however, optimizer is no instanced before that. It means kwargs record the lr from dummyoptimizer which will be used to initialize scheduler. I think this may cause the problem.
To avoid that, I choose customize the scheduler by myself without accelerate.prepare() and I delete the config about scheduler in .json file. Finally, it looks like scheduler works well in two params groups with different lr as we want. Here is the code I fix follow these steps.
However, The result shows muon perform worse than adamW in any learning rate as the figue I showed before. The dark and yellow curve are adamW with or without accelerate.prepare() (without or with my customize scheduler).

Thanks for your patience, I wanna know whether I make muon work exactly.

test.py
test_acc.log

[delock]

Hi @Zazexy which model are you using? I see in test.py there is only a very simple model. Could it be the model is too small for Muon to be better than Adam? Have you tried to use something like Qwen2 0.5B?

Hi @delock , I understood what you said, since I face these problems as well. I tried to find out the process of accelerate. Finally, I found some weired behaviours in https://github.com/huggingface/accelerate/blob/4826b66d4d4b7c60a2722e83dfe90347fd971502/src/accelerate/accelerator.py#L2391C12-L2391C73. It looks like that accelerate fills the kwargs as super parameters table first, however, optimizer is no instanced before that. It means kwargs record the lr from dummyoptimizer which will be used to initialize scheduler. I think this may cause the problem. To avoid that, I choose customize the scheduler by myself without accelerate.prepare() and I delete the config about scheduler in .json file. Finally, it looks like scheduler works well in two params groups with different lr as we want. Here is the code I fix follow these steps. However, The result shows muon perform worse than adamW in any learning rate as the figue I showed before. The dark and yellow curve are adamW with or without accelerate.prepare() (without or with my customize scheduler).

Thanks for your patience, I wanna know whether I make muon work exactly.

test.py test_acc.log