Fix node calculation in replication for World object

[snarayan21]

Description of changes:

A customer was hitting an issue where the number of nodes was not being calculated correctly in the replication function of the World object. This is because we did not properly account and test for cases where the number of TP blocks was not cleanly divisible by the original number of ranks per node. For example, with 37 nodes, 8 GPUs per node, and replication of 8, the World object returned by the replication function would have 5 nodes and 7 ranks per node, which did not account for the last 2 TP blocks (which in this case, were the last 2 nodes).

To address this, we check if the number of TP blocks is divisible by the number of ranks per node, and if it's not, we set the number of nodes to 1. So, using the previous example, we now return a World object with 1 node and 37 ranks in the node, corresponding to the 37 TP blocks. Download optimality is still guaranteed with 1 node since canonical nodes will still be partitioned among these 37 ranks of the new World object, ensuring that no node will have to download more samples than intended.

I also expanded the test cases to include previously failing cases, since the test runs very fast anyways.

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.

[j316chuck]

Thanks for unblocking the customer quickly!