(1/2) Implementing Compiler Pass for AutoTP

deepspeed/compile/config.py

@@ -6,7 +6,7 @@
 from typing import List, Optional, Literal
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
 
-PassName = Literal["z1", "z3", "autosp"]
+PassName = Literal["z1", "z3", "autosp", "autotp"]
 
 
 class CompileConfig(DeepSpeedConfigModel):

deepspeed/compile/custom_ops/__init__.py

@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 from .all_to_all import all_to_all
+from .tp_collectives import copy_to_tp_region, reduce_from_tp_region
 from . import sp_dp_registry
 
-__all__ = ["all_to_all", "sp_dp_registry", "sp_compat"]
+__all__ = ["all_to_all", "copy_to_tp_region", "reduce_from_tp_region", "sp_dp_registry", "sp_compat"]

deepspeed/compile/custom_ops/tp_collectives.py

@@ -0,0 +1,72 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed.comm as dist
+from deepspeed.utils import groups
+
+
+def get_tp_group():
+    """Return the tensor-parallel group created by the existing AutoTP setup.
+
+    The AutoTP pass reuses the groups that ``TpTrainingManager`` already builds, so the compiled
+    collectives always communicate over the same group as the module-level ones they replace.
+    """
+    return groups.get_tensor_model_parallel_group()
+
+
+@torch.library.custom_op("autotp::copy_to_tp_region", mutates_args=())
+def copy_to_tp_region(input: torch.Tensor) -> torch.Tensor:
+    """Identity in the forward pass, all-reduce in the backward pass.
+
+    This is Megatron's ``f``. It is inserted before a column-parallel matmul: the activation is
+    already replicated across the tensor-parallel group, so nothing has to happen in the forward
+    pass, while each rank contributes a partial gradient that must be summed in the backward pass.
+    """
+    return input.clone()
+
+
+@torch.library.register_fake("autotp::copy_to_tp_region")
+def copy_to_tp_region_fake(input: torch.Tensor):
+    return torch.empty_like(input)
+
+
+@torch.library.custom_op("autotp::reduce_from_tp_region", mutates_args=())
+def reduce_from_tp_region(input: torch.Tensor) -> torch.Tensor:
+    """All-reduce in the forward pass, identity in the backward pass.
+
+    This is Megatron's ``g``. It is inserted after a row-parallel matmul, whose output is only a
+    partial sum because each rank holds a slice of the input dimension.
+    """
+    output = input.contiguous().clone()
+    dist.all_reduce(output, group=get_tp_group())
+    return output
+
+
+@torch.library.register_fake("autotp::reduce_from_tp_region")
+def reduce_from_tp_region_fake(input: torch.Tensor):
+    return torch.empty_like(input)
+
+
+def _copy_to_tp_region_backward(ctx, grad):
+    # f and g are duals, so f's backward is simply g.
+    return reduce_from_tp_region(grad.contiguous())
+
+
+def _reduce_from_tp_region_backward(ctx, grad):
+    return grad
+
+
+def _setup_context_without_saved_tensors(ctx, inputs, output):
+    # Both collectives are shape-preserving and stateless, so their backwards need nothing saved.
+    pass
+
+
+torch.library.register_autograd("autotp::copy_to_tp_region",
+                                _copy_to_tp_region_backward,
+                                setup_context=_setup_context_without_saved_tensors)
+torch.library.register_autograd("autotp::reduce_from_tp_region",
+                                _reduce_from_tp_region_backward,
+                                setup_context=_setup_context_without_saved_tensors)

deepspeed/compile/init_tp.py

@@ -0,0 +1,35 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch.fx import GraphModule
+
+from deepspeed.utils.torch import required_torch_version
+
+from .passes.tp_compile import apply_autotp, defer_collectives_to_compiler
+
+AUTOTP_MIN_TORCH_VERSION = 2.6
+
+
+def _check_autotp_compatibility():
+    if not required_torch_version(min_version=AUTOTP_MIN_TORCH_VERSION):
+        raise RuntimeError(f"The AutoTP compile pass requires PyTorch >= {AUTOTP_MIN_TORCH_VERSION}, found "
+                           f"{torch.__version__}.")
+
+
+def init_autotp(model):
+    """Hand the tensor-parallel collectives of an AutoTP-partitioned model over to the compiler.
+
+    The model is expected to have been partitioned already by the regular AutoTP path, so this only
+    suppresses the module-level collectives and returns a backend that emits them as graph nodes.
+    """
+    _check_autotp_compatibility()
+    defer_collectives_to_compiler(model)
+
+    def backend_fn(gm: GraphModule, real_inputs):
+        apply_autotp(gm, real_inputs)
+        return torch._inductor.compile(gm, real_inputs)
+
+    return backend_fn

deepspeed/compile/passes/tp_compile.py

@@ -0,0 +1,138 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Dict, List
+
+import torch
+from torch.fx import GraphModule, Node
+
+from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer
+
+from ..custom_ops import tp_collectives  # noqa: F401
+
+COLUMN_PARALLEL_OP = torch.ops.autotp.copy_to_tp_region.default
+ROW_PARALLEL_OP = torch.ops.autotp.reduce_from_tp_region.default
+
+# AutoTP replaces nn.Linear with these layers and shards their weights, so the injected layer type
+# already records the partitioning decision the pass needs. Reading it back is more robust than
+# re-deriving column/row from parameter-name patterns.
+COLUMN_PARALLEL_LAYER = LinearLayer
+ROW_PARALLEL_LAYER = LinearAllreduce
+
+# The injected layers compute their matmul with torch.matmul; the plain nn.Linear spelling is
+# accepted too so the pass keeps working if a layer is lowered differently.
+_MATMUL_TARGETS = {
+    torch.matmul,
+    torch.ops.aten.matmul.default,
+    torch.ops.aten.linear.default,
+    torch._C._nn.linear,
+}
+
+
+def defer_collectives_to_compiler(model) -> int:
+    """Suppress the module-level TP collectives on layers this pass will handle in the graph.
+
+    Returns the number of layers handed over to the pass. Layers the pass does not rewrite (the
+    fused sub-param variants, conv and embedding layers) keep their module-level collectives and
+    stay correct as-is.
+    """
+    deferred = 0
+    for name, module in model.named_modules():
+        is_row_parallel = type(module) is ROW_PARALLEL_LAYER
+        is_column_parallel = type(module) is COLUMN_PARALLEL_LAYER
+        if not (is_row_parallel or is_column_parallel):
+            continue
+        if module.mp_group is None:
+            continue
+        if type(module).tp_overlap_comm:
+            raise NotImplementedError("AutoTP compile pass does not support tp_overlap_comm. Set "
+                                      "'tp_overlap_comm': false to emit the collectives into the graph.")
+        # GatherFromTensorParallelRegion reads the gathered shard sizes back into Python, which the
+        # full graph this pass needs cannot capture. Leaving such a layer on the module-level path
+        # is not an option either: the pass identifies column-parallel layers by type, so it would
+        # add a second collective on top of the module's own and reduce the input gradient twice.
+        if is_column_parallel and module.gather_output:
+            raise NotImplementedError(
+                f"AutoTP compile pass does not support gather_output layers, but '{name}' is one. Partition it "
+                "without gather_output, or drop 'autotp' from the DeepCompile passes for this model.")
+        module.defer_collectives_to_compiler = True
+        deferred += 1
+    return deferred
+
+
+def _originating_layer_type(node: Node):
+    """Return the innermost nn.Module type a node was traced from, or None."""
+    module_stack = node.meta.get("nn_module_stack")
+    if not module_stack:
+        return None
+    _, module_type = list(module_stack.values())[-1]
+    return module_type
+
+
+def _insert_row_collective(gm: GraphModule, matmul: Node) -> Node:
+    """Insert g after a row-parallel matmul.
+
+    Every consumer has to read the reduced value, which is also what the module-level
+    RowParallel.apply this replaces produces.
+    """
+    with gm.graph.inserting_after(matmul):
+        collective_node = gm.graph.call_function(ROW_PARALLEL_OP, args=(matmul, ))
+    collective_node.meta["val"] = matmul.meta.get("val")
+    matmul.replace_all_uses_with(collective_node)
+    collective_node.update_arg(0, matmul)
+    return collective_node
+
+
+def _insert_column_collective(gm: GraphModule, activation: Node, consumers: List[Node]) -> Node:
+    """
+    Insert f in front of the column-parallel matmuls that share activation.
+    """
+    with gm.graph.inserting_before(consumers[0]):
+        collective_node = gm.graph.call_function(COLUMN_PARALLEL_OP, args=(activation, ))
+    collective_node.meta["val"] = activation.meta.get("val")
+    for consumer in consumers:
+        consumer.replace_input_with(activation, collective_node)
+    return collective_node
+
+
+def pass_insert_tp_collectives(gm: GraphModule, real_inputs):
+    """Insert the tensor-parallel collectives around the matmuls of the injected AutoTP layers."""
+    column_consumers: Dict[Node, List[Node]] = {}
+
+    for node in list(gm.graph.nodes):
+        if node.op != "call_function" or node.target not in _MATMUL_TARGETS:
+            continue
+
+        layer_type = _originating_layer_type(node)
+        if layer_type is ROW_PARALLEL_LAYER:
+            _insert_row_collective(gm, node)
+        elif layer_type is COLUMN_PARALLEL_LAYER:
+            activation = node.args[0]
+            column_consumers.setdefault(activation, []).append(node)
+
+    for activation, consumers in column_consumers.items():
+        _insert_column_collective(gm, activation, consumers)
+
+
+def pass_canonicalize(gm: GraphModule, real_inputs):
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()
+
+
+AUTOTP_PASSES = [
+    pass_insert_tp_collectives,
+    pass_canonicalize,
+]
+
+
+def apply_autotp(gm: GraphModule, real_inputs, passes=None):
+    """Apply the AutoTP transformation passes to the graph.
+
+    The collectives are shape-preserving, so unlike AutoSP this needs no shape re-propagation.
+    """
+    for opt_pass in passes or AUTOTP_PASSES:
+        opt_pass(gm, real_inputs)
+    return gm

deepspeed/module_inject/layers.py

@@ -305,6 +305,10 @@ def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any):
         """
         super().__init__()
         self.support_training: bool = False
+        # DeepCompile's AutoTP pass emits the tensor-parallel collectives as graph nodes so the
+        # scheduler and profiler can see them. The module-level collectives are suppressed in that
+        # mode, but mp_group is still needed for parameter gathering and checkpointing.
+        self.defer_collectives_to_compiler: bool = False
         self.mp_group = mp_group
         if mp_group is not None:
             self.tp_world_size: int = dist.get_world_size(self.mp_group)
@@ -638,7 +642,8 @@ def __init__(self, module, mp_group, **kwargs):
 
     def forward(self, input):
         output = torch.matmul(input, self.weight.transpose(-1, -2))
-        output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
+        if not self.defer_collectives_to_compiler:
+            output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
         if self.bias is not None:
             output = add_bias(output, self.bias)
         return output
@@ -734,7 +739,7 @@ def __init__(self, module, mp_group=None, skip_partition=False, gather_output=Fa
 
     def forward(self, input):
         if not self.__class__.tp_overlap_comm:
-            if getattr(self, 'mp_group', None) is not None:
+            if getattr(self, 'mp_group', None) is not None and not self.defer_collectives_to_compiler:
                 input = ColumnParallel.apply(self.mp_group, input)
             output = torch.matmul(input, self.weight.transpose(-1, -2))
             if self.bias is not None:

deepspeed/runtime/engine.py

@@ -150,6 +150,7 @@
 from deepspeed.compile.init_z3 import init_z3
 from deepspeed.compile.z3_eager_fallback import deepcompile_z3_forward_context
 from deepspeed.compile.init_sp import init_autosp
+from deepspeed.compile.init_tp import init_autotp
 
 MEMORY_OPT_ALLREDUCE_SIZE = 500000000
 
@@ -1213,6 +1214,19 @@ def compile_autosp(self):
         """Determines if AutoSP is set in deepcompile's passes attributes."""
         return "autosp" in (getattr(self._config.compile_config, "passes", None) or [])
 
+    def compile_autotp(self):
+        """Determines if AutoTP is set in deepcompile's passes attributes."""
+        return "autotp" in (getattr(self._config.compile_config, "passes", None) or [])
+
+    def uses_parallelization_pass_only(self):
+        """Determines if the compiled graph comes from a parallelization pass rather than ZeRO.
+
+        AutoSP and AutoTP rewrite the graph and then rely on regular autograd, so a run using only
+        those passes must keep the standard gradient reduction instead of the one the z1/z3 passes
+        install.
+        """
+        return self.compile_autosp() or self.compile_autotp()
+
     def mics_shard_size(self):
         return self._config.mics_shard_size
 
@@ -2808,7 +2822,7 @@ def print_forward_breakdown(self, fwd_time):
     def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
         # Skip gradient reduction when DeepCompile is enabled
         # DeepCompile handles its own gradient reduction through compiled graph operations
-        if self.is_deepcompile_active() and not self.compile_autosp():
+        if self.is_deepcompile_active() and not self.uses_parallelization_pass_only():
             return
 
         # Pass (PP) gas boundary flag to optimizer (required for zero)
@@ -2867,7 +2881,9 @@ def _backward_prologue(self):
         assert not self.eigenvalue_enabled(), "Eigenvalue is not supported with non-scalar backward"
         assert not self.amp_enabled(), "Apex AMP is not supported with non-scalar backward"
 
-        if self.is_deepcompile_active():
+        # The AutoTP pass installs no backward hooks and keeps no DeepCompile state, so the
+        # prologue would only force the DeepCompile native extension to load for nothing.
+        if self.is_deepcompile_active() and not self.compile_autotp():
             deepcompile_backward_prologue(self.is_gradient_accumulation_boundary())
 
         if isinstance(self.optimizer, ZeROOptimizer):
@@ -2902,7 +2918,7 @@ def _backward_epilogue(self):
                 self.optimizer.backward_epilogue()
             self.optimizer.exit_backward()
 
-        if self.is_deepcompile_active():
+        if self.is_deepcompile_active() and not self.compile_autotp():
             deepcompile_backward_epilogue()
 
         see_memory_usage("Engine after backward", force=self.memory_breakdown())
@@ -5476,6 +5492,26 @@ def get_autosp_backend(self, compile_kwargs):
         compile_kwargs['fullgraph'] = True
         return init_autosp(self._config)
 
+    def get_autotp_backend(self, compile_kwargs):
+        if self.autotp_size() <= 1:
+            logger.info("AutoTP compile pass requires tensor_parallel.autotp_size > 1. "
+                        "Falling back to the torch compiler.")
+            return None
+
+        # The one-shot dataloader consistency check broadcasts Python objects, which cannot be
+        # captured in a full graph, so it has to go before the module is compiled.
+        if self.first_dataloader_check is not None:
+            self.first_dataloader_check.remove()
+            self.first_dataloader_check = None
+            logger.warning("Skipping the TP dataloader consistency check because the AutoTP compile pass "
+                           "requires a full graph. Ensure the dataloader yields identical inputs on every "
+                           "rank of the TP group.")
+
+        # A graph break would leave part of the model without the collectives the pass inserts,
+        # which is silently wrong rather than slow, so the whole module must be captured.
+        compile_kwargs['fullgraph'] = True
+        return init_autotp(self.module)
+
     def get_deepcompile_backend(self, backend, compile_kwargs, schedule):
         if self.zero_optimization_stage() != ZeroStageEnum.optimizer_states \
                 and self.zero_optimization_stage() != ZeroStageEnum.weights \
@@ -5514,8 +5550,14 @@ def passes_name_to_fn(passes):
 
         assert backend in ['inductor', 'eager'], f"Backend {backend} is not supported for DeepCompile."
 
+        if self.compile_autotp() and (self.compile_autosp() or self.compile_zero_optimization_stage()):
+            raise NotImplementedError("The AutoTP compile pass cannot yet be combined with AutoSP or the ZeRO "
+                                      "passes. Run 'autotp' on its own until the passes are made composable.")
+
         if self.compile_autosp():
             resolved_backend = self.get_autosp_backend(compile_kwargs)
+        elif self.compile_autotp():
+            resolved_backend = self.get_autotp_backend(compile_kwargs)
         else:
             resolved_backend = self.get_deepcompile_backend(backend, compile_kwargs, schedule)