Make Trainer compatible with sharded checkpoints
# What does this PR do?

The `Trainer` is currently incompatible with the new sharded checkpoint feature in two places:
- resuming from a checkpoint
- loading the best model at the end of training

In both cases, the model state dict is loaded back inside the model **but** there is no model save file if the model was above the default size for sharding, resulting in errors (as was pointed out by #16976 ).

This PR addresses this by:
1. Creating a new function `load_sharded_checkpoint` that does the same thing as `model.load_state_dict` for regular model files, but loads a sharded checkpoint (and errors in case of missing/unexpected keys when `strict=True`).
2. Use that function inside the Trainer in the two places mentioned above.

A test is added to make sure resuming works from a sharded checkpoint.

Fixes #16976 
_The documentation is not available anymore as the PR was closed or merged._