Use state_dict Torchmetrics Serialization

[nik-mosaic]

What does this PR do?

Our metrics objects were previously serialized using pickle. The pickling includes many fields which are unnecessary and may change from version to version, causing a mismatch when we stop and start training with slightly different configurations or upgrade Torchmetrics.

This PR changes the saving and loading of Composer state objects so that state.train_metrics and state.eval_metrics use Torchmetrics' built in state_dict() method instead of Pickle. The state_dict() method essentially returns a dictionary of <metric_name, metric_value> <key, value> pairs, without any of the other data.

When loading a state dict from its serialized version, we recreate the metric as follows:.
(1) Get the state.model
(2) Call model.get_metrics(), which creates a default version of the metrics
(3) Loop over the dictionary of <metric_name, metric_value> pairs from the serialization, matching metric names to the metrics from model.get_metrics() and populating values with the saved metric_value. If a saved metric isn't present in the model, we skip it.

Other changes in this PR:

• Gate _ensure_metrics_device_and_dtype, which is only necessary for Deepspeed models, behind an is_model_deepspeed() check. This should fix CO-1910.

Testing:

• We have modified the test_checkpoint.py and test_sharded_checkpoint.py files to include a metric equality check wherever we have a weight equality or an optimizer equality check.
• Added test_load_remote_checkpoint, a backwards compatibility checkpoint test. I have uploaded a checkpoint saved with Composer 0.13.5 and default dependenices (torchmetrics 0.11.3, etc.). This test downloads the checkpoint and ensures equivalency with a currently trained version. As we continue to push to Composer, this test will become more useful because our remote checkpoint will remain frozen while our local one will include all future trainer/model/metrics changes.
• Locally, all tests pass on torchmetrics 0.11.4, but we do not bump the version in this PR.
• @coryMosaicML has verified that CO-1910 has been fixed with this PR.

What issue(s) does this change relate to?

CO-1918, CO-1907, CO-1853.

[mvpatel2000]

This adds a backwards compatibility test -- very nice! But is this change backwards compatible? The checkpoint appears to be off dev and not 0.13.5

Also, please show screenshot of test manually passing as it is only run on daily (since its remote)

[dakinggg]

Thanks for digging through this @nik-mosaic! Left a few comments and questions, mostly around the usage of private attributes.

[nik-mosaic]

Here it is passing pytest tests/trainer/test_checkpoint.py -m remote. The remote checkpoint file has now changed, it is now one created off the Composer dev branch. To get this test to pass, we add a case to the state.py serialization/deserialization section to support deserializing a Torchmetrics object. Old composer checkpoints will have Torchmetrics objects, new ones will have dictionaries with metric tensors.

[mvpatel2000]

Why do we not need the backwards compatibiltiy hacks with this? because wont they still be the old format?

[nik-mosaic]

Why do we not need the backwards compatibility hacks with this? because wont they still be the old format?

I have added a comment to address this directly in the code: The explanation is as follows.

Given the rest of a Composer checkpoint, the state_dict() and _computed attributes of a Torchmetrics object are enough information to recreate it upon serialization. We only serialize the minimum metric information to maximize backwards compatibility --- old checkpoints will continue to be compatible even if other Torchmetrics attributes have changed.

[mvpatel2000]

LGTM, thank you for taking this on!