Fix deepcopy bottleneck in state_dict()

[yuxin00j]

Fix #8393
This PR significantly resolves a major CPU bottleneck in IterableDataset dataloading when buffer_size and max_buffer_input_shards are > 1.

The Bug:
Currently, _BaseExamplesIterable.state_dict() falls back to performing a full recursive copy.deepcopy(self._state_dict).
In nested iterator chains (like BufferShuffledExamplesIterable wrapping CyclingMultiSourcesExamplesIterable wrapping multiple ArrowExamplesIterable shards), the state_dict() is queried and updated on every yielded row.

Because copy.deepcopy relies on Python's pickle mechanics for object introspection and memoization, it incurs significant overhead. Worse, pickle often crashes when it encounters unpicklable objects like generators. When mbis is scaled up and the dictionary size grows, the deepcopy completely blocks the CPU, hiding underneath the shuffle buffer processing.

The Fix:
We replaced the expensive deepcopy() logic with a dynamic, lightweight _fast_copy() helper directly inside _BaseExamplesIterable.state_dict().

Because dataset state_dicts strictly consist of standard collections and primitive state objects (and importantly, contain no circular references), _fast_copy() manually recursively walks dict and list structures. For any other object, it natively returns the object (or uses its .copy() method if available), bypassing pickle and deepcopy entirely.

This allows us to dynamically serialize state without explicitly overriding state_dict in every single dataset subclass, restoring maintainability and drastically improving speed.

Impact:
We measured workloads using a high number of concurrent shards (max_buffer_input_shards=10):

import datasets
from torch.utils.data import DataLoader
from datasets.distributed import split_dataset_by_node

ds = datasets.load_dataset(
    "parquet",
    data_files="gs://your-bucket/benchmark_data/shard_*.parquet",
    split="train",
    streaming=True
)

ds = ds.shuffle(seed=0, buffer_size=10000, max_buffer_input_shards=10)
ds = split_dataset_by_node(ds, rank=0, world_size=8)
ds = ds.with_format("torch")

loader = DataLoader(
    ds,
    batch_size=8,
    num_workers=4,
    prefetch_factor=2,
    persistent_workers=True,
)

for epoch in range(3):
    ds.set_epoch(epoch)
    for batch in loader:
        pass

• Before this PR: ~103.5 seconds per epoch.
• After this PR: ~65.8 seconds per epoch.

This single PR shaves ~38 seconds of pure CPU overhead off the workload, generating a 36% overall speedup, entirely restoring the expected linear scaling of shuffle buffers when reading across many shards!

[zhixiangli]

How about moving this func to module level for reusability?

[zhixiangli]

We need to throw TypeError for unexpected types, such as those outside of the expected dict, list, tuple, primitive, numpy scalar, and numpy array types.

[zhixiangli]

_fast_copy?

[zhixiangli]

Could you confirm are there any other deepcopy needed to be replaced with _fast_copy?