fix rank of megatron dp

c0cfcc47 · Rick Ho · 593feab2 · c0cfcc47
Commit c0cfcc47 authored Feb 26, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

fmoe/megatron.py fmoe/megatron.py +4 -1

No files found.
--- a/fmoe/megatron.py
+++ b/fmoe/megatron.py
@@ -84,7 +84,10 @@ class MegatronMLP(FMoETransformerMLP):
                world_size=world_size, mp_group=group,
                expert_dp_comm='none' if args.distributed_experts else 'dp')
        self.hidden_size = args.hidden_size
-        self.rank = args.rank
+        if args.distributed_experts:
+            self.rank = args.rank
+        else:
+            self.rank = 0
        self.sigma = args.init_method_std
        self.num_layers = args.num_layers
        self.reset_parameters()