Adds Ability to Sub-Sample Data for Data Constrained Scaling Law Experiments

[Helw150]

Allows mixture datasets to specify a target budget and a experiment budget. This then computes what percentage of the data to sample overall in order to enable data constrained experiments like the above figure.

[dlwh]

what do you think about making a custom asyncdataset that is basically just a slice and leave this logic out of the mixture dataset?

[Helw150]

what do you think about making a custom asyncdataset that is basically just a slice and leave this logic out of the mixture dataset?

With the idea of having it wrap the mixture dataset in the simulation case? Or each of the subdatasets of a mixture would be of this type?

I think the prior makes a lot of sense! I just need to wrap my head around the impl. that would lead to consistent slices of each sub-dataset as well.

The latter seems like it would require more configuration wiring that might not be worth it.

[Helw150]

Looking at it more, I don't know if there's actually a clean way to make it work for the prior.

The latter feels fine as well though and likely has less footguns with respect to the effective dataset size being different than all the different utilities so I'll get that in.

[dlwh]

i was thinking the latter yeah. basically just "if subsample: datasets = map(datasets, slice_dataset)" but I guess your point is we don't know how many samples we're going to go through until we're in the mixture dataset

[Helw150]

I've got a rework I'll push after NLP lunch that I feel like is probably cleaner.

[dlwh]

actually no i don't understand why isn't it just:

class SimulatedDataRatioDataset:
  async def wait_until_len_at_least(len):
          target_for_len = len * self.ratio_target
          actual = await self.ds.wait_until_len_at_least(int(target_for_len))

          return int(actual * self.ratio_target))

[Helw150]

Ok - how does this look?

The test is probably in the wrong spot though realistically?

[Helw150]

actually no i don't understand why isn't it just:

class SimulatedDataRatioDataset:
  async def wait_until_len_at_least(len):
          target_for_len = len * self.ratio_target
          actual = await self.ds.wait_until_len_at_least(int(target_for_len))

          return int(actual * self.ratio_target))

Just saw this! My one concern with this v.s. what I have in the current one is that it makes more assumptions about the implementation of the underlying get_batch code.

(edit: for example, you couldn't wrap a mixture in this class right?)

[dlwh]

i don't see why not?

[Helw150]

actually no i don't understand why isn't it just:

class SimulatedDataRatioDataset:
  async def wait_until_len_at_least(len):
          target_for_len = len * self.ratio_target
          actual = await self.ds.wait_until_len_at_least(int(target_for_len))

          return int(actual * self.ratio_target))

Maybe I'm missing something, but this assumes the underlying dataset uses get_batch logic which depends on it's own wait_until_len_at_least to determine the dataset is exhausted.

For a mixture dataset, this isn't quite true - it instead uses the wait_until_len_at_least of each of the sub-datasets. Since overriding the parents call doesn't affect these, it wouldn't change the epoching behavior right?

[Helw150]

Oop just saw the slice version I have on my branch isn't pushed here. I think that'll be less mixed concerns than the current PR

[Helw150]

Ok - pushed. I realized I hadn't pushed before since this currently has a pre-commit type failure which I'll fix assuming the core logic makes sense to you!

The core difference between this and your proposal is overriding get_batch rather than overriding the length waiting function.

    async def get_batch(self, indices: Sequence[int]) -> Sequence[T_co]:
        shifted_indices = [(index + self.start_index) for index in indices]
        max_index = max(shifted_indices)

        if self.end_index is not None and max_index > self.end_index:
            raise ValueError("Requested indices beyond the end of the dataset")

        return await self.dataset.get_batch(indices)

I prefer this because the effects of calling the slice (to me) are clearer for the above code without checking anything else, whereas I feel like understanding the effects of the wait_until_len_at_least override feels like it would require knowing how the get_batch of dataset you are wrapping is implemented for it to make sense.

LMK what you think!

[dlwh]

this is the thing i was trying to work around but I agree it's not worth fixing

[Helw150]

SG!

[dlwh]

can you fix the mypy errors then good to merge

[Helw150]

can you fix the mypy errors then good to merge

Fixed!