returnn 1.20240903.205823__tar.gz → 1.20240905.105440__tar.gz

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/PKG-INFO

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

returnn-1.20240905.105440/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240905.105440'
+long_version = '1.20240905.105440+git.d79c168'

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/returnn/datasets/distrib_files.py

@@ -317,15 +317,36 @@ class DistributeFilesDataset(CachedDataset2):
         return True
 
     def _get_sub_dataset_dict(self, files: List[FileTree]) -> Dict[str, Any]:
+        import tree
+
         dataset_dict = self.get_sub_epoch_dataset(files)
         dataset_dict = extend_dataset_dict_from_parent_dataset(dataset_dict, parent_dataset=self)
-        if dataset_dict.get("partition_epoch", 1) != 1:
-            raise ValueError(f"{self}: sub dataset should not have partition_epoch, got: {dataset_dict}")
-        if "seq_ordering" not in dataset_dict and "seq_order_control_dataset" not in dataset_dict:
+
+        flat_sub_dset = tree.flatten_with_path(dataset_dict)
+
+        part_epoch_cfg = next(
+            ((path, v) for path, v in flat_sub_dset if path[-1] == "partition_epoch" and v != 1), None
+        )
+        if part_epoch_cfg is not None:
+            path, subeps = part_epoch_cfg
             raise ValueError(
-                f"{self}: sub dataset should have explicit seq_ordering "
-                f"(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
+                f"{self}: sub dataset should not have partition_epoch, "
+                f'but got "partition_epoch": {subeps} at {".".join(path)} in {dataset_dict}.'
             )
+
+        # Heuristic check for well-definedness of seq ordering. Might need to be extended in the
+        # future if there are other ways of defining a seq order than the ones below.
+        if (
+            not any(path[-1] == "seq_ordering" for path, _ in flat_sub_dset)
+            and not any(path[-1] == "seq_order_control_dataset" for path, _ in flat_sub_dset)
+            and not any(path[-1] == "map_seq_stream" for path, _ in flat_sub_dset)
+        ):
+            raise ValueError(
+                f"{self}: there should be an explicit seq_ordering somewhere in the sub dataset "
+                f"(or seq_order_control_dataset for MetaDataset or map_seq_stream for PostprocessingDataset), "
+                f"but found none in {dataset_dict}."
+            )
+
         return dataset_dict
 
     @staticmethod

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/returnn/torch/engine.py

@@ -3,13 +3,15 @@ Main engine for PyTorch
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Union, Callable, Dict
+from typing import Optional, Any, Union, Callable, Dict, Set
 from contextlib import nullcontext
 
 import gc
 import os
 import time
 import socket
+import fnmatch
+import re
 
 import torch
 import torch.distributed
@@ -41,6 +43,7 @@ from .data import extern_data as extern_data_util
 from .data.queued_data_iter import QueuedDataIter
 from .frontend.bridge import rf_module_to_pt_module
 from .util import diagnose_gpu
+from .util import module as util_module
 from .distributed import DistributedContext, get_ctx as dist_get_ctx
 
 
@@ -83,6 +86,7 @@ class Engine(EngineBase):
             self._forward_step_expected_outputs = TensorDict()
             self._forward_step_expected_outputs.update(self.config.typed_value("model_outputs"), auto_convert=True)
         self._save_model_epoch_interval = 1
+        self._ignore_param_set: Set[str] = set()  # for the updater and for saving the model checkpoint
         self._updater = None  # type: Optional[Updater]
 
         self._use_autocast = False
@@ -721,6 +725,7 @@ class Engine(EngineBase):
 
         self._create_model(epoch=epoch, step=step)
 
+        self._ignore_param_set.clear()
         loaded_state_keys = set()
         missing_keys = set()
         unexpected_keys = set()
@@ -745,6 +750,7 @@ class Engine(EngineBase):
 
         preload_from_files = self.config.typed_value("preload_from_files", {})
         if preload_from_files:
+            model_state_keys_set = set(self._pt_model.state_dict().keys())
             # see `preload_from_files` in tf engine and `returnn.tf.network.CustomCheckpointLoader`
             # We use the reversed sorted order here to achieve consistent behavior with the TF engine.
             # There, the keys are used in sorted order but if a variable is loaded,
@@ -754,12 +760,39 @@ class Engine(EngineBase):
             # In order to get consistent behavior, we use the reversed order.
             for preload_key, opts in reversed(sorted(preload_from_files.items())):
                 assert isinstance(opts, dict) and "filename" in opts
-                if opts.get("init_for_train", False):
-                    if not is_first_train_epoch:
-                        continue
+                init_for_train = opts.get("init_for_train", False)
+                if init_for_train:
+                    if isinstance(init_for_train, str) and init_for_train == "always":
+                        # No matter if this is the first train epoch
+                        # or training with loading some prev epoch,
+                        # those parameters will always be loaded via preload_from_files,
+                        # and thus also not stored in our own checkpoint.
+                        pass
+                    elif isinstance(init_for_train, bool) and init_for_train:
+                        if not is_first_train_epoch:
+                            continue
+                    else:
+                        raise ValueError(
+                            f"preload key {preload_key}:"
+                            f" invalid init_for_train value {init_for_train!r} (type {type(init_for_train).__name__})"
+                        )
                 else:  # default: init for recog
                     if is_training:
                         continue
+                if opts["filename"] is None:
+                    print(f"Pre-load (initialize) weights for key '{preload_key}'", file=log.v3)
+                    pattern = opts["pattern"]
+                    match = re.compile(fnmatch.translate(pattern)).match
+                    remove = []
+                    for name in self._pt_model.state_dict().keys():
+                        if match(name) and name in missing_keys:
+                            remove.append(name)
+                    if remove:
+                        print(f"Randomly initialize params: {remove}", file=log.v3)
+                        missing_keys.difference_update(remove)
+                    else:
+                        print("(No relevant parameters matching.)", file=log.v3)
+                    continue
                 print(f"Pre-load weights for key '{preload_key}' from {opts['filename']}", file=log.v3)
                 preload_model_state = torch.load(opts["filename"])
                 if opts.get("checkpoint_key", "model") is not None:
@@ -786,6 +819,8 @@ class Engine(EngineBase):
                         preload_model_state.pop(key)
                 for new_name, name_in_checkpoint in opts.get("var_name_mapping", {}).items():
                     preload_model_state[new_name] = preload_model_state.pop(name_in_checkpoint)
+                if init_for_train == "always":
+                    self._ignore_param_set.update(set(preload_model_state.keys()).intersection(model_state_keys_set))
                 missing_keys_preload, unexpected_keys_preload = self._pt_model.load_state_dict(
                     preload_model_state, strict=False
                 )
@@ -797,7 +832,7 @@ class Engine(EngineBase):
                 if opts.get("prefix", ""):
                     prefix_keys = [key for key in self._pt_model.state_dict() if key.startswith(opts.get("prefix", ""))]
                 else:
-                    prefix_keys = self._pt_model.state_dict().keys()
+                    prefix_keys = model_state_keys_set
                 missing_keys_preload = (
                     set(prefix_keys).intersection(set(missing_keys_preload)).difference(loaded_state_keys)
                 )
@@ -816,6 +851,9 @@ class Engine(EngineBase):
                     )
                     unexpected_keys.update(unexpected_keys_preload)
 
+        if self._ignore_param_set:
+            util_module.convert_parameters_to_buffers(self._pt_model, self._ignore_param_set, persistent=False)
+
         if missing_keys:
             raise Exception(
                 "\n".join(
@@ -913,9 +951,16 @@ class Engine(EngineBase):
         tmp_filename = filename + ".tmp_write"
         if os.path.exists(tmp_filename):
             os.unlink(tmp_filename)
+        state_dict = self._pt_model.state_dict()
+        if self._ignore_param_set:
+            # Do some extra check that we don't save the ignored parameters.
+            # Should not be in the state_dict anymore because we should have converted them to buffers
+            # via util_module.convert_parameters_to_buffers before.
+            remaining = set(state_dict.keys()).intersection(self._ignore_param_set)
+            assert not remaining, f"_save_model: found remaining params in state_dict to ignore: {remaining}"
         torch.save(
             {
-                "model": self._pt_model.state_dict(),
+                "model": state_dict,
                 "epoch": self.epoch,
                 "step": self.global_train_step,
                 "effective_learning_rate": self._updater.get_effective_learning_rate() if self._updater else None,

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/returnn/torch/frontend/bridge.py

@@ -149,13 +149,16 @@ class RFModuleAsPTModule(torch.nn.Module):
         for name, rf_param in self._rf_module.named_parameters(recurse=False):
             pt_param = getattr(self, name)
             if rf_param.auxiliary and self._aux_params_as_buffers:
-                assert isinstance(pt_param, torch.Tensor)  # but not torch.nn.Parameter
-                # See similar logic in torch.nn.Module._apply.
-                pt_param = torch.nn.Parameter(pt_param, pt_param.requires_grad)
-            else:
-                assert isinstance(pt_param, torch.nn.Parameter), (
-                    f"{self}.{name} is not a Parameter" f" but {type(pt_param).__name__}"
-                )
+                if not isinstance(pt_param, torch.nn.Parameter):
+                    assert isinstance(pt_param, torch.Tensor)  # but not torch.nn.Parameter
+                    # See similar logic in torch.nn.Module._apply.
+                    pt_param = torch.nn.Parameter(pt_param, pt_param.requires_grad)
+            # Otherwise, we do not care whether it is a torch.nn.Parameter or not.
+            # Its type might have changed due to convert_parameters_to_buffers.
+            # Just make sure it is a tensor.
+            assert isinstance(pt_param, torch.Tensor)
+            # noinspection PyProtectedMember
+            rf_param.dtype = rf_param._raw_backend.get_dtype_name_raw(pt_param)  # dtype might have changed
             rf_param.raw_tensor = pt_param
 
     def register_parameter(self, name: str, param: Optional[torch.nn.Parameter]) -> None:

returnn-1.20240905.105440/returnn/torch/util/module.py

@@ -0,0 +1,43 @@
+"""
+Utils for modules
+"""
+
+from __future__ import annotations
+from typing import Collection
+import torch
+
+
+def convert_parameters_to_buffers(
+    module: torch.nn.Module, parameter_names: Collection[str], *, deep: bool = True, persistent: bool
+):
+    """
+    :param module:
+    :param parameter_names:
+    :param deep: parameter_name can contain '.' to access submodules
+    :param persistent: whether the buffer is persistent. if True, the buffer will be saved to the state_dict.
+        passed to module.register_buffer.
+    """
+    for parameter_name in parameter_names:
+        convert_parameter_to_buffer(module, parameter_name, deep=deep, persistent=persistent)
+
+
+def convert_parameter_to_buffer(module: torch.nn.Module, parameter_name: str, *, deep: bool = True, persistent: bool):
+    """
+    :param module:
+    :param parameter_name:
+    :param deep: parameter_name can contain '.' to access submodules
+    :param persistent: whether the buffer is persistent. if True, the buffer will be saved to the state_dict.
+        passed to module.register_buffer.
+    """
+    if "." in parameter_name:
+        if not deep:
+            raise ValueError("parameter_name can't contain '.' when deep is False")
+        module_path, _, parameter_name = parameter_name.rpartition(".")
+        module = module.get_submodule(module_path)
+
+    parameter = getattr(module, parameter_name)
+    if not isinstance(parameter, torch.nn.Parameter):
+        raise ValueError(f"{parameter_name} is not a torch.nn.Parameter, got type {type(parameter).__name__}")
+    delattr(module, parameter_name)
+    parameter.requires_grad = False
+    module.register_buffer(parameter_name, parameter, persistent=persistent)

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/returnn.egg-info/SOURCES.txt

@@ -302,6 +302,7 @@ returnn/torch/util/__init__.py
 returnn/torch/util/array_.py
 returnn/torch/util/diagnose_gpu.py
 returnn/torch/util/gradient_checkpoint.py
+returnn/torch/util/module.py
 returnn/torch/util/scaled_gradient.py
 returnn/util/__init__.py
 returnn/util/basic.py

{returnn-1.20240903.205823 → returnn-1.20240905.105440}/tests/test_torch_frontend.py

@@ -1005,6 +1005,29 @@ def test_forward_hook():
     assert call_count == 5 and hook_call_count == 2
 
 
+def test_convert_parameter_to_buffer():
+    from returnn.torch.frontend.bridge import rf_module_to_pt_module
+    from returnn.torch.util.module import convert_parameters_to_buffers
+
+    # Replicate the case of Engine preload_from_files with init_for_train == "always".
+    in_dim = Dim(2, name="in")
+    out_dim = Dim(3, name="out")
+    mod_rf = rf.Linear(in_dim, out_dim)
+    mod_pt = rf_module_to_pt_module(mod_rf)
+    convert_parameters_to_buffers(mod_pt, ["weight", "bias"], persistent=False)
+    assert dict(mod_pt.named_parameters()) == {}
+    assert set(dict(mod_pt.named_buffers()).keys()) == {"weight", "bias"}
+    # The following type checks are implementation details,
+    # and maybe we might change this behavior in the future.
+    assert type(mod_pt.weight) is torch.nn.Parameter  # still same instance
+    # After mod_pt.to(...), buffers will not be torch.nn.Parameter anymore.
+    # Note, the test here is actually to run the to() method,
+    # which runs our RFModuleAsPTModule._apply, because there was a bug, leading to:
+    #   Linear[RF→PT]().weight is not a Parameter but Tensor
+    mod_pt.to(torch.bfloat16)
+    assert type(mod_pt.weight) is torch.Tensor
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20240903.205823/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240903.205823'
-long_version = '1.20240903.205823+git.eb0f22e'