Dist-optim fixes

[akoumpa]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Changelog

• Please update the CHANGELOG.md under next version with high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? Make sure to also update the NeMo Framework User Guide which contains the tutorials

Checklist when contributing a new algorithm

• Does the trainer resume and restore model state all states?
• Does the trainer support all parallelism techniques(PP, TP, DP)?
• Does the trainer support max_steps=-1 and validation?
• Does the trainer only call APIs defined in alignable_interface.py?
• Does the trainer have proper logging?

Additional Information

• Related to # (issue)

[terrykong]

@akoumpa Could you PR against #380 ? I caught a few of the changes missing for mcore dist opt over there

[terrykong]

Closing so you can re-create a PR targetting that branch

[akoumpa]

synced with @terrykong offline; will pull 392 in 380.