gather on rank 0 to avoid oom

OpenMOSS · Aug 23, 2023 · ac6eed4 · ac6eed4
1 parent f434afa
commit ac6eed4
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/collie/controller/trainer.py b/collie/controller/trainer.py
@@ -556,10 +556,14 @@ def save_model(self, path: str, process_exclusion: bool = False,
                 self._checkpoint_prologue()
                 for name, param in self.engine.module.named_parameters():
                     with deepspeed.zero.GatheredParameters(param):
-                        state_dict[name] = param.detach().cpu()
+                        if env.dp_rank == 0:
+                            state_dict[name] = param.detach().cpu()
                 self._checkpoint_epilogue()
             else:
-                state_dict = self.engine.module.state_dict()
+                if env.dp_rank == 0:
+                    state_dict = self.engine.module.state_dict()
+                else:
+                    state_dict = {}
             self.engine.module.save_parallel_state_dict(
                 state_dict=state_dict,
                 path=path,