[v0ropaev]

registry.get_configured_optimizer_cls("muon_hyperball", lr=0.01, hyperball_radius=1.0) raises TypeError: Optimizer 'MuonHyperball' does not accept arguments: {'lr'}. The same happens for momentum, nesterov and every other Muon argument. validate_optimizer_args only looks at signature(opt_cls), and MuonHyperball.__init__ is (*args, hyperball_radius, hyperball_eps=1e-15, **kwargs), so the only names it accepts are args, hyperball_radius, hyperball_eps and kwargs. The nanogpt example works around this by hand in build_optimizer by merging in Muon's signature when the name is muon_hyperball.

The validation now walks the class MRO and collects the named parameters of each __init__, moving on to the parent only while the current __init__ takes **kwargs. The *args and **kwargs entries are no longer counted as argument names. MuonHyperball is the only registered optimizer with **kwargs in its signature, so every other optimizer stops at its own __init__ and accepts the same names as before. For MuonHyperball the walk stops at Muon, so names that only exist further up, like scaled_orthogonalize_fn on OrthogonalizedOptimizer, are still rejected.

One side effect is that weight_update_hook now passes validation for muon_hyperball, since Muon's __init__ takes it. MuonHyperball still raises KeyError for it when the optimizer is constructed, so the error shows up there instead of at validation time. I didn't special case it in the registry. I also left the nanogpt workaround alone. That example looks up classes through the registry but does its own argument check and never calls validate_optimizer_args, so it can be simplified separately.

I added two tests to tests/test_registry.py. One builds muon_hyperball through get_configured_optimizer_cls with lr, momentum, nesterov and hyperball_radius and constructs it on a unit norm parameter. The other checks that nonexistent_arg, args, kwargs and scaled_orthogonalize_fn are rejected for MuonHyperball. On main the first test errors with TypeError and the args and kwargs cases fail, and with the change all of them pass. I ran python tests/test_registry.py --device=cpu --seed=42 (18 tests OK) and python tests/test_orthogonalized_optimizer.py (145 tests OK) on CPU with torch 2.14.1 and Python 3.12. I also looped over all 27 registered optimizers and compared the accepted names before and after. They are identical for all of them except muon_hyperball, which now accepts every Muon argument and no longer accepts args or kwargs. ruff 0.9.9 (check, isort and format) is clean on both changed files and mypy 1.19.1 is clean on registry.py.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

The PR lets muon_hyperball accept arguments forwarded to its parent constructor.

• Muon Hyperball accepts the options it forwards to Muon.

_{Reviews (2) · Last reviewed commit: "Use a plain loop to collect init params ..." · Reviewed by Greptile}

[skyw]

Largely LGTM. Couple of small comments, will approve after.

Thanks for the contribution. It has been on my list to redesign the registry, it is very limited and not definitely not bullet proof at the moment.

[skyw]

class?

[v0ropaev]

Renamed it to cls, class itself is a keyword.

[skyw]

hard to read comprehension, favor plain for loop.

[v0ropaev]

Done, it's a plain loop now.

[skyw]

/ok to test 521873f