py1e randomized

[snarayan21]

Description of changes:

Previously, the range of each shard in the py1e shuffle algorithm was not randomized. So when training on a large number of nodes, each node would need to download NCN/PN number of shards at the same time. By randomizing the range of each shard (within 0.75xSBS to 1.25xSBS), we ensure that nodes need to download shards at different times during training. If there is a cluster-level internet bandwidth limit, this PR will help alleviate potential bottlenecks.

In this first image, we see that using the py1br shuffle algorithm can cause throughput drops when node internet bandwidth is very low. However, downloads are relatively balanced across nodes due to randomized shuffle block size.

In this next image, we see that using the py1e shuffle algorithm eliminates throughput drops since downloads inside each node are much more balanced. However, overall downloads show a stepwise behavior since all nodes are downloading shards at the same time. This may cause issues with cluster-level internet bandwidth limits and cause unforeseen throughput drops due to this bottleneck.

As a result, bringing the randomized shuffle block size from py1br to py1e will ensure that the balanced intra-node downloading of the original py1e is combined with the balanced inter-node downloading of py1br.

Issue #, if available:

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.