Per Stream Batching

[snarayan21]

Description of changes:

Enables batching where each batch consists of samples from just one stream! Users can access this functionality by setting batching_method to per_stream.

How it works:

• make partition per stream
• shuffle each partition
• reshape to get global batches per stream, in order of traversal
• discard batches with any padding at all
• construct new overall batch ordering by selecting from stream batches
• pad new batches if necessary
• reshape/transpose back to get final partition

This is a useful feature requested mainly by advanced vision users. For example, each stream can be images of a different aspect ratio and resolution. Using per_stream batching means that pytorch now does not throw errors when trying to train with data of different dimensions now, as long as data of the same dimension is in its own stream.

Ran multi-GPU tests with toy datasets to make sure that samples in the same batch were only from a single stream. Ran regression tests over 200 batches using different text datasets to ensure that elastic determinism and resumption still work when using per_stream 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-124

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]

Overall looks good to me. Can you also update the description with what all integration testing that you have run to validate this PR ?