fix: adjust layer params in source

torch-ext/megablocks/layers.py

@@ -683,26 +683,23 @@ def moe_forward(
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        moe_top_k = getattr(self, "moe_top_k", 4)
-        moe_num_experts = getattr(self, "moe_num_experts", 128)
-        gradient_scale = getattr(self, "gradient_scale", None)
-        alpha = getattr(self, "alpha", 1.702)
-        moe_capacity_factor = getattr(self, "moe_capacity_factor", 1.0)
-        moe_jitter_eps = getattr(self, "moe_jitter_eps", None)
-        moe_normalize_expert_weights = getattr(
-            self, "moe_normalize_expert_weights", None
-        )
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-
+        
         has_parallel = hasattr(self, "expert_parallel_group")
-        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        expert_parallel_group = torch.distributed.group.WORLD
         forward_fn = parallel_forward_once if has_parallel else forward_once
-        sort_end_bit = max(
-            int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1
-        )
-        mlp_impl = getattr(self, "mlp_impl", "grouped")  # or sparse
-
-        output, expert_weights_out, _ = moe_forward(
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward(
             x=x,
             router_weight=self.router.weight,
             moe_top_k=moe_top_k,
@@ -725,4 +722,4 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
-        return output, expert_weights_out
+        return output, expert_weights_out