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