Amuse optimizer

docs/configuration_reference/optimizer_settings.rst

@@ -142,6 +142,12 @@ optimizer
     A dictionary with a ``class`` entry for the optimizer.
     Other keys are passed as parameters to the constructor of the optimizer class.
 
+    With ``class: "multi"``, multiple optimizers can be combined over disjoint parameter subsets,
+    e.g. ``torch.optim.Muon`` (which only accepts 2D parameters) for the hidden matrix weights
+    together with ``torch.optim.AdamW`` for all remaining parameters.
+    See the module docstring of :mod:`returnn.torch.optim.multi` for the config interface
+    (PyTorch backend only).
+
 relative_error_div_by_old
     If true the relative error is computed by dividing the error difference by the old error value instead of the
     current error value.

returnn/torch/engine.py

@@ -474,6 +474,8 @@ def train_epoch(self):
         self._reset_dev_memory_stats()
 
         self._on_epoch_start(dataset_name="train")
+        # Schedule-free optimizers: switch to the training iterate (see set_optimizer_training_mode).
+        self._updater.set_optimizer_training_mode(train=True)
 
         if self.config.bool("debug_shell_before_train_loop", False):
             print("debug_shell_before_train_loop", file=log.v1)
@@ -752,6 +754,8 @@ def _debug_func() -> torch.Tensor:
         self._maybe_report_dev_memory_stats()
 
         self._on_epoch_end(dataset_name="train")
+        # Schedule-free optimizers: switch to the averaged weights (see set_optimizer_training_mode).
+        self._updater.set_optimizer_training_mode(train=False)
 
         if self.epoch % self._save_model_epoch_interval == 0 or self.epoch == self._final_epoch:
             if self.model_filename:

returnn/torch/optim/README.md

@@ -1,5 +1,6 @@
 Here we can put some arbitrary external optimizers.
 It might be copied from some existing code, or our own implementation.
 It might also happen that some of these will be added to later versions of PyTorch.
-So, regarding the user config, the optimizers here should be differentiated
-by having the full module name, e.g. like ``returnn.torch.optim.lion.Lion``.
+The optimizers here can be referenced in the user config
+by short name (e.g. ``lion``, ``amuse``, ``multi``; ``torch.optim`` names take precedence)
+or by full module name, e.g. ``returnn.torch.optim.lion.Lion``.