Relaxing divisibility constraints on num_canonical_nodes and num_physical_nodes

[snarayan21]

Description of changes:

Relaxes constraints on num_canonical_nodes being divisible by physical_nodes or vice versa, opening up many more numbers of nodes to deterministically train and resume on. This can be enabled by setting partition_algo='reserved' in StreamingDataset instantiation, or by setting partition_algo: reserved in yamls for each dataset. The only assumptions we make are that:

• num_canonical_nodes plays nicely with physical_nodes only in the initial run.
• global batch size stays constant during resumption.
• global batch size is divisible by the total number of devices during resumption.
  This reserved partition algo also preserves earlier functionality with determinism and should replace the orig algorithm as the default, since orig functionality is a subset of this partitioning algo (in an upcoming PR).

Testing:

• Added unit tests proving deterministic sample ordering with multiple numbers of nodes that were not possible with the orig partition algorithm
• Conducted local CPU tests with saving and loading dataset state dict to ensure that values were being set and read in correctly
• Conducted resumption tests on interactive instance with 4 GPUs to ensure that state dict was being set and loaded correctly
• Conducted multi-stream elastic deterministic resumption tests going from 2 to 3 nodes with num_canonical_nodes set to 2 and 64. This would not have been possible earlier, and results are shown below. Some numerical instability on newer clusters is likely responsible for slight deviations.
• Conducted tests going from 1->2 nodes, to make sure behavior was same as orig partitioning.
• Conducted tests going from 2 -> 3 -> 4 nodes, to make sure the initial physical nodes were being persisted in the state and partitions were still deterministic across multiple resumptions.

How it works:

• if partitioning for the first time, use the orig partition method. Save the initial number of physical nodes in the state.
• if resuming, partition over the initial number of physical nodes and the initial device batch size, then repartition this across the new number of physical nodes and the new device batch size. orig partition is still used if NCN and PN play nicely to preserve all earlier functionality and have nicer downloading for nice numbers of PN.

In the worst case, a node during the resumed run will have to do twice as much shard downloads as a node from the initial run. This is because of the way a global batch has to be split across the new number of nodes. When combined with better shuffling algorithms and batching methods, this should be an acceptable tradeoff.

Test with num_canonical_nodes=2.

• Light blue: full 200 step run on 2 nodes.
• Brown: first 100 steps on 2 nodes.
• Gray: last 100 steps on 3 nodes.

Test with num_canonical_nodes=64.

• Blue: full 200 step run on 2 nodes.
• Red: first 100 steps on 2 nodes.
• Yellow: last 100 steps on 3 nodes.

Test from 1->2 nodes, preserving orig partitioning behavior:

• Blue: full 200 step run on 1 node.
• Red: first 100 steps on 2 nodes.
• Green: last 100 steps on 2 nodes.

Test from 2->3->4 nodes, with NCN=2, to test multiple resumptions:

• Pink: full 200 step run on 2 nodes.
• Purple: first 100 steps on 2 nodes.
• Orange: middle 100 steps on 3 nodes.
• Light Blue: last 100 steps on 4 nodes.

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.

[karan6181]

LGTM. Thank You!