[BUG] datasets < 4.0.0 makes `IterableDataset.repeat(None)` unusable with `split_dataset_by_node`

[dhruvdcoder]

Describe the bug

We pin datasets < 4.0.0 (see wiki/compatibility_issues.md, §1) because huggingface/datasets >= 4.0.0 drops support for the script-based billion-word-benchmark/lm1b dataset that we use for training. A side-effect of that pin is that IterableDataset.repeat(None) cannot be combined with split_dataset_by_node — the two raise a TypeError on the very first iteration in any DDP setup. This is why xlm.datamodule ships its own Python-level _CycleDataset workaround (enabled via DatasetManager(make_infinite=True)) instead of using the upstream repeat() API.

Root cause

Every _BaseExamplesIterable subclass in datasets is expected to implement:

def shard_data_sources(self, num_shards: int, index: int, contiguous: bool = True) -> "_BaseExamplesIterable": ...

@property
def num_shards(self) -> int: ...

But in datasets < 4.0.0 (see src/datasets/iterable_dataset.py, class RepeatExamplesIterable), RepeatExamplesIterable — the wrapper that IterableDataset.repeat() returns — instead exposes:

def shard_data_sources(self, worker_id: int, num_workers: int) -> "RepeatExamplesIterable":
    return RepeatExamplesIterable(
        self.ex_iterable.shard_data_sources(worker_id, num_workers),
        num_times=self.num_times,
    )

@property
def n_shards(self) -> int:                    # legacy name, not the one HF actually calls
    return self.ex_iterable.n_shards

Two things are wrong:

1. The shard_data_sources parameter names (worker_id, num_workers) and arity (no contiguous) do not match the protocol the rest of the library uses.
2. The shards property is named n_shards instead of num_shards, so callers that expect num_shards fall through to the (no-op / NotImplementedError) default on the base class.

Why this breaks our pipeline

IterableDataset._prepare_ex_iterable_for_iteration is invoked on every iteration over a distributed iterable dataset and unconditionally calls (in datasets/iterable_dataset.py):

if ex_iterable.num_shards % world_size == 0:
    ex_iterable = ex_iterable.shard_data_sources(
        num_shards=world_size, index=rank, contiguous=False,
    )

If the outer _ex_iterable is a RepeatExamplesIterable (i.e. repeat() was called after the dataset was made iterable / distributed), Python raises:

TypeError: shard_data_sources() got an unexpected keyword argument 'num_shards'

Reversing the order — calling repeat(None) before split_dataset_by_node — does not help either: _split_by_node_iterable_dataset keeps the same outer _ex_iterable, so the same call site still explodes on iteration.

To Reproduce

Minimal repro (DDP, 2 ranks, requires datasets < 4.0.0):

from datasets import Dataset
from datasets.distributed import split_dataset_by_node

def gen():
    for i in range(8):
        yield {"x": i}

ds = Dataset.from_generator(gen).to_iterable_dataset(num_shards=4)
ds = ds.repeat(None)                       # wraps in RepeatExamplesIterable
ds = split_dataset_by_node(ds, rank=0, world_size=2)

for _ in ds:                               # raises on first iteration
    break
# TypeError: shard_data_sources() got an unexpected keyword argument 'num_shards'

Expected behavior

repeat(None) composed with split_dataset_by_node should produce an infinite, per-rank iterable that survives sharding without raising, exactly as it does in datasets >= 4.0.0.

The fix in 4.0.0 (already upstream)

Upstream PR huggingface/datasets#7581 ("Add missing property on RepeatExamplesIterable", merged Jun 5, 2025) renames n_shards to num_shards and rewrites shard_data_sources to match the standard protocol:

def shard_data_sources(self, num_shards: int, index: int, contiguous: bool = True) -> "RepeatExamplesIterable":
    return RepeatExamplesIterable(
        self.ex_iterable.shard_data_sources(num_shards, index, contiguous=contiguous),
        num_times=self.num_times,
    )

@property
def num_shards(self) -> int:
    return self.ex_iterable.num_shards

That patch first shipped in datasets >= 4.0.0.

Current workaround in xlm-core

Because we cannot rely on repeat(), the codebase implements its own Python-level _CycleDataset wrapper (see src/xlm/datamodule.py, class _CycleDataset) that gives the same "infinite iterable that survives sharding" semantics without going through RepeatExamplesIterable. It is enabled via DatasetManager(make_infinite=True).

The trade-off is that _CycleDataset.state_dict() deliberately returns {}, so StatefulDataLoader cannot mid-cycle resume — epoch boundaries are effectively encoded in max_steps rather than in dataset state. RepeatExamplesIterable already wires up _init_state_dict correctly (it tracks repeat_index and re-initialises the inner state dict at every cycle), so once we relax the pin we get proper mid-cycle resume support for free.

Resolution path

Unpinning datasets is blocked on the lm1b script-based-dataset issue (wiki/compatibility_issues.md, §1). Once we either (a) validate dvruette/lm1b as a non-script replacement, or (b) drop lm1b as a dependency for training, we can:

• bump the pin to datasets >= 4.0.0,
• delete _CycleDataset,
• replace make_infinite=True with IterableDataset.repeat(None) in DatasetManager.setup,
• and gain mid-cycle resume support from StatefulDataLoader for free.

Environment

• OS: Linux (Ubuntu 22.04 / kernel 6.8)
• PyTorch: 2.x (any version compatible with datasets < 4.0.0)
• Lightning: as pinned by requirements.txt
• Transformers: as pinned by requirements.txt
• Hydra: as pinned by requirements.txt
• datasets: < 4.0.0 (the pin is the bug)

Additional context

• Wiki entry with full explanation: wiki/compatibility_issues.md §1.
• Workaround class: src/xlm/datamodule.py → class _CycleDataset (gated by DatasetManager(make_infinite=True)).
• Upstream fix: huggingface/datasets#7581, shipped in datasets >= 4.0.0.