Amuse optimizer by polycrit · Pull Request #1830 · rwth-i6/returnn

polycrit · 2026-08-03T13:49:57Z

Added an implementation of AMUSE optimizer (mirrored from authors' github repository and adapted for RETURNN as returnn.torch.optim.amuse)
Strip callables from the saved optimizer_opts metadata in updater.py so optimizer checkpoints are loadable under the torch >= 2.6 weights_only default

…lity

albertz · 2026-08-03T15:28:28Z

So currently we don't even have the much more popular Muon optimizer in RETURNN. I don't think adding AMUSE makes sense when we don't have Muon. But maybe we should add both then.

See the failing code style errors also:

  returnn/torch/optim/amuse.py:18: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:19: WEAK WARNING PyPep8NamingInspection: Argument name should be lowercase
  returnn/torch/optim/amuse.py:25: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:28: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:31: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:34: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:35: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:36: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:39: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:44: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:184: WEAK WARNING PyPep8NamingInspection: Variable in function should be lowercase
  returnn/torch/optim/amuse.py:201: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:212: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:223: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:257: WEAK WARNING PyAttributeOutsideInitInspection: Instance attribute beta1 defined outside __init__
  returnn/torch/optim/amuse.py:441: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:443: WARNING PyUnresolvedReferencesInspection: Cannot find reference 'train' in 'None'
  returnn/torch/optim/amuse.py:446: WEAK WARNING PyMissingOrEmptyDocstringInspection: Missing docstring
  returnn/torch/optim/amuse.py:448: WARNING PyUnresolvedReferencesInspection: Cannot find reference 'eval' in 'None'
  returnn/torch/optim/amuse.py:464: ERROR PyUnresolvedReferencesInspection: Unresolved reference 'CodeWrapper'
  returnn/torch/optim/amuse.py:464: ERROR PyUnresolvedReferencesInspection: Unresolved reference 'i6_core'
  returnn/torch/optim/amuse.py:492: ERROR PyUnresolvedReferencesInspection: Unresolved reference 'CallImport'
  returnn/torch/optim/amuse.py:492: ERROR PyUnresolvedReferencesInspection: Unresolved reference 'Import'
  returnn/torch/optim/amuse.py:492: ERROR PyUnresolvedReferencesInspection: Unresolved reference 'i6_core'

polycrit · 2026-08-03T15:31:08Z

So currently we don't even have the much more popular Muon optimizer in RETURNN. I don't think adding AMUSE makes sense when we don't have Muon. But maybe we should add both then.

was also thinking about this, makes total sense, I'll add it as well

albertz · 2026-08-03T15:37:02Z

Re Muon, I always use i6_experiments.users.zeyer.experiments.exp2024_04_23_baselines.optim_ext.muon.Muon. Maybe you can copy that code?

polycrit · 2026-08-03T15:50:42Z

Re Muon, I always use i6_experiments.users.zeyer.experiments.exp2024_04_23_baselines.optim_ext.muon.Muon. Maybe you can copy that code?

I just found oud there's an already implemented Muon from pytorch (https://docs.pytorch.org/docs/2.13/generated/torch.optim.Muon.html). So I think maybe we don't need it separately implemented after all? Just a small adapter.

albertz · 2026-08-03T15:53:41Z

Ah Muon seems to exist since PyTorch 2.10. Hm so good question. Maybe you are right.

polycrit · 2026-08-03T16:00:57Z

Ah Muon seems to exist since PyTorch 2.10. Hm so good question. Maybe you are right.

The problem is that in RETURNN you can't do something like:

muon = torch.optim.Muon(hidden_matrix_params, ...)
adamw = torch.optim.AdamW(other_params, ...)

loss.backward()
muon.step()
adamw.step()

RETURNN’s PyTorch updater currently manages one optimizer object. So what we could implement there is a compact Muon/AdamW adapter that presents Muon and AdamW as that single object.

albertz · 2026-08-03T16:04:38Z

Doesn't the Torch Muon already handle this automatically? At least i6_experiments.users.zeyer.experiments.exp2024_04_23_baselines.optim_ext.muon.Muon does. (Hybrid Muon (2-D hidden weights) + AdamW (everything else))

albertz · 2026-08-03T16:07:53Z

Btw, I think it was discussed maybe in the past whether we somehow should have a possibility to use multiple optimizers for different params. Not sure?

polycrit · 2026-08-03T16:09:01Z

Doesn't the Torch Muon already handle this automatically? At least i6_experiments.users.zeyer.experiments.exp2024_04_23_baselines.optim_ext.muon.Muon does. (Hybrid Muon (2-D hidden weights) + AdamW (everything else))

As far I understand, it doesn't, you have to explicitly call two separate optimizers.

polycrit · 2026-08-03T16:09:53Z

Btw, I think it was discussed maybe in the past whether we somehow should have a possibility to use multiple optimizers for different params. Not sure?

Can't remember this, but I think it's generally a very good idea, given we want to experiment with such multi-level optimizers. I will look into it.

polycrit · 2026-08-03T18:36:26Z

@albertz I imagine we can create an interface that does something like this:

optimizer = {
    "class": "multi",
    "optimizers": [
        {
            "class": "amuse",
            "update_type": "muon",
            "params_filter": make_hidden_matrix_filter(),
            "momentum": 0.95,
            "warmup_steps": 10000,
        },
        {
            "class": "amuse",
            "update_type": "adamw",
            "learning_rate_multiplier": 0.015,
            "warmup_steps": 10000,
        },
    ],
}

or:

optimizer = {
    "class": "multi",
    "optimizers": [
        {
            "class": "muon",
            "params_filter": make_hidden_matrix_filter(),
            "momentum": 0.95,
            "adjust_lr_fn": "match_rms_adamw",
            "weight_decay": 1e-2,
        },
        {
            "class": "adamw",
            "learning_rate_multiplier": 0.015, 
            "weight_decay": 1e-2,
            "epsilon": 1e-8,
        },
    ],
}

params_filter defines parameter ownership here, which accepts a predicate function like:

def my_filter(*, full_param_name: str, param: torch.nn.Parameter,
              module: torch.nn.Module, rf_module, **_kwargs) -> bool:
    ...

similar to RETURNN's existing weight_decay_custom_include_check (updater.py:619-623). Example for Muon:

def muon_filter(*, full_param_name, param, **_kw):
    lname = full_param_name.lower()
    return param.dim() == 2 and not any(s in lname for s in ("embed", "logit", "head"))

albertz · 2026-08-03T19:22:50Z

Yea looks good. Feel free to implement that.

…er-sub weight decay split

Copilot

Pull request overview

This PR adds two new Torch optimizers to RETURNN—AMUSE (schedule-free) and MultiOptimizer (composite optimizer over disjoint parameter subsets)—and updates optimizer checkpoint metadata handling so checkpoints remain loadable with PyTorch’s weights_only=True default (>= 2.6). It also wires schedule-free optimizer train()/eval() switching into the Torch engine epoch boundaries and adds extensive tests + documentation.

Changes:

Add returnn.torch.optim.amuse.AMUSE and returnn.torch.optim.multi.MultiOptimizer (plus helper filter) for mixed-optimizer setups.
Extend optimizer resolution + updater logic (multi-optimizer construction, param-group splitting on subsets, schedule-free mode switching, callable stripping in saved metadata).
Add engine hooks, tests, and docs for the new optimizers and config interface.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_torch_engine.py	Adds coverage for multi-optimizer behavior, schedule-free hooks, AMUSE integration, and save/load contracts.
returnn/util/basic.py	Extends `collect_class_init_kwargs` to include keyword-only args/defaults.
returnn/torch/updater.py	Adds short-name resolution for shipped optimizers; implements `multi` optimizer construction; adds schedule-free mode switching; strips callables from saved optimizer metadata.
returnn/torch/optim/README.md	Documents short-name referencing for shipped optimizers.
returnn/torch/optim/multi.py	Introduces `MultiOptimizer` composite optimizer and a picklable hidden-matrix filter helper.
returnn/torch/optim/amuse.py	Adds AMUSE schedule-free optimizer implementation with selectable inner update type.
returnn/torch/engine.py	Switches schedule-free optimizers to train/eval modes at train epoch boundaries.
docs/configuration_reference/optimizer_settings.rst	Documents `optimizer: {class: "multi", ...}` in the config reference.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+        optimizer_opts_to_save = self._optimizer_opts
+        if isinstance(optimizer_opts_to_save, dict):
+            optimizer_opts_to_save = _drop_callables_deep(optimizer_opts_to_save)


+    def __init__(self, sub_optimizers: Sequence[torch.optim.Optimizer]):
+        self._sub_optimizers = list(sub_optimizers)
+
+    def _owning_sub(self, param: torch.nn.Parameter) -> Optional[torch.optim.Optimizer]:
+        for sub in self._sub_optimizers:
+            for group in sub.param_groups:
+                for other in group["params"]:
+                    if other is param:
+                        return sub
+        return None


+        c_warmup = group.get("c_warmup", 1.0 / self.warmup_steps)
+        s_t = (ckp1 * (1.0 - c_warmup)) / (c_warmup * (1.0 - ckp1))
+        return 1.0 - (s_t**self.rho) * (1.0 - self.beta1_init)


…configs

… view

polycrit added 2 commits August 3, 2026 15:31

add AMUSE optimizer

01cbf09

fix optimizer checkpoint saving for the torch>=2.6 weights_only default

fe9dded

polycrit requested review from a team and albertz as code owners August 3, 2026 13:49

polycrit added 2 commits August 3, 2026 15:54

apply ruff formatting

c77eb5c

conditionally apply add_safe_globals for torch < 2.4 import compatibi…

ce30c8a

…lity

polycrit force-pushed the amuse-optimizer branch from e41c0a2 to ce30c8a Compare August 3, 2026 13:57

polycrit marked this pull request as draft August 3, 2026 19:40

polycrit added 7 commits August 5, 2026 17:28

fix collect_class_init_kwargs to also collect keyword-only args

224b0f8

add MultiOptimizer composing sub-optimizers over disjoint param subsets

55bc603

resolve returnn custom optimizers by short name (multi, lion, amuse)

c9068f4

create multi optimizer from config with params_filter partition and p…

3f006ba

…er-sub weight decay split

call schedule-free optimizer train()/eval() at train epoch boundaries

ead1878

restructure AMUSE to one update type per instance, composed via multi

0f331d0

Merge remote-tracking branch 'origin/master' into amuse-optimizer

eee94f4

polycrit requested a lite review from Copilot August 5, 2026 20:24

Copilot started reviewing on behalf of polycrit August 5, 2026 20:25 View session

Copilot AI reviewed Aug 5, 2026

View reviewed changes

polycrit added 5 commits August 5, 2026 22:43

save valid optimizer checkpoint metadata also for callable optimizer …

0f3a28e

…configs

cache the param to sub-optimizer mapping in the multi optimizer state…

bc7541e

… view

guard AMUSE beta1 computation against a degenerate averaging weight

ce15546

guard AMUSE beta1 ramp against zero lr at the warmup boundary

c128e7e

validate AMUSE constructor arguments

7182473

Navigation Menu

Uh oh!

Amuse optimizer - #1830

Amuse optimizer#1830
polycrit wants to merge 16 commits into
masterfrom
amuse-optimizer

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026 •

edited

Loading

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

polycrit commented Aug 3, 2026

Uh oh!

albertz commented Aug 3, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

albertz commented Aug 3, 2026 •

edited

Loading