enable arg enabled normalization of routing weights (#58)

* enable arg enabled nomalization of routing weights * pr comment
databricks · Dec 13, 2023 · a6ac3ce · a6ac3ce
1 parent 5897cd6
commit a6ac3ce
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/megablocks/layers/arguments.py b/megablocks/layers/arguments.py
@@ -24,6 +24,7 @@ class Arguments:
     moe_num_experts : int = 1
     moe_top_k : int = 1
     moe_capacity_factor : int = 1
+    moe_normalize_expert_weights: bool = False
     moe_loss_weight : float = 0.1
     moe_jitter_eps : Optional[float] = None
     moe_lbl_in_fp32 : bool = False

diff --git a/megablocks/layers/router.py b/megablocks/layers/router.py
@@ -55,6 +55,8 @@ def forward(self, x):
 
         scores = self.layer(x.view(-1, x.shape[-1])).softmax(dim=-1)
         expert_weights, expert_indices = self._top_k(scores)
+        if self.args.moe_normalize_expert_weights:
+            expert_weights /= expert_weights.sum(dim=-1, keepdim=True)
 
         expert_indices = (
             _uniform_expert_assignment(expert_indices, self.args.moe_num_experts)