Expose num_ns_steps

[janEbert]

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the
scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

[greptile-apps]

Greptile Summary

Exposes Muon’s construction-time num_ns_steps value through a read-only property.

• Stores the validated Newton–Schulz step count on each Muon instance.
• Prevents post-initialization assignment because the orthogonalization closure retains the original value.

Confidence Score: 5/5

The PR appears safe to merge with no actionable defects identified.

The new property is initialized from the same validated constructor argument used by the orthogonalization closure, and current subclass and serialization paths preserve that value consistently.

Important Files Changed

Filename	Overview
emerging_optimizers/orthogonalized_optimizers/muon.py	Adds a read-only property that consistently reports the same validated step count captured by Muon’s orthogonalization closure.

_{Reviews (1): Last reviewed commit: "Expose `num_ns_steps`" | Re-trigger Greptile}

[skyw]

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

It can be useful to obtain the num_ns_steps in certain cases, e.g., to balance compute around existing instances' num_ns_steps.

Could also refactor this to allow people to set num_ns_steps post-init by either (a) creating a new scaled_orthogonalize_fn when setting num_ns_steps or (b) simply passing self into the scaled_orthogonalize_fn closure, so that the instance's current num_ns_steps are always accessed.

It is very dangerous to modify scaled_orthogonalize_fn post initialization, so none of its arguments is maintained. What use cases do you have in mind, Jan? For simply use, inspect signiture can provide that information.

Having said that, we designed for relatively simple customization, like QKV split. If more complex customization is needed, we should revisit the entire design of OrthogonalizedOptimizer.

[janEbert]

It is very dangerous to modify scaled_orthogonalize_fn post initialization, so none of its arguments is maintained. What use cases do you have in mind, Jan? For simply use, inspect signiture can provide that information.

Enabling also setting num_ns_steps post-init was more a suggestion for potential build on top of this. Whether it's useful is up to you, but it isn't really related to this PR. It was just an idea that I could follow this PR up with.

For getting num_ns_steps from the closure: yes, it's possible with inspect.getclosurevars, but it's super hacky and I'd prefer to have a "simple" public-facing API to get the value.

The value is necessary to balance parallel work around the compute cost of the optimizer. I'm currently using this formula to approximate NS orthogonalization work:

cost = numel * (short_dim * num_steps + 1)

[skyw]

It is very dangerous to modify scaled_orthogonalize_fn post initialization, so none of its arguments is maintained. What use cases do you have in mind, Jan? For simply use, inspect signiture can provide that information.

Enabling also setting num_ns_steps post-init was more a suggestion for potential build on top of this. Whether it's useful is up to you, but it isn't really related to this PR. It was just an idea that I could follow this PR up with.

For getting num_ns_steps from the closure: yes, it's possible with inspect.getclosurevars, but it's super hacky and I'd prefer to have a "simple" public-facing API to get the value.

The value is necessary to balance parallel work around the compute cost of the optimizer. I'm currently using this formula to approximate NS orthogonalization work:

cost = numel * (short_dim * num_steps + 1)

The design was against post initialization change of orthogonalization method, which would open too many doors and needs to be carefully guarded and tested. I don't have any case that we need to change that inflight yet.

Caller has all the information and cam maintain them if necessary. We can revisit when we have a case that Muon optimizer itself needs it.

Alternative is storing and exposing all arguments to orthogonalization, which can be justified.

[janEbert]

It's fair to not allow post-initialization changes.

As for the num_ns_steps information being available to callers: not necessarily. If I get a Muon optimizer instance passed from elsewhere and need to know its num_ns_steps as an outside observer, I currently cannot get this information.

What's the issue with adding the read-only attribute? Can I address it somehow?

[skyw]

It's fair to not allow post-initialization changes.

As for the num_ns_steps information being available to callers: not necessarily. If I get a Muon optimizer instance passed from elsewhere and need to know its num_ns_steps as an outside observer, I currently cannot get this information.

What's the issue with adding the read-only attribute? Can I address it somehow?

It's weird to have only num_ns_steps expose but not others, for example the coefficient type is highly coupled with number of steps. Understand calculating only steps matters for calculating cost.

I'd say do one of the following for short term:

• single private attribute _num_ns_steps, basically remove the property and setter guard. implies for internal (and adult) use only.
• make num_ns_steps, coefficient_type, scale_mode, extra_scale_factor property. basically saving everything need for orthogonalization.

Long term, there is a pending item on my end to make orthogonalizer class that should store those attribute. We started with very simple function but the level of customization we later added has beyond originally designed for .