[moe][compile] turn on capture_scalar_outputs and assert fullgraph on grouped experts

torchtitan/experiments/llama4/infra/parallelize.py

@@ -510,16 +510,17 @@ def apply_compile(model: nn.Module, compile_config: CompileConfig):
     """
     # NOTE: This flag is needed for torch.compile to avoid graph breaking on dynamic shapes in token-choice MoE
     # but it is experimental.
-    # torch._dynamo.config.capture_scalar_outputs = True
+    torch._dynamo.config.capture_scalar_outputs = True
     for layer_id, transformer_block in model.layers.named_children():
         # TODO: remove when torch.compile supports fullgraph=True for MoE
-        fullgraph = True
         if transformer_block.moe_enabled:
-            fullgraph = False
+            transformer_block.moe.experts._forward = torch.compile(transformer_block.moe.experts._forward, backend=compile_config.backend, fullgraph=True)
+            transformer_block.moe.experts.forward = torch.compiler.disable(transformer_block.moe.experts.forward)
+
         transformer_block = torch.compile(
             transformer_block,
             backend=compile_config.backend,
-            fullgraph=fullgraph,
+            fullgraph=not transformer_block.moe_enabled,
         )
         model.layers.register_module(layer_id, transformer_block)
 

torchtitan/models/moe.py

@@ -143,6 +143,13 @@ def forward(
         self,
         x: torch.Tensor,
         num_tokens_per_expert: torch.Tensor,
+    ) -> torch.Tensor:
+        return self._forward(x, num_tokens_per_expert)
+
+    def _forward(
+        self,
+        x: torch.Tensor,
+        num_tokens_per_expert: torch.Tensor,
     ) -> torch.Tensor:
         if isinstance(self.w1, DTensor):
             # Convert parameters from DTensors to plain Tensors, to work with