Inconsistent p.ndim usage in Muon optimizer

[N-damo]

There's an inconsistency in how p.ndim is used in the optimize.py file. In the get_optimizer function, parameters are selected with the condition p.ndim >= 2, but in the Muon class constructor, there's an assertion assert p.ndim == 2, p.ndim which only allows 2D parameters.

This inconsistency can cause issues when using the Muon optimizer with parameters that have more than 2 dimensions.

In get_optimizer
muon_params = [
    p
    for name, p in model.named_parameters()
    if p.ndim >= 2 and "classifiers" not in name and "embedding" not in name
]

In Muon class constructor
for p in muon_params:
    # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer
    assert p.ndim == 2, p.ndim
    self.state[p]["use_muon"] = True

could you give any suggestions?

[hyhsjd2]

Same question.

[DiMarzioBian]

Have the same question. The report states that parameters larger than 2D are flattened into 2D, but they hardcode the assertion that p.ndim == 2. Torch's implementation has the same problem (link).