(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,9 @@
 # DeepSpeed Team
 
 from .all_to_all import all_to_all
+from .tp_collectives import copy_to_tp_region, gather_from_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", "gather_from_tp_region", "reduce_from_tp_region", "sp_dp_registry", "sp_compat"
+]

deepspeed/compile/custom_ops/tp_collectives.py

@@ -0,0 +1,108 @@
+# 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."""
+    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.
+
+    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.
+
+    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)
+
+
+@torch.library.custom_op("autotp::gather_from_tp_region", mutates_args=())
+def gather_from_tp_region(input: torch.Tensor) -> torch.Tensor:
+    """All-gather the last dimension in the forward pass, take this rank's slice in the backward.
+
+    Inserted after a column-parallel matmul whose layer asks for gather_output, so that
+    every rank leaves the layer holding the full output rather than its own shard.
+    """
+    group = get_tp_group()
+    world_size = dist.get_world_size(group=group)
+    if world_size == 1:
+        return input.clone()
+
+    local_shard = input.contiguous()
+    flat_gathered = torch.empty((world_size * local_shard.shape[0], *local_shard.shape[1:]),
+                                dtype=local_shard.dtype,
+                                device=local_shard.device)
+    dist.all_gather_into_tensor(flat_gathered, local_shard, group=group)
+    shards = flat_gathered.view(world_size, *local_shard.shape)
+    return torch.cat(shards.unbind(0), dim=-1)
+
+
+@torch.library.register_fake("autotp::gather_from_tp_region")
+def gather_from_tp_region_fake(input: torch.Tensor):
+    world_size = dist.get_world_size(group=get_tp_group())
+    return input.new_empty((*input.shape[:-1], input.shape[-1] * world_size))
+
+
+def _copy_to_tp_region_backward(ctx, grad):
+    # copy_to_tp_region and reduce_from_tp_region are duals,
+    # so copy_to_tp_region's backward is simply reduce_from_tp_region.
+    return reduce_from_tp_region(grad.contiguous())
+
+
+def _reduce_from_tp_region_backward(ctx, grad):
+    return grad
+
+
+def _gather_from_tp_region_backward(ctx, grad):
+    group = get_tp_group()
+    world_size = dist.get_world_size(group=group)
+    if world_size == 1:
+        return grad
+    shard_width = grad.shape[-1] // world_size
+    shard_start = dist.get_rank(group=group) * shard_width
+    return grad.narrow(-1, shard_start, shard_width).contiguous()
+
+
+def _setup_context_without_saved_tensors(ctx, inputs, output):
+    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)
+torch.library.register_autograd("autotp::gather_from_tp_region",
+                                _gather_from_tp_region_backward,
+                                setup_context=_setup_context_without_saved_tensors)

deepspeed/compile/init_tp.py

@@ -0,0 +1,61 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from packaging.version import Version
+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
+BROKEN_TRANSFORMERS_MOE_VERSIONS = ("5.8.0", "5.10.1")
+
+
+def _check_autotp_compatibility(model):
+    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__}.")
+    _check_broken_transformers_moe(model)
+
+
+def _check_broken_transformers_moe(model):
+    """Reject models whose experts forward the installed transformers cannot capture in a graph."""
+
+    try:
+        import transformers
+    except ImportError:
+        return
+    first_broken, first_fixed = BROKEN_TRANSFORMERS_MOE_VERSIONS
+    if not (Version(first_broken) <= Version(transformers.__version__) < Version(first_fixed)):
+        return
+    experts_modules = [
+        name for name, module in model.named_modules()
+        if hasattr(module, "_apply_gate") and hasattr(module, "is_concatenated")
+        and getattr(getattr(module, "config", None), "_experts_implementation", None) == "batched_mm"
+    ]
+    if experts_modules:
+        raise RuntimeError(
+            f"transformers {transformers.__version__} mutates the MoE routing tensor in place inside "
+            "batched_mm_experts_forward (huggingface/transformers#45621, fixed by #45634), and this model "
+            f"routes through it (e.g. '{experts_modules[0]}'), so the AutoTP compile pass cannot capture a "
+            f"full graph. Upgrade to transformers >= {first_fixed}.")
+
+
+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(model)
+    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,148 @@
+# 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, LmHeadLinearAllreduce,
+                                            SubParamLinearAllreduce, SubParamLinearLayer, TensorParallel_Layer)
+
+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
+GATHER_OUTPUT_OP = torch.ops.autotp.gather_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.
+
+COLUMN_PARALLEL_LAYERS = (LinearLayer, SubParamLinearLayer)
+ROW_PARALLEL_LAYERS = (LinearAllreduce, SubParamLinearAllreduce)
+# LmHeadLinearAllreduce subclasses LinearAllreduce but slices its own input and reduces with
+# inference_all_reduce rather than going through RowParallel, so the pass cannot stand in for it.
+UNSUPPORTED_LAYERS = (LmHeadLinearAllreduce, )
+
+_MATMUL_TARGETS = {
+    torch.matmul,
+    torch.ops.aten.matmul.default,
+    torch.ops.aten.linear.default,
+    torch._C._nn.linear,
+}
+
+
+def _is_column_parallel(layer_type) -> bool:
+    return _in_family(layer_type, COLUMN_PARALLEL_LAYERS)
+
+
+def _is_row_parallel(layer_type) -> bool:
+    return _in_family(layer_type, ROW_PARALLEL_LAYERS)
+
+
+def _in_family(layer_type, family) -> bool:
+    if not isinstance(layer_type, type) or issubclass(layer_type, UNSUPPORTED_LAYERS):
+        return False
+    return issubclass(layer_type, family)
+
+
+def defer_collectives_to_compiler(model) -> None:
+    """Suppress the module-level TP collectives on layers this pass will handle in the graph."""
+    tp_modules = []
+    for name, module in model.named_modules():
+        if not isinstance(module, TensorParallel_Layer) or module.mp_group is None:
+            continue
+
+        layer_type = type(module)
+        is_column_parallel = _is_column_parallel(layer_type)
+        if not (is_column_parallel or _is_row_parallel(layer_type)):
+            raise NotImplementedError(
+                f"AutoTP compile pass cannot rewrite '{name}' ({layer_type.__name__}), and leaving it on the "
+                "module-level path under a full graph would silently drop its backward collective. Drop "
+                "'autotp' from the DeepCompile passes for this model.")
+        if layer_type.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.")
+
+        tp_modules.append(module)
+
+    for module in tp_modules:
+        module.defer_collectives_to_compiler = True
+
+
+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.
+
+    Only f and g are inserted here. The output gather of a gather_output layer is emitted by the
+    layer's own forward (see LinearLayer.forward): it changes the activation's width, so the ops
+    downstream of it only trace correctly if it is already present during graph capture.
+    """
+    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 _is_row_parallel(layer_type):
+            _insert_row_collective(gm, node)
+        elif _is_column_parallel(layer_type):
+            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."""
+    for opt_pass in passes or AUTOTP_PASSES:
+        opt_pass(gm, real_inputs)
+    return gm

deepspeed/module_inject/auto_tp.py

@@ -596,9 +596,55 @@ def _slice_embedding(self, child, name, conv_linear_layer):
         setattr(child, "replaced", True)
         return new_embedding
 
+    def register_replicated_grad_hooks(self, model):
+        """Sum the gradients of replicated parameters marked grad_allreduce across the TP group.
+
+        A parameter left replicated but applied to sharded activations (Qwen3's q_norm/k_norm
+        normalize the local attention heads) receives a different partial gradient on every
+        tensor-parallel rank. Without this reduction the ranks silently drift apart after the
+        first optimizer step. The hook runs when the gradient is computed, before accumulation,
+        so it holds under gradient accumulation and under the DeepCompile pass, whose compiled
+        backward still accumulates leaf gradients through the autograd engine.
+        """
+        if self.partition_config is None or self.mp_group is None or self.mp_size <= 1:
+            return
+
+        def make_grad_allreduce_hook(group):
+
+            def hook(grad):
+                grad = grad.contiguous()
+                dist.all_reduce(grad, group=group)
+                return grad
+
+            return hook
+
+        model_type = self._get_model_type()
+        registered = []
+        for param_name, param in model.named_parameters():
+            spec = self.partition_config.find_matching_spec(param_name, model_type)
+            if spec is None or not spec.grad_allreduce:
+                continue
+            if getattr(param, "_ds_grad_allreduce_registered", False):
+                continue
+            param.register_hook(make_grad_allreduce_hook(self.mp_group))
+            param._ds_grad_allreduce_registered = True
+            registered.append(param_name)
+        if registered:
+            print_dist(
+                f"AutoTP: registered tensor-parallel grad all-reduce for {len(registered)} replicated "
+                f"parameters, e.g. {registered[0]!r}",
+                ranks=[0])
+
     def update_mp_params(self, child):
         if getattr(child, "replaced", False) == True:
             return
+        # Fused-expert containers (Mixtral/Llama4/Qwen-MoE style) hold their weights as 3D
+        # parameters that AutoTP does not shard, so their dimension attributes must stay whole.
+        # Halving e.g. Llama4TextExperts.hidden_size while its weights keep the full size breaks
+        # the experts' batched matmul.
+        if any(param.dim() >= 3 for param in child.parameters(recurse=False)):
+            setattr(child, "replaced", True)
+            return
         param_list = [
             "n_heads", "inner_dim", "num_heads", "num_kv", "num_attention_heads", "num_attn_heads", "all_head_size",
             "embed_dim", "hidden_size", "num_key_value_heads", "num_kv_heads", "kv_n_heads", "d_model",