Sample Range for Metadataset

[lvoegtle]

This feature could offer multiple use-cases:

• Limiting the training dataset to a small subset for debugging
• Post-splitting a dataset into val/train/test
• Changing sampling weights for a subset

This would be implemented like:

__module__: megatron.energon
__class__: Metadataset
splits:
  train:
    datasets:
      - sample_range: [0, 1000]
        path: ...
      - ...

Setting the range of the dataset to consider to the first 1k sample. Sample indexes are computed by concatenating the webdataset shards (from the split.yaml).

This feature only works for the innermost level, i.e. not for nested metadatasets (counting for that is not clear, because nested datasets would need to be concatenated, which is not the case for training).

Currently, for get_val_dataset, there already exists a limit option, which is different: That only limits the iterations (i.e. batches) to that number, which gives different samples than setting the sample range (range is based on sample index, the iterations may interleave different portions of the dataset).

[voegtlel]

There also was a request to allow splitting the dataset by samples for train/val/test (not only by whole shards).