Add support for torch 2.0

[dakinggg]

What does this PR do?

This PR upgrades the torch pin to support torch 2.0. It includes related fixes, mostly resulting from the use of use_orig_params=True with FSDP, which is necessary to support compile.

Changes:

• Different way of loading FSDP state dicts for torch 2
• _LRScheduler -> LRScheduler
• using summon_full_params to get HF generate to work with FSDP
• fixing the way we compute optimizer metrics to allow for the fact that not all ranks have all params
• assorted minor fixes and test fixes

Manual tests:

• sharded multinode autoresume

• local autoresume

• full resume

• optimizer monitor (https://wandb.ai/mosaic-ml/daniel-debug/runs/34o29laf?workspace=user-danielking)
• base 1b run (https://wandb.ai/mosaic-ml/daniel-debug/runs/37e2g60h?workspace=user-danielking)
• hf generate for gpt neo (https://wandb.ai/mosaic-ml/daniel-debug/groups/torch2-evals/workspace?workspace=user-danielking)
• hf generate for mosaicgpt (https://wandb.ai/mosaic-ml/daniel-debug/runs/10lknf0v?workspace=user-danielking)
• added remote tests pass

What issue(s) does this change relate to?

Closes CO-2029
Closes #2147

Before submitting

• Have you read the contributor guidelines?
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[mvpatel2000]

Why is there a gap in the resumption tests?

Mostly LGTM / minor comments. will do one more pass after comments are resolved before approval because PR is massive

[dakinggg]

@mvpatel2000 the resumptions with a gap are because run 1 went for 10 batches, and then run 2 was run with autoresume true and an increased max duration, rather than deleting some checkpoints or something. Lmk if that makes sense.

[karan6181]

Some minor comments. Overall looks good. I liked the detailed comments that you have added at multiple places to get better understanding of a code.

Minor nit: More descriptive PR header name ?

[mvpatel2000]

LGTM. Let's get more approvals though before merging. Only outstanding is adding GPU daily tests

[karan6181]

LGTM. Thanks!

[nik-mosaic]

Should we add helpful error messages to the ONNX export method if a user runs into an error and is running PyTorch 2.0? We could suggest they try downgrading PyTorch versions if one of their model operators is not supported.

This is not a blocking suggestion --- we can merge without this.

[dakinggg]

I think the message you get from ONNX directly is about as clear as it gets...since we don't know which operator they're having trouble with and how that correspond to opset version and torch version. but open to another suggestion