add sharded loading for safetensors in AutoTP (#4854)

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>

deepspeed/module_inject/replace_module.py

@@ -566,7 +566,12 @@ def replace_module(model, orig_class, replace_fn, _replace_policy, checkpoint=No
     """
     sd = None
     if checkpoint is not None:
-        sd = torch.load(checkpoint, map_location='cpu')
+        if checkpoint.endswith(".safetensors"):
+            from safetensors.torch import load_file
+            sd = load_file(checkpoint)
+        else:
+            sd = torch.load(checkpoint, map_location='cpu')
+
     policy = {}
     if orig_class is not None:
         policy.update({orig_class: (replace_fn, _replace_policy)})

requirements/requirements-inf.txt

@@ -1,5 +1,6 @@
 google
 lm-eval==0.3.0
 protobuf
+safetensors
 transformers>=4.32.1
 transformers[sentencepiece]