Upgrade to latest Megatron-lm version
There were some small breaking changes when using the latest Megatron-lm version.

Fixing those here.

# TODO
~~We need to cut a new pip package from Megatron and update NeMo requirements.~~
