Add support for auto packing ratio

[irenedea]

Adds 'auto' packing ratio for finetuning.

If 'auto' is specified, packing will be profiled for the dataloader configuration across a maximum of 20 packing ratios until the highest packing ratio with 0 waste is found.

If there are multiple ranks, we take the minimum 'auto' packing ratio across all ranks.

Testing

• Manual tests
  • 1 epoch on mosaicml/dolly_hhrlhf, auto packing had 1.3% waste overall
    • finetune-auto-pack-test-CHlwSQ https://wandb.ai/mosaic-ml/irene-test/runs/7ccl7k7a
      • time/token: 9390338, time/sample: 4896, time/batch: 51, time/epoch: 1
    • finetune-auto-pack-test-baseline2-zptNpN https://wandb.ai/mosaic-ml/irene-test/runs/b0xqxmlj
      • time/token: 9518016, time/sample: 59288, time/batch: 618, time/epoch: 1
  • No packing run are the same before and after PR
    - before: finetune-auto-pack-test-baseline2-zptNpN , after: finetune-ratio-no-pack-branch-test-8N5I1y
  • Packing ratio of 3.0 is same before and after PR
    - before: finetune-ratio-3-pack-test-kW8gp6 , after: finetune-ratio-3-pack-with-auto-test-YVP00E
• Unit tests
  • Tests packing on small tensors
  • Tests packing on small tensors with leftovers and waste calculation
  • Tests auto packing ratio selection for single and multi ranks
  • Tests auto packing with dataloader
  • Adds 'auto' packing_ratio parameterization to existing dataloader tests

[alextrott16]

LGTM. Thanks for adding this!!
Will hold off on approval so this can get another set of eyes!

Also left a couple comments (nothing major) along with a suggestion for an alternative way to search for the optimal packing ratio. I think re-using my brute force approach might not be the best way to go. If you agree, that will require more of an overhaul of this code, unfortunately. But potentially better to do that than simply inherit my hasty decision :)

[dakinggg]

Looking good!

Two more manual tests I would like to see before merging:
(1) a short training run without packing, before and after this PR. These should be identical.
(2) a short training run with a set packing ratio, before and after this PR. These should be identical.

[dakinggg]

LGTM except for one determinism related comment

[dakinggg]

LGTM, with one more question relating to the randomness