Add an opt-in AutoEP non-MoE regional compile mode

deepspeed/runtime/compiler.py

@@ -3,11 +3,17 @@
 
 # DeepSpeed Team
 
+from __future__ import annotations
+
 import torch
 import contextlib
 import functools
+from collections import OrderedDict
+
+import torch.nn as nn
 from deepspeed.utils.torch import required_torch_version
 from deepspeed.accelerator import get_accelerator
+from deepspeed.utils import logger
 
 try:
     from torch.compiler import is_compiling as torch_is_compiling
@@ -111,3 +117,63 @@ def compile():
         return torch.compile
     else:
         return dummy_decorator
+
+
+def compile_autoep_non_moe_regions(model: nn.Module, backend, compile_kwargs: dict) -> list[str]:
+    """Compile decoder regions around AutoEP layers while keeping AutoEP eager."""
+    from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer
+
+    named_modules = dict(model.named_modules())
+    autoep_modules = [(name, module) for name, module in named_modules.items() if isinstance(module, AutoEPMoELayer)]
+    if not autoep_modules:
+        raise ValueError("compile_mode='autoep_non_moe' requires at least one AutoEPMoELayer. "
+                         "Enable expert_parallel and call compile() after deepspeed.initialize().")
+
+    if "fullgraph" in compile_kwargs and compile_kwargs["fullgraph"] is not False:
+        raise ValueError("compile_mode='autoep_non_moe' requires fullgraph=False because AutoEP is an eager graph "
+                         "break.")
+    if "dynamic" in compile_kwargs and compile_kwargs["dynamic"] is not False:
+        raise ValueError("compile_mode='autoep_non_moe' currently requires dynamic=False.")
+
+    resolved_compile_kwargs = {
+        "fullgraph": False,
+        "dynamic": False,
+        **compile_kwargs,
+        "backend": backend,
+    }
+    regions: OrderedDict[str, nn.Module] = OrderedDict()
+    for module_name, _ in autoep_modules:
+        parent_name, separator, _ = module_name.rpartition(".")
+        if not separator:
+            raise ValueError("compile_mode='autoep_non_moe' cannot compile an AutoEPMoELayer at the model root.")
+        parent = named_modules[parent_name]
+        if type(parent).forward is nn.Module.forward:
+            raise ValueError(f"AutoEP compile region '{parent_name}' has no forward implementation. "
+                             "The MoE layer must be a direct child of a callable decoder block.")
+        if getattr(parent, "_compiled_call_impl", None) is not None:
+            raise ValueError(f"AutoEP compile region '{parent_name}' is already compiled.")
+        regions.setdefault(parent_name, parent)
+
+    original_forwards = {}
+    original_compiled_calls = {}
+    try:
+        for module_name, module in autoep_modules:
+            original_forwards[module] = module.__dict__.get("forward")
+            module.forward = disable(module.forward)
+            logger.debug("AutoEP regional compile: disabled compiler tracing for '%s'.", module_name)
+
+        for region_name, region in regions.items():
+            original_compiled_calls[region] = getattr(region, "_compiled_call_impl", None)
+            region.compile(**resolved_compile_kwargs)
+            logger.info("AutoEP regional compile: compiled '%s' with backend=%s.", region_name, backend)
+    except BaseException:
+        for module, original_forward in original_forwards.items():
+            if original_forward is None:
+                del module.forward
+            else:
+                module.forward = original_forward
+        for region, compiled_call in original_compiled_calls.items():
+            region._compiled_call_impl = compiled_call
+        raise
+
+    return list(regions)

deepspeed/runtime/engine.py

@@ -557,6 +557,8 @@ def __init__(self,
         self.unflatten = _unflatten_dense_tensors
 
         self._is_compiled = False
+        self._compile_mode = None
+        self._compiled_regions = []
         if is_deepcompile_supported():
             # Predefined compile passes
             self.register_compile_pass(zero_1_and_2_compile.NAME_Z1, zero_1_and_2_compile.add_z1_reduce,
@@ -5802,23 +5804,61 @@ def compile(self,
                 backend=get_accelerator().get_compile_backend(),
                 compile_kwargs={},
                 schedule=None,
-                compiled_autograd_enabled=False) -> None:
+                compiled_autograd_enabled=False,
+                compile_mode="model") -> None:
         """Compile the module using the specified backend and kwargs.
-        If a compiler_fn is set, it will be used instead of torch.compile().
+
+        ``compile_mode="model"`` compiles the full module. ``compile_mode="autoep_non_moe"``
+        compiles each decoder block containing an AutoEP layer while keeping the AutoEP
+        router, token movement, expert compute, and collectives eager.
         """
         # Avoid graph breaks
         deepspeed.utils.nvtx.enable_nvtx = False
 
         if not is_compile_supported():
             raise RuntimeError("compile is not supported in your version of PyTorch.")
 
+        if compile_mode not in ("model", "autoep_non_moe"):
+            raise ValueError(f"Unknown compile_mode={compile_mode!r}; expected 'model' or 'autoep_non_moe'.")
+
         if self.is_compiled:
-            return
+            if self._compile_mode == compile_mode:
+                return
+            raise RuntimeError(f"Engine is already compiled with compile_mode={self._compile_mode!r}.")
 
         if 'backend' in compile_kwargs:
             logger.warning("The `backend` in `compile_kwargs` will be overridden. Use the `backend` argument instead.")
 
-        logger.info(f"Compiling deepcompile={self.is_deepcompile_enabled()} backend={backend}")
+        logger.info(f"Compiling mode={compile_mode} deepcompile={self.is_deepcompile_enabled()} backend={backend}")
+
+        if compile_mode == "autoep_non_moe":
+            if self.is_deepcompile_enabled():
+                raise ValueError("compile_mode='autoep_non_moe' uses vanilla torch.compile and cannot be combined "
+                                 "with DeepCompile.")
+            autoep_config = getattr(self._config, "expert_parallel_config", None)
+            if getattr(autoep_config, "comm_backend", "comm") != "comm":
+                raise ValueError("compile_mode='autoep_non_moe' supports only expert_parallel.comm_backend='comm'.")
+            if self.autotp_size() > 1:
+                raise ValueError("compile_mode='autoep_non_moe' does not support AutoEP+AutoTP folding yet.")
+            if self._autoep_sequence_parallel_world_size() > 1:
+                raise ValueError("compile_mode='autoep_non_moe' does not support sequence parallelism yet.")
+            folding_spec = getattr(self, "_autoep_folding_spec", None)
+            if getattr(self, "pipeline_parallelism", False) or getattr(folding_spec, "pp_size", 1) > 1:
+                raise ValueError("compile_mode='autoep_non_moe' does not support pipeline parallelism yet.")
+            if self.zero_optimization_partition_weights():
+                raise ValueError("compile_mode='autoep_non_moe' does not support ZeRO Stage 3 yet.")
+            if self.zero_offload_optimizer() is not None or self.zero_offload_param() is not None:
+                raise ValueError("compile_mode='autoep_non_moe' does not support optimizer or parameter offload yet.")
+            if schedule is not None:
+                raise ValueError("compile_mode='autoep_non_moe' does not support DeepCompile schedules.")
+            if compiled_autograd_enabled:
+                raise ValueError("compile_mode='autoep_non_moe' does not support compiled autograd yet.")
+            from .compiler import compile_autoep_non_moe_regions
+            self._compiled_regions = compile_autoep_non_moe_regions(self.module, backend, compile_kwargs)
+            self._is_compiled = True
+            self._compile_mode = compile_mode
+            self._compile_kwargs = compile_kwargs
+            return
 
         resolved_backend = None
         if self.is_deepcompile_enabled():
@@ -5842,6 +5882,7 @@ def compile(self,
             raise
 
         self._is_compiled = True
+        self._compile_mode = compile_mode
         self._compile_kwargs = compile_kwargs
         if compiled_autograd_enabled:
             if not self._deepcompile_active:

docs/code-docs/source/autoep.rst

@@ -74,6 +74,38 @@ Weights-only/module-only Universal Checkpoint loads use the converted
         }
     }
 
+Experimental regional ``torch.compile``
+----------------------------------------
+
+AutoEP can keep its router, token movement, expert computation, and collectives
+in eager mode while compiling the surrounding decoder blocks with vanilla
+``torch.compile``. This targets fragmented attention, normalization, residual,
+and dense backward work without capturing AutoEP communication in the graph.
+The path is opt-in and does not change the default eager execution:
+
+.. code-block:: python
+
+    engine, optimizer, _, _ = deepspeed.initialize(
+        model=model,
+        model_parameters=model.parameters(),
+        config=ds_config,
+    )
+    engine.compile(compile_mode="autoep_non_moe")
+
+The call must happen after ``deepspeed.initialize()`` so AutoEP replacement is
+complete. DeepSpeed discovers each ``AutoEPMoELayer`` and regionally compiles
+its direct parent decoder block with ``fullgraph=False`` and ``dynamic=False``.
+The AutoEP layer is an explicit compiler-disabled graph break, so routing,
+AllToAll dispatch/combine, and expert execution remain eager.
+
+The initial experimental path supports vanilla ``torch.compile`` with the
+standard ``comm`` backend, sequence and pipeline parallel sizes of one, and
+ZeRO stages 0, 1, and 2. Distributed performance and parity validation currently
+target ZeRO stage 1. It rejects DeepEP, DeepCompile, AutoEP+AutoTP folding,
+sequence or pipeline parallelism, ZeRO stage 3, optimizer or parameter offload,
+compiled autograd, DeepCompile schedules, and any ``fullgraph`` or ``dynamic``
+value other than ``False`` instead of silently changing the requested behavior.
+
 **How it works:**
 
 1. During ``deepspeed.initialize()``, AutoEP scans the model for MoE layers