Stratified Batching

[snarayan21]

Description of changes:

This PR lets users control the composition of each batch. Previously, users could set proportions for streams, which would be fulfilled in aggregate over the whole training run. Now, users can have each batch also consistently be stratified between streams in specified proportions as well. This will help customers training multimodal models and can also make some fancier dataset mixing strategies possible. Users can access this functionality by setting batching_method to stratified.

How it works:

• make partition per stream, with 1 physical node and batch size 1
• shuffle each partition
• reshape to get global batch parts per stream, in order of traversal
• discard batches with any padding at all
• construct new overall batches by concatenating batch parts from each stream.

This is a useful feature that's good for multimodal applications and for some dataset mixing strategies. For example, we can combine text datasets with code datasets in definite proportions in every single batch, as long as the data is of the same dimension in each batch.

Ran multi-GPU tests with toy datasets to make sure that batches always had the same number of samples from each stream. Ran regression tests over 200 batches from different text datasets to ensure that elastic determinism and resumption still work when using stratified batching. These ran without issue and the results are here: https://wandb.ai/mosaic-ml/batching_methods/reports/new-batching-methods-testing--Vmlldzo1MzIxNjI1

Issue #, if available:

https://mosaicml.atlassian.net/browse/STR-130

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

• I have read the contributor guidelines
• This is a documentation change or typo fix. If so, skip the rest of this checklist.
• I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
• I have updated any necessary documentation, including README and API docs (if appropriate).

Tests

• I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
• I have added tests that prove my fix is effective or that my feature works (if appropriate).
• I ran the tests locally to make sure it pass. (check out testing)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.

[karan6181]

Can you also update the description with what all integration testing that you have run to validate this PR ?

[karan6181]

LGTM. Thank You!