Add support for automatically registering models to UC at the end of training

[dakinggg]

Adds support for directly logging models to mlflow during training for easy deployment on databricks model serving. Tested and was able to deploy both an MPT model and a llama model succesfully.

• Confirm can log MPT and llama and deploy
• Measure time of final checkpoint for different sizes
• wait for next composer patch release with mlflow model logging changes
• write tests
• add support for using UC
• refactor to skip the extra load to memory.

along for the ride: fixing the hf checkpointer to work with non fsdp models as well

• redo manual tests once ready for review and unit tests pass (successfully logged and deployed mpt and llama models)

[mvpatel2000]

What scale of models can this take? What happens if I try shoving in a 30b?

Should this be here instead of the mlflow logger as a public fn that this callback calls?

[dakinggg]

@mvpatel2000 large models work. at least 30b, haven't had cluster space to run 70b yet. this does call the mlflow logger log_model from composer. the work in this callback is just preparing the inputs to that function call.

[eracah]

ok lgtm