Allow per-variable optimizer, add DispatchOptimizer.

- Adds a property `variable.optimizer` that defaults to `None`
- Adds a `DispatchOptimizer` that scans the list of trainable variables during build,
  collects all unique per-variable optimizers, then dispatches the apply/stateless_apply
  function to the correct optimizer if applicable.
- Modifies `trainer` so that during the optimizer build stage, checks if any variables
  have a custom optimizer attached, and if so inserts a `DispatchOptimizer` to properly
  handle them.  This allows usage to be hidden from the user.

Context: for large embedding tables, we need special optimizers to be used so that
the tables can be updated in-place, rather than returning large gradients.  The layer
will handle setting of the custom optimizers, but we need the trainer to be aware
of them and dispatch the embedding tables to different optimizers appropriately.

Author: cantonios

keras/src/backend/common/variables.py

@@ -154,6 +154,8 @@ def __init__(
         # whether this variable should be overwritten by the computed gradient.
         # Ref: https://github.com/google/flax/blob/main/flax/linen/fp8_ops.py
         self._overwrite_with_gradient = False
+        # Per-variable optimizer.
+        self._optimizer = None
         if isinstance(initializer, str):
             from keras.src import initializers
 
@@ -372,6 +374,26 @@ def regularizer(self, value):
             )
         self._regularizer = value
 
+    @property
+    def optimizer(self):
+        """Per-variable custom optimizer."""
+        return self._optimizer
+
+    @optimizer.setter
+    def optimizer(self, value):
+        from keras.src import optimizers
+
+        if isinstance(value, str):
+            value = optimizers.get(value)
+
+        if value is not None and not isinstance(value, optimizers.Optimizer):
+            raise ValueError(
+                "Invalid value for attribute `optimizer`. Expected an "
+                "instance of `keras.optimizers.Optimizer`, or `None`. "
+                f"Received: regularizer={value}"
+            )
+        self._optimizer = value
+
     @property
     def constraint(self):
         return self._constraint

keras/src/backend/common/variables_test.py

@@ -8,6 +8,7 @@
 from keras.src import backend
 from keras.src import initializers
 from keras.src import ops
+from keras.src import optimizers
 from keras.src.backend.common import dtypes
 from keras.src.backend.common.variables import AutocastScope
 from keras.src.backend.common.variables import shape_equal
@@ -419,6 +420,18 @@ def test_deferred_initialize_within_stateless_scope(self):
             ):
                 v._deferred_initialize()
 
+    def test_optimizer_setter(self):
+        v = backend.Variable(
+            initializer=initializers.RandomNormal(),
+            shape=(2, 2),
+        )
+        self.assertIsNone(v.optimizer)
+        v.optimizer = "sgd"
+        self.assertTrue(isinstance(v.optimizer, optimizers.Optimizer))
+
+        with self.assertRaisesRegex(ValueError, "Invalid value"):
+            v.optimizer = True
+
 
 class VariableDtypeShapeNdimRepr(test_case.TestCase):
     """tests for dtype, shape, ndim, __repr__"""

keras/src/backend/tensorflow/optimizer.py

@@ -71,9 +71,25 @@ def assign_sub(self, variable, value):
         else:
             variable.assign_sub(value)
 
-    def _var_key(self, variable):
+    def _convert_to_tf_variable(self, variable):
         if isinstance(variable, backend.Variable):
-            variable = variable.value  # Convert to tf.Variable
+            tf_variable = variable.value
+            # Copy additional properties.
+            if getattr(variable, "optimizer", None) is not None:
+                tf_variable.optimizer = variable.optimizer
+            if getattr(variable, "overwrite_with_gradient", False):
+                tf_variable.overwrite_with_gradient = True
+            return tf_variable
+        elif isinstance(variable, tf.Variable):
+            return variable
+        else:
+            raise ValueError(
+                f"Variable {variable} must be of type keras.Variable or "
+                f"tf.Variable, received {value.__class__.__name__}."
+            )
+
+    def _var_key(self, variable):
+        variable = self._convert_to_tf_variable(variable)
         if hasattr(variable, "_distributed_container"):
             variable = variable._distributed_container()
         elif (
@@ -98,8 +114,7 @@ def weight_decay_fn(variable):
                     variable.assign_sub(variable * wd * lr)
 
             for variable in variables:
-                if isinstance(variable, backend.Variable):
-                    variable = variable.value  # Convert to tf.Variable
+                variable = self._convert_to_tf_variable(variable)
                 distribution.extended.update(
                     variable, weight_decay_fn, group=False
                 )

keras/src/optimizers/__init__.py

@@ -5,6 +5,7 @@
 from keras.src.optimizers.adam import Adam
 from keras.src.optimizers.adamax import Adamax
 from keras.src.optimizers.adamw import AdamW
+from keras.src.optimizers.dispatch_optimizer import DispatchOptimizer
 from keras.src.optimizers.ftrl import Ftrl
 from keras.src.optimizers.lion import Lion
 from keras.src.optimizers.loss_scale_optimizer import LossScaleOptimizer

keras/src/optimizers/base_optimizer.py

@@ -204,13 +204,27 @@ def iterations(self):
     def _track_variable(self, variable):
         self._tracker.add_to_store("variables", variable)
 
+    def _has_custom_optimizer(self, variable):
+        return (
+            hasattr(variable, "optimizer")
+            and variable.optimizer is not None
+            and variable.optimizer != self
+        )
+
     @tracking.no_automatic_dependency_tracking
     def build(self, variables):
         if self.use_ema:
             self._model_variables_moving_average = []
         if self.gradient_accumulation_steps:
             self._accumulated_gradients = []
         for i, variable in enumerate(variables):
+            if self._has_custom_optimizer(variable):
+                warnings.warn(
+                    f"Variable {variable} has a custom optimizer "
+                    f"{variable.optimizer} that is being ignored. "
+                    "See `keras.optimizers.DispatchOptimizer` to allow "
+                    "dispatching to the correct per-variable optimizer."
+                )
             self._trainable_variables_indices[self._var_key(variable)] = i
             if self.use_ema:
                 self._model_variables_moving_average.append(
@@ -568,7 +582,7 @@ def stateless_apply(self, optimizer_variables, grads, trainable_variables):
             )
         if len(trainable_variables) != len(self._trainable_variables):
             raise ValueError(
-                "Argument `optimizer_variables` must be a list of tensors "
+                "Argument `trainable_variables` must be a list of tensors "
                 "corresponding 1:1 to the trainable variables list that "
                 "the optimizer was built with. Received "
                 f"len(trainable_variables) == {len(trainable_variables)} "