returnn 1.20250204.4017__tar.gz → 1.20250206.144022__tar.gz

{returnn-1.20250204.4017/returnn.egg-info → returnn-1.20250206.144022}/PKG-INFO

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

returnn-1.20250206.144022/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250206.144022'
+long_version = '1.20250206.144022+git.550e757'

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/returnn/datasets/util/vocabulary.py

@@ -78,6 +78,8 @@ class Vocabulary:
         :param seq_postfix: labels will be added to the seq in self.get_seq
         :param labels:
         """
+        if vocab_file and not isinstance(vocab_file, str):  # sometimes it is a Path
+            vocab_file = str(vocab_file)
         self.vocab_file = vocab_file
 
         if special_symbols_via_file:

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/returnn/frontend/encoder/transformer.py

@@ -35,6 +35,7 @@ class TransformerEncoder(rf.Module):
         layer: Optional[Union[TransformerEncoderLayer, rf.Module, type, Dict[str, Any], Any]] = None,
         layer_opts: Optional[Dict[str, Any]] = None,
         embed_dim: Optional[Dim] = None,
+        input_embedding: Union[None, rf.Module, type, Dict[str, Any]] = rf.Embedding,
         input_embedding_scale: float = None,
         input_dropout: float = None,
         sequential=rf.Sequential,
@@ -53,6 +54,7 @@ class TransformerEncoder(rf.Module):
         :param layer: an instance of :class:`TransformerEncoderLayer` or similar
         :param layer_opts: options for the encoder layer
         :param embed_dim: if given, will first have an embedding [vocab,embed] and then a linear [embed,model].
+        :param input_embedding:
         :param input_embedding_scale:
         :param input_dropout:
         :param sequential:
@@ -77,9 +79,15 @@ class TransformerEncoder(rf.Module):
         self.model_dim = model_dim
         self.embed_dim = embed_dim
 
-        # We could make this optional or configurable if we ever need to.
-        # Or maybe you would just have another separate implementation of this module then...
-        self.input_embedding = rf.Embedding(vocab_dim, embed_dim or model_dim)
+        if input_embedding is None or isinstance(input_embedding, rf.Module):
+            pass
+        elif isinstance(input_embedding, type):
+            input_embedding: rf.Embedding = input_embedding(vocab_dim, embed_dim or model_dim)
+        elif isinstance(input_embedding, dict):
+            input_embedding = rf.build_from_dict(input_embedding, vocab_dim, embed_dim or model_dim)
+        else:
+            raise TypeError(f"unexpected input_embedding {input_embedding!r} type {type(input_embedding)}")
+        self.input_embedding = input_embedding
 
         self.input_embedding_proj = None
         if embed_dim:
@@ -88,17 +96,13 @@ class TransformerEncoder(rf.Module):
         if pos_enc is None:
             pass
         elif isinstance(pos_enc, dict):
-            pos_enc = rf.build_from_dict(
-                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
-            )
+            pos_enc = rf.build_from_dict(pos_enc, feat_dim=embed_dim or model_dim, dtype=rf.get_default_float_dtype())
         elif isinstance(pos_enc, rf.Module):
             pass
         elif isinstance(pos_enc, FunctionType):
-            pos_enc = functools.partial(
-                pos_enc, feat_dim=embed_dim or model_dim, dtype=self.input_embedding.weight.dtype
-            )
+            pos_enc = functools.partial(pos_enc, feat_dim=embed_dim or model_dim, dtype=rf.get_default_float_dtype())
         else:
-            raise TypeError(f"unexpected pos_enc type {pos_enc!r}")
+            raise TypeError(f"unexpected pos_enc {pos_enc!r} type {type(pos_enc)}")
         self.pos_enc = pos_enc
         if input_embedding_scale is None:
             input_embedding_scale = model_dim.dimension**0.5
@@ -157,7 +161,11 @@ class TransformerEncoder(rf.Module):
         :param collected_outputs:
         :return: final encoder output, after final layer norm
         """
-        decoded = self.input_embedding(source) * self.input_embedding_scale
+        if self.input_embedding is not None:
+            decoded = self.input_embedding(source) * self.input_embedding_scale
+        else:
+            assert self.model_dim in source.dims
+            decoded = source
         if self.pos_enc is not None:
             decoded = decoded + self.pos_enc(spatial_dim=spatial_dim)
         decoded = rf.dropout(decoded, self.input_dropout)

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/returnn/tensor/tensor_dict.py

@@ -139,7 +139,7 @@ class TensorDict:
         """
         visited_dims = set()
         for key, value in self.data.items():
-            assert key in raw_tensor_dict
+            assert key in raw_tensor_dict, f"key {key} not in raw_tensor_dict {list(raw_tensor_dict.keys())}"
             value.raw_tensor = raw_tensor_dict[key]
             for i, dim in enumerate(value.dims):
                 dim: Dim

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/returnn/torch/engine.py

@@ -505,12 +505,35 @@ class Engine(EngineBase):
                             file=log.v1,
                         )
 
+                        print("Checking for inf/nan in model parameters...", file=log.v1)
+                        count_nan_inf_params = 0
+                        for name, param in self._pt_model.named_parameters():
+                            got_nan_inf_t = torch.stack([torch.isnan(param).any(), torch.isinf(param).any()]).cpu()
+                            got_nan = got_nan_inf_t[0].item()
+                            got_inf = got_nan_inf_t[1].item()
+                            if got_nan or got_inf:
+                                s = "/".join([s_ for s_, b in [("nan", got_nan), ("inf", got_inf)] if b])
+                                print(f"  {name} {param}: {s}", file=log.v1)
+                                count_nan_inf_params += 1
+                        if count_nan_inf_params == 0:
+                            print("(No inf/nan in model parameters.)", file=log.v1)
+
                         def _debug_func() -> torch.Tensor:
                             self._run_step(extern_data, train_flag=True, train_func=True)
-                            return rf.get_run_ctx().total_loss()
+                            loss = rf.get_run_ctx().total_loss()
+                            assert isinstance(loss, Tensor)
+                            return loss.raw_tensor
 
                         print("Running debug_inf_nan...", file=log.v1)
                         debug_inf_nan(_debug_func, with_grad=True)
+                        if count_nan_inf_params > 0 and self.global_train_step == 1:
+                            print(
+                                "This was the second step, so likely the first step grad was broken."
+                                " Try again with reset model...",
+                                file=log.v1,
+                            )
+                            self._load_model()
+                            debug_inf_nan(_debug_func, with_grad=True)
                         raise Exception(f"Inf/nan score in step {step_idx}.")
 
                 step_idx += 1

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/returnn/torch/util/debug_inf_nan.py

@@ -52,6 +52,7 @@ def debug_inf_nan(
     *,
     with_grad: bool = False,
     report_every_op_call: bool = True,
+    stop_reporting_after_first_inf_nan: bool = True,
     file: Optional[Union[TextIO, TextIOBase]] = None,
 ):
     """
@@ -61,6 +62,7 @@ def debug_inf_nan(
         and we will call `loss = func(); loss.backward()`.
     :param with_grad: whether to compute and debug gradients for inf/nan.
     :param report_every_op_call: whether to report every op call.
+    :param stop_reporting_after_first_inf_nan: whether to stop reporting after the first inf/nan.
     :param file: where to write the output to. Default is stdout.
     """
 
@@ -69,13 +71,18 @@ def debug_inf_nan(
 
     # noinspection PyUnresolvedReferences,PyProtectedMember
     cur_frame: FrameType = sys._getframe()
-    trace_ops = _TraceOps(root_frame=cur_frame, file=file, report_every_op_call=report_every_op_call)
+    trace_ops = _TraceOps(
+        root_frame=cur_frame,
+        file=file,
+        report_every_op_call=report_every_op_call,
+        stop_reporting_after_first_inf_nan=stop_reporting_after_first_inf_nan,
+    )
 
     if with_grad:
-
         with torch.autograd.detect_anomaly():
             with trace_ops:  # currently only for forward (but we might want to trace the backward too)
                 loss = func()
+            file.flush()  # the backward detect_anomaly might screw up the output otherwise
             try:
                 loss.backward()
             except RuntimeError as exc:
@@ -89,23 +96,46 @@ def debug_inf_nan(
 
 # For efficiency, and to be less spammy
 _TraceFuncNameBlacklist = {
-    "aten::detach",
     "aten::zeros_like",
     "aten::ones_like",
+    "aten::full",
+    "aten::scalar_tensor",  # when we deliberately create a scalar inf tensor
+    "aten::_local_scalar_dense",
+    "aten::where.self",  # when we intentionally mask with inf
+    "aten::detach",
+    "aten::_to_copy",
+    "aten::clone",
+    "aten::stack",
+    "aten::view",
+    "aten::_unsafe_view",
+    "aten::permute",
+    "aten::t",
+    "aten::split_with_sizes",
+    "aten::slice.Tensor",
+    "aten::select.int",
 }
 
 
 class _TraceOps(TorchDispatchMode):
-    def __init__(self, *, root_frame: FrameType, file: Union[TextIO, TextIOBase], report_every_op_call: bool = True):
+    def __init__(
+        self,
+        *,
+        root_frame: FrameType,
+        file: Union[TextIO, TextIOBase],
+        report_every_op_call: bool = True,
+        stop_reporting_after_first_inf_nan: bool = True,
+    ):
         super().__init__()
         self.root_frame = root_frame
         self.file = file
+        self.enabled = True
         self.report_every_op_call = report_every_op_call
+        self.stop_reporting_after_first_inf_nan = stop_reporting_after_first_inf_nan
 
     def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         if kwargs is None:
             kwargs = {}
-        if func.name() in _TraceFuncNameBlacklist:
+        if not self.enabled or func.name() in _TraceFuncNameBlacklist:
             return func(*args, **kwargs)
         if self.report_every_op_call:
             print(f"--- op {func.name()}", file=self.file)
@@ -121,6 +151,8 @@ class _TraceOps(TorchDispatchMode):
                     traceback.print_list(
                         _extract_stack_up_to(skip_top_num_frames=1, root_frame=self.root_frame), file=self.file
                     )
+                    if self.stop_reporting_after_first_inf_nan:
+                        self.enabled = False
         return out
 
 

{returnn-1.20250204.4017 → returnn-1.20250206.144022/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/tests/rf_utils.py

@@ -55,6 +55,7 @@ def run_model(
     dyn_dim_max_sizes: Optional[Dict[Dim, int]] = None,
     dyn_dim_min_sizes: Optional[Dict[Dim, int]] = None,
     test_tensorflow: bool = True,
+    allow_inf_nan_in_output: bool = False,
 ) -> TensorDict:
     """run"""
     print(f"* run_model with dyn_dim_max_sizes={dyn_dim_max_sizes!r}")
@@ -70,6 +71,10 @@ def run_model(
         # get the values now because dims might get overwritten
         out_pt_raw = out_pt.as_raw_tensor_dict(include_const_sizes=True)
 
+    if not allow_inf_nan_in_output:
+        for k, v in out_pt.data.items():
+            assert numpy.isfinite(v.raw_tensor).all(), f"output {k!r} has non-finite values: {v.raw_tensor}"
+
     if not test_tensorflow:
         return out_pt
     if disable_tf:
@@ -86,8 +91,10 @@ def run_model(
     random_journal: RandomJournal
     assert random_journal.reached_end()
 
-    print(out_pt, out_tf)
-    assert set(out_pt.data.keys()) == set(out_tf.data.keys())
+    print("Output PT/TF:", out_pt, out_tf)
+    assert set(out_pt.data.keys()) == set(
+        out_tf.data.keys()
+    ), f"PT output {list(out_pt.data.keys())} vs TF output {list(out_tf.data.keys())}"
     for k, v_pt in out_pt.data.items():
         v_tf = out_tf[k]
         # We cannot really check the dims directly for equality,
@@ -123,11 +130,31 @@ def _run_model_torch(extern_data: TensorDict, get_model: rf.GetModelFunc, forwar
     # We recover extern_data in the end.
     tensor_dict_numpy_to_torch_(extern_data)
 
+    for k, v in extern_data.data.items():
+        if v.raw_tensor.dtype.is_floating_point:
+            v.raw_tensor.requires_grad = True
+
     model = get_model(epoch=1, step=0)
     rf.init_forward_step_run_ctx(epoch=1, step=0)
     forward_step(model=model, extern_data=extern_data)
     outputs = rf.get_run_ctx().outputs
     assert outputs.data
+
+    if "loss" in outputs.data:
+        loss = outputs.data["loss"]
+        assert isinstance(loss, Tensor)
+        assert loss.raw_tensor.dtype.is_floating_point
+        loss = rf.reduce_sum(loss, axis=loss.dims)
+        print("loss:", loss.raw_tensor.detach().numpy().item())
+        loss.raw_tensor.backward()
+        for k, v in list(extern_data.data.items()):
+            if v.raw_tensor.dtype.is_floating_point:
+                assert v.raw_tensor.grad is not None, f"no grad for {k}"
+                v_grad = v.copy_template()
+                v_grad.raw_tensor = v.raw_tensor.grad
+                assert f"{k}_grad" not in outputs.data
+                outputs.data[f"{k}_grad"] = v_grad
+
     tensor_dict_torch_to_numpy_(outputs)
 
     extern_data.assign_from_raw_tensor_dict_(extern_data_raw)
@@ -242,6 +269,20 @@ def _run_model_net_dict_tf(
             layer = net.get_layer(layer_name)
             outputs_tf.data[k] = layer.output.copy()
 
+        if "loss" in outputs_tf.data:
+            data_ = {name: data for name, data in net.extern_data.data.items() if data.dtype.startswith("float")}
+            loss = outputs_tf.data["loss"]
+            assert isinstance(loss, Tensor)
+            assert loss.dtype.startswith("float")
+            loss = rf.reduce_sum(loss, axis=loss.dims)
+            d_grads = tf.gradients(loss.raw_tensor, [d.raw_tensor for d in data_.values()])
+            for (name, data), d_grad_tf in zip(data_.items(), d_grads):
+                assert isinstance(data, Tensor)
+                assert isinstance(d_grad_tf, tf.Tensor)
+                d_grad = data.copy_template()
+                d_grad.raw_tensor = d_grad_tf
+                outputs_tf.data[f"{name}_grad"] = d_grad
+
         fetches = outputs_tf.as_raw_tensor_dict(expected_value_type=tf.Tensor)
         assert set(extern_data.data.keys()) == set(net.extern_data.data.keys())
         extern_data_tf_placeholders = net.extern_data.as_raw_tensor_dict(expected_value_type=tf.Tensor)

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/tests/test_rf_array.py

@@ -627,6 +627,27 @@ def test_gather_time_static_clip_to_valid():
     run_model(extern_data_template, lambda *, epoch, step: rf.Module(), _forward_step)
 
 
+def test_gather_3d_embed():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    num_embeddings_dim = Dim(2, name="num_embeddings")
+    embed_dim = Dim(11, name="embed")
+    extern_data_template = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim], sparse_dim=in_dim, dtype="int32"),
+            "embed": Tensor("embed", [in_dim, num_embeddings_dim, embed_dim], dtype="float32"),
+        }
+    )
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        x, embed = extern_data["data"], extern_data["embed"]
+        out = rf.gather(embed, indices=x)
+        out.mark_as_default_output(shape=(batch_dim, time_dim, num_embeddings_dim, embed_dim))
+        rf.reduce_sum(out, axis=out.dims).mark_as_output("loss")
+
+    run_model(extern_data_template, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
 def test_scatter_fill_inf():
     batch_dim_ = Dim(3, name="batch")
     states_dim = Dim(7, name="states")
@@ -645,7 +666,13 @@ def test_scatter_fill_inf():
         )  # [S], per state
         scores.mark_as_default_output(shape=[states_dim])
 
-    res = run_model(TensorDict(), lambda *, epoch, step: rf.Module(), _forward_step, test_tensorflow=False)
+    res = run_model(
+        TensorDict(),
+        lambda *, epoch, step: rf.Module(),
+        _forward_step,
+        test_tensorflow=False,
+        allow_inf_nan_in_output=True,
+    )
     batch_size = res["batch_size"].raw_tensor.item()
     assert res["start_states"].raw_tensor.shape == (batch_size,)
     assert res["output"].raw_tensor.shape == (states_dim.dimension,)

{returnn-1.20250204.4017 → returnn-1.20250206.144022}/tests/test_torch_util.py

@@ -286,12 +286,12 @@ def test_debug_inf_nan():
 
     # Run directly, to just test that it goes through without exception.
     # For some reason, the detect_anomaly does not print the forward op?
-    debug_inf_nan(func, with_grad=True)
+    debug_inf_nan(func, with_grad=True, stop_reporting_after_first_inf_nan=False)
 
     from io import StringIO
 
     out = StringIO()
-    debug_inf_nan(func, file=out)
+    debug_inf_nan(func, file=out, stop_reporting_after_first_inf_nan=False)
     assert "inf in aten.exp" in out.getvalue()
     assert "nan in aten.div" in out.getvalue()
     assert "mod5" in out.getvalue()

returnn-1.20250204.4017/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250204.004017'
-long_version = '1.20250204.004017+git.43f54b8'