Repository: huggingface/datasets

Issue:
Epochs shortened after resuming mid-epoch with Iterable dataset+StatefulDataloader(persistent_workers=True)
### Describe the bug

When `torchdata.stateful_dataloader.StatefulDataloader(persistent_workers=True)` the epochs after resuming only iterate through the examples that were left in the epoch when the training was interrupted. For example, in the script below training is interrupted on step 124 (epoch 1) when 3 batches are left. Then after resuming, the rest of epochs (2 and 3) only iterate through these 3 batches.

### Steps to reproduce the bug

Run the following script with and with PERSISTENT_WORKERS=true.

```python
# !/usr/bin/env python3
# torch==2.5.1
# datasets==3.3.2
# torchdata>=0.9.0
import datasets
import pprint
from torchdata.stateful_dataloader import StatefulDataLoader

import os

PERSISTENT_WORKERS = (
    os.environ.get("PERSISTENT_WORKERS", "False").lower() == "true"
)

# PERSISTENT_WORKERS = True  # Incorrect resume


# ds = datasets.load_from_disk("dataset").to_iterable_dataset(num_shards=4)
def generator():
    for i in range(128):
        yield {"x": i}


ds = datasets.Dataset.from_generator(
    generator, features=datasets.Features({"x": datasets.Value("int32")})
).to_iterable_dataset(num_shards=4)

dl = StatefulDataLoader(
    ds, batch_size=2, num_workers=2, persistent_workers=PERSISTENT_WORKERS
)
global_step = 0
epoch = 0
ds_state_dict = None
state_dict = None
resumed = False
while True:
    if epoch >= 3:
        break
    if state_dict is not None:
        dl.load_state_dict(state_dict)
        state_dict = None
        ds_state_dict = None
        resumed = True
        print("resumed")
    for i, batch in enumerate(dl):
        print(f"epoch: {epoch}, global_step: {global_step}, batch: {batch}")
        global_step += 1  # consume datapoint
        # simulate error
        if global_step == 124 and not resumed:
            ds_state_dict = ds.state_dict()
            state_dict = dl.state_dict()
            print("checkpoint")
            print("ds_state_dict")
            pprint.pprint(ds_state_dict)
            print("dl_state_dict")
            pprint.pprint(state_dict)
            break

    if state_dict is None:
        ds.set_epoch(epoch)
        epoch += 1
```

The script checkpoints when there are three batches left in the second epoch. After resuming, only the last three batches are repeated in the rest of the epochs.

If it helps, following are the two state_dicts for the dataloader save at the same step with the two settings. The left one is for `PERSISTENT_WORKERS=False`

![Image](https://github.com/user-attachments/assets/c97d6502-d7bd-4ef4-ae2d-66fe1a9732b1)

### Expected behavior

All the elements in the dataset should be iterated through in the epochs following the one where we resumed. The expected behavior can be seen by setting `PERSISTENT_WORKERS=False`. 

### Environment info

torch==2.5.1
datasets==3.3.2
torchdata>=0.9.0

Hints from issue comments:
@lhoestq: Thanks for reporting ! Maybe we should store the epoch in the state_dict, and then when the dataset is iterated on again after setting a new epoch it should restart from scratch instead of resuming ? wdyt ?@dhruvdcoder: But why does this only happen when `persistent_workers=True`? I would expect it to work correctly even without storing the epoch number in the state_dict of the iterable dataset. @lhoestq: I think persistent_workers=False simply ignores the dataset state_dict when it starts a new epoch, that's why the issue doesn't appear in that case@lhoestq: I opened https://github.com/huggingface/datasets/pull/7451 to fix the issue, let me know if it works for you@lhoestq: I just released `datasets` 3.4 that includes the fix :)

PS: in your script you probably want to set the epoch like this, otherwise it's still set to 0 after the first epoch:

```diff
    if state_dict is None:
-       ds.set_epoch(epoch)
        epoch += 1
+       ds.set_epoch(epoch)
```@tianyu-l: @lhoestq 
If I understand correctly, the issue was:
when training saves a checkpoint of dataloader in epoch 1, the resumed training only consumes partial data in epoch 2, 3, etc.

However, with the fix we are facing the issue that:
when training saves a checkpoint of dataloader in epoch 2, the resumed training starts from scratch instead of consuming remaining partial data in epoch 2.

This makes training inconsistent between resuming from a checkpoint vs. original training if continued without a checkpoint.

Solution patch:
diff --git a/docs/source/audio_load.mdx b/docs/source/audio_load.mdx
index 2d4fed9a2..c9db4fd56 100644
--- a/docs/source/audio_load.mdx
+++ b/docs/source/audio_load.mdx
@@ -96,3 +96,27 @@ For more information about creating your own `AudioFolder` dataset, take a look
 </Tip>
 
 For a guide on how to load any type of dataset, take a look at the <a class="underline decoration-sky-400 decoration-2 font-semibold" href="./loading">general loading guide</a>.
+
+## Audio decoding
+
+By default, audio files are decoded sequentially as NumPy arrays when you iterate on a dataset.
+However it is possible to speed up the dataset significantly using multithreaded decoding:
+
+```python
+>>> import os
+>>> num_threads = num_threads = min(32, (os.cpu_count() or 1) + 4)
+>>> dataset = dataset.decode(num_threads=num_threads)
+>>> for example in dataset:  # up to 20 times faster !
+...     ...
+```
+
+You can enable multithreading using `num_threads`. This is especially useful to speed up remote data streaming.
+However it can be slower than `num_threads=0` for local data on fast disks.
+
+If you are not interested in the images decoded as NumPy arrays and would like to access the path/bytes instead, you can disable decoding:
+
+```python
+>>> dataset = dataset.decode(False)
+```
+
+Note: [`IterableDataset.decode`] is only available for streaming datasets at the moment.
diff --git a/docs/source/image_load.mdx b/docs/source/image_load.mdx
index ead0d2d72..33afd5217 100644
--- a/docs/source/image_load.mdx
+++ b/docs/source/image_load.mdx
@@ -138,3 +138,27 @@ You can load a WebDataset like this:
 
 >>> dataset = load_dataset("webdataset", data_dir="/path/to/folder", streaming=True)
 ```
+
+## Image decoding
+
+By default, images are decoded sequentially as `PIL.Images` when you iterate on a dataset.
+However it is possible to speed up the dataset significantly using multithreaded decoding:
+
+```python
+>>> import os
+>>> num_threads = num_threads = min(32, (os.cpu_count() or 1) + 4)
+>>> dataset = dataset.decode(num_threads=num_threads)
+>>> for example in dataset:  # up to 20 times faster !
+...     ...
+```
+
+You can enable multithreading using `num_threads`. This is especially useful to speed up remote data streaming.
+However it can be slower than `num_threads=0` for local data on fast disks.
+
+If you are not interested in the images decoded as `PIL.Images` and would like to access the path/bytes instead, you can disable decoding:
+
+```python
+>>> dataset = dataset.decode(False)
+```
+
+Note: [`IterableDataset.decode`] is only available for streaming datasets at the moment.
diff --git a/docs/source/package_reference/main_classes.mdx b/docs/source/package_reference/main_classes.mdx
index 62dc9127d..86f1788bb 100644
--- a/docs/source/package_reference/main_classes.mdx
+++ b/docs/source/package_reference/main_classes.mdx
@@ -163,6 +163,7 @@ The base class [`IterableDataset`] implements an iterable Dataset backed by pyth
     - select_columns
     - cast_column
     - cast
+    - decode
     - __iter__
     - iter
     - map
diff --git a/docs/source/video_load.mdx b/docs/source/video_load.mdx
index 428c8aadf..ed9702cf6 100644
--- a/docs/source/video_load.mdx
+++ b/docs/source/video_load.mdx
@@ -169,3 +169,29 @@ You can load a WebDataset like this:
 
 >>> dataset = load_dataset("webdataset", data_dir="/path/to/folder", streaming=True)
 ```
+
+## Video decoding
+
+By default, videos are decoded sequentially as torchvision `VideoReaders` when you iterate on a dataset.
+It sequentially decodes the metadata of the videos, and doesn't read the video frames until you access them.
+
+However it is possible to speed up the dataset significantly using multithreaded decoding:
+
+```python
+>>> import os
+>>> num_threads = num_threads = min(32, (os.cpu_count() or 1) + 4)
+>>> dataset = dataset.decode(num_threads=num_threads)
+>>> for example in dataset:  # up to 20 times faster !
+...     ...
+```
+
+You can enable multithreading using `num_threads`. This is especially useful to speed up remote data streaming.
+However it can be slower than `num_threads=0` for local data on fast disks.
+
+If you are not interested in the images decoded as torchvision `VideoReaders` and would like to access the path/bytes instead, you can disable decoding:
+
+```python
+>>> dataset = dataset.decode(False)
+```
+
+Note: [`IterableDataset.decode`] is only available for streaming datasets at the moment.
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index cf41688c3..cc39ffb56 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -3551,7 +3551,7 @@ class Dataset(DatasetInfoMixin, IndexableMixin, TensorflowDatasetMixin):
                         task.cancel(msg="KeyboardInterrupt")
                     try:
                         loop.run_until_complete(asyncio.gather(*tasks))
-                    except asyncio.CancelledError:
+                    except (asyncio.CancelledError, ValueError):
                         logger.debug("Tasks canceled.")
                 raise
 
diff --git a/src/datasets/commands/convert.py b/src/datasets/commands/convert.py
index f8a706af9..f82547efd 100644
--- a/src/datasets/commands/convert.py
+++ b/src/datasets/commands/convert.py
@@ -167,7 +167,7 @@ class ConvertCommand(BaseDatasetsCLICommand):
                 output_file = os.path.join(output_dir, f_name)
                 os.makedirs(output_dir, exist_ok=True)
                 self._logger.info(f"Adding directory {output_dir}")
-                imports_to_builder_map.update({imp: output_dir for imp in tfds_imports})
+                imports_to_builder_map.update(dict.fromkeys(tfds_imports, output_dir))
             else:
                 # Utilities will be moved at the end
                 utils_files.append(output_file)
diff --git a/src/datasets/dataset_dict.py b/src/datasets/dataset_dict.py
index c743dc90f..9393d0960 100644
--- a/src/datasets/dataset_dict.py
+++ b/src/datasets/dataset_dict.py
@@ -931,7 +931,7 @@ class DatasetDict(dict):
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
 
         dataset_dict = {}
         for split, dataset in self.items():
@@ -1051,7 +1051,7 @@ class DatasetDict(dict):
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.filter(
@@ -1109,7 +1109,7 @@ class DatasetDict(dict):
         """
         self._check_values_type()
         if cache_file_names is None:
-            cache_file_names = {k: None for k in self}
+            cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.flatten_indices(
@@ -1176,7 +1176,7 @@ class DatasetDict(dict):
         """
         self._check_values_type()
         if indices_cache_file_names is None:
-            indices_cache_file_names = {k: None for k in self}
+            indices_cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.sort(
@@ -1254,13 +1254,13 @@ class DatasetDict(dict):
             raise ValueError("Please specify seed or seeds, but not both")
         seeds = seed if seed is not None else seeds
         if seeds is None:
-            seeds = {k: None for k in self}
+            seeds = dict.fromkeys(self)
         elif not isinstance(seeds, dict):
-            seeds = {k: seeds for k in self}
+            seeds = dict.fromkeys(self, seeds)
         if generators is None:
-            generators = {k: None for k in self}
+            generators = dict.fromkeys(self)
         if indices_cache_file_names is None:
-            indices_cache_file_names = {k: None for k in self}
+            indices_cache_file_names = dict.fromkeys(self)
         return DatasetDict(
             {
                 k: dataset.shuffle(
@@ -1326,7 +1326,7 @@ class DatasetDict(dict):
         fs, _ = url_to_fs(dataset_dict_path, **(storage_options or {}))
 
         if num_shards is None:
-            num_shards = {k: None for k in self}
+            num_shards = dict.fromkeys(self)
         elif not isinstance(num_shards, dict):
             raise ValueError(
                 "Please provide one `num_shards` per dataset in the dataset dictionary, e.g. {{'train': 128, 'test': 4}}"
@@ -1696,7 +1696,7 @@ class DatasetDict(dict):
         ```
         """
         if num_shards is None:
-            num_shards = {k: None for k in self}
+            num_shards = dict.fromkeys(self)
         elif not isinstance(num_shards, dict):
             raise ValueError(
                 "Please provide one `num_shards` per dataset in the dataset dictionary, e.g. {{'train': 128, 'test': 4}}"
diff --git a/src/datasets/formatting/formatting.py b/src/datasets/formatting/formatting.py
index c07c9e2c1..165f178f3 100644
--- a/src/datasets/formatting/formatting.py
+++ b/src/datasets/formatting/formatting.py
@@ -270,7 +270,7 @@ class LazyDict(MutableMapping):
         self.pa_table = pa_table
         self.formatter = formatter
 
-        self.data = {key: None for key in pa_table.column_names}
+        self.data = dict.fromkeys(pa_table.column_names)
         self.keys_to_format = set(self.data.keys())
 
     def __len__(self):
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index 7fab8bcc9..423abeeb5 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -2,6 +2,7 @@ import asyncio
 import copy
 import inspect
 import itertools
+import multiprocessing.pool
 import sys
 from collections import Counter
 from collections.abc import Iterable, Iterator
@@ -24,6 +25,7 @@ from .features.features import (
     Value,
     _align_features,
     _check_if_features_can_be_aligned,
+    _visit,
     cast_to_python_objects,
 )
 from .formatting import (
@@ -210,7 +212,7 @@ class ExamplesIterable(_BaseExamplesIterable):
         self.kwargs = kwargs
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -248,7 +250,7 @@ class ShuffledDataSourcesExamplesIterable(ExamplesIterable):
         self.generator = deepcopy(generator)
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -288,7 +290,7 @@ class ArrowExamplesIterable(_BaseExamplesIterable):
         return self._iter_arrow
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -355,7 +357,7 @@ class ShuffledDataSourcesArrowExamplesIterable(ArrowExamplesIterable):
         self.generator = deepcopy(generator)
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0}
+        self._state_dict = {"shard_idx": 0, "shard_example_idx": 0, "type": self.__class__.__name__}
         return self._state_dict
 
     def __iter__(self):
@@ -435,11 +437,12 @@ class RebatchedArrowExamplesIterable(_BaseExamplesIterable):
 
     def _init_state_dict(self) -> dict:
         self._state_dict = {
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
             "previous_state": None,
             "batch_idx": 0,
             "num_chunks_since_previous_state": 0,
             "cropped_chunk_length": 0,
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -678,6 +681,7 @@ class CyclingMultiSourcesExamplesIterable(_BaseExamplesIterable):
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
             "previous_states": [None] * len(self.ex_iterables),
             "is_exhausted": [False] * len(self.ex_iterables),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -776,6 +780,7 @@ class VerticallyConcatenatedMultiSourcesExamplesIterable(_BaseExamplesIterable):
         self._state_dict = {
             "ex_iterable_idx": 0,
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -856,7 +861,10 @@ class HorizontallyConcatenatedMultiSourcesExamplesIterable(_BaseExamplesIterable
         return self.ex_iterables[0].features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables]}
+        self._state_dict = {
+            "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -958,6 +966,7 @@ class RandomlyCyclingMultiSourcesExamplesIterable(CyclingMultiSourcesExamplesIte
             "ex_iterables": [ex_iterable._init_state_dict() for ex_iterable in self.ex_iterables],
             "previous_states": [None] * len(self.ex_iterables),
             "is_exhausted": [False] * len(self.ex_iterables),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1010,6 +1019,7 @@ class MappedExamplesIterable(_BaseExamplesIterable):
         fn_kwargs: Optional[dict] = None,
         formatting: Optional["FormattingConfig"] = None,
         features: Optional[Features] = None,
+        max_num_running_async_map_functions_in_parallel: Optional[int] = None,
     ):
         super().__init__()
         self.ex_iterable = ex_iterable
@@ -1023,6 +1033,9 @@ class MappedExamplesIterable(_BaseExamplesIterable):
         self.fn_kwargs = fn_kwargs or {}
         self.formatting = formatting  # required for iter_arrow
         self._features = features
+        self.max_num_running_async_map_functions_in_parallel = (
+            max_num_running_async_map_functions_in_parallel or config.MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL
+        )
         # sanity checks
         if formatting and formatting.is_table:
             # batch_size should match for iter_arrow
@@ -1036,6 +1049,8 @@ class MappedExamplesIterable(_BaseExamplesIterable):
                     f"The {formatting.format_type.capitalize()}-formatted {type(self).__name__} has batch_size={batch_size if batched else 1} which is"
                     f"different from {ex_iterable.batch_size=} from its underlying iterable."
                 )
+        # to enable graceful ends
+        self._owned_loops_and_tasks: list[tuple[asyncio.AbstractEventLoop, list[asyncio.Task]]] = []
 
     @property
     def iter_arrow(self):
@@ -1052,10 +1067,11 @@ class MappedExamplesIterable(_BaseExamplesIterable):
 
     def _init_state_dict(self) -> dict:
         self._state_dict = {
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
             "previous_state": None,
             "num_examples_since_previous_state": 0,
             "previous_state_example_idx": 0,
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1174,6 +1190,7 @@ class MappedExamplesIterable(_BaseExamplesIterable):
                 loop = asyncio.get_running_loop()
             except RuntimeError:
                 loop = asyncio.new_event_loop()
+            self._owned_loops_and_tasks.append((loop, tasks))
         else:
             loop = None
 
@@ -1191,15 +1208,15 @@ class MappedExamplesIterable(_BaseExamplesIterable):
                     indices.append(i)
                     tasks.append(loop.create_task(async_apply_function(key_example, i)))
                     # keep the total active tasks under a certain number
-                    if len(tasks) >= config.MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL:
+                    if len(tasks) >= self.max_num_running_async_map_functions_in_parallel:
                         done, pending = loop.run_until_complete(
                             asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
                         )
-                        while tasks and len(pending) >= config.MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL:
+                        while tasks and len(pending) >= self.max_num_running_async_map_functions_in_parallel:
                             done, pending = loop.run_until_complete(
                                 asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
                             )
-                    if len(tasks) >= 10 * config.MAX_NUM_RUNNING_ASYNC_MAP_FUNCTIONS_IN_PARALLEL:
+                    if len(tasks) >= 10 * self.max_num_running_async_map_functions_in_parallel:
                         loop.run_until_complete(tasks[0])
                     # yield finished tasks
                     while tasks and tasks[0].done():
@@ -1257,7 +1274,7 @@ class MappedExamplesIterable(_BaseExamplesIterable):
                     task.cancel(msg="KeyboardInterrupt")
                 try:
                     loop.run_until_complete(asyncio.gather(*tasks))
-                except asyncio.CancelledError:
+                except (asyncio.CancelledError, ValueError):
                     logger.debug("Tasks canceled.")
             raise
 
@@ -1347,6 +1364,7 @@ class MappedExamplesIterable(_BaseExamplesIterable):
             fn_kwargs=self.fn_kwargs,
             formatting=self.formatting,
             features=self.features,
+            max_num_running_async_map_functions_in_parallel=self.max_num_running_async_map_functions_in_parallel,
         )
 
     def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> "MappedExamplesIterable":
@@ -1363,6 +1381,7 @@ class MappedExamplesIterable(_BaseExamplesIterable):
             fn_kwargs=self.fn_kwargs,
             formatting=self.formatting,
             features=self.features,
+            max_num_running_async_map_functions_in_parallel=self.max_num_running_async_map_functions_in_parallel,
         )
 
     @property
@@ -1567,7 +1586,11 @@ class SkipExamplesIterable(_BaseExamplesIterable):
         return self.ex_iterable.features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"skipped": False, "ex_iterable": self.ex_iterable._init_state_dict()}
+        self._state_dict = {
+            "skipped": False,
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -1631,7 +1654,8 @@ class RepeatExamplesIterable(_BaseExamplesIterable):
     def _init_state_dict(self) -> dict:
         self._state_dict = {
             "repeat_index": 0,
-            "ex_iterable": self.ex_iterable._init_state_dict(),
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
         }
         return self._state_dict
 
@@ -1644,7 +1668,7 @@ class RepeatExamplesIterable(_BaseExamplesIterable):
             repeat_index += 1
             if self._state_dict:
                 self._state_dict["repeat_index"] = repeat_index
-                self._state_dict["ex_iterable"] = self.ex_iterable._init_state_dict()
+                self._state_dict["examples_iterable"] = self.ex_iterable._init_state_dict()
 
     def shuffle_data_sources(self, generator: np.random.Generator) -> "RepeatExamplesIterable":
         """Shuffle the underlying iterable, then repeat."""
@@ -1686,7 +1710,11 @@ class TakeExamplesIterable(_BaseExamplesIterable):
         return self.ex_iterable.features
 
     def _init_state_dict(self) -> dict:
-        self._state_dict = {"num_taken": 0, "ex_iterable": self.ex_iterable._init_state_dict()}
+        self._state_dict = {
+            "num_taken": 0,
+            "examples_iterable": self.ex_iterable._init_state_dict(),
+            "type": self.__class__.__name__,
+        }
         return self._state_dict
 
     def __iter__(self):
@@ -1945,9 +1973,8 @@ class IterableDataset(DatasetInfoMixin):
         self._token_per_repo_id: dict[str, Union[str, bool, None]] = token_per_repo_id or {}
         self._epoch: Union[int, "torch.Tensor"] = _maybe_share_with_torch_persistent_workers(0)
         self._starting_state_dict: Optional[dict] = None
-        self._prepared_ex_iterable = self._prepare_ex_iterable_for_iteration()
-        self._state_dict = self._prepared_ex_iterable._init_state_dict()
-        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)
+        self._prepare_ex_iterable_for_iteration()  # set state_dict
+        _maybe_add_torch_iterable_dataset_parent_class(self.__class__)  # subclass of torch IterableDataset
 
     def state_dict(self) -> dict:
         """Get the current state_dict of the dataset.
@@ -2050,7 +2077,6 @@ class IterableDataset(DatasetInfoMixin):
         >>> dataloader.load_state_dict(state_dict)  # uses ds.load_state_dict() under the hood
         ```
         """
-        self._prepared_ex_iterable.load_state_dict(state_dict)
         self._starting_state_dict = state_dict
 
     def __repr__(self):
@@ -2125,9 +2151,12 @@ class IterableDataset(DatasetInfoMixin):
             ex_iterable = ex_iterable.shard_data_sources(
                 num_shards=worker_info.num_workers, index=worker_info.id, contiguous=False
             )
-            self._state_dict = ex_iterable._init_state_dict()
-            if self._starting_state_dict:
-                ex_iterable.load_state_dict(self._starting_state_dict)
+            self._state_dict = {
+                "examples_iterable": ex_iterable._init_state_dict(),
+                "epoch": self.epoch,
+            }
+            if self._starting_state_dict and self.epoch == self._starting_state_dict["epoch"]:
+                ex_iterable.load_state_dict(self._starting_state_dict["examples_iterable"])
 
             if self._formatting and (ex_iterable.iter_arrow or self._formatting.is_table):
                 formatter = get_formatter(self._formatting.format_type, features=self.features)
@@ -2205,9 +2234,12 @@ class IterableDataset(DatasetInfoMixin):
                 token_per_repo_id=self._token_per_repo_id,
             )
 
-        self._state_dict = ex_iterable._init_state_dict()
-        if self._starting_state_dict:
-            ex_iterable.load_state_dict(self._starting_state_dict)
+        self._state_dict = {
+            "examples_iterable": ex_iterable._init_state_dict(),
+            "epoch": self.epoch,
+        }
+        if self._starting_state_dict and self.epoch == self._starting_state_dict["epoch"]:
+            ex_iterable.load_state_dict(self._starting_state_dict["examples_iterable"])
         return ex_iterable
 
     def __iter__(self):
@@ -3189,6 +3221,99 @@ class IterableDataset(DatasetInfoMixin):
             token_per_repo_id=self._token_per_repo_id,
         )
 
+    def decode(self, enable: bool = True, num_threads: int = 0) -> "IterableDataset":
+        """
+        Enable or disable the dataset features decoding for audio, image, video.
+
+        When enabled (default), media types are decoded:
+
+        * audio -> dict of "array" and "sampling_rate" and "path"
+        * image -> PIL.Image
+        * video -> torchvision.io.VideoReader
+
+        You can enable multithreading using `num_threads`. This is especially useful to speed up remote
+        data streaming. However it can be slower than `num_threads=0` for local data on fast disks.
+
+        Disabling decoding is useful if you want to iterate on the paths or bytes of the media files
+        without actually decoding their content. To disable decoding you can use `.decode(False)`, which
+        is equivalent to calling `.cast()` or `.cast_column()` with all the Audio, Image and Video types
+        set to `decode=False`.
+
+        Args:
+            enable (`bool`, defaults to `True`):
+                Enable or disable features decoding.
+            num_threads (`int`, defaults to `0`):
+                Enable multithreading for features decoding.
+
+        Returns:
+            `IterableDataset`: A copy of the dataset with casted features.
+
+        Examples:
+
+        Disable decoding:
+
+        ```py
+        >>> from datasets import load_dataset
+        >>> ds = load_dataset("sshh12/planet-textures", split="train", streaming=True)
+        >>> next(iter(ds))
+        {'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=2048x1024>,
+        'text': 'A distant celestial object with an icy crust, displaying a light blue shade, covered with round pits and rugged terrains.'}
+        >>> ds = ds.decode(False)
+        >>> ds.features
+        {'image': Image(mode=None, decode=False, id=None),
+        'text': Value(dtype='string', id=None)}
+        >>> next(iter(ds))
+        {
+          'image': {
+            'path': 'hf://datasets/sshh12/planet-textures@69dc4cef7a5c4b2cfe387727ec8ea73d4bff7302/train/textures/0000.png',
+            'bytes': None
+          },
+          'text': 'A distant celestial object with an icy crust, displaying a light blue shade, covered with round pits and rugged terrains.'
+        }
+        ```
+
+        Speed up streaming with multithreading:
+
+        ```py
+        >>> import os
+        >>> from datasets import load_dataset
+        >>> from tqdm import tqdm
+        >>> ds = load_dataset("sshh12/planet-textures", split="train", streaming=True)
+        >>> num_threads = min(32, (os.cpu_count() or 1) + 4)
+        >>> ds = ds.decode(num_threads=num_threads)
+        >>> for _ in tqdm(ds):  # 20 times faster !
+        ...     ...
+        ```
+        """
+        if not self.features:
+            raise ValueError(
+                "Features decoding is only available for datasets with known features, but features are Unknown. "
+                "Please set the datasets features with `ds = ds.cast(features)`."
+            )
+        ds = self
+
+        def set_decoding(decode: bool, feature):
+            if hasattr(feature, "decode"):
+                feature.decode = decode
+
+        if enable and num_threads > 0:
+            disabled_decoding_features = self.features.copy()
+            enabled_decoding_features = self.features.copy()
+
+            _visit(disabled_decoding_features, partial(set_decoding, False))
+            _visit(enabled_decoding_features, partial(set_decoding, True))
+            ds = ds.cast(disabled_decoding_features)
+            pool = multiprocessing.pool.ThreadPool(num_threads)
+            func = partial(_apply_async, pool, enabled_decoding_features.decode_example)
+            ds = ds.map(func, features=enabled_decoding_features)
+            assert isinstance(ds._ex_iterable, MappedExamplesIterable)
+            ds._ex_iterable.max_num_running_async_map_functions_in_parallel = 2 * num_threads
+        else:
+            features = ds.features.copy()
+            _visit(features, partial(set_decoding, enable))
+            ds = ds.cast(features)
+        return ds
+
     def _step(self, step: int, offset: int) -> "IterableDataset":
         ex_iterable = StepExamplesIterable(self._ex_iterable, step=step, offset=offset)
         return IterableDataset(
@@ -3407,3 +3532,12 @@ def _split_by_node_iterable_dataset(dataset: IterableDataset, rank: int, world_s
         distributed=distributed,
         token_per_repo_id=dataset._token_per_repo_id,
     )
+
+
+async def _apply_async(pool, func, x):
+    future = pool.apply_async(func, (x,))
+    while True:
+        if future.ready():
+            return future.get()
+        else:
+            await asyncio.sleep(0)


Test patch:
diff --git a/tests/test_iterable_dataset.py b/tests/test_iterable_dataset.py
index ea3c9d850..fd3168416 100644
--- a/tests/test_iterable_dataset.py
+++ b/tests/test_iterable_dataset.py
@@ -1581,6 +1581,17 @@ def test_iterable_dataset_set_epoch(dataset: IterableDataset):
     assert dataset._epoch == 42
 
 
+def test_iterable_dataset_set_epoch_resuming(dataset: IterableDataset):
+    dataset_length = len(list(dataset))
+    assert len(list(dataset)) == dataset_length > 0
+    dataset.load_state_dict(dataset.state_dict())
+    assert len(list(dataset)) == 0
+    dataset.set_epoch(1)
+    assert len(list(dataset)) == dataset_length > 0
+    dataset.load_state_dict(dataset.state_dict())
+    assert len(list(dataset)) == 0
+
+
 @pytest.mark.parametrize("seed", [None, 42, 1337])
 @pytest.mark.parametrize("epoch", [None, 0, 1, 10])
 def test_iterable_dataset_set_epoch_of_shuffled_dataset(dataset: IterableDataset, seed, epoch):
@@ -2474,3 +2485,31 @@ def test_iterable_dataset_batch():
         assert len(batch["text"]) == 3
         assert batch["id"] == [3 * i, 3 * i + 1, 3 * i + 2]
         assert batch["text"] == [f"Text {3 * i}", f"Text {3 * i + 1}", f"Text {3 * i + 2}"]
+
+
+class DecodableFeature:
+    decode_example_num_calls = 0
+
+    def __init__(self):
+        self.decode = True
+
+    def decode_example(self, example, token_per_repo_id=None):
+        type(self).decode_example_num_calls += 1
+        return "decoded" if self.decode else example
+
+
+def test_decode():
+    data = [{"i": i} for i in range(10)]
+    features = Features({"i": DecodableFeature()})
+    ds = IterableDataset.from_generator(lambda: (x for x in data), features=features)
+    assert next(iter(ds)) == {"i": "decoded"}
+    assert DecodableFeature.decode_example_num_calls == 1
+    ds = ds.decode(False)
+    assert next(iter(ds)) == {"i": 0}
+    assert DecodableFeature.decode_example_num_calls == 1
+    ds = ds.decode(True)
+    assert next(iter(ds)) == {"i": "decoded"}
+    assert DecodableFeature.decode_example_num_calls == 2
+    ds = ds.decode(num_threads=1)
+    assert next(iter(ds)) == {"i": "decoded"}
+    assert DecodableFeature.decode_example_num_calls == 4