Disable mcore saving

Signed-off-by: Igor Gitman <igitman@nvidia.com>
Kipok · Oct 17, 2023 · 7d96a34 · 7d96a34
1 parent b6a648c
commit 7d96a34
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -1315,7 +1315,7 @@ def on_save_checkpoint(self, checkpoint) -> None:
         """
 
         # mcore uses distributed checkpointing
-        if self.mcore_gpt:
+        if False:
             checkpoint['sharded_state_dict'] = self.sharded_state_dict()
 
         # legacy checkpointing for interleaved
@@ -1332,7 +1332,7 @@ def on_load_checkpoint(self, checkpoint) -> None:
         """
 
         # mcore uses distributed checkpointing
-        if self.mcore_gpt:
+        if False:
             if 'state_dict' in checkpoint and checkpoint['state_dict']:
                 for index, module in enumerate(self.get_gpt_module_list()):
                     if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
@@ -1366,7 +1366,7 @@ def sharded_state_dict(self, prefix: str = '') -> Dict[str, Any]:
         The sharded tensor mapping is defined in the GPTModel class from mcore.
         """
 
-        if self.mcore_gpt:
+        if False:
             module_prefix = f'{prefix}model.'
             sharded_state_dict = {}
             for index, module in enumerate(self.get_gpt_module_list()):