From a748e56622615cfd07679ced2331121f1a350a4d Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Tue, 28 Jul 2026 20:21:54 -0500 Subject: [PATCH] Memoize optimizer state per rule and parameter --- pytensor_ml/optim/alias.py | 17 +++++++- pytensor_ml/optim/base.py | 86 +++++++++++++++++++++++++++++++------- tests/optim/test_rules.py | 53 +++++++++++++++++++++++ 3 files changed, 141 insertions(+), 15 deletions(-) diff --git a/pytensor_ml/optim/alias.py b/pytensor_ml/optim/alias.py index d3917d1..da9b15e 100644 --- a/pytensor_ml/optim/alias.py +++ b/pytensor_ml/optim/alias.py @@ -1,6 +1,12 @@ from collections.abc import Callable, Sequence -from pytensor_ml.optim.base import LossOrGradients, Parameter, UpdateRule, Updates +from pytensor_ml.optim.base import ( + LossOrGradients, + Parameter, + UpdateRule, + Updates, + reuses_state, +) from pytensor_ml.optim.rules import ( adadelta_updates, adagrad_updates, @@ -29,6 +35,7 @@ def sgd(learning_rate: float = 0.01, momentum: float = 0.0, nesterov: bool = Fal Use Nesterov momentum. Ignored when ``momentum`` is 0. Default False. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: if not momentum: return sgd_updates(loss_or_gradients, parameters, learning_rate=learning_rate) @@ -50,6 +57,7 @@ def adam( Adam optimizer. See :func:`~pytensor_ml.optim.rules.adam_updates` for the update rule. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return adam_updates( loss_or_gradients, @@ -78,6 +86,7 @@ def adamw( :func:`~pytensor_ml.optim.rules.adamw_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return adamw_updates( loss_or_gradients, @@ -105,6 +114,7 @@ def nadam( :func:`~pytensor_ml.optim.rules.nadam_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return nadam_updates( loss_or_gradients, @@ -129,6 +139,7 @@ def adamax( :func:`~pytensor_ml.optim.rules.adamax_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return adamax_updates( loss_or_gradients, @@ -153,6 +164,7 @@ def rprop( Rprop optimizer (resilient backpropagation). See :func:`~pytensor_ml.optim.rules.rprop_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return rprop_updates( loss_or_gradients, @@ -178,6 +190,7 @@ def rmsprop( RMSProp optimizer. See :func:`~pytensor_ml.optim.rules.rmsprop_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return rmsprop_updates( loss_or_gradients, @@ -197,6 +210,7 @@ def adagrad(learning_rate: float = 0.01, epsilon: float = 1e-8) -> UpdateRule: AdaGrad optimizer. See :func:`~pytensor_ml.optim.rules.adagrad_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return adagrad_updates( loss_or_gradients, parameters, learning_rate=learning_rate, epsilon=epsilon @@ -210,6 +224,7 @@ def adadelta(learning_rate: float = 1.0, rho: float = 0.9, epsilon: float = 1e-8 AdaDelta optimizer. See :func:`~pytensor_ml.optim.rules.adadelta_updates`. """ + @reuses_state def rule(loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter]) -> Updates: return adadelta_updates( loss_or_gradients, parameters, learning_rate=learning_rate, rho=rho, epsilon=epsilon diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py index 4000373..8b06100 100644 --- a/pytensor_ml/optim/base.py +++ b/pytensor_ml/optim/base.py @@ -1,4 +1,6 @@ from collections.abc import Callable, Sequence +from contextvars import ContextVar +from functools import wraps import numpy as np import pytensor @@ -52,13 +54,67 @@ def get_gradients( return grad(rewrite_pregrad(loss_or_gradients), list(parameters)) # type: ignore[return-value] +# A per-parameter slot, or the bare name of a rule-wide counter. +type _StateKey = tuple[Parameter, str] | str + +# Bound by reuses_state for the duration of one rule invocation; None means "allocate fresh". +_state_buffers: ContextVar[dict[_StateKey, Parameter] | None] = ContextVar( + "optimizer_state_buffers", default=None +) + + +def reuses_state(rule: UpdateRule) -> UpdateRule: + """ + Give ``rule`` a private set of optimizer-state buffers, reused on every invocation. + + A configured rule such as ``adam(1e-3)`` reads as a value, so it is natural to compile two training + functions from one. Without this, each invocation allocates fresh momentum under the *same* derived + name: the two steps then share parameters but not optimizer state, which is silently wrong at runtime + and raises only later when both are checkpointed together. + + The buffers are keyed per rule rather than globally so two independently configured optimizers stay + independent. They are bound dynamically because :func:`state_for` is reached through the update-rule + functions, several call layers below, and threading a cache down would touch every one of them. + + Parameters + ---------- + rule : UpdateRule + The rule to wrap. Its buffers live as long as the wrapper does. + """ + buffers: dict[_StateKey, Parameter] = {} + + @wraps(rule) + def rule_with_persistent_state( + loss_or_gradients: LossOrGradients, parameters: Sequence[Parameter] + ) -> Updates: + token = _state_buffers.set(buffers) + try: + return rule(loss_or_gradients, parameters) + finally: + _state_buffers.reset(token) + + return rule_with_persistent_state + + +def _reuse_or_allocate(key: _StateKey, allocate: Callable[[], Parameter]) -> Parameter: + buffers = _state_buffers.get() + if buffers is None: + return allocate() + if key not in buffers: + buffers[key] = allocate() + return buffers[key] + + def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Parameter: """ - Allocate an optimizer-state shared variable shaped and typed like ``parameter``. + Return the optimizer-state shared variable shaped and typed like ``parameter``. The variable is named ``"{parameter.name}/{slot}"`` so it can be matched by name at serialization boundaries. The name is never used to *find* the variable at runtime — callers hold the returned object - directly. + directly, and reuse within a rule is keyed on the parameter object, so two same-named parameters still + get distinct buffers and collide loudly at save time rather than silently sharing. + + Allocates unless the enclosing rule was wrapped in :func:`reuses_state` and already holds this slot. Parameters ---------- @@ -68,19 +124,26 @@ def state_for(parameter: Parameter, slot: str, fill_value: float = 0.0) -> Param A short role tag for the slot, e.g. ``"adam/first_moment"`` or ``"trace/velocity"``. fill_value : float Constant to initialize the state with. Default 0.0. - - Returns - ------- - shared tensor variable - A freshly allocated state variable. """ if parameter.name is None: raise ValueError( f"Cannot allocate optimizer state {slot!r} for an unnamed parameter. Stateful optimizers rely on " "parameter names to identify their state at serialization boundaries; give the parameter a name." ) - value = parameter.get_value(borrow=True) - return pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}") + + def allocate() -> Parameter: + value = parameter.get_value(borrow=True) + return pytensor.shared(np.full_like(value, fill_value), name=f"{parameter.name}/{slot}") + + return _reuse_or_allocate((parameter, slot), allocate) + + +def counter(name: str) -> Parameter: + """Return an int64 step counter shared variable initialized to zero, reused across invocations of a + rule wrapped in :func:`reuses_state` so its step count keeps advancing.""" + return _reuse_or_allocate( + name, lambda: pytensor.shared(np.asarray(0, dtype="int64"), name=name) + ) def require_unique_state_names(updates: Updates) -> None: @@ -109,11 +172,6 @@ def require_unique_state_names(updates: Updates) -> None: seen.add(name) -def counter(name: str) -> Parameter: - """Allocate an int64 step counter shared variable initialized to zero.""" - return pytensor.shared(np.asarray(0, dtype="int64"), name=name) - - def chain(*transforms: Transform) -> Transform: """ Compose transforms left to right into a single transform. diff --git a/tests/optim/test_rules.py b/tests/optim/test_rules.py index 838bd1d..0ff19cf 100644 --- a/tests/optim/test_rules.py +++ b/tests/optim/test_rules.py @@ -152,6 +152,59 @@ def test_adam_updates_keyed_by_object_with_named_state(): assert state_names == {"adam/step_count", "w/adam/first_moment", "w/adam/second_moment"} +@pytest.mark.parametrize( + "make_rule", + [lambda: adam(learning_rate=1e-2), lambda: sgd(learning_rate=1e-2, momentum=0.9)], + ids=["state_from_a_rule", "state_from_a_transform"], +) +def test_reused_rule_shares_its_optimizer_state(make_rule): + """A configured rule reads as a value, so compiling two training functions from one is natural. Both + must drive the same buffers: separate ones under the same derived name are silently wrong at runtime, + and collide only later when both are checkpointed together. Momentum SGD is included because its + velocity comes from a transform rather than the rule, which is a separate allocation path.""" + p = trainable(np.zeros(3), name="w") + loss = (p**2).sum() + rule = make_rule() + + first = {key for key in rule(loss, [p]) if key is not p} + second = {key for key in rule(loss, [p]) if key is not p} + + assert first and first == second + + +def test_two_functions_from_one_rule_continue_the_same_momentum(): + """What the shared buffers buy: the second function continues the first's trajectory instead of + restarting it. Under a constant gradient, momentum SGD's step at iteration ``t`` is + ``lr * g * (1 - m**t) / (1 - m)``, so a continued second step is 1.9x a restarted one at ``m = 0.9``.""" + p = trainable(np.zeros(2), name="w") + gradient = np.array([2.0, -0.5]) + loss = (pt.constant(gradient) * p).sum() # constant gradient, independent of p + learning_rate, momentum = 0.1, 0.9 + rule = sgd(learning_rate=learning_rate, momentum=momentum) + + step_once = function([], loss, updates=rule(loss, [p])) + step_again = function([], loss, updates=rule(loss, [p])) + + step_once() + before = p.get_value().copy() + step_again() + + continued = -learning_rate * gradient * (1 - momentum**2) / (1 - momentum) + np.testing.assert_allclose(p.get_value() - before, continued, rtol=1e-6) + + +def test_separately_configured_rules_keep_independent_state(): + """Buffers are memoized per rule, not globally, so two optimizers over the same parameter do not + quietly train through each other's momentum.""" + p = trainable(np.zeros(3), name="w") + loss = (p**2).sum() + + first = {key for key in adam(learning_rate=1e-2)(loss, [p]) if key is not p} + second = {key for key in adam(learning_rate=1e-2)(loss, [p]) if key is not p} + + assert not first & second + + def test_adamw_first_step_applies_decoupled_decay(): """AdamW adds a decoupled decay term to Adam's sign-descent step: the first-step displacement is ``-lr * (sign(g) + weight_decay * p)``. At t = 1 bias correction makes the Adam part ``sign(g)`` (the