returnn 1.20241020.5643__tar.gz → 1.20241022.173306__tar.gz

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/LICENSE

@@ -176,4 +176,11 @@ Contains code from PyTorch:
  - Copyright 2016..2023 various developers
  - Various code snippets
 
+Contains code from lion-pytorch:
+ - https://github.com/google/automl/blob/master/lion/lion_pytorch.py
+ - https://github.com/lucidrains/lion-pytorch/
+ - MIT License / Apache License 2.0
+ - Copyright 2023 Google Research, Phil Wang
+ - torch/optim/lion.py
+
 Various code snippets from StackOverflow, which are under Creative Commons / Public Domain.

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241020.5643
+Version: 1.20241022.173306
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20241022.173306/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241022.173306'
+long_version = '1.20241022.173306+git.f1877f3'

returnn-1.20241022.173306/returnn/torch/optim/README.md

@@ -0,0 +1,5 @@
+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``.

returnn-1.20241022.173306/returnn/torch/optim/__init__.py

@@ -0,0 +1,3 @@
+"""
+Any custom optimizer
+"""

returnn-1.20241022.173306/returnn/torch/optim/lion.py

@@ -0,0 +1,205 @@
+"""
+Lion optimizer <https://arxiv.org/abs/2302.06675>
+
+Code adapted from https://github.com/lucidrains/lion-pytorch/,
+which is adapted from https://github.com/google/automl/blob/master/lion/lion_pytorch.py.
+"""
+
+from __future__ import annotations
+from typing import Optional, Tuple, Callable
+import inspect
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class Lion(Optimizer):
+    """
+    Lion (Evolved Sign Momentum (Evo_l_ved S_i_gn M_o_me_n_tum)) optimizer <https://arxiv.org/abs/2302.06675>
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-4,
+        betas: Tuple[float, float] = (0.9, 0.99),
+        weight_decay: float = 0.0,
+        use_triton: Optional[bool] = None,
+        decoupled_weight_decay: bool = False,
+    ):
+        assert lr > 0.0
+        assert all([0.0 <= beta <= 1.0 for beta in betas])
+
+        self._init_lr = lr
+        self.decoupled_wd = decoupled_weight_decay
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+
+        super().__init__(params, defaults)
+
+        if use_triton is None:
+            use_triton = bool(triton_update_fn)
+        self.use_triton = use_triton
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable] = None):
+        """update step"""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                beta1, beta2 = group["betas"]
+                grad, lr, wd, state, decoupled_wd, init_lr = (
+                    p.grad,
+                    group["lr"],
+                    group["weight_decay"],
+                    self.state[p],
+                    self.decoupled_wd,
+                    self._init_lr,
+                )
+
+                # maybe decoupled weight decay
+
+                if decoupled_wd:
+                    wd /= init_lr
+
+                # init state - exponential moving average of gradient values
+
+                if len(state) == 0:
+                    state["exp_avg"] = torch.zeros_like(p)
+
+                exp_avg = state["exp_avg"]
+
+                if self.use_triton and p.is_cuda:
+                    triton_update_fn(p, grad, exp_avg, lr, wd, beta1, beta2)
+                else:
+                    update_fn(p, grad, exp_avg, lr, wd, beta1, beta2)
+
+        return loss
+
+
+# update functions
+
+
+def update_fn(p, grad, exp_avg, lr, wd, beta1, beta2):
+    """
+    Lion update function
+    """
+    # stepweight decay
+
+    p.data.mul_(1.0 - lr * wd)
+
+    # weight update
+
+    update = exp_avg.clone().mul_(beta1).add(grad, alpha=1.0 - beta1).sign_()
+    p.add_(update, alpha=-lr)
+
+    # decay the momentum running average coefficient
+
+    exp_avg.mul_(beta2).add_(grad, alpha=1.0 - beta2)
+
+
+try:
+    # noinspection PyPackageRequirements
+    import triton
+
+    # noinspection PyPackageRequirements
+    import triton.language as tl
+except ImportError as e:
+    triton = None
+    tl = None
+
+
+# restore_value is not available in older versions of triton
+if triton and "restore_value" in inspect.signature(triton.autotune).parameters:
+    # triton cuda kernel
+
+    # noinspection PyPep8Naming,PyArgumentList
+    @triton.autotune(
+        configs=[
+            triton.Config({"BLOCK_SIZE": 128}, num_warps=4),
+            triton.Config({"BLOCK_SIZE": 1024}, num_warps=8),
+        ],
+        key=["n_elements"],
+        restore_value=["p_ptr", "exp_avg_ptr"],
+    )
+    @triton.jit
+    def _triton_update_fn_kernel(
+        p_ptr,
+        grad_ptr,
+        exp_avg_ptr,
+        lr,
+        wd,
+        beta1,
+        beta2,
+        n_elements,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        pid = tl.program_id(axis=0)
+
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+        mask = offsets < n_elements
+
+        # offsetted pointers
+
+        offset_p_ptr = p_ptr + offsets
+        offset_grad_ptr = grad_ptr + offsets
+        offset_exp_avg_ptr = exp_avg_ptr + offsets
+
+        # load
+
+        p = tl.load(offset_p_ptr, mask=mask)
+        grad = tl.load(offset_grad_ptr, mask=mask)
+        exp_avg = tl.load(offset_exp_avg_ptr, mask=mask)
+
+        # stepweight decay
+
+        p = p * (1 - lr * wd)
+
+        # diff between momentum running average and grad
+
+        diff = exp_avg - grad
+
+        # weight update
+
+        update = diff * beta1 + grad
+
+        # torch.sign
+
+        can_update = update != 0
+        update_sign = tl.where(update > 0, -lr, lr)
+
+        p = p + update_sign * can_update
+
+        # decay the momentum running average coefficient
+
+        exp_avg = diff * beta2 + grad
+
+        # store new params and momentum running average coefficient
+
+        tl.store(offset_p_ptr, p, mask=mask)
+        tl.store(offset_exp_avg_ptr, exp_avg, mask=mask)
+
+    def triton_update_fn(
+        p: torch.Tensor, grad: torch.Tensor, exp_avg: torch.Tensor, lr: float, wd: float, beta1: float, beta2: float
+    ):
+        """
+        Lion update function using triton kernel
+        """
+        assert all([t.is_cuda for t in (p, grad, exp_avg)])
+        n_elements = p.numel()
+
+        def _grid(meta):
+            return tuple((triton.cdiv(n_elements, meta["BLOCK_SIZE"]),))
+
+        _triton_update_fn_kernel[_grid](p, grad, exp_avg, lr, wd, beta1, beta2, n_elements)
+
+else:
+    triton_update_fn = None

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/returnn/torch/updater.py

@@ -5,7 +5,7 @@ and model param update logic in general.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Any, Type, Sequence, Iterable, Set, Dict, List, Tuple
+from typing import Optional, Union, Any, Type, Callable, Sequence, Iterable, Set, Dict, List, Tuple
 import os
 import gc
 import torch
@@ -38,28 +38,40 @@ def _init_optimizer_classes_dict():
         _OptimizerClassesDict[name.lower()] = cls
 
 
-def get_optimizer_class(class_name) -> Type[torch.optim.Optimizer]:
+def get_optimizer_class(
+    class_name: Union[str, Type[torch.optim.Optimizer], Callable[[], Type[torch.optim.Optimizer]]]
+) -> Type[torch.optim.Optimizer]:
     """
-    :param str|()->torch.optim.Optimizer|type[torch.optim.Optimizer] class_name:
-        Optimizer data, e.g. "adam", torch.optim.Adam...
-    :return: Optimizer class
+    :param class_name: Optimizer class, either as str (e.g. "adam"), as type (torch.optim.Adam) or callable.
+        If str, we support all torch.optim optimizers (ignoring case) (e.g. "adam"),
+        or class names with full module path (e.g. "returnn.torch.optim.lion.Lion").
+    :return: Optimizer class, e.g. torch.optim.Adam
     """
     _init_optimizer_classes_dict()
     if isinstance(class_name, type):
         assert issubclass(class_name, torch.optim.Optimizer)
+        return class_name
     elif callable(class_name):
-        class_name = class_name()
-    else:
-        assert isinstance(class_name, str)
-        assert (
-            class_name.lower() in _OptimizerClassesDict
-        ), "%s not found in the available torch optimizers list: %s." % (
-            class_name.lower(),
-            ", ".join("'%s'" % key for key in _OptimizerClassesDict),
-        )
-        class_name = _OptimizerClassesDict[class_name.lower()]
+        return class_name()
+    elif isinstance(class_name, str):
+        if "." in class_name:
+            import importlib
 
-    return class_name
+            mod_name, class_name_ = class_name.rsplit(".", 1)
+            mod = importlib.import_module(mod_name)
+            return getattr(mod, class_name_)
+
+        if class_name.lower() not in _OptimizerClassesDict:
+            raise ValueError(
+                "Optimizer %r not found in the available torch optimizers list: %s."
+                % (
+                    class_name.lower(),
+                    ", ".join("'%s'" % key for key in _OptimizerClassesDict),
+                )
+            )
+        return _OptimizerClassesDict[class_name.lower()]
+    else:
+        raise TypeError(f"Invalid optimizer class_name {class_name!r} type {type(class_name).__name__}")
 
 
 def _get_class_init_kwargs(optim_class):
@@ -411,7 +423,7 @@ class Updater:
         # If the user specified it as epsilon, parse it as eps for the optimizer
         if "eps" in optim_class_init_kwargs and "epsilon" in opt_kwargs:
             opt_kwargs["eps"] = opt_kwargs.pop("epsilon")
-        if "learning_rate" in opt_kwargs:
+        if "learning_rate" in opt_kwargs or "lr" in opt_kwargs:
             raise ValueError("'learning_rate' should be set outside of the 'optimizer' dict.")
         lr = lr * opt_kwargs.pop("learning_rate_multiplier", 1.0)
         opt_kwargs["lr"] = lr

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241020.5643
+Version: 1.20241022.173306
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/returnn.egg-info/SOURCES.txt

@@ -298,6 +298,9 @@ returnn/torch/frontend/_backend.py
 returnn/torch/frontend/_rand.py
 returnn/torch/frontend/bridge.py
 returnn/torch/frontend/raw_ops.py
+returnn/torch/optim/README.md
+returnn/torch/optim/__init__.py
+returnn/torch/optim/lion.py
 returnn/torch/util/README.md
 returnn/torch/util/__init__.py
 returnn/torch/util/array_.py

{returnn-1.20241020.5643 → returnn-1.20241022.173306}/tests/test_torch_engine.py

@@ -584,6 +584,28 @@ def test_dynamic_learning_rate():
     assert len(epoch_continuous_diffs) == engine.global_train_step
 
 
+def test_torch_engine_train_lion_optimizer():
+    config = Config(
+        dict(
+            task="train",
+            device="cpu",
+            extern_data={"data": {"dim": 9}, "classes": {"dim": 2, "sparse": True}},
+            get_model=TrainTestModel,
+            train_step=TrainTestModel.train_step,
+            batch_size=500,
+            optimizer={"class": "returnn.torch.optim.lion.Lion"},
+            num_epochs=1,
+        )
+    )
+    dataset = init_dataset({"class": "Task12AXDataset", "num_seqs": 10, "name": "train"})
+    dataset.init_seq_order(epoch=1)
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_train_from_config(train_data=dataset)
+        engine.train()
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20241020.5643/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241020.005643'
-long_version = '1.20241020.005643+git.cc3b047'