From 4442abfa22f4279d66516f3433916780bef2f759 Mon Sep 17 00:00:00 2001 From: DEV TYAGI Date: Mon, 5 Oct 2026 23:15:49 +0530 Subject: [PATCH 1/2] fix ExpertTensorParallel Dynamo crash --- .../moe_training/reference_parallel_styles.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/test/prototype/moe_training/reference_parallel_styles.py b/test/prototype/moe_training/reference_parallel_styles.py index c3d540e952a..75227635e5f 100644 --- a/test/prototype/moe_training/reference_parallel_styles.py +++ b/test/prototype/moe_training/reference_parallel_styles.py @@ -165,12 +165,12 @@ def _token_dispatch(self, mod, inputs, device_mesh): routed_input, num_tokens_per_expert = inputs routed_input = DTensor.from_local( - routed_input, device_mesh["tp"], (Replicate(),) + routed_input, self.tp_mesh, (Replicate(),) ).to_local(grad_placements=(Partial(),)) inputs = (routed_input, num_tokens_per_expert) - return super()._token_dispatch(mod, inputs, device_mesh["ep"]) + return super()._token_dispatch(mod, inputs, self.ep_mesh) def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> None: from torch.distributed.tensor import distribute_tensor @@ -189,15 +189,9 @@ def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> N ) def _token_combine(self, mod, routed_output, device_mesh): - return super()._token_combine(mod, routed_output, device_mesh["ep"]) + return super()._token_combine(mod, routed_output, self.ep_mesh) - def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: - from torch.distributed.tensor import distribute_module - - return distribute_module( - module, - device_mesh, - partition_fn=self._partition_fn, - input_fn=self._token_dispatch, - output_fn=self._token_combine, - ) + def _apply(self, mod, device_mesh): + self.tp_mesh = device_mesh["tp"] + self.ep_mesh = device_mesh["ep"] + return super()._apply(mod, device_mesh) From 6686c7bc458f1607ae82bcc630f3060b9ac33f6d Mon Sep 17 00:00:00 2001 From: DEV TYAGI Date: Mon, 5 Oct 2026 23:23:54 +0530 Subject: [PATCH 2/2] fix: restore _apply type hints and keyword signature compatibility --- test/prototype/moe_training/reference_parallel_styles.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/prototype/moe_training/reference_parallel_styles.py b/test/prototype/moe_training/reference_parallel_styles.py index 75227635e5f..1064b88a2ce 100644 --- a/test/prototype/moe_training/reference_parallel_styles.py +++ b/test/prototype/moe_training/reference_parallel_styles.py @@ -191,7 +191,7 @@ def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> N def _token_combine(self, mod, routed_output, device_mesh): return super()._token_combine(mod, routed_output, self.ep_mesh) - def _apply(self, mod, device_mesh): + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: self.tp_mesh = device_mesh["tp"] self.ep_mesh = device_mesh["ep"] - return super()._apply(mod, device_mesh) + return super()._apply(module, device_mesh)