Fix for concat map dataset

[1-800-BAD-CODE]

What does this PR do ?

Fix for concat map dataset #5117

Only caveat I know of is that it will not reshuffle each epoch when using round robin. I'm not sure if anyone actually uses round robin in practice, though.

Usage

from collections import defaultdict

import torch
from pytorch_lightning import seed_everything

from nemo.collections.common.data import ConcatMapDataset

seed_everything(42)


class IntegerDataset(torch.utils.data.Dataset):
    def __init__(self, n: int, dataset_id: int):
        self._ints = list(range(n))
        self._id = torch.tensor(dataset_id)

    def __len__(self):
        return len(self._ints)

    def __getitem__(self, idx):
        worker_info = torch.utils.data.get_worker_info()
        worker_id = worker_info.id
        return self._id, worker_id, self._ints[idx]


for sampling_technique in ["temperature", "random", "round-robin"]:
    probs = [0.1, 0.2, 0.3]
    concat_dataset = ConcatMapDataset(
        datasets=[
            IntegerDataset(5, 0),
            IntegerDataset(10, 1),
            IntegerDataset(20, 2),
        ],
        seed=123456,
        sampling_temperature=2,
        sampling_probabilities=probs,
        sampling_technique=sampling_technique
    )

    dataloader = torch.utils.data.DataLoader(
        dataset=concat_dataset,
        batch_size=1,
        num_workers=2,
        shuffle=False  # Should be False to preserve round robin
    )

    print(f"Item-by-item results for technique '{sampling_technique}':")
    collected = defaultdict(list)
    for batch in dataloader:
        batch_dataset_id, batch_worker_id, batch_value = batch
        print(
            f"Dataset {batch_dataset_id.item()} worker {batch_worker_id.item()} fetched {batch_value.item()}"
        )
        collected[batch_dataset_id.item()].append(batch_value.item())

    print("*" * 80)
    print(f"Summary of returned results for each dataset for technique '{sampling_technique}'")
    if sampling_technique == "random":
        print(f"\t{probs=}")
    for dataset_id, values in sorted(collected.items(), key=lambda x: x[0]):
        dataset_raw_length = len(concat_dataset.datasets[dataset_id])
        print(f"\tDataset #{dataset_id} (raw length {dataset_raw_length}): {values}")
    print("*" * 80)
    print()

Output:

Item-by-item results for technique 'temperature':
Global seed set to 42
Dataset 2 worker 0 fetched 0
Dataset 2 worker 1 fetched 12
Dataset 1 worker 0 fetched 3
Dataset 1 worker 1 fetched 1
Dataset 0 worker 0 fetched 0
Dataset 2 worker 1 fetched 18
Dataset 1 worker 0 fetched 5
Dataset 0 worker 1 fetched 4
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 1
Dataset 0 worker 0 fetched 2
Dataset 0 worker 1 fetched 1
Dataset 2 worker 0 fetched 3
Dataset 1 worker 1 fetched 6
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 2
Dataset 2 worker 0 fetched 5
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 9
Dataset 0 worker 1 fetched 0
Dataset 2 worker 0 fetched 9
Dataset 2 worker 1 fetched 8
Dataset 2 worker 0 fetched 13
Dataset 2 worker 1 fetched 7
Dataset 1 worker 0 fetched 4
Dataset 2 worker 1 fetched 19
Dataset 2 worker 0 fetched 4
Dataset 2 worker 1 fetched 16
Dataset 2 worker 0 fetched 10
Dataset 2 worker 1 fetched 17
Dataset 0 worker 0 fetched 2
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 7
Dataset 1 worker 1 fetched 0
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 8
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 14
Dataset 2 worker 1 fetched 2
Dataset 1 worker 0 fetched 0
Dataset 2 worker 1 fetched 6
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 15
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'temperature'
	Dataset #0; raw length 5; num drawn = 13; values = [0, 4, 3, 2, 1, 1, 3, 0, 2, 4, 1, 2, 3]
	Dataset #1; raw length 10; num drawn = 13; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 5, 0, 7]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

Item-by-item results for technique 'random':
Dataset 2 worker 0 fetched 0
Dataset 2 worker 1 fetched 12
Dataset 2 worker 0 fetched 18
Dataset 1 worker 1 fetched 3
Dataset 0 worker 0 fetched 0
Dataset 2 worker 1 fetched 1
Dataset 1 worker 0 fetched 1
Dataset 0 worker 1 fetched 4
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 3
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 1 worker 0 fetched 6
Dataset 2 worker 1 fetched 5
Dataset 1 worker 0 fetched 2
Dataset 2 worker 1 fetched 9
Dataset 2 worker 0 fetched 8
Dataset 2 worker 1 fetched 13
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 9
Dataset 1 worker 0 fetched 4
Dataset 0 worker 1 fetched 0
Dataset 2 worker 0 fetched 7
Dataset 2 worker 1 fetched 19
Dataset 2 worker 0 fetched 4
Dataset 2 worker 1 fetched 16
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 10
Dataset 2 worker 0 fetched 17
Dataset 2 worker 1 fetched 14
Dataset 2 worker 0 fetched 2
Dataset 2 worker 1 fetched 6
Dataset 0 worker 0 fetched 3
Dataset 0 worker 1 fetched 2
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 0
Dataset 1 worker 0 fetched 8
Dataset 1 worker 1 fetched 8
Dataset 0 worker 0 fetched 4
Dataset 1 worker 1 fetched 7
Dataset 2 worker 0 fetched 15
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'random'
	probs=[0.1, 0.2, 0.3]
	Dataset #0; raw length 5; num drawn = 10; values = [0, 4, 3, 2, 1, 0, 3, 2, 1, 4]
	Dataset #1; raw length 10; num drawn = 12; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 8, 7]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

Item-by-item results for technique 'round-robin':
Dataset 0 worker 0 fetched 0
Dataset 1 worker 1 fetched 3
Dataset 2 worker 0 fetched 0
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 1
Dataset 2 worker 1 fetched 12
Dataset 0 worker 0 fetched 3
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 18
Dataset 0 worker 1 fetched 2
Dataset 1 worker 0 fetched 6
Dataset 2 worker 1 fetched 1
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 2
Dataset 2 worker 0 fetched 3
Dataset 0 worker 1 fetched 0
Dataset 1 worker 0 fetched 9
Dataset 2 worker 1 fetched 5
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 4
Dataset 2 worker 0 fetched 9
Dataset 0 worker 1 fetched 1
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 8
Dataset 0 worker 0 fetched 3
Dataset 1 worker 1 fetched 0
Dataset 2 worker 0 fetched 13
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 8
Dataset 2 worker 1 fetched 7
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 6
Dataset 2 worker 0 fetched 19
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 3
Dataset 2 worker 1 fetched 4
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 16
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 9
Dataset 2 worker 1 fetched 10
Dataset 0 worker 0 fetched 0
Dataset 1 worker 1 fetched 8
Dataset 2 worker 0 fetched 17
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 14
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 1
Dataset 2 worker 0 fetched 2
Dataset 0 worker 1 fetched 0
Dataset 1 worker 0 fetched 2
Dataset 2 worker 1 fetched 6
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 4
Dataset 2 worker 0 fetched 15
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 0
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'round-robin'
	Dataset #0; raw length 5; num drawn = 20; values = [0, 4, 3, 2, 1, 0, 2, 1, 3, 4, 1, 3, 2, 4, 0, 3, 2, 0, 1, 4]
	Dataset #1; raw length 10; num drawn = 20; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 6, 3, 5, 9, 8, 7, 1, 2, 4, 0]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[okuchaiev]

@1-800-BAD-CODE is this still WIP or ready for review?

[1-800-BAD-CODE]

@okuchaiev ready for review; I previously changed the status but didn't update the title.

I'll add some of my own critical comments:

• I believe that round-robin should shuffle every epoch, but this PR does not.
• For large datasets, the non-round-robin techniques don't re-sample from the dataset each epoch, and some samples from each dataset are used more than the others (the ones chosen during the initial oversampling).

These are slightly harder to solve for a map-style dataset, since when the data loader says "give me item 113" we should always return the same item "113", whereas with the iterable-style datasets we simply need to return an arbitrary unique item.

[MaximumEntropy]

Thanks for the fixes!