mcore distOpt restore fix

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

nemo/collections/nlp/parts/nlp_overrides.py

@@ -444,6 +444,10 @@ 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: if it contains an mcore dist opt, param_groups is under ['optimizer'].data['param_groups']
+        # with ['optimizer'] being a ShardedObject
+        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)