Additional logging from Model definition

[CYHSM]

What does this PR do?

Allows additonal logging of model metrics directly from the model definition through the training loop to tracking platforms (e.g. W&B). In this case, the model returns a dictionary of custom metrics during the forward pass which are accumulated, reduced across ranks and formatte for W&B. What I used in my branch is a bit more complex but this is general enough to work across any potential upcoming models.

General Changes

• Support for models to return custom metrics (example metrics also added to gpt2_model.py)
• New logging file (could also integrate into trainer.py) that flattens, unpacks and formats metrics
• Metrics are accumulated in trainer.py

Breaking Changes

• No breaking changes

Checklist before submitting final PR

• My PR is minimal and addresses one issue in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have checked that all tests run through (python tests/tests.py)
• I have updated the internal changelog (CHANGELOG_DEV.md)