Skip to content

GitLab

"...git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "a6314a8d4e1c301bce4e45c10f325f594220617f"
Commit 039b4550 authored by Rick Ho's avatar Rick Ho
Browse files

add an option of init sync

parent 6fe415e9
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
fmoe/distributed.py
View file @ 039b4550
...@@ -26,6 +26,7 @@ class DistributedGroupedDataParallel(nn.Module): ...@@ -26,6 +26,7 @@ class DistributedGroupedDataParallel(nn.Module):
self, self,
module, module,
auto_allreduce=False, auto_allreduce=False,
need_sync=True,
**kwargs **kwargs
): ):
assert not auto_allreduce, "Automatic all-reduce is not implemented yet" assert not auto_allreduce, "Automatic all-reduce is not implemented yet"
...@@ -75,7 +76,8 @@ class DistributedGroupedDataParallel(nn.Module): ...@@ -75,7 +76,8 @@ class DistributedGroupedDataParallel(nn.Module):
g.copy_(s) g.copy_(s)
self.allreduce_params = allreduce_params self.allreduce_params = allreduce_params
self._sync_params() if need_sync:
self._sync_params()
def _sync_params(self): def _sync_params(self):
groups = dict() groups = dict()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment