Skip to content

GitLab

Commit 039b4550 authored by Rick Ho's avatar Rick Ho
Browse files

add an option of init sync

parent 6fe415e9
Show 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,6 +76,7 @@ class DistributedGroupedDataParallel(nn.Module): ...@@ -75,6 +76,7 @@ class DistributedGroupedDataParallel(nn.Module):
g.copy_(s) g.copy_(s)
self.allreduce_params = allreduce_params self.allreduce_params = allreduce_params
if need_sync:
self._sync_params() self._sync_params()
def _sync_params(self): def _sync_params(self):
......
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