Custom Optimizers¶

MuonWithAuxAdam¶

MuonWithAuxAdam is a custom optimizer that combines Muon updates with an auxiliary Adam optimizer. This Muon implementation is based on the original Keller Jordan GitHub repository: https://github.com/KellerJordan/Muon/blob/master/muon.py

class nequip.train.MuonWithAuxAdam(params)[source]¶

Non-distributed variant of MuonWithAuxAdam (Originally SingleDeviceMuonWithAuxAdam).

Config example with NequIP parameter groups:

_target_: nequip.train.MuonWithAuxAdam
param_groups:
  _target_: nequip.model.MuonParamGroups
  muon:
    lr: 0.01
    weight_decay: 1e-5
  adam:
    lr: 0.01
    weight_decay: 1e-5