Refactor of pre/post weight update hook functions by mkhona-nvidia · Pull Request #224 · NVIDIA-NeMo/Emerging-Optimizers

mkhona-nvidia · 2026-06-04T18:10:07Z

Weight Update Hooks

This change adds a small weight_update_hooks library for reusable behavior around an optimizer's final in-place parameter update. A hook is a configured object passed into an optimizer, for example:

from emerging_optimizers.weight_update_hooks import RadialBrake

optimizer = Muon(
    params,
    weight_decay=0.0,
    weight_update_hook=RadialBrake(outward_scale_factor=0.5),
)

The optimizer owns only one hook object. Hook-specific arguments live on the hook constructor, not on the optimizer constructor. The base update flow is:

pre_update_state = weight_update_hook.pre_weight_update_inplace(p, update)
p.add_(update, alpha=-lr)
post_weight_update_fn_inplace(p)
weight_update_hook.post_weight_update_inplace(p, pre_update_state)

pre_update_state is a transient value returned by the pre hook and immediately consumed by the post hook for the same parameter. It is not stored in optimizer.state, so it does not pollute checkpoints.

The library currently provides three hook implementations, including the newly added RadialBrake based on https://nilin.github.io/radial-brake/:

NoOpWeightUpdateHook: default no-op behavior.
Hyperball: normalizes the update to a target radius before the update, then projects the updated weight back to that radius.
RadialBrake: applies the normal optimizer update, then rescales the updated weight so radial norm changes are damped. For pre-update weight w_prev and updated weight w, it sets:

$$\|w_{\text{brake}}\| = \|w_{\text{prev}}\| + s(\|w\| - \|w_{\text{prev}}\|)$$

where s = outward_scale_factor if the update increases the norm, otherwise s = inward_scale_factor.

MuonHyperball now uses this shared hook machinery by passing Hyperball(radius=hyperball_radius, eps=hyperball_eps) into Muon, instead of implementing custom pre/post update methods and storing temporary values in optimizer state.

Signed-off-by: mikail <mkhona@nvidia.com>

copy-pr-bot · 2026-06-04T18:10:10Z

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 · 2026-06-04T18:14:57Z

Greptile Summary

The PR introduces reusable pre/post weight-update hooks and migrates MuonHyperball onto the shared hook mechanism.

Adds Hyperball, RadialBrake, RelativeUpdate, and no-op hook implementations.
Integrates hook execution into the shared orthogonalized optimizer step.
Adds validation and focused tests for hook behavior.
The Hyperball zero-norm fix can still leave a parameter at zero rather than its configured radius.

Confidence Score: 4/5

The PR is not yet safe to merge because Hyperball can leave an updated parameter at zero instead of restoring its configured radius.

The shared Hyperball scaling helper selects a zero multiplier when the post-update parameter norm is below eps, so an exactly cancelling update violates the optimizer's fixed-radius invariant.

Files Needing Attention: emerging_optimizers/weight_update_hooks/hyperball.py

Important Files Changed

Filename	Overview
emerging_optimizers/weight_update_hooks/hyperball.py	Adds fixed-radius normalization, but its zero-norm post-update branch leaves parameters at zero and violates the radius constraint.
emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py	Integrates transient pre/post hook execution around the final in-place parameter update.
emerging_optimizers/orthogonalized_optimizers/muon_hyperball.py	Migrates MuonHyperball to the shared Hyperball hook and retains construction-time radius validation.
emerging_optimizers/weight_update_hooks/radial_brake.py	Adds bounded radial norm damping and now prevents amplifying scale factors.
emerging_optimizers/weight_update_hooks/relative_update.py	Adds pre-update scaling that matches update norm to the current weight norm.
tests/test_weight_update_hooks.py	Covers normal, boundary, and validation behavior for the new hooks.

Sequence Diagram

sequenceDiagram
    participant Step as OrthogonalizedOptimizer.step
    participant Hook as WeightUpdateHook
    participant Param as Parameter
    Step->>Hook: pre_weight_update_inplace(p, update)
    Hook-->>Step: transient hook state
    Step->>Param: "p.add_(update, alpha=-lr)"
    Step->>Step: optimizer-specific post hook
    Step->>Hook: post_weight_update_inplace(p, state)
    Hook->>Param: project/rescale updated parameter

_{Reviews (3): Last reviewed commit: "Merge branch 'NVIDIA-NeMo:main' into mkh..." | Re-trigger Greptile}

greptile-apps · 2026-06-04T18:15:00Z

+        current_norm = torch.linalg.vector_norm(p.detach().to(torch.float32))
+        if current_norm.item() == 0:
+            raise ValueError("Hyperball requires all parameters to have non-zero norm.")
+
+        radius = (
+            torch.as_tensor(self.radius, device=p.device, dtype=torch.float32)
+            if self.radius is not None
+            else current_norm
+        )


The zero-norm guard fires even when self.radius is not None, but in that branch current_norm is never used as the radius — it is only assigned to radius when self.radius is None. A fixed-radius Hyperball applied to a parameter that starts at zero (e.g. a zero-initialized bias) will therefore raise a ValueError on the very first optimizer step, even though the fixed-radius path has no mathematical dependency on the pre-update parameter norm.

Suggested change

current_norm = torch.linalg.vector_norm(p.detach().to(torch.float32))

if current_norm.item() == 0:

raise ValueError("Hyperball requires all parameters to have non-zero norm.")

radius = (

torch.as_tensor(self.radius, device=p.device, dtype=torch.float32)

if self.radius is not None

else current_norm

)

current_norm = torch.linalg.vector_norm(p.detach().to(torch.float32))

if self.radius is not None:

radius = torch.as_tensor(self.radius, device=p.device, dtype=torch.float32)

else:

if current_norm.item() == 0:

raise ValueError("Hyperball requires all parameters to have non-zero norm when radius is not fixed.")

radius = current_norm

Signed-off-by: mikail <mkhona@nvidia.com>

skyw

Code logic LGTM. A lot of test names can be improved, not super critical.

Before merge, please get familiar with protocol, i.e. being able to explain what purpose it serves here and why not use subclass.

One critical thing is 0 handling, i.e. epsilon.

at least, it should be consistently applied, for vector_norm for example.
Comparing a floating point number directly against zero (a == 0 for example) is actually testing underflow, not numerical 0. Same eps apply, if numbers smaller than eps are considered 0, logical test to determine a numerical 0 should be abs(a) < eps
Magnitude aware handling maybe out of scope, but something to consider as we running into this more and more.

skyw · 2026-06-22T20:07:44Z


                # perform weight update with pre and post weight update functions for subclass customization
                self.pre_weight_update_fn_inplace(p, orth_grad)
+                weight_update_hook_pre_update_state = self.weight_update_hook.pre_weight_update_inplace(p, orth_grad)


name is too long, two update in the name.

skyw · 2026-06-22T20:10:07Z

+__all__ = ["NoOpWeightUpdateHook", "WeightUpdateHook"]
+
+
+class WeightUpdateHook(Protocol):


@mkhona-nvidia I don't oppose using Protocol. But I think need to get yourself comfortable with PEP544 before this can be merged.

skyw · 2026-06-22T20:12:19Z

+        p: torch.Tensor,
+        update: torch.Tensor,
+    ) -> torch.Tensor:
+        current_norm = torch.linalg.vector_norm(p.detach().to(torch.float32))


detach is not necessary as all of our optimizers are wrapped in no_grad.

Actually, vector_norm has an dtype argument, explicit to fp32 is not necessary.

It technically not dtype but compute type though, don't know who added it to pytorch.

skyw · 2026-06-22T20:16:40Z

+        if self.radius is not None:
+            radius = torch.as_tensor(self.radius, device=p.device, dtype=torch.float32)
+        else:
+            if current_norm.item() == 0:


This triggers synced device to host copy. Can't have it in every call. We may have it already, but it needs to be fixed. Otherwise can't use in any pratical runs.

skyw · 2026-06-22T20:17:06Z

+    def __init__(
+        self,
+        radius: float | None = None,
+        eps: float = 1e-8,


Maybe reduce this to 1e-15?

skyw · 2026-06-22T20:21:43Z

+        pre_norm = pre_update_state
+        post_norm = torch.linalg.vector_norm(p.detach().to(torch.float32))
+        norm_delta = post_norm - pre_norm
+        scale_factor = self.outward_scale_factor if norm_delta.item() > 0 else self.inward_scale_factor


same as before, don't use .item().

skyw · 2026-06-22T20:22:16Z

+        """MuonHyperball manages its own Hyperball hook internally."""
+        test_param = nn.Parameter(torch.randn((5, 7), dtype=torch.float32, device=self.device))
+
+        with self.assertRaisesRegex(TypeError, "does not accept a 'weight_update_hook' argument"):


This is better to be a KeyError,

skyw · 2026-06-22T20:23:15Z

+        pre_update_state = hook.pre_weight_update_inplace(param, update)
+        hook.post_weight_update_inplace(param, pre_update_state)
+
+        torch.testing.assert_close(param, param_before, atol=0.0, rtol=0.0)


Good to test exact match.

We should probably have a partial function called assert_equal given it is so widely used.

skyw · 2026-06-22T20:26:26Z

+        torch.testing.assert_close(param, param_before, atol=0.0, rtol=0.0)
+        torch.testing.assert_close(update, update_before, atol=0.0, rtol=0.0)
+
+    def test_radial_brake_dampens_outward_norm_change(self) -> None:


Suggestion: always state expected behavior in test name. e.g. in this case, what close to what. "change" is too vague of a behavior.

skyw · 2026-06-22T20:33:27Z

        weight_decay_method: opt_mixin.WeightDecayT,
        fp32_matmul_prec: FP32MatmulPrecT,
        scaled_orthogonalize_fn: Callable | None = None,
+        weight_update_hook: WeightUpdateHook | None = None,


Let's make it private, _weight_update_hook, to suggest it is not supposed to be modified after initialization.

Hook easily invites abuse. making it private at least making people aware of the abusing. A setter can be provided if we want to support properly change the hook inflight.

Signed-off-by: mkhona <mkhona@nvidia.com>

…factors

greptile-apps · 2026-08-11T15:50:47Z

+        scale = torch.where(is_numerical_zero, torch.zeros_like(norm), radius / norm.clamp_min(self.eps))
+        tensor.mul_(scale.to(dtype=tensor.dtype))


Zero projection breaks radius constraint

When the normalized update exactly cancels the parameter, the post-update norm falls below eps and this branch multiplies the parameter by zero. The hook therefore leaves the parameter at zero instead of projecting it back to the configured nonzero radius.

Knowledge Base Used: Orthogonalized Optimizers

skyw

Bottom line is it can be merged after trimming the tests.

skyw · 2026-08-11T16:30:26Z

+
+
+__all__ = [
+    "Hyperball",


Q: why some has Hook while others don't?
If there is no particular reason, make naming consistent.

skyw · 2026-08-11T16:32:06Z

+class WeightUpdateHook(Protocol[HookStateT]):
+    """Static structural contract for behavior around the final weight update.
+
+    PEP 544 structural typing lets third-party hooks satisfy this interface without


too verbose and not always accurate, remove.

skyw · 2026-08-11T16:36:30Z

        scale_mode: MuonScaleT = "spectral",
        extra_scale_factor: float = 1.0,
        use_syrk: bool = False,
+        weight_update_hook: WeightUpdateHook[Any] | None = None,


Q: what's the reason to use Any while a HookStateT is defined?
can set bound to HookStateT to accept known types it needs to handle.

skyw · 2026-08-11T16:40:24Z

+class RadialBrake:
+    """Dampen radial norm changes after an optimizer update.
+
+    The optimizer first applies its usual update ``w = w_prev + dw``. This hook then rescales ``w`` so that


Add argument section.

skyw · 2026-08-11T16:41:56Z

+    """
+
+    def __init__(self, eps: float = 1e-15) -> None:
+        if not math.isfinite(eps) or eps <= 0.0:


This is overkill, check magnitude on eps would be enough.
It would be strange to check if values are finite at one place while not everywhere else

skyw · 2026-08-11T16:42:43Z

+        param.add_(update)
+        hook.post_weight_update_inplace(param, pre_update_state)
+
+        torch.testing.assert_close(torch.linalg.vector_norm(param), torch.tensor(7.5, device=self.device))


Comment why it should be 7.5

skyw · 2026-08-11T16:43:10Z

+        param.add_(update)
+        hook.post_weight_update_inplace(param, pre_update_state)
+
+        torch.testing.assert_close(torch.linalg.vector_norm(param), torch.tensor(9.0, device=self.device))


Same as above. always comment hardcoded test target.

skyw · 2026-08-11T16:47:42Z

@@ -0,0 +1,234 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.


Too many hardcoded simple tests don't do a lot different, please trim.

Also think about how to improve tests coverage, includes but not limits to:

Coverage, a lot of tests only test shape 2 input.

Could it be exact match? can manipulate to value to have exactly representable norm for example.

refactor of pre/post weight update hook functions

a49944b

Signed-off-by: mikail <mkhona@nvidia.com>

mkhona-nvidia requested a review from skyw June 4, 2026 18:10

mkhona-nvidia changed the title ~~refactor of pre/post weight update hook functions~~ Refactor of pre/post weight update hook functions Jun 4, 2026

greptile-apps Bot reviewed Jun 4, 2026

View reviewed changes

fixed greptile caught errors

16d9c70

Signed-off-by: mikail <mkhona@nvidia.com>

skyw requested changes Jun 22, 2026

View reviewed changes

mkhona-nvidia added 3 commits August 10, 2026 16:51

Merge main and address review feedback

a463197

Signed-off-by: mkhona <mkhona@nvidia.com>

Add relative update hook

3e2753e

Signed-off-by: mkhona <mkhona@nvidia.com>

Merge branch 'NVIDIA-NeMo:main' into mkhona/pre_post_weight_update_re…

3164484

…factors

greptile-apps Bot reviewed Aug 11, 2026

View reviewed changes

skyw requested changes Aug 11, 2026

View reviewed changes

		__all__ = ["NoOpWeightUpdateHook", "WeightUpdateHook"]


		class WeightUpdateHook(Protocol):

		scale = torch.where(is_numerical_zero, torch.zeros_like(norm), radius / norm.clamp_min(self.eps))
		tensor.mul_(scale.to(dtype=tensor.dtype))

		@@ -0,0 +1,234 @@
		# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Uh oh!

Conversation

mkhona-nvidia commented Jun 4, 2026

Weight Update Hooks

Uh oh!

copy-pr-bot Bot commented Jun 4, 2026

Uh oh!

greptile-apps Bot commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

skyw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

skyw Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

skyw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

greptile-apps Bot commented Jun 4, 2026 •

edited

Loading

skyw Jun 22, 2026 •

edited

Loading