returnn 1.20251106.185107__tar.gz → 1.20260118.2517__tar.gz

{returnn-1.20251106.185107/returnn.egg-info → returnn-1.20260118.2517}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251106.185107
+Version: 1.20260118.2517
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer
@@ -36,7 +36,7 @@ Welcome to RETURNN
 `RETURNN paper 2018 <https://arxiv.org/abs/1805.05225>`_.
 
 RETURNN - RWTH extensible training framework for universal recurrent neural networks,
-is a Theano/TensorFlow-based implementation of modern recurrent neural network architectures.
+is a PyTorch/TensorFlow-based implementation of modern recurrent neural network architectures.
 It is optimized for fast and reliable training of recurrent neural networks in a multi-GPU environment.
 
 The high-level features and goals of RETURNN are:

{returnn-1.20251106.185107 → returnn-1.20260118.2517}/README.rst

@@ -7,7 +7,7 @@ Welcome to RETURNN
 `RETURNN paper 2018 <https://arxiv.org/abs/1805.05225>`_.
 
 RETURNN - RWTH extensible training framework for universal recurrent neural networks,
-is a Theano/TensorFlow-based implementation of modern recurrent neural network architectures.
+is a PyTorch/TensorFlow-based implementation of modern recurrent neural network architectures.
 It is optimized for fast and reliable training of recurrent neural networks in a multi-GPU environment.
 
 The high-level features and goals of RETURNN are:

returnn-1.20260118.2517/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20260118.002517'
+long_version = '1.20260118.002517+git.cfcf6ca'

{returnn-1.20251106.185107 → returnn-1.20260118.2517}/pyproject.toml

@@ -16,6 +16,9 @@ extend-exclude = '''
 [tool.ruff]
 line-length = 120
 target-version = "py38"  # https://github.com/rwth-i6/returnn/issues/1326
+extend-exclude = [
+    ".github/workflows/_pytorch_collect_env.py",
+]
 
 [build-system]
 requires = ["setuptools", "numpy"]

{returnn-1.20251106.185107 → returnn-1.20260118.2517}/returnn/__old_mod_loader__.py

@@ -17,7 +17,7 @@ This is supported as well.
 import sys
 import os
 import types
-import typing
+from typing import Any, Dict
 import importlib
 
 old_to_new_mod_mapping = {
@@ -122,7 +122,7 @@ class _LazyLoader(types.ModuleType):
             fn = "%s/%s/__init__.py" % (_base_dir, full_mod_name.replace(".", "/"))
             assert os.path.exists(fn), "_LazyLoader: mod %r not found in %r" % (full_mod_name, _base_dir)
         self.__file__ = fn
-        self._lazy_mod_config = dict(full_mod_name=full_mod_name, **kwargs)  # type: typing.Dict[str]
+        self._lazy_mod_config: Dict[str, Any] = dict(full_mod_name=full_mod_name, **kwargs)
 
     def _load(self):
         full_mod_name = self.__name__
@@ -172,6 +172,30 @@ class _LazyLoader(types.ModuleType):
         return super(_LazyLoader, self).__getattribute__(item)
 
     def __getattr__(self, item):
+        if item == "torch":
+            # torch.compile Dynamo hashing can trigger this, when it uses pickle to serialize some object state,
+            # which iterates through sys.modules and does getattr on each module.
+            # In this case, it searches for torch.
+            #   File ".../torch/_inductor/codecache.py", line 607 in dumps
+            #   File ".../torch/_inductor/codecache.py", line 622 in get_hash
+            #   File ".../torch/_inductor/codecache.py", line 961 in compiled_fx_graph_hash
+            #   ...
+            # Unfortunately, Pickler.dump is native code, so we cannot easily check whether that is the parent frame.
+            # The C stacktrace looks like:
+            # ...
+            # 7   Python                              0x0000000102e7d504 call_attribute + 80
+            # 8   Python                              0x0000000102e7d400 _Py_slot_tp_getattr_hook + 576
+            # 9   Python                              0x0000000102e507a0 PyObject_GetOptionalAttr + 248
+            # 10  _pickle.cpython-313-darwin.so       0x0000000102d24fb4 get_deep_attribute + 104
+            # 11  _pickle.cpython-313-darwin.so       0x0000000102d250b8 _checkmodule + 88
+            # 12  _pickle.cpython-313-darwin.so       0x0000000102d22588 save_global + 3024
+            # 13  _pickle.cpython-313-darwin.so       0x0000000102d1eddc save + 3424
+            # ...
+            # Right now, we just check for `item == "torch"` as a heuristic,
+            # which should never exist for any of the old-style wrapped modules here.
+            # We could maybe also check sys._getframe(1).f_code or so and add some other heuristics...
+            raise AttributeError(f"module {self.__name__} has no attribute {item} (lazy loading skipped)")
+
         module = self._load()
         return getattr(module, item)
 

{returnn-1.20251106.185107 → returnn-1.20260118.2517}/returnn/datasets/lm.py

@@ -86,6 +86,7 @@ class LmDataset(CachedDataset2):
         delayed_seq_data_start_symbol="[START]",
         dtype: Optional[str] = None,
         tag_prefix: Optional[str] = None,
+        _debug_limit_line_count: Optional[int] = None,
         **kwargs,
     ):
         """
@@ -138,6 +139,8 @@ class LmDataset(CachedDataset2):
           delayed_seq_data_start_symbol + original_sequence[:-1].
         :param str delayed_seq_data_start_symbol: used for add_delayed_seq_data.
         :param dtype: explicit dtype. if not given, automatically determined based on the number of labels.
+        :param tag_prefix: prefix for sequence tags. by default "line-".
+        :param _debug_limit_line_count:
         """
         super(LmDataset, self).__init__(**kwargs)
 
@@ -316,6 +319,10 @@ class LmDataset(CachedDataset2):
         self.num_skipped = 0
         self.num_unknown = 0
 
+        if _debug_limit_line_count is None:
+            _debug_limit_line_count = _get_debug_limit_line_count()
+        self._debug_limit_line_count = _debug_limit_line_count
+
     def _lazy_init(self):
         if self._orths_offsets_and_lens is not None:
             return
@@ -340,6 +347,9 @@ class LmDataset(CachedDataset2):
         lens_per_corpus_file = []
         start_time = time.time()
         last_print_time = start_time
+        debug_limit_line_count = self._debug_limit_line_count
+        debug_limit_est_total = 0
+        debug_limit_hit = False
 
         def _init_tmp_file():
             nonlocal tmp_file, tmp_file_orth_files_index
@@ -368,13 +378,16 @@ class LmDataset(CachedDataset2):
 
             if time.time() - last_print_time > 10:
                 print(
-                    f"  ... loaded {len(self._orths_offsets_and_lens)} sequences,"
+                    f"  ... loaded {len(orths)} sequences,"
                     f" {human_bytes_size(total_bytes_read)},"
                     f" after {hms(time.time() - start_time)}",
                     file=log.v4,
                 )
                 last_print_time = time.time()
 
+            if debug_limit_line_count is not None and len(orths) - prev_orth_len >= debug_limit_line_count:
+                raise _ReachedDebugLimitLineCount()
+
         # If a list of files is provided, concatenate all.
         if isinstance(corpus_file, str):
             corpus_file = [corpus_file]
@@ -383,37 +396,46 @@ class LmDataset(CachedDataset2):
         for file_name in corpus_file:
             if self._use_cache_manager:
                 file_name = cf(file_name)
-            if _is_bliss(file_name):
-                _init_tmp_file()
-                _iter_bliss(filename=file_name, callback=_tmp_file_add_line, decode=False)
-            elif file_name.endswith(".gz"):
-                _init_tmp_file()
-                _iter_txt(
-                    filename=file_name,
-                    callback=_tmp_file_add_line,
-                    skip_empty_lines=self._skip_empty_lines,
-                    decode=False,
-                )
-            else:  # Raw txt file
-                # Directly mmap the file.
-                # We just need to scan once through it to find line offsets.
-                file = open(file_name, "rb")
-                file_mmap = mmap.mmap(file.fileno(), 0, flags=mmap.MAP_PRIVATE)
-                file_index = len(self._orth_files)
-                self._orth_files.append(file)
-                self._orth_mmaps.append(file_mmap)
-
-                pos = 0
-                while True:
-                    next_new_line = file_mmap.find(b"\n", pos)
-                    if next_new_line == -1:
-                        break
-                    line_len = next_new_line - pos
-                    if line_len or not self._skip_empty_lines:
-                        orths.append((file_index, pos, line_len))
-                    total_bytes_read += line_len + 1
-                    pos = next_new_line + 1
-                    _maybe_report_status()
+
+            try:
+                if _is_bliss(file_name):
+                    _init_tmp_file()
+                    _iter_bliss(filename=file_name, callback=_tmp_file_add_line, decode=False)
+                elif file_name.endswith(".gz"):
+                    _init_tmp_file()
+                    _iter_txt(
+                        filename=file_name,
+                        callback=_tmp_file_add_line,
+                        skip_empty_lines=self._skip_empty_lines,
+                        decode=False,
+                    )
+                else:  # Raw txt file
+                    # Directly mmap the file.
+                    # We just need to scan once through it to find line offsets.
+                    file = open(file_name, "rb")
+                    file_mmap = mmap.mmap(file.fileno(), 0, flags=mmap.MAP_PRIVATE)
+                    file_index = len(self._orth_files)
+                    self._orth_files.append(file)
+                    self._orth_mmaps.append(file_mmap)
+
+                    pos = 0
+                    while True:
+                        next_new_line = file_mmap.find(b"\n", pos)
+                        if next_new_line == -1:
+                            break
+                        line_len = next_new_line - pos
+                        if line_len or not self._skip_empty_lines:
+                            orths.append((file_index, pos, line_len))
+                        total_bytes_read += line_len + 1
+                        pos = next_new_line + 1
+                        _maybe_report_status()
+
+            except _ReachedDebugLimitLineCount as exc:
+                assert exc.estimated_total_num_seqs is not None  # currently only for _iter_txt implemented
+                debug_limit_est_total += exc.estimated_total_num_seqs
+                debug_limit_hit = True
+            else:  # iteration completed without hitting debug limit
+                debug_limit_est_total += len(orths) - prev_orth_len
 
             lens_per_corpus_file.append(len(orths) - prev_orth_len)
             prev_orth_len = len(orths)
@@ -447,6 +469,18 @@ class LmDataset(CachedDataset2):
             file=log.v4,
         )
 
+        if debug_limit_hit:
+            est_frac_loaded = len(self._orths_offsets_and_lens) / debug_limit_est_total
+            new_partition_epoch = max(int(self.partition_epoch * est_frac_loaded), 1)
+            print(
+                f"LmDataset: debug limit of {debug_limit_line_count} lines (per file) hit,"
+                f" estimated total num seqs {debug_limit_est_total},"
+                f" loaded {len(self._orths_offsets_and_lens)}, {est_frac_loaded:.2%},"
+                f" adjusting partition_epoch from {self.partition_epoch} to {new_partition_epoch}",
+                file=log.v4,
+            )
+            self.partition_epoch = new_partition_epoch
+
         # It's only estimated because we might filter some out or so.
         self._estimated_num_seqs = len(self._orths_offsets_and_lens) // self.partition_epoch
 
@@ -784,19 +818,34 @@ def _iter_txt(
     :param decode:
     """
     f = open(filename, "rb")
+    f_ = f
     if filename.endswith(".gz"):
         f = gzip.GzipFile(fileobj=f)
 
-    for line in f:
-        if decode:
-            try:
-                line = line.decode("utf8")
-            except UnicodeDecodeError:
-                line = line.decode("latin_1")  # or iso8859_15?
-        line = line.strip()
-        if skip_empty_lines and not line:
-            continue
-        callback(line)
+    count = 0
+    try:
+        for line in f:
+            if decode:
+                try:
+                    line = line.decode("utf8")
+                except UnicodeDecodeError:
+                    line = line.decode("latin_1")  # or iso8859_15?
+            line = line.strip()
+            if skip_empty_lines and not line:
+                continue
+            count += 1
+            callback(line)
+
+    except _ReachedDebugLimitLineCount as exc:
+        print(f"Reached debug limit line count for {filename}, stopping early", file=log.v4)
+        pos = f_.tell()
+        f_.seek(0, os.SEEK_END)
+        size = f_.tell()
+        print(f"  stopped at byte {human_bytes_size(pos)} / {human_bytes_size(size)}", file=log.v4)
+        estimated_num_seqs = int(count * (size / pos))
+        print(f"  estimated total num seqs: {estimated_num_seqs}", file=log.v4)
+        exc.estimated_total_num_seqs = estimated_num_seqs
+        raise
 
 
 def iter_corpus(
@@ -2517,6 +2566,25 @@ def get_post_processor_function(opts):
     return chained_post_processors
 
 
+def _get_debug_limit_line_count() -> Optional[int]:
+    """
+    :return: if set, limit to this many lines for debugging
+    """
+    from returnn.config import get_global_config
+
+    config = get_global_config(raise_exception=False)
+    if not config:
+        return None
+
+    return config.int("lm_dataset_debug_limit_line_count", None)
+
+
+class _ReachedDebugLimitLineCount(Exception):
+    """internal exception to signal reached debug limit line count"""
+
+    estimated_total_num_seqs: Optional[int] = None
+
+
 def _main():
     from returnn.util import better_exchook
 

{returnn-1.20251106.185107 → returnn-1.20260118.2517}/returnn/datasets/meta.py

@@ -253,22 +253,12 @@ class MetaDataset(CachedDataset2):
         }
 
         self._seq_list_file = seq_list_file
-        self.seq_list_original = self._load_seq_list(seq_list_file)
-        self.num_total_seqs = len(self.seq_list_original[self.default_dataset_key])
-        for key in self.dataset_keys:
-            assert len(self.seq_list_original[key]) == self.num_total_seqs
-
-        self.tag_idx = {tag: idx for (idx, tag) in enumerate(self.seq_list_original[self.default_dataset_key])}
+        self.seq_list_original: Optional[Dict[str, List[str]]] = None
+        self.tag_idx: Optional[Dict[str, int]] = None
 
         self._seq_lens: Optional[Dict[str, NumbersDict]] = None
         self._num_timesteps: Optional[NumbersDict] = None
         self._seq_lens_file = seq_lens_file
-        if seq_lens_file:
-            seq_lens = load_json(filename=seq_lens_file)
-            assert isinstance(seq_lens, dict)
-            # dict[str,NumbersDict], seq-tag -> data-key -> len
-            self._seq_lens = {tag: NumbersDict(l) for (tag, l) in seq_lens.items()}
-            self._num_timesteps = sum([self._seq_lens[s] for s in self.seq_list_original[self.default_dataset_key]])
 
         if data_dims:
             data_dims = convert_data_dims(data_dims)
@@ -290,19 +280,20 @@ class MetaDataset(CachedDataset2):
         self.num_outputs = self.data_dims
 
         self.orig_seq_order_is_initialized = False
+        self._current_seq_order: List[int] = []
         self.seq_list_ordered: Optional[Dict[str, List[str]]] = None
 
-    def _load_seq_list(self, seq_list_file: Optional[Union[str, Dict[str, str]]] = None) -> Dict[str, List[str]]:
-        """
-        :param seq_list_file:
-        :return: dict: dataset key -> seq list
-        """
-        if not seq_list_file:
+    def _lazy_init_seq_list(self):
+        if self.seq_list_original is not None:
+            return
+
+        if not self._seq_list_file:
             # We create a sequence list from all the sequences of the default dataset
             # and hope that it also applies to the
             # other datasets.
             # This can only work if all datasets have the same tag format and the sequences in the other
             # datasets are a subset of those in the default dataset.
+            # (But the order does not matter.)
             default_dataset = self.datasets[self.default_dataset_key]
             assert isinstance(default_dataset, Dataset)
             print(
@@ -349,17 +340,18 @@ class MetaDataset(CachedDataset2):
                         break  # only print one
                 del seq_list_set
                 raise Exception("Dataset %r is missing seqs." % key)
-        elif isinstance(seq_list_file, str):
-            seq_list = Dataset._load_seq_list_file(seq_list_file, expect_list=False)
-        elif isinstance(seq_list_file, dict):
+        elif isinstance(self._seq_list_file, str):
+            seq_list = Dataset._load_seq_list_file(self._seq_list_file, expect_list=False)
+        elif isinstance(self._seq_list_file, dict):
             for key in self.dataset_keys:
-                if key not in seq_list_file:
+                if key not in self._seq_list_file:
                     raise ValueError(f"seq_list_file does not contain all datasets, missing {key}")
-            seq_list = {key: Dataset._load_seq_list_file(seq_list_file[key]) for key in self.dataset_keys}
+            seq_list = {key: Dataset._load_seq_list_file(self._seq_list_file[key]) for key in self.dataset_keys}
         else:
-            raise TypeError(f"unexpected seq_list_file type {type(seq_list_file)}")
+            raise TypeError(f"unexpected seq_list_file type {type(self._seq_list_file)}")
 
         if isinstance(seq_list, list):
+            # Use same seq list for all datasets
             seq_list = {key: seq_list for key in self.dataset_keys}
         elif isinstance(seq_list, dict):
             for key in self.dataset_keys:
@@ -368,10 +360,29 @@ class MetaDataset(CachedDataset2):
         else:
             raise TypeError(f"unexpected seq_list type {type(seq_list)}")
 
-        return seq_list
+        for key in self.dataset_keys:
+            assert len(seq_list[key]) == len(seq_list[self.default_dataset_key])
+
+        self.seq_list_original = seq_list
+
+    def _lazy_init_tag_idx(self):
+        if self.tag_idx is not None:
+            return
+        self._lazy_init_seq_list()
+        self.tag_idx = {tag: idx for (idx, tag) in enumerate(self.seq_list_original[self.default_dataset_key])}
+
+    def _lazy_init_seq_lens(self):
+        if self._seq_lens is not None:
+            return
+        assert self._seq_lens_file
+        seq_lens = load_json(filename=self._seq_lens_file)
+        assert isinstance(seq_lens, dict)
+        # dict[str,NumbersDict], seq-tag -> data-key -> len
+        self._seq_lens = {tag: NumbersDict(lens) for (tag, lens) in seq_lens.items()}
 
     def _get_dataset_seq_length(self, seq_idx: int):
         if not self.orig_seq_order_is_initialized:
+            self._lazy_init_seq_list()
             # To use get_seq_length() we first have to init the sequence order once in original order.
             # If sequence lengths are not needed by get_seq_order_for_epoch this is never executed.
             self.datasets[self.default_dataset_key].init_seq_order(
@@ -379,6 +390,9 @@ class MetaDataset(CachedDataset2):
             )
             self.orig_seq_order_is_initialized = True
 
+        # Warning: This is not correct in the general case.
+        # get_seq_length needs to have load_seqs called beforehand per API contract.
+        # For some datasets, it might anyway work.
         return self.datasets[self.default_dataset_key].get_seq_length(seq_idx)["data"]
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
@@ -392,6 +406,7 @@ class MetaDataset(CachedDataset2):
             self.epoch is None
             or self.epoch != epoch
             or self.seq_list_ordered is None
+            or not self._current_seq_order
             or seq_list is not None
             or seq_order is not None
             or self.expected_load_seq_start > 0
@@ -401,16 +416,17 @@ class MetaDataset(CachedDataset2):
             # This is called via initialize() with epoch=None, just to init some other things.
             # We are not expected to have prepared any real epoch here.
             self._num_seqs = 0
+            self._current_seq_order = []
             return True
 
         if not need_reinit:
-            self._num_seqs = len(self.seq_list_ordered[self.default_dataset_key])
             return False
 
         seq_order_dataset = None
         if seq_order is not None:
             seq_index = seq_order
         elif seq_list is not None:
+            self._lazy_init_tag_idx()
             seq_index = [self.tag_idx[tag] for tag in seq_list]
         elif self.seq_order_control_dataset:
             seq_order_dataset = self.datasets[self.seq_order_control_dataset]
@@ -418,13 +434,15 @@ class MetaDataset(CachedDataset2):
             seq_order_dataset.init_seq_order(epoch=epoch)
             seq_index = seq_order_dataset.get_current_seq_order()
         else:
-            if self._seq_lens:
+            if self._seq_lens_file:
 
                 def get_seq_len(s):
                     """
                     :param int s:
                     :rtype: int
                     """
+                    self._lazy_init_seq_list()
+                    self._lazy_init_seq_lens()
                     return self._seq_lens[self.seq_list_original[self.default_dataset_key][s]]["data"]
 
             elif self._seq_order_seq_lens_file:
@@ -432,8 +450,10 @@ class MetaDataset(CachedDataset2):
             else:
                 self.orig_seq_order_is_initialized = False
                 get_seq_len = self._get_dataset_seq_length
-            seq_index = self.get_seq_order_for_epoch(epoch, self.num_total_seqs, get_seq_len)
+            seq_index = self.get_seq_order_for_epoch(epoch, self.get_total_num_seqs(), get_seq_len)
         self._num_seqs = len(seq_index)
+        self._current_seq_order = seq_index
+        self._lazy_init_seq_list()
         self.seq_list_ordered = {key: [ls[s] for s in seq_index] for (key, ls) in self.seq_list_original.items()}
 
         for dataset_key, dataset in self.datasets.items():
@@ -447,7 +467,7 @@ class MetaDataset(CachedDataset2):
         """supports sorting"""
         if self.seq_order_control_dataset:
             return self.datasets[self.seq_order_control_dataset].supports_seq_order_sorting()
-        if self._seq_lens or self._seq_order_seq_lens_file:
+        if self._seq_lens_file or self._seq_order_seq_lens_file:
             return True
         return False
 
@@ -464,20 +484,40 @@ class MetaDataset(CachedDataset2):
         :return: current seq order for the current epoch, after self.init_seq_order was called.
         :rtype: list[int]
         """
-        return [self.tag_idx[tag] for tag in self.seq_list_ordered[self.default_dataset_key]]
+        return self._current_seq_order
 
     def get_all_tags(self):
         """
         :return: list of all seq tags, of the whole dataset, without partition epoch
         :rtype: list[str]
         """
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_all_tags()
+        self._lazy_init_seq_list()
+        assert self.seq_list_original is not None
         return self.seq_list_original[self.default_dataset_key]
 
     def get_total_num_seqs(self, *, fast: bool = False) -> int:
         """
+        :param fast: if True, might raise an exception if not possible to get fast.
         :return: total number of seqs, without partition epoch
         """
-        return self.num_total_seqs
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_total_num_seqs(fast=fast)
+        if fast and self.seq_list_original is None:
+            raise OptionalNotImplementedError(f"{self} get_total_num_seqs, seq list not loaded yet")
+        self._lazy_init_seq_list()
+        assert self.seq_list_original is not None
+        return len(self.seq_list_original[self.default_dataset_key])
+
+    def get_num_timesteps(self):
+        """num timesteps"""
+        if self._num_timesteps is None and self._seq_lens_file:
+            self._lazy_init_seq_lens()
+            self._num_timesteps = sum([self._seq_lens[s] for s in self.get_all_tags()], start=NumbersDict())
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_num_timesteps()
+        return super().get_num_timesteps()
 
     def finish_epoch(self, *, free_resources: bool = False):
         """
@@ -503,8 +543,9 @@ class MetaDataset(CachedDataset2):
         if start_ < end:
             for dataset_key in self.dataset_keys:
                 self.datasets[dataset_key].load_seqs(start_, end)
-                for seq_idx in range(start_, end):
-                    self._check_dataset_seq(dataset_key, seq_idx)
+                if self.seq_list_ordered is not None:
+                    for seq_idx in range(start_, end):
+                        self._check_dataset_seq(dataset_key, seq_idx)
         super(MetaDataset, self)._load_seqs(start=start, end=end)
 
     def _check_dataset_seq(self, dataset_key, seq_idx):
@@ -531,7 +572,7 @@ class MetaDataset(CachedDataset2):
         :type seq_idx: int
         :rtype: DatasetSeq
         """
-        seq_tag = self.seq_list_ordered[self.default_dataset_key][seq_idx]
+        seq_tag = self.get_tag(seq_idx)
         features = {data_key: self._get_data(seq_idx, data_key) for data_key in self.data_keys}
         return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=features)
 
@@ -540,8 +581,9 @@ class MetaDataset(CachedDataset2):
         :param int sorted_seq_idx:
         :rtype: NumbersDict
         """
-        if self._seq_lens:
-            return self._seq_lens[self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]]
+        if self._seq_lens_file:
+            self._lazy_init_seq_lens()
+            return self._seq_lens[self.get_tag(sorted_seq_idx)]
         return super(MetaDataset, self).get_seq_length(sorted_seq_idx)
 
     def get_tag(self, sorted_seq_idx):
@@ -549,7 +591,10 @@ class MetaDataset(CachedDataset2):
         :param int sorted_seq_idx:
         :rtype: str
         """
-        return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
+        if self.seq_list_ordered is not None:
+            return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
+        else:
+            return self.datasets[self.default_dataset_key].get_tag(sorted_seq_idx)
 
     def get_complete_frac(self, sorted_seq_idx: int, **kwargs) -> Optional[float]:
         """
@@ -961,6 +1006,7 @@ class CombinedDataset(CachedDataset2):
         super(CombinedDataset, self).__init__(**kwargs)
         assert self.shuffle_frames_of_nseqs == 0  # not implemented. anyway only for non-recurrent nets
 
+        self.data_map = data_map
         self.dataset_keys = set([m[0] for m in data_map.keys()])  # type: typing.Set[str]
         self.dataset_idx2key_map = dict(enumerate(sorted(self.dataset_keys)))  # idx -> dataset-key
         self.data_keys = set(data_map.values())  # type: typing.Set[str]
@@ -1248,6 +1294,10 @@ class CombinedDataset(CachedDataset2):
         # Cur meaning for the next sequence to be added to dataset_sorted_seq_idx_list.
         seq_idx = self.used_num_seqs_per_subset[dataset_idx]
         cur_start, cur_end = self._sub_dataset_cur_loaded_seq_range[dataset_idx]
+
+        if not self.datasets[self.dataset_idx2key_map[dataset_idx]].is_less_than_num_seqs(seq_idx):
+            return False
+
         if seq_idx >= cur_end:
             self._sub_dataset_load_seqs(dataset_idx, cur_start, seq_idx + 1)
             return True
@@ -1294,10 +1344,12 @@ class CombinedDataset(CachedDataset2):
                 complete_fracs_and_ds_idx = [
                     (
                         self.datasets[self.dataset_idx2key_map[j]].get_complete_frac(
-                            self.used_num_seqs_per_subset[j] - 1, allow_only_lr_suitable=True
+                            self.used_num_seqs_per_subset[j], allow_only_lr_suitable=True
+                        )
+                        if self.datasets[self.dataset_idx2key_map[j]].is_less_than_num_seqs(
+                            self.used_num_seqs_per_subset[j]
                         )
-                        if self.used_num_seqs_per_subset[j] > 0
-                        else 0.0,
+                        else float("inf"),
                         j,
                     )
                     for j in range(len(self.datasets))
@@ -1309,9 +1361,7 @@ class CombinedDataset(CachedDataset2):
                 # Sort by complete frac, i.e. datasets with the lowest complete frac first.
                 complete_fracs_and_ds_idx.sort()
                 for complete_frac, dataset_idx in complete_fracs_and_ds_idx:
-                    if self.datasets[self.dataset_idx2key_map[dataset_idx]].is_less_than_num_seqs(
-                        self.used_num_seqs_per_subset[dataset_idx]
-                    ):
+                    if complete_frac < float("inf"):
                         break
                 else:
                     return False  # No dataset has remaining data