Stop ConcatDataset from re-sharding itself on every epoch

[VaggelisGian]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

ConcatDataset.__iter__ wrote its rank's torch.utils.data.Subset shards back into self.datasets, so every epoch re-sharded already-sharded data: under world_size > 1 each rank received fewer samples every successive epoch, and the temperature/random samplers computed their weights from shrinking lengths. This PR stops the mutation.

Collection: common

Changelog

• ConcatDataset.__iter__ now builds a local list of Subset shards per iteration and leaves self.datasets untouched. The local list feeds get_iterable, the index generator (temperature/random/round-robin), item lookup, and the mid-epoch exhaustion handler.
• Iterable sources are unaffected semantically (the local name aliases the same objects).
• Four new unit tests in tests/collections/common/data/test_dataset_concat.py.

Observed before the fix with two map-style datasets of 10 and 7 items on 2 ranks: epoch 1 gave rank 0 [a0..a4] and rank 1 [a5..a9] halves; iterating again sharded those halves a second time (5 -> 2 items for rank 0) and duplicated samples once a shard shrank below the rank slice. Sampler weights drifted between epochs because generator probabilities were derived from the mutated lengths.

Usage

No API or config change. Existing training setups get stable per-rank data across epochs:

from nemo.collections.common.data.dataset import ConcatDataset

concat = ConcatDataset(datasets=[ds_a, ds_b], sampling_technique='round-robin', shuffle=False,
                       global_rank=0, world_size=2)
epoch1 = list(iter(concat))
assert epoch1 == list(iter(concat))   # failed before, holds now

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

Trusted PRs run automatically through copy-pr-bot. For an untrusted PR, a maintainer can trigger CI by commenting
/ok to test <head-sha>; repeat this after a new push if the PR remains untrusted.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? Yes: four tests covering cross-epoch repeatability, expected rank shard contents, cross-rank disjointness over repeated iterations, and the world_size=1 path. Two of them fail on unfixed code.
• Did you add or update any necessary documentation? No documentation change needed: bug fix only, no public API, config key, CLI behavior, or user workflow changed.
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc) No optional dependency is touched.
  • Reviewer: Does the PR have correct import guards for all optional libraries? N/A.

PR Type:

• New Feature
• Bugfix
• Documentation

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to none (found while reading the codebase)
• Note for reviewers: CodeSwitchedDataset.prep_underlying_datasets in the same file has the same write-back pattern into self.datasets[lang] and is re-invoked from its __iter__; it is a separate class with separate sampling machinery and is intentionally left out here to keep this PR single-purpose. Happy to follow up separately.

Exact checks run locally (Python 3.12.6 via dedicated venv, torch available):

python -m pytest tests/collections/common/data/test_dataset_concat.py -q -p no:cacheprovider -m unit
  -> 4 passed (0.84s)
same suite with only the fix reverted (dataset.py at origin/main):
  -> 2 failed: test_iter_is_repeatable_across_epochs, test_ranks_get_disjoint_shards;
     2 passed (the two first-epoch-only tests)
python -m black --check --diff <changed files>      -> clean ("2 files would be left unchanged")
python -m isort --check-only <changed files>        -> clean
python -m flake8 <changed files>                    -> only pre-existing findings on lines
                                                       untouched by this diff (F821 line 356,
                                                       E741 lines 385/390)
git diff --check                                    -> clean

Intentionally skipped: full common+asr suites (blast radius is one method plus its new tests), docs build (uv unavailable on this machine), pre-commit binary absent locally; black/isort/flake8 were run directly against the repo's pinned configs instead.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.