fix large checkpoints in pipe parallel (EleutherAI#33)

sweinbach · Samuel Weinbach · web-flow · commit 87fbb8fbea8f · 2021-07-29T18:40:21.000+02:00
Co-authored-by: Samuel Weinbach &lt;samuel.weinbach@gmail.com&gt;
diff --git a/deepspeed/runtime/pipe/module.py b/deepspeed/runtime/pipe/module.py
@@ -553,7 +553,10 @@ def save_state_dict(self, save_dir):
             model_ckpt_path = self.ckpt_layer_path(save_dir, idx)
             if not hasattr(layer, 'state_dict'):
                 continue
-            torch.save(layer.state_dict(), model_ckpt_path)
+            
+            # for some reason torch is saving lots of unnecessary information
+            # this line forces a copy of the tensor in order to get rid of temporary stuff
+            torch.save({k: (v.cpu() if torch.is_tensor(v) else v) for k, v in layer.state_dict().items()}, model_ckpt_path)
 
     def load_state_dir(self, load_dir, strict=True):
         rank = dist.get_rank()
