Update nemo/core/optim/mcore_optim.py

Co-authored-by: mikolajblaz <[email protected]> Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · May 22, 2024 · 0a9cd71 · 0a9cd71
1 parent d7cf7f0
commit 0a9cd71
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/core/optim/mcore_optim.py b/nemo/core/optim/mcore_optim.py
@@ -59,7 +59,8 @@ def sharded_state_dict(
  self, model_sharded_state_dict, optimizer_state_dict=None, is_loading=False, dist_ckpt_parallel_save=False
  ):
  # TODO(@akoumparouli, @mikolajblaz): switch to sharding_type once support for fully_sharded_model_space merged in mcore.
- sharding_type = 'fully_sharded_model_space' if dist_ckpt_parallel_save else 'dp_zero_gather_scatter'
+ # sharding_type = 'fully_sharded_model_space' if dist_ckpt_parallel_save else 'dp_zero_gather_scatter'
+ sharding_type = 'dp_zero_gather_scatter'
  return self.mcore_optimizer.sharded_state_dict(
  model_sharded_state_dict, is_loading=is_loading, sharding_type=sharding_type
  )