mcore distOpt restore fix (NVIDIA#9421)

Signed-off-by: Alexandros Koumparoulis <[email protected]>

nemo/collections/nlp/parts/nlp_overrides.py

@@ -444,6 +444,9 @@ def _check_param_groups_mismatch(self, checkpoint_path: Union[str, Path], sharde
             bool: True if the number of param groups does not match
         """
         common_state_dict = dist_checkpointing.load_common_state_dict(checkpoint_path)
+        # @akoumparouli: check if it contains an mcore dist opt
+        if common_state_dict.get('optimizer_states', [{}])[0].get('param_groups', None) is None:
+            return False
         model_param_groups = self._get_param_group(common_state_dict)
         checkpoint_param_groups = self._get_param_group(sharded_state_dict)
         return len(model_param_groups) != len(checkpoint_param_groups)