Parameter Groups for Optimizers

nequip.model.MuonParamGroups(model: Module, muon: dict, adam: dict)[source]

Build optimizer parameter groups, splitting parameters between a Muon-based optimizer and Adam (or Adam-like) optimizer.

Assigned to Adam group:

  • Any parameter whose name does not contain the substring "layer".

  • Any parameter not matching the Muon-specific rules below.

Assigned to Muon group:

  • Edge MLP weights: parameters whose name contains "edge_mlp" and that are 2D tensors (i.e., matrix weights).

  • e3nn convolution linear weights: parameters whose name contains "conv.linear".

For e3nn Linear layers, the returned Muon parameter group includes an e3nn_reshaping dictionary mapping the index of the parameter within the Muon group to the module’s weight_index_slices. This metadata will be used by the to reshape or operate on corresponding matrix weights.

Parameters:

  • model (torch.nn.Module) – The model to optimize.

  • muon (dict) – Muon config parameters.

  • adam (dict) – Adam config parameters.