fix: extract expert device mesh for group from unused prehook

torch-ext/megablocks/layers.py

@@ -680,6 +680,17 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +702,9 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         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)
-        
-        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+
+        device_mesh = get_device_mesh(self)
+        expert_parallel_group = device_mesh.get_group() if device_mesh else None
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once