returnn 1.20250204.160236__tar.gz → 1.20250206.151011__tar.gz

{returnn-1.20250204.160236/returnn.egg-info → returnn-1.20250206.151011}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250204.160236
+Version: 1.20250206.151011
 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.151011/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250206.151011'
+long_version = '1.20250206.151011+git.6fa4b38'

{returnn-1.20250204.160236 → returnn-1.20250206.151011}/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.160236 → returnn-1.20250206.151011}/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.160236 → returnn-1.20250206.151011}/returnn/torch/engine.py

@@ -3,7 +3,7 @@ Main engine for PyTorch
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Union, Callable, Dict, Set
+from typing import Optional, Any, Union, Callable, Dict, Set, Tuple
 from contextlib import nullcontext, ExitStack, contextmanager
 
 import gc
@@ -371,6 +371,7 @@ class Engine(EngineBase):
         total_data_size_packed = NumbersDict()
         total_data_size_padded = NumbersDict()
 
+        report_prefix = f"ep {self.epoch} train"
         try:
             while True:
                 with torch.no_grad():
@@ -398,21 +399,13 @@ class Engine(EngineBase):
                     {k: int(util.prod(extern_data_raw[k].shape[:2])) for k in keys_w_seq_len},
                 )
 
-                num_seqs_ = (
-                    int(extern_data_raw["num_seqs"]) if extern_data_raw.get("num_seqs", None) is not None else -1
+                num_seqs, last_seq_idx = _get_num_seqs_last_seq_idx(
+                    report_prefix=report_prefix,
+                    extern_data_raw=extern_data_raw,
+                    step_idx=step_idx,
+                    prev_num_seqs=num_seqs,
+                    prev_last_seq_idx=last_seq_idx,
                 )
-                # Note: The batches might have been shuffled,
-                # thus we cannot really assert that the seq_idx is always increasing.
-                last_seq_idx = max(int(extern_data_raw["seq_idx"].max()), last_seq_idx)
-                if step_idx == 0:
-                    if num_seqs_ >= 0:
-                        print(f"Epoch {self.epoch} num_seqs: {num_seqs_}", file=log.v5)
-                        num_seqs = num_seqs_
-                elif num_seqs_ >= 0:
-                    assert num_seqs_ == num_seqs
-                del num_seqs_
-                if num_seqs is not None:
-                    assert last_seq_idx < num_seqs
                 epoch_continuous = (self.epoch - 1 + (last_seq_idx + 1) / num_seqs) if num_seqs is not None else None
 
                 # clear the gradients when every gradient accumulation loop starts
@@ -485,7 +478,7 @@ class Engine(EngineBase):
                 accumulated_inv_norm_factors_dict += inv_norm_factors_dict
                 eval_info = self._maybe_extend_losses_info(losses_dict / inv_norm_factors_dict)
                 _print_process(
-                    f"ep {self.epoch} train",
+                    report_prefix,
                     step=step_idx,
                     eval_info=dict(eval_info),
                     step_duration=step_duration,
@@ -505,12 +498,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
@@ -1253,6 +1269,8 @@ class Engine(EngineBase):
             new_dim.dyn_size_ext = _get_tensor_wo_batch_numpy(dim.dyn_size_ext)
             return new_dim
 
+        num_seqs = None
+        last_seq_idx = 0
         report_prefix = f"ep {self.epoch} {dataset.name} forward"
         with torch.no_grad():
             callback.init(model=self._orig_model)
@@ -1260,6 +1278,15 @@ class Engine(EngineBase):
             step_idx = 0
             for extern_data_raw in data_loader:
                 step_begin_time = time.monotonic()
+
+                num_seqs, last_seq_idx = _get_num_seqs_last_seq_idx(
+                    report_prefix=report_prefix,
+                    extern_data_raw=extern_data_raw,
+                    step_idx=step_idx,
+                    prev_num_seqs=num_seqs,
+                    prev_last_seq_idx=last_seq_idx,
+                )
+
                 if self._forward_step_expected_outputs:
                     # Also resets any dyn dims, which might have been set in the prev step.
                     self._forward_step_expected_outputs.reset_content()
@@ -1296,11 +1323,19 @@ class Engine(EngineBase):
                         model_outputs_per_batch.data[k] = _get_tensor_wo_batch_numpy(v)
                     callback.process_seq(seq_tag=seq_tag, outputs=model_outputs_per_batch)
 
-                elapsed_computation_time += time.monotonic() - step_begin_time
+                step_end_time = time.monotonic()
+                step_duration = step_end_time - step_begin_time
+                elapsed_computation_time += step_duration
+
                 _print_process(
                     report_prefix,
                     step=step_idx,
                     eval_info=None,
+                    step_duration=step_duration,
+                    start_elapsed=step_end_time - epoch_start_time,
+                    seq_idx=last_seq_idx,
+                    num_seqs=num_seqs,
+                    batch_size_info=_get_batch_size_info(extern_data) if self._log_batch_size else None,
                     log_memory_usage_device=self._device if self._log_memory_usage else None,
                 )
                 step_idx += 1
@@ -1578,3 +1613,27 @@ def _get_total_grad_norm(model: torch.nn.Module, p: float) -> float:
             p=p,
         ).item()
     )
+
+
+def _get_num_seqs_last_seq_idx(
+    *,
+    report_prefix: str,
+    extern_data_raw: Dict[str, Any],
+    step_idx: int,
+    prev_num_seqs: Optional[int],
+    prev_last_seq_idx: int,
+) -> Tuple[Optional[int], int]:
+    num_seqs = prev_num_seqs
+    num_seqs_ = int(extern_data_raw["num_seqs"]) if extern_data_raw.get("num_seqs", None) is not None else -1
+    # Note: The batches might have been shuffled,
+    # thus we cannot really assert that the seq_idx is always increasing.
+    last_seq_idx = max(int(extern_data_raw["seq_idx"].max()), prev_last_seq_idx)
+    if step_idx == 0:
+        if num_seqs_ >= 0:
+            print(f"{report_prefix} num_seqs: {num_seqs_}", file=log.v5)
+            num_seqs = num_seqs_
+    elif num_seqs_ >= 0:
+        assert num_seqs_ == num_seqs
+    if num_seqs is not None:
+        assert last_seq_idx < num_seqs
+    return num_seqs, last_seq_idx

{returnn-1.20250204.160236 → returnn-1.20250206.151011}/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.160236 → returnn-1.20250206.151011/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250204.160236
+Version: 1.20250206.151011
 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.160236 → returnn-1.20250206.151011}/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.160236 → returnn-1.20250206.151011}/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.160236 → returnn-1.20250206.151011}/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.160236/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250204.160236'
-long_version = '1.20250204.160236+git.e147886'