fix loading with torch dist ckpt

Signed-off-by: jiemingz <jiemingz@nvidia.com>
JimmyZhang12 · Mar 23, 2024 · 7fbf9e2 · 7fbf9e2
1 parent a44253d
commit 7fbf9e2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py
@@ -442,9 +442,13 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path]) -> Dict[str, Any]:
  # after dist_checkpointing.load, sharded tensors will be replaced with tensors
  checkpoint['state_dict'] = sharded_state_dict
  checkpoint['optimizer_states'] = [self.optimizer_sharded_state_dict()]
- strategy = tensorstore.TensorStoreLoadShardedStrategy(load_directly_on_device=True)
+
+ if self.torch_dist_ckpt:
+ sharded_strategy = ('torch_dist', 1)
+ else:
+ sharded_strategy = tensorstore.TensorStoreLoadShardedStrategy(load_directly_on_device=True)
  checkpoint = dist_checkpointing.load(
- sharded_state_dict=checkpoint, checkpoint_dir=checkpoint_path, sharded_strategy=strategy
+ sharded_state_dict=checkpoint, checkpoint_dir=checkpoint_path, sharded_strategy=sharded_strategy
  )
 
  return checkpoint