Muon optimizer

[ankitpatnala]

Description

Implemented Muon optimizer

• Currently, it is only implemented for forecasting configs. config/config_forecasting.yml
• Just adapted the code to log multiple lr, to keep track of lr of AdamW and Muon separately (needs to be tested)

Issue Number

Closes #2638

[ankitpatnala]

@sophie-xhonneux
@florianscheidl

[clessig]

@ankitpatnala : could you outline the hyperparams that are involved. How should the lr change between muon and AdamW?

[ankitpatnala]

The name field inside optimizer selects which optimizer is used: adamw or muon. If it is set to adamw, the muon config block is ignored.

  optimizer:
   # which optimizer to use: "adamw" or "muon"
   # muon only optimizes 2D weight matrices of hidden layers; all other
   # parameters (biases, norms, ...) are still optimized with adamw
   name: muon
   grad_clip: 1.0
   weight_decay: 0.1
   log_grad_norms: False
   adamw :
     # parameters are scaled by number of DDP workers
     beta1 : 0.98125 # == 0.85 on 2 nodes x 4 gpus
     beta2 : 0.9875  # == 0.90 on 2 nodes x 4 gpus
     eps : 2e-08
   muon:
     # peak lr override for muon; null shares adamw's lr schedule (relies on
     # adjust_lr_fn=match_rms_adamw's per-matrix scaling instead of a fixed multiplier)
     lr_max: null
     momentum: 0.95
     nesterov: True
     ns_steps: 5
     eps: 1e-7
     # "original" | "match_rms_adamw"
     adjust_lr_fn: match_rms_adamw

When it is set to muon, the general hyperparameters stay the same:

grad_clip: 1.0
weight_decay: 0.1
log_grad_norms: False

The optimizer then explicitly checks whether each parameter is 2D (optimizer.py#L97):

• Non-2D parameters (biases, norms, etc.) use the base AdamW config:

adamw:
  # parameters are scaled by number of DDP workers
  beta1: 0.98125  # == 0.85 on 2 nodes x 4 gpus
  beta2: 0.9875   # == 0.90 on 2 nodes x 4 gpus
  eps: 2e-08

• 2D parameters are updated by Muon, using this config:

muon:
  # peak lr override for muon; null shares adamw's lr schedule (relies on
  # adjust_lr_fn=match_rms_adamw's per-matrix scaling instead of a fixed multiplier)
  lr_max: null
  momentum: 0.95
  nesterov: True
  ns_steps: 5
  eps: 1e-7
  # "original" | "match_rms_adamw"
  adjust_lr_fn: match_rms_adamw

If lr_max is defined, Muon's learning rate is scaled between this value and the lr_max set in learning_rate_scheduling optimizer.py#37. It is typically set to 1e-2, but that did not work for me:

learning_rate_scheduling:
  lr_start: 1e-6
  lr_max: 5e-5
  lr_final_decay: 2e-6
  lr_final: 0.0
  num_steps_warmup: 256
  num_steps_cooldown: 512
  policy_warmup: "cosine"
  policy_decay: "constant"
  policy_cooldown: "linear"
  parallel_scaling_policy: "sqrt"

All other Muon parameters take their defaults from torch.optim.Muon.

I tested roughly lr_max=1e-2, and training failed — I'm not yet sure why. Leaving lr_max: null together with adjust_lr_fn: match_rms_adamw (optimizer.py#L56) instead gives stable training. Because Muon optimizer now handles the LR scaling internally, it logs the same learning rate for both optimizers.

[sophie-xhonneux]

Only tangentially related to this PR, but do we want to remove this scaling of beta parameters with batch size because I think it makes setting those variables so much more complicated

[ankitpatnala]

Should we then hardcode it for the 2-node, 4-GPU setup?

[clessig]

We would need at least some evidence that it's neutral and doesn't affect skill.

[florianscheidl]

Hey Ankit,
I ran some memory and pytorch profiler traces with and without the Muon optimizer on config_forecasting and a fixed seed.

• Little impact on memory: qualitatively, the memory traces look similar, the peak memory is in the same ballpark.

Adam:

Muon:

• Computational overhead: The Muon optimizer launches many small tensor ops. We could investigate if this can be optimized. We see a sum in the gradient scaler, which launches a sum and a synchronization. We can probably remove gradient scaling (since we're training with fp32 and bf16) and thereby the sync.

Muon

Adam: The optimizer step is nearly negligible, so not showing it here.

Note: these are only qualitative, not quantitative results. For the latter, we could compare the peak memory consumption and throughput for a longer run.

[clessig]

Hey Ankit, I ran some memory and pytorch profiler traces with and without the Muon optimizer on config_forecasting and a fixed seed.

• Little impact on memory: qualitatively, the memory traces look similar, the peak memory is in the same ballpark.

Adam:

Muon:

• Computational overhead: The Muon optimizer launches many small tensor ops. We could investigate if this can be optimized. We see a sum in the gradient scaler, which launches a sum and a synchronization. We can probably remove gradient scaling (since we're training with fp32 and bf16) and thereby the sync.

Muon

Adam: The optimizer step is nearly negligible, so not showing it here.

Note: these are only qualitative, not quantitative results. For the latter, we could compare the peak memory consumption and throughput for a longer run.

Thanks! Could we also try with the config with the GEOs that is more demanding in terms of memory to make sure we don't break configs that are currently working.