From 8abe7ed428e31378132894da764f969ef0ca90b7 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Sat, 1 Aug 2026 22:56:12 +0000
Subject: [PATCH 01/11] added support for autotp

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/config.py                   |   2 +-
 deepspeed/compile/custom_ops/__init__.py      |   3 +-
 .../compile/custom_ops/tp_collectives.py      |  72 +++++++
 deepspeed/compile/init_tp.py                  |  23 +++
 deepspeed/compile/passes/tp_compile.py        | 116 +++++++++++
 deepspeed/module_inject/layers.py             |   9 +-
 deepspeed/runtime/engine.py                   |  48 ++++-
 tests/unit/compile/test_tp_compile.py         | 183 ++++++++++++++++++
 8 files changed, 449 insertions(+), 7 deletions(-)
 create mode 100644 deepspeed/compile/custom_ops/tp_collectives.py
 create mode 100644 deepspeed/compile/init_tp.py
 create mode 100644 deepspeed/compile/passes/tp_compile.py
 create mode 100644 tests/unit/compile/test_tp_compile.py

diff --git a/deepspeed/compile/config.py b/deepspeed/compile/config.py
index 2137b94722f2..5bb249450448 100644
--- a/deepspeed/compile/config.py
+++ b/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):
diff --git a/deepspeed/compile/custom_ops/__init__.py b/deepspeed/compile/custom_ops/__init__.py
index e5fc593a2e7e..d183eaa11344 100644
--- a/deepspeed/compile/custom_ops/__init__.py
+++ b/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"]
diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py
new file mode 100644
index 000000000000..ec48ce37ed1a
--- /dev/null
+++ b/deepspeed/compile/custom_ops/tp_collectives.py
@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation.
+# 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)
diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py
new file mode 100644
index 000000000000..dd0218f29bd5
--- /dev/null
+++ b/deepspeed/compile/init_tp.py
@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch.fx import GraphModule
+from .passes.tp_compile import apply_autotp, defer_collectives_to_compiler
+
+
+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.
+    """
+    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
diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
new file mode 100644
index 000000000000..b88088f5ee89
--- /dev/null
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -0,0 +1,116 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+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 (a
+    column-parallel layer that gathers its output, the fused sub-param variants, conv and
+    embedding layers) keep their module-level collectives and stay correct as-is.
+    """
+    deferred = 0
+    for module in model.modules():
+        is_row_parallel = type(module) is ROW_PARALLEL_LAYER
+        # gather_output adds a further collective that this pass does not emit yet, so leave those
+        # layers to the module-level path.
+        is_column_parallel = type(module) is COLUMN_PARALLEL_LAYER and not module.gather_output
+        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.")
+        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_after(gm: GraphModule, node: Node, op) -> Node:
+    """Insert ``op(node)`` and re-point every consumer of ``node`` at the new node."""
+    with gm.graph.inserting_after(node):
+        collective_node = gm.graph.call_function(op, args=(node, ))
+    collective_node.meta["val"] = node.meta.get("val")
+    # Steal every consumer first, then hand the original back as this node's own input; doing it in
+    # the other order would leave the new node feeding itself.
+    node.replace_all_uses_with(collective_node)
+    collective_node.update_arg(0, 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."""
+    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_after(gm, node, ROW_PARALLEL_OP)
+        elif layer_type is COLUMN_PARALLEL_LAYER:
+            activation = node.args[0]
+            # Column-parallel layers that share an activation (q/k/v, gate/up) need only one
+            # collective. Inserting it already re-pointed the sibling matmuls at the new node, so
+            # finding one here means this activation has been handled.
+            if activation.op == "call_function" and activation.target is COLUMN_PARALLEL_OP:
+                continue
+            _insert_after(gm, activation, COLUMN_PARALLEL_OP)
+
+
+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
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 33b1fbe3dbd0..6b74da06ab72 100644
--- a/deepspeed/module_inject/layers.py
+++ b/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:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index b1b99c305f92..2ee87728c9a0 100755
--- a/deepspeed/runtime/engine.py
+++ b/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)
 
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
new file mode 100644
index 000000000000..79c2972a270b
--- /dev/null
+++ b/tests/unit/compile/test_tp_compile.py
@@ -0,0 +1,183 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+import deepspeed
+import deepspeed.comm as dist
+from deepspeed.accelerator import get_accelerator
+from deepspeed.utils import groups
+from deepspeed.utils.torch import required_torch_version
+
+from unit.common import DistributedTest
+
+pytestmark = pytest.mark.skipif(not required_torch_version(min_version=2.9),
+                                reason="The AutoTP compile pass requires PyTorch >= 2.9")
+
+HIDDEN_DIM = 64
+INTERMEDIATE_DIM = 128
+
+
+class MLPBlock(torch.nn.Module):
+    """Llama-style MLP: gate/up are column-parallel and down is row-parallel."""
+
+    def __init__(self):
+        super().__init__()
+        self.gate_proj = torch.nn.Linear(HIDDEN_DIM, INTERMEDIATE_DIM, bias=False)
+        self.up_proj = torch.nn.Linear(HIDDEN_DIM, INTERMEDIATE_DIM, bias=False)
+        self.down_proj = torch.nn.Linear(INTERMEDIATE_DIM, HIDDEN_DIM, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MLPModel(torch.nn.Module):
+
+    def __init__(self, nlayers=2):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([MLPBlock() for _ in range(nlayers)])
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def build_config(tp_size, use_compile_pass):
+    config = {
+        "train_micro_batch_size_per_gpu": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-6
+            }
+        },
+        "tensor_parallel": {
+            "autotp_size": tp_size,
+            "partition_config": {
+                "use_default_specs":
+                False,
+                "layer_specs": [{
+                    "patterns": [".*\\.gate_proj\\.weight$", ".*\\.up_proj\\.weight$"],
+                    "partition_type": "column",
+                }, {
+                    "patterns": [".*\\.down_proj\\.weight$"],
+                    "partition_type": "row",
+                }],
+            },
+        },
+        "zero_optimization": {
+            "stage": 0,
+        },
+    }
+    if use_compile_pass:
+        config["compile"] = {"deepcompile": True, "passes": ["autotp"]}
+    return config
+
+
+def build_engine(tp_size, use_compile_pass):
+    # Both engines are built from the same seed so they hold identical shards, which lets the
+    # gradients be compared directly without gathering them first.
+    torch.manual_seed(42)
+    model = MLPModel()
+    engine, _, _, _ = deepspeed.initialize(model=model,
+                                           model_parameters=model.parameters(),
+                                           config=build_config(tp_size, use_compile_pass))
+    if use_compile_pass:
+        engine.compile()
+    return engine
+
+
+class TestAutoTPCompileEquivalence(DistributedTest):
+    """The compile pass must reproduce the module-injection AutoTP path exactly.
+
+    Both paths shard the weights the same way, so the compiled model is compared against the
+    module-level collectives it replaces rather than against a single-device run.
+    """
+
+    world_size = 2
+    non_daemonic_procs = True
+
+    @pytest.mark.sequential
+    def test_matches_module_injection(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU does not support this test yet")
+
+        device = torch.device(get_accelerator().current_device_name())
+        reference_engine = build_engine(self.world_size, use_compile_pass=False)
+        compiled_engine = build_engine(self.world_size, use_compile_pass=True)
+
+        # The TP group must see identical inputs on every rank.
+        torch.manual_seed(1234)
+        x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32)
+
+        reference_out = reference_engine(x)
+        compiled_out = compiled_engine(x)
+        assert torch.allclose(reference_out, compiled_out, atol=1e-5), \
+            "AutoTP compile pass changed the forward result"
+
+        reference_engine.backward(reference_out.sum())
+        compiled_engine.backward(compiled_out.sum())
+
+        # A missing or duplicated collective usually leaves the forward pass intact and only
+        # corrupts gradients, so the gradients are what this test really checks.
+        for (name, reference_param), (_, compiled_param) in zip(reference_engine.module.named_parameters(),
+                                                                compiled_engine.module.named_parameters()):
+            assert torch.allclose(reference_param.grad, compiled_param.grad, atol=1e-5), \
+                f"AutoTP compile pass changed the gradient of {name}"
+
+
+class TestAutoTPCompileDataParallelGradients(DistributedTest):
+    """Gradients must still be reduced across data-parallel replicas.
+
+    The engine skips its own gradient reduction when DeepCompile is active because the ZeRO passes
+    emit that reduction into the graph. The AutoTP pass does not: its collectives only sum partial
+    results inside a TP group and never touch the DP axis. Only a run with more than one
+    data-parallel replica shows whether the reduction still happens.
+    """
+
+    world_size = 4
+    non_daemonic_procs = True
+
+    @pytest.mark.sequential
+    def test_gradients_are_reduced_across_dp_group(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU does not support this test yet")
+
+        tp_size = 2
+        device = torch.device(get_accelerator().current_device_name())
+        engine = build_engine(tp_size, use_compile_pass=True)
+
+        dp_group = groups.get_data_parallel_group()
+        assert dist.get_world_size(group=dp_group) == self.world_size // tp_size
+
+        # Every data-parallel replica gets different data, so an unreduced gradient differs between
+        # replicas. Ranks inside a TP group must still agree, hence seeding on the replica index.
+        replica_index = dist.get_rank() // tp_size
+        torch.manual_seed(1234 + replica_index)
+        x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32)
+
+        out = engine(x)
+        engine.backward(out.sum())
+
+        for name, param in engine.module.named_parameters():
+            gathered = [torch.empty_like(param.grad) for _ in range(dist.get_world_size(group=dp_group))]
+            dist.all_gather(gathered, param.grad.contiguous(), group=dp_group)
+            assert torch.allclose(gathered[0], gathered[-1], atol=1e-5), \
+                f"Gradient of {name} was not reduced across the data-parallel group"
+
+
+class TestAutoTPCompileRejectsUnsupportedCombinations(DistributedTest):
+
+    world_size = 1
+
+    def test_autotp_with_zero_pass_raises(self):
+        model = MLPModel()
+        config = build_config(tp_size=1, use_compile_pass=True)
+        config["compile"]["passes"] = ["autotp", "z1"]
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
+        with pytest.raises(NotImplementedError, match="cannot yet be combined"):
+            engine.compile()

From d4420016acea2ef39f905db781e897f57b495f28 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Wed, 5 Aug 2026 19:50:06 -0400
Subject: [PATCH 02/11] fix for 2x residuals

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 .../compile/custom_ops/tp_collectives.py      |  2 +-
 deepspeed/compile/init_tp.py                  |  2 +-
 deepspeed/compile/passes/tp_compile.py        | 70 ++++++++++++-------
 tests/unit/compile/test_tp_compile.py         | 55 ++++++++++++---
 4 files changed, 92 insertions(+), 37 deletions(-)

diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py
index ec48ce37ed1a..65e7a3070110 100644
--- a/deepspeed/compile/custom_ops/tp_collectives.py
+++ b/deepspeed/compile/custom_ops/tp_collectives.py
@@ -1,4 +1,4 @@
-# Copyright (c) Microsoft Corporation.
+# Copyright (c) DeepSpeed Team.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py
index dd0218f29bd5..0963e9b528d5 100644
--- a/deepspeed/compile/init_tp.py
+++ b/deepspeed/compile/init_tp.py
@@ -1,4 +1,4 @@
-# Copyright (c) Microsoft Corporation.
+# Copyright (c) DeepSpeed Team.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
index b88088f5ee89..db82f5ac4c64 100644
--- a/deepspeed/compile/passes/tp_compile.py
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -1,8 +1,10 @@
-# Copyright (c) Microsoft Corporation.
+# 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
 
@@ -32,16 +34,14 @@
 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 (a
-    column-parallel layer that gathers its output, the fused sub-param variants, conv and
-    embedding layers) keep their module-level collectives and stay correct as-is.
+    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 module in model.modules():
+    for name, module in model.named_modules():
         is_row_parallel = type(module) is ROW_PARALLEL_LAYER
-        # gather_output adds a further collective that this pass does not emit yet, so leave those
-        # layers to the module-level path.
-        is_column_parallel = type(module) is COLUMN_PARALLEL_LAYER and not module.gather_output
+        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:
@@ -49,6 +49,14 @@ def defer_collectives_to_compiler(model) -> int:
         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
@@ -63,35 +71,49 @@ def _originating_layer_type(node: Node):
     return module_type
 
 
-def _insert_after(gm: GraphModule, node: Node, op) -> Node:
-    """Insert ``op(node)`` and re-point every consumer of ``node`` at the new node."""
-    with gm.graph.inserting_after(node):
-        collective_node = gm.graph.call_function(op, args=(node, ))
-    collective_node.meta["val"] = node.meta.get("val")
-    # Steal every consumer first, then hand the original back as this node's own input; doing it in
-    # the other order would leave the new node feeding itself.
-    node.replace_all_uses_with(collective_node)
-    collective_node.update_arg(0, node)
+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_after(gm, node, ROW_PARALLEL_OP)
+            _insert_row_collective(gm, node)
         elif layer_type is COLUMN_PARALLEL_LAYER:
             activation = node.args[0]
-            # Column-parallel layers that share an activation (q/k/v, gate/up) need only one
-            # collective. Inserting it already re-pointed the sibling matmuls at the new node, so
-            # finding one here means this activation has been handled.
-            if activation.op == "call_function" and activation.target is COLUMN_PARALLEL_OP:
-                continue
-            _insert_after(gm, activation, COLUMN_PARALLEL_OP)
+            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):
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index 79c2972a270b..b42e65b17cb4 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -1,4 +1,4 @@
-# Copyright (c) Microsoft Corporation.
+# Copyright (c) DeepSpeed Team.
 # SPDX-License-Identifier: Apache-2.0
 
 # DeepSpeed Team
@@ -9,13 +9,14 @@
 import deepspeed
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
+from deepspeed.compile.init_tp import AUTOTP_MIN_TORCH_VERSION
 from deepspeed.utils import groups
 from deepspeed.utils.torch import required_torch_version
 
 from unit.common import DistributedTest
 
-pytestmark = pytest.mark.skipif(not required_torch_version(min_version=2.9),
-                                reason="The AutoTP compile pass requires PyTorch >= 2.9")
+pytestmark = pytest.mark.skipif(not required_torch_version(min_version=AUTOTP_MIN_TORCH_VERSION),
+                                reason=f"The AutoTP compile pass requires PyTorch >= {AUTOTP_MIN_TORCH_VERSION}")
 
 HIDDEN_DIM = 64
 INTERMEDIATE_DIM = 128
@@ -31,7 +32,10 @@ def __init__(self):
         self.down_proj = torch.nn.Linear(INTERMEDIATE_DIM, HIDDEN_DIM, bias=False)
 
     def forward(self, x):
-        return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+        # The residual is what makes this block interesting for the pass: x feeds the two
+        # column-parallel matmuls and the addition, and only the matmuls may be routed through the
+        # backward all-reduce. Reducing the residual gradient too would scale it by the TP size.
+        return x + self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
 
 
 class MLPModel(torch.nn.Module):
@@ -39,14 +43,20 @@ class MLPModel(torch.nn.Module):
     def __init__(self, nlayers=2):
         super().__init__()
         self.layers = torch.nn.ModuleList([MLPBlock() for _ in range(nlayers)])
+        self.head = torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM, bias=False)
 
     def forward(self, x):
         for layer in self.layers:
             x = layer(x)
-        return x
+        return self.head(x)
 
 
-def build_config(tp_size, use_compile_pass):
+def build_config(tp_size, use_compile_pass, gather_output_head=False):
+    head_spec = {
+        "patterns": [".*\\.head\\.weight$"],
+        "partition_type": "column",
+        "gather_output": gather_output_head,
+    }
     config = {
         "train_micro_batch_size_per_gpu": 1,
         "optimizer": {
@@ -66,7 +76,7 @@ def build_config(tp_size, use_compile_pass):
                 }, {
                     "patterns": [".*\\.down_proj\\.weight$"],
                     "partition_type": "row",
-                }],
+                }, head_spec],
             },
         },
         "zero_optimization": {
@@ -78,14 +88,14 @@ def build_config(tp_size, use_compile_pass):
     return config
 
 
-def build_engine(tp_size, use_compile_pass):
+def build_engine(tp_size, use_compile_pass, gather_output_head=False):
     # Both engines are built from the same seed so they hold identical shards, which lets the
     # gradients be compared directly without gathering them first.
     torch.manual_seed(42)
     model = MLPModel()
     engine, _, _, _ = deepspeed.initialize(model=model,
                                            model_parameters=model.parameters(),
-                                           config=build_config(tp_size, use_compile_pass))
+                                           config=build_config(tp_size, use_compile_pass, gather_output_head))
     if use_compile_pass:
         engine.compile()
     return engine
@@ -112,10 +122,11 @@ def test_matches_module_injection(self):
 
         # The TP group must see identical inputs on every rank.
         torch.manual_seed(1234)
-        x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32)
+        x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32, requires_grad=True)
+        compiled_x = x.detach().clone().requires_grad_(True)
 
         reference_out = reference_engine(x)
-        compiled_out = compiled_engine(x)
+        compiled_out = compiled_engine(compiled_x)
         assert torch.allclose(reference_out, compiled_out, atol=1e-5), \
             "AutoTP compile pass changed the forward result"
 
@@ -129,6 +140,9 @@ def test_matches_module_injection(self):
             assert torch.allclose(reference_param.grad, compiled_param.grad, atol=1e-5), \
                 f"AutoTP compile pass changed the gradient of {name}"
 
+        assert torch.allclose(x.grad, compiled_x.grad, atol=1e-5), \
+            "AutoTP compile pass changed the gradient reaching the model input"
+
 
 class TestAutoTPCompileDataParallelGradients(DistributedTest):
     """Gradients must still be reduced across data-parallel replicas.
@@ -170,6 +184,25 @@ def test_gradients_are_reduced_across_dp_group(self):
                 f"Gradient of {name} was not reduced across the data-parallel group"
 
 
+class TestAutoTPCompileRejectsGatherOutput(DistributedTest):
+    """gather_output layers must be rejected instead of silently losing a collective.
+
+    Their gather reads shard sizes back into Python, which the full graph the pass needs cannot
+    capture, so the pass can neither emit the collectives nor leave them to the module.
+    """
+
+    world_size = 2
+    non_daemonic_procs = True
+
+    @pytest.mark.sequential
+    def test_gather_output_raises(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU does not support this test yet")
+
+        with pytest.raises(NotImplementedError, match="gather_output"):
+            build_engine(self.world_size, use_compile_pass=True, gather_output_head=True)
+
+
 class TestAutoTPCompileRejectsUnsupportedCombinations(DistributedTest):
 
     world_size = 1

From 478c1e9d99891943c2020d1decf87115fadf7347 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Wed, 5 Aug 2026 19:53:32 -0400
Subject: [PATCH 03/11] restore AutoTP torch version check

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/init_tp.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py
index 0963e9b528d5..9da3be3f9174 100644
--- a/deepspeed/compile/init_tp.py
+++ b/deepspeed/compile/init_tp.py
@@ -5,8 +5,19 @@
 
 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.
@@ -14,6 +25,7 @@ def init_autotp(model):
     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):

From e7762c3bdeaa992c55bb48b6503170274375a672 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Fri, 7 Aug 2026 06:28:42 +0000
Subject: [PATCH 04/11] autotp works now

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/custom_ops/__init__.py      |   6 +-
 .../compile/custom_ops/tp_collectives.py      |  49 +++-
 deepspeed/compile/passes/tp_compile.py        |  83 +++---
 deepspeed/module_inject/auto_tp.py            |  46 ++++
 deepspeed/module_inject/autotp_config.py      |   7 +
 deepspeed/module_inject/layers.py             |   9 +-
 deepspeed/module_inject/tp_plan_converter.py  |  34 ++-
 deepspeed/runtime/engine.py                   |   2 +
 docs/_pages/config-json.md                    |  28 ++-
 tests/unit/compile/test_tp_compile.py         | 237 ++++++++++++++++--
 .../model_parallelism/test_tp_plan_e2e.py     |  79 ++++++
 .../module_inject/test_tp_plan_converter.py   |  29 ++-
 12 files changed, 538 insertions(+), 71 deletions(-)

diff --git a/deepspeed/compile/custom_ops/__init__.py b/deepspeed/compile/custom_ops/__init__.py
index d183eaa11344..5885328f1fd8 100644
--- a/deepspeed/compile/custom_ops/__init__.py
+++ b/deepspeed/compile/custom_ops/__init__.py
@@ -4,7 +4,9 @@
 # DeepSpeed Team
 
 from .all_to_all import all_to_all
-from .tp_collectives import copy_to_tp_region, reduce_from_tp_region
+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", "copy_to_tp_region", "reduce_from_tp_region", "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"
+]
diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py
index 65e7a3070110..e937cc4881f6 100644
--- a/deepspeed/compile/custom_ops/tp_collectives.py
+++ b/deepspeed/compile/custom_ops/tp_collectives.py
@@ -50,6 +50,37 @@ 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.
+
+    This is 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. AutoTP only
+    builds such a layer when the output dimension divides evenly by the TP size, so every shard has
+    the same width and the sizes are known statically.
+    """
+    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)
+    # The gather stacks whole shards along dim 0, but the partitioning split the last dimension,
+    # so the shards are re-joined there in rank order to rebuild the unpartitioned output.
+    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):
     # f and g are duals, so f's backward is simply g.
     return reduce_from_tp_region(grad.contiguous())
@@ -59,8 +90,21 @@ def _reduce_from_tp_region_backward(ctx, grad):
     return grad
 
 
+def _gather_from_tp_region_backward(ctx, grad):
+    # The forward concatenated the shards in rank order, so each rank owns a contiguous slice of
+    # the gradient and no communication is needed to recover it.
+    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):
-    # Both collectives are shape-preserving and stateless, so their backwards need nothing saved.
+    # The collectives are stateless and their shapes are fixed by the TP size, so their backwards
+    # need nothing saved.
     pass
 
 
@@ -70,3 +114,6 @@ def _setup_context_without_saved_tensors(ctx, inputs, output):
 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)
diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
index db82f5ac4c64..46a79b7e345b 100644
--- a/deepspeed/compile/passes/tp_compile.py
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -8,18 +8,29 @@
 import torch
 from torch.fx import GraphModule, Node
 
-from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer
+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. Reading it back is more robust than
 # re-deriving column/row from parameter-name patterns.
-COLUMN_PARALLEL_LAYER = LinearLayer
-ROW_PARALLEL_LAYER = LinearAllreduce
+#
+# The families are matched by subclass rather than by exact type. AutoTP injects several variants
+# per family (fused QKV, conv, packed gate/up, Yuan), and they either inherit the base forward or
+# repeat its shape, so the pass rewrites all of them identically. Matching exact types instead
+# leaves the variants on the module-level path, which is not merely unoptimized but wrong: see the
+# comment in defer_collectives_to_compiler.
+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, )
 
 # 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.
@@ -31,35 +42,45 @@
 }
 
 
-def defer_collectives_to_compiler(model) -> int:
+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.
 
-    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.
+    Any tensor-parallel layer the pass cannot rewrite is rejected rather than left on the
+    module-level path. That path looks like the safe fallback but is not: the pass compiles with
+    fullgraph=True, and ColumnParallel's forward is a plain identity, so tracing folds it away and
+    takes its backward all-reduce with it. The forward still matches and only the gradients are
+    wrong, which is the worst way for this to fail.
     """
-    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:
+        if not isinstance(module, TensorParallel_Layer) or module.mp_group is None:
             continue
-        if type(module).tp_overlap_comm:
+
+        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.")
-        # 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):
@@ -98,7 +119,12 @@ def _insert_column_collective(gm: GraphModule, activation: Node, consumers: List
 
 
 def pass_insert_tp_collectives(gm: GraphModule, real_inputs):
-    """Insert the tensor-parallel collectives around the matmuls of the injected AutoTP layers."""
+    """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):
@@ -106,9 +132,9 @@ def pass_insert_tp_collectives(gm: GraphModule, real_inputs):
             continue
 
         layer_type = _originating_layer_type(node)
-        if layer_type is ROW_PARALLEL_LAYER:
+        if _is_row_parallel(layer_type):
             _insert_row_collective(gm, node)
-        elif layer_type is COLUMN_PARALLEL_LAYER:
+        elif _is_column_parallel(layer_type):
             activation = node.args[0]
             column_consumers.setdefault(activation, []).append(node)
 
@@ -131,7 +157,8 @@ def pass_canonicalize(gm: GraphModule, real_inputs):
 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.
+    The inserted collectives are shape-preserving (the shape-changing gather is emitted by the
+    layer forwards during capture), so unlike AutoSP this needs no shape re-propagation.
     """
     for opt_pass in passes or AUTOTP_PASSES:
         opt_pass(gm, real_inputs)
diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py
index 783a00d08ab6..81fa9cdb2324 100755
--- a/deepspeed/module_inject/auto_tp.py
+++ b/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",
diff --git a/deepspeed/module_inject/autotp_config.py b/deepspeed/module_inject/autotp_config.py
index 36896d452a9e..a1fbf998a05a 100644
--- a/deepspeed/module_inject/autotp_config.py
+++ b/deepspeed/module_inject/autotp_config.py
@@ -127,6 +127,12 @@ class TPLayerSpec:
     # Gather column-parallel output shards so every TP rank receives the full output
     gather_output: bool = False
 
+    # For SKIP specs only: the parameter stays replicated, but each tensor-parallel rank computes
+    # only its shard's contribution to the gradient, so the gradients must be summed across the
+    # tensor-parallel group. This is HuggingFace's 'replicated_with_grad_allreduce' style (e.g.
+    # Qwen3's q_norm/k_norm, which normalize sharded attention heads).
+    grad_allreduce: bool = False
+
     def __post_init__(self):
         if isinstance(self.partition_type, str):
             self.partition_type = PartitionType(self.partition_type.lower())
@@ -296,6 +302,7 @@ def from_dict(cls, config_dict: dict) -> "AutoTPConfig":
                     patterns=spec_dict.get("patterns", []),
                     partition_type=partition_type,
                     gather_output=spec_dict.get("gather_output", False),
+                    grad_allreduce=spec_dict.get("grad_allreduce", False),
                     shape=shape,
                     partition_dim=spec_dict.get("partition_dim"),
                     model_types=spec_dict.get("model_types"),
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 6b74da06ab72..95d22fb5b2d9 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -748,7 +748,14 @@ def forward(self, input):
             output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
 
         if self.gather_output:
-            output = GatherFromTensorParallelRegion.apply(self.mp_group, output)
+            if self.defer_collectives_to_compiler:
+                # The gather changes the activation's width, so downstream ops (e.g. a depthwise
+                # conv sized for the full width) only trace correctly if it happens inline. The
+                # custom op is graph-capturable, unlike GatherFromTensorParallelRegion, which
+                # reads gathered shard sizes back into Python.
+                output = torch.ops.autotp.gather_from_tp_region(output)
+            else:
+                output = GatherFromTensorParallelRegion.apply(self.mp_group, output)
 
         return output
 
diff --git a/deepspeed/module_inject/tp_plan_converter.py b/deepspeed/module_inject/tp_plan_converter.py
index c972764b2e33..761f43ce8e82 100644
--- a/deepspeed/module_inject/tp_plan_converter.py
+++ b/deepspeed/module_inject/tp_plan_converter.py
@@ -9,7 +9,7 @@
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise"}
+SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce"}
 # `colwise_rep` was renamed to `colwise_gather_output` in huggingface/transformers#42809.
 
 
@@ -20,28 +20,46 @@ class TPPlanConverter:
     def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
         """Convert HF tp_plan to DeepSpeed layer specs.
 
-        Returns None if the plan contains any unsupported partition styles,
-        allowing the caller to fall back to the existing AutoTP path.
+        Entries whose style is not supported are converted to SKIP specs instead of invalidating
+        the whole plan. Discarding the plan used to send models like Llama4 or Qwen3 down the
+        heuristic path, which shards by name patterns alone and has no notion of the modules the
+        plan deliberately excluded — on Llama4 it wraps the MoE router, whose forward returns a
+        tuple, and breaks the model. A SKIP spec keeps such layers untouched on purpose while the
+        supported entries are still applied.
+
+        Returns None only when no entry is convertible, so the caller can fall back to the
+        existing AutoTP path for models whose plan gives us nothing to work with.
         """
         unsupported = {style for style in hf_tp_plan.values() if style.lower() not in SUPPORTED_STYLES}
         if unsupported:
             logger.warning(
                 "HuggingFace tp_plan contains unsupported partition style(s): %s. "
-                "Falling back to AutoTP preset-based partitioning.", sorted(unsupported))
-            return None
+                "Layers with these styles are left unpartitioned; the remaining entries are still applied.",
+                sorted(unsupported))
 
         layer_specs = []
+        convertible_entries = 0
 
         for pattern, partition in hf_tp_plan.items():
             regex_pattern = TPPlanConverter._wildcard_to_regex(pattern)
             partition_style = partition.lower()
             gather_output = False
+            grad_allreduce = False
 
             if partition_style in ("colwise", "colwise_rep", "colwise_gather_output"):
                 partition_type = PartitionType.COLUMN
                 gather_output = partition_style != "colwise"
+                convertible_entries += 1
             elif partition_style == "rowwise":
                 partition_type = PartitionType.ROW
+                convertible_entries += 1
+            elif partition_style == "replicated_with_grad_allreduce":
+                # The parameter stays whole; only its gradient needs summing across the group.
+                partition_type = PartitionType.SKIP
+                grad_allreduce = True
+                convertible_entries += 1
+            else:
+                partition_type = PartitionType.SKIP
 
             # Only add .weight suffix if not already present
             if not regex_pattern.endswith(r"\.weight"):
@@ -54,8 +72,14 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
                     patterns=[regex_pattern],
                     partition_type=partition_type,
                     gather_output=gather_output,
+                    grad_allreduce=grad_allreduce,
                 ))
 
+        if convertible_entries == 0:
+            logger.warning("HuggingFace tp_plan has no convertible entries; styles=%s.",
+                           sorted(set(hf_tp_plan.values())))
+            return None
+
         return layer_specs
 
     @staticmethod
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 2ee87728c9a0..20dc63ff32b2 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -757,6 +757,7 @@ def lm_head_entries(tp_plan):
             autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
             autotp.update_linear_policies()
             autotp._replace_module(model)
+            autotp.register_replicated_grad_hooks(model)
             setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model))
             setattr(model, "ds_autotp_parsed", True)
             return
@@ -795,6 +796,7 @@ def lm_head_entries(tp_plan):
                 autotp.set_tensor_parallel_config(tp_size, tp_config.tensor_parallel.tp_group)
                 autotp.update_linear_policies()
                 autotp._replace_module(model)
+                autotp.register_replicated_grad_hooks(model)
                 setattr(model, UNIVERSAL_CHECKPOINT_INFO, collect_autotp_universal_checkpoint_info(model))
                 setattr(model, "ds_autotp_parsed", True)
                 return
diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md
index fb4294d1e72b..16e8c4c0326c 100755
--- a/docs/_pages/config-json.md
+++ b/docs/_pages/config-json.md
@@ -2165,9 +2165,31 @@ DeepSpeed provides compiler-based optimization passes through the `compile` conf
 
 <i>**passes**</i>: [array of strings]
 
-| Description                                                              | Default |
-| ------------------------------------------------------------------------ | ------- |
-| List of compiler passes to apply. Currently supported: `["autosp"]`.     | `[]`    |
+| Description                                                                       | Default |
+| ----------------------------------------------------------------------------------- | ------- |
+| List of compiler passes to apply. Currently supported: `["autosp", "autotp"]`.    | `[]`    |
+
+### AutoTP options
+
+The `autotp` pass emits AutoTP's tensor-parallel collectives into the compiled graph instead of
+running them from inside the injected `LinearLayer` / `LinearAllreduce` modules. The model is
+partitioned by the regular AutoTP path, so `tensor_parallel.autotp_size` must be greater than 1
+and the pass reuses the same tensor-parallel group.
+
+```json
+{
+    "zero_optimization": {"stage": 0},
+    "tensor_parallel": {"autotp_size": 4},
+    "compile": {
+        "deepcompile": true,
+        "passes": ["autotp"],
+    }
+}
+```
+
+The pass compiles the module with `fullgraph=True`, because a graph break would leave part of the
+model without the collectives it suppressed at the module level. It cannot yet be combined with the
+ZeRO passes (`z1`, `z3`) or with `autosp`, and it does not support `tensor_parallel.tp_overlap_comm`.
 
 ### Data Type options
 
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index b42e65b17cb4..0da686b798a8 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -53,7 +53,7 @@ def forward(self, x):
 
 def build_config(tp_size, use_compile_pass, gather_output_head=False):
     head_spec = {
-        "patterns": [".*\\.head\\.weight$"],
+        "patterns": ["(.*\\.)?head\\.weight$"],
         "partition_type": "column",
         "gather_output": gather_output_head,
     }
@@ -112,13 +112,14 @@ class TestAutoTPCompileEquivalence(DistributedTest):
     non_daemonic_procs = True
 
     @pytest.mark.sequential
-    def test_matches_module_injection(self):
+    @pytest.mark.parametrize("gather_output_head", [False, True])
+    def test_matches_module_injection(self, gather_output_head):
         if get_accelerator().device_name() == "cpu":
             pytest.skip("CPU does not support this test yet")
 
         device = torch.device(get_accelerator().current_device_name())
-        reference_engine = build_engine(self.world_size, use_compile_pass=False)
-        compiled_engine = build_engine(self.world_size, use_compile_pass=True)
+        reference_engine = build_engine(self.world_size, use_compile_pass=False, gather_output_head=gather_output_head)
+        compiled_engine = build_engine(self.world_size, use_compile_pass=True, gather_output_head=gather_output_head)
 
         # The TP group must see identical inputs on every rank.
         torch.manual_seed(1234)
@@ -130,6 +131,12 @@ def test_matches_module_injection(self):
         assert torch.allclose(reference_out, compiled_out, atol=1e-5), \
             "AutoTP compile pass changed the forward result"
 
+        # Comparing the two paths cannot catch a gather that both of them dropped, so the width of
+        # the head output is checked against the partitioning it was configured with.
+        expected_head_width = HIDDEN_DIM if gather_output_head else HIDDEN_DIM // self.world_size
+        assert compiled_out.shape[-1] == expected_head_width, \
+            f"Expected a head output of width {expected_head_width}, got {compiled_out.shape[-1]}"
+
         reference_engine.backward(reference_out.sum())
         compiled_engine.backward(compiled_out.sum())
 
@@ -144,6 +151,209 @@ def test_matches_module_injection(self):
             "AutoTP compile pass changed the gradient reaching the model input"
 
 
+class FusedQKVBlock(torch.nn.Module):
+    """Attention-shaped block whose column-parallel projection is a shaped sub-param layer.
+
+    AutoTP injects a fused QKV projection as SubParamLinearLayer rather than LinearLayer, which is
+    the case an exact-type check in the pass would miss.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.qkv_proj = torch.nn.Linear(HIDDEN_DIM, 3 * HIDDEN_DIM, bias=False)
+        self.o_proj = torch.nn.Linear(HIDDEN_DIM, HIDDEN_DIM, bias=False)
+
+    def forward(self, x):
+        query, key, value = self.qkv_proj(x).chunk(3, dim=-1)
+        return x + self.o_proj(query * key + value)
+
+
+class FusedQKVModel(torch.nn.Module):
+
+    def __init__(self, nlayers=2):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([FusedQKVBlock() for _ in range(nlayers)])
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+def build_fused_qkv_engine(tp_size, use_compile_pass):
+    config = {
+        "train_micro_batch_size_per_gpu": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-6
+            }
+        },
+        "tensor_parallel": {
+            "autotp_size": tp_size,
+            "partition_config": {
+                "use_default_specs":
+                False,
+                "layer_specs": [{
+                    "patterns": [".*\\.qkv_proj\\.weight$"],
+                    "partition_type": "column",
+                    "shape": [3, -1],
+                    "partition_dim": 0,
+                }, {
+                    "patterns": [".*\\.o_proj\\.weight$"],
+                    "partition_type": "row",
+                }],
+            },
+        },
+        "zero_optimization": {
+            "stage": 0
+        },
+    }
+    if use_compile_pass:
+        config["compile"] = {"deepcompile": True, "passes": ["autotp"]}
+
+    torch.manual_seed(42)
+    model = FusedQKVModel()
+    engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config)
+    if use_compile_pass:
+        engine.compile()
+    return engine
+
+
+class TestAutoTPCompileLayerVariants(DistributedTest):
+    """Every injected tensor-parallel layer variant has to be rewritten, not just the base classes.
+
+    AutoTP injects a family of variants per partitioning style (fused QKV, conv, packed gate/up,
+    Yuan, shaped sub-params). A variant left on the module-level path is not merely unoptimized:
+    the pass compiles with fullgraph=True, and ColumnParallel's forward is an identity, so tracing
+    folds it away and drops its backward all-reduce. The forward still matches and only the
+    gradients are wrong.
+    """
+
+    world_size = 2
+    non_daemonic_procs = True
+
+    @pytest.mark.sequential
+    def test_fused_qkv_matches_module_injection(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU does not support this test yet")
+
+        device = torch.device(get_accelerator().current_device_name())
+        reference_engine = build_fused_qkv_engine(self.world_size, use_compile_pass=False)
+        compiled_engine = build_fused_qkv_engine(self.world_size, use_compile_pass=True)
+
+        from deepspeed.module_inject.layers import SubParamLinearLayer
+        qkv = compiled_engine.module.layers[0].qkv_proj
+        assert isinstance(qkv, SubParamLinearLayer), f"expected a shaped sub-param layer, got {type(qkv).__name__}"
+        assert qkv.defer_collectives_to_compiler, "the fused QKV layer was left on the module-level path"
+
+        torch.manual_seed(1234)
+        x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32, requires_grad=True)
+        compiled_x = x.detach().clone().requires_grad_(True)
+
+        reference_out = reference_engine(x)
+        compiled_out = compiled_engine(compiled_x)
+        assert torch.allclose(reference_out, compiled_out, atol=1e-5)
+
+        reference_engine.backward(reference_out.sum())
+        compiled_engine.backward(compiled_out.sum())
+
+        for (name, reference_param), (_, compiled_param) in zip(reference_engine.module.named_parameters(),
+                                                                compiled_engine.module.named_parameters()):
+            assert torch.allclose(reference_param.grad, compiled_param.grad, atol=1e-5), \
+                f"AutoTP compile pass changed the gradient of {name}"
+
+        assert torch.allclose(x.grad, compiled_x.grad, atol=1e-5), \
+            "AutoTP compile pass changed the gradient reaching the model input"
+
+
+class TestAutoTPCompileMoE(DistributedTest):
+    """A mixture-of-experts model must survive the pass.
+
+    MoE models are the reason the tp_plan converter has to skip unknown styles per entry rather
+    than discard the plan: their expert and router entries use styles AutoTP does not implement,
+    while their attention entries are ordinary colwise/rowwise. The experts themselves stay
+    replicated, so the pass only has to leave them alone and rewrite the attention projections.
+    """
+
+    world_size = 2
+    non_daemonic_procs = True
+
+    @pytest.mark.sequential
+    def test_mixtral_matches_module_injection(self):
+        if get_accelerator().device_name() == "cpu":
+            pytest.skip("CPU does not support this test yet")
+        transformers = pytest.importorskip("transformers")
+        if not hasattr(transformers, "MixtralForCausalLM"):
+            pytest.skip("transformers build has no Mixtral")
+
+        def build(use_compile_pass):
+            config = transformers.MixtralConfig(vocab_size=256,
+                                                hidden_size=HIDDEN_DIM,
+                                                intermediate_size=INTERMEDIATE_DIM,
+                                                num_hidden_layers=2,
+                                                num_attention_heads=8,
+                                                num_key_value_heads=8,
+                                                num_local_experts=4,
+                                                num_experts_per_tok=2,
+                                                max_position_embeddings=64,
+                                                use_cache=False,
+                                                tie_word_embeddings=False)
+            config._attn_implementation = "sdpa"
+            # The default eager experts route tokens with data-dependent indexing, which a full
+            # graph cannot capture; batched_mm is the static-shape implementation.
+            config._experts_implementation = "batched_mm"
+            torch.manual_seed(42)
+            model = transformers.MixtralForCausalLM(config)
+            ds_config = {
+                "train_micro_batch_size_per_gpu": 1,
+                "optimizer": {
+                    "type": "Adam",
+                    "params": {
+                        "lr": 1e-6
+                    }
+                },
+                "tensor_parallel": {
+                    "autotp_size": self.world_size
+                },
+                "zero_optimization": {
+                    "stage": 0
+                },
+            }
+            if use_compile_pass:
+                ds_config["compile"] = {"deepcompile": True, "passes": ["autotp"]}
+            engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
+            if use_compile_pass:
+                engine.compile()
+            return engine
+
+        device = torch.device(get_accelerator().current_device_name())
+        reference_engine = build(False)
+        compiled_engine = build(True)
+
+        torch.manual_seed(1234)
+        input_ids = torch.randint(0, 256, (1, 16), device=device)
+        labels = torch.randint(0, 256, (1, 16), device=device)
+
+        def step(engine):
+            logits = engine(input_ids=input_ids, use_cache=False).logits
+            loss = torch.nn.functional.cross_entropy(logits.reshape(-1, 256).float(), labels.reshape(-1))
+            engine.backward(loss)
+            return loss
+
+        reference_loss = step(reference_engine)
+        compiled_loss = step(compiled_engine)
+        assert torch.allclose(reference_loss, compiled_loss, atol=1e-5), \
+            f"MoE loss changed: {reference_loss.item()} vs {compiled_loss.item()}"
+
+        for (name, reference_param), (_, compiled_param) in zip(reference_engine.module.named_parameters(),
+                                                                compiled_engine.module.named_parameters()):
+            if reference_param.grad is None or compiled_param.grad is None:
+                continue
+            assert torch.allclose(reference_param.grad, compiled_param.grad, atol=1e-5), \
+                f"AutoTP compile pass changed the gradient of {name}"
+
+
 class TestAutoTPCompileDataParallelGradients(DistributedTest):
     """Gradients must still be reduced across data-parallel replicas.
 
@@ -184,25 +394,6 @@ def test_gradients_are_reduced_across_dp_group(self):
                 f"Gradient of {name} was not reduced across the data-parallel group"
 
 
-class TestAutoTPCompileRejectsGatherOutput(DistributedTest):
-    """gather_output layers must be rejected instead of silently losing a collective.
-
-    Their gather reads shard sizes back into Python, which the full graph the pass needs cannot
-    capture, so the pass can neither emit the collectives nor leave them to the module.
-    """
-
-    world_size = 2
-    non_daemonic_procs = True
-
-    @pytest.mark.sequential
-    def test_gather_output_raises(self):
-        if get_accelerator().device_name() == "cpu":
-            pytest.skip("CPU does not support this test yet")
-
-        with pytest.raises(NotImplementedError, match="gather_output"):
-            build_engine(self.world_size, use_compile_pass=True, gather_output_head=True)
-
-
 class TestAutoTPCompileRejectsUnsupportedCombinations(DistributedTest):
 
     world_size = 1
diff --git a/tests/unit/model_parallelism/test_tp_plan_e2e.py b/tests/unit/model_parallelism/test_tp_plan_e2e.py
index 9ad2ee81368e..a5573678339e 100644
--- a/tests/unit/model_parallelism/test_tp_plan_e2e.py
+++ b/tests/unit/model_parallelism/test_tp_plan_e2e.py
@@ -303,3 +303,82 @@ def test_tp_plan_with_zero2(self):
         loss = output.mean()
         engine.backward(loss)
         engine.step()
+
+
+class TestReplicatedGradAllreduce(DistributedTest):
+    """Replicated parameters marked replicated_with_grad_allreduce must not drift across TP ranks.
+
+    Qwen3's q_norm/k_norm normalize the locally sharded attention heads, so each rank computes only
+    its shard's contribution to their gradient. Before the style was implemented the whole tp_plan
+    was discarded and the partial gradients were silently left unsummed: training ran, the loss
+    went down, and the ranks held different weights after the first optimizer step.
+    """
+
+    world_size = 2
+    non_daemonic_procs = True
+
+    def test_qwen3_norm_grads_match_across_ranks(self):
+        if get_accelerator().device_name() == "cpu":
+            import pytest
+            pytest.skip("CPU does not support this test yet")
+        import pytest
+        transformers = pytest.importorskip("transformers")
+        if not hasattr(transformers, "Qwen3ForCausalLM"):
+            pytest.skip("transformers build has no Qwen3")
+
+        config = transformers.Qwen3Config(vocab_size=256,
+                                          hidden_size=64,
+                                          intermediate_size=128,
+                                          num_hidden_layers=2,
+                                          num_attention_heads=8,
+                                          num_key_value_heads=8,
+                                          head_dim=8,
+                                          max_position_embeddings=64,
+                                          use_cache=False,
+                                          tie_word_embeddings=False)
+        config._attn_implementation = "sdpa"
+        torch.manual_seed(42)
+        model = transformers.Qwen3ForCausalLM(config)
+
+        ds_config = {
+            "train_micro_batch_size_per_gpu": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-6
+                }
+            },
+            "tensor_parallel": {
+                "autotp_size": self.world_size
+            },
+            "zero_optimization": {
+                "stage": 0
+            },
+        }
+        engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=ds_config)
+
+        norm_params = [name for name, _ in engine.module.named_parameters() if "q_norm" in name or "k_norm" in name]
+        assert norm_params, "the model should expose q_norm/k_norm parameters"
+        registered = [
+            name for name, param in engine.module.named_parameters()
+            if getattr(param, "_ds_grad_allreduce_registered", False)
+        ]
+        assert set(norm_params) <= set(registered), \
+            f"grad all-reduce hooks missing for {sorted(set(norm_params) - set(registered))}"
+
+        device = torch.device(get_accelerator().current_device_name())
+        torch.manual_seed(1234)
+        input_ids = torch.randint(0, 256, (1, 16), device=device)
+        logits = engine(input_ids=input_ids, use_cache=False).logits
+        engine.backward(logits.float().pow(2).mean())
+
+        tp_group = groups.get_tensor_model_parallel_group()
+        tp_world = dist.get_world_size(group=tp_group)
+        for name, param in engine.module.named_parameters():
+            if name not in norm_params:
+                continue
+            gathered = [torch.empty_like(param.grad) for _ in range(tp_world)]
+            dist.all_gather(gathered, param.grad.contiguous(), group=tp_group)
+            for other in gathered[1:]:
+                assert torch.allclose(gathered[0], other, atol=1e-6), \
+                    f"gradient of {name} differs across tensor-parallel ranks"
diff --git a/tests/unit/module_inject/test_tp_plan_converter.py b/tests/unit/module_inject/test_tp_plan_converter.py
index f04b5be00019..c2b4ea2a161d 100644
--- a/tests/unit/module_inject/test_tp_plan_converter.py
+++ b/tests/unit/module_inject/test_tp_plan_converter.py
@@ -69,10 +69,9 @@ def test_pattern_weight_suffix_already_present(self):
         assert specs[0].patterns[0].endswith(r"\.weight$")
 
     def test_empty_plan(self):
-        hf_plan = {}
-        specs = TPPlanConverter.convert(hf_plan)
-
-        assert len(specs) == 0
+        """An empty plan has nothing to convert, so the caller should fall back rather than
+        silently partition nothing."""
+        assert TPPlanConverter.convert({}) is None
 
     def test_multiple_patterns(self):
         hf_plan = {
@@ -109,11 +108,25 @@ def test_pattern_matches_param_name(self):
 
         assert re.match(down_pattern.patterns[0], "model.layers.5.mlp.down_proj.weight")
 
-    def test_unsupported_style_returns_none(self):
-        """Unsupported styles cause convert() to return None for fallback."""
+    def test_unsupported_style_is_skipped_not_fatal(self):
+        """One unknown style must not discard the entries next to it.
+
+        Dropping the whole plan sends the model down the heuristic path, which shards by name
+        alone and has no notion of the modules the plan deliberately excluded. On Llama4 that
+        wraps the MoE router, whose forward returns a tuple, and breaks the model outright.
+        """
         hf_plan = {"layers.*.q_proj": "local_colwise", "layers.*.o_proj": "rowwise"}
-        result = TPPlanConverter.convert(hf_plan)
-        assert result is None
+        specs = TPPlanConverter.convert(hf_plan)
+
+        assert specs is not None
+        by_type = {spec.partition_type for spec in specs}
+        assert PartitionType.ROW in by_type, "the supported entry should still be applied"
+        assert PartitionType.SKIP in by_type, "the unsupported entry should be left unpartitioned"
+
+    def test_all_styles_unsupported_returns_none(self):
+        """With nothing convertible there is no plan to apply, so fall back."""
+        hf_plan = {"layers.*.q_proj": "local_colwise", "layers.*.experts": "moe_tp_experts"}
+        assert TPPlanConverter.convert(hf_plan) is None
 
     def test_alternate_prefixes(self):
         """Test tp_plan with non-layers prefix"""

From 4738778e9f1fc604d0e964b077f70899520ebc3b Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Sun, 9 Aug 2026 00:26:41 +0000
Subject: [PATCH 05/11] cleaning up

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 .../compile/custom_ops/tp_collectives.py      | 27 ++++++-------------
 deepspeed/compile/passes/tp_compile.py        | 27 ++++---------------
 deepspeed/module_inject/layers.py             |  3 ---
 deepspeed/runtime/engine.py                   | 13 +--------
 docs/_pages/config-json.md                    | 16 +++++------
 5 files changed, 21 insertions(+), 65 deletions(-)

diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py
index e937cc4881f6..a71509759bec 100644
--- a/deepspeed/compile/custom_ops/tp_collectives.py
+++ b/deepspeed/compile/custom_ops/tp_collectives.py
@@ -9,11 +9,7 @@
 
 
 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 the tensor-parallel group created by the existing AutoTP setup."""
     return groups.get_tensor_model_parallel_group()
 
 
@@ -21,7 +17,7 @@ def get_tp_group():
 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
+    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.
     """
@@ -36,8 +32,8 @@ def copy_to_tp_region_fake(input: torch.Tensor):
 @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
+    
+    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()
@@ -54,10 +50,8 @@ def reduce_from_tp_region_fake(input: torch.Tensor):
 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.
 
-    This is 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. AutoTP only
-    builds such a layer when the output dimension divides evenly by the TP size, so every shard has
-    the same width and the sizes are known statically.
+    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)
@@ -69,8 +63,6 @@ def gather_from_tp_region(input: torch.Tensor) -> torch.Tensor:
                                 dtype=local_shard.dtype,
                                 device=local_shard.device)
     dist.all_gather_into_tensor(flat_gathered, local_shard, group=group)
-    # The gather stacks whole shards along dim 0, but the partitioning split the last dimension,
-    # so the shards are re-joined there in rank order to rebuild the unpartitioned output.
     shards = flat_gathered.view(world_size, *local_shard.shape)
     return torch.cat(shards.unbind(0), dim=-1)
 
@@ -82,7 +74,8 @@ def gather_from_tp_region_fake(input: torch.Tensor):
 
 
 def _copy_to_tp_region_backward(ctx, grad):
-    # f and g are duals, so f's backward is simply g.
+    # 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())
 
 
@@ -91,8 +84,6 @@ def _reduce_from_tp_region_backward(ctx, grad):
 
 
 def _gather_from_tp_region_backward(ctx, grad):
-    # The forward concatenated the shards in rank order, so each rank owns a contiguous slice of
-    # the gradient and no communication is needed to recover it.
     group = get_tp_group()
     world_size = dist.get_world_size(group=group)
     if world_size == 1:
@@ -103,8 +94,6 @@ def _gather_from_tp_region_backward(ctx, grad):
 
 
 def _setup_context_without_saved_tensors(ctx, inputs, output):
-    # The collectives are stateless and their shapes are fixed by the TP size, so their backwards
-    # need nothing saved.
     pass
 
 
diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
index 46a79b7e345b..3152f59afddc 100644
--- a/deepspeed/compile/passes/tp_compile.py
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -18,22 +18,14 @@
 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. Reading it back is more robust than
-# re-deriving column/row from parameter-name patterns.
-#
-# The families are matched by subclass rather than by exact type. AutoTP injects several variants
-# per family (fused QKV, conv, packed gate/up, Yuan), and they either inherit the base forward or
-# repeat its shape, so the pass rewrites all of them identically. Matching exact types instead
-# leaves the variants on the module-level path, which is not merely unoptimized but wrong: see the
-# comment in defer_collectives_to_compiler.
+# 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, )
 
-# 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,
@@ -60,10 +52,7 @@ def defer_collectives_to_compiler(model) -> None:
     """Suppress the module-level TP collectives on layers this pass will handle in the graph.
 
     Any tensor-parallel layer the pass cannot rewrite is rejected rather than left on the
-    module-level path. That path looks like the safe fallback but is not: the pass compiles with
-    fullgraph=True, and ColumnParallel's forward is a plain identity, so tracing folds it away and
-    takes its backward all-reduce with it. The forward still matches and only the gradients are
-    wrong, which is the worst way for this to fail.
+    module-level path.
     """
     for name, module in model.named_modules():
         if not isinstance(module, TensorParallel_Layer) or module.mp_group is None:
@@ -107,9 +96,7 @@ def _insert_row_collective(gm: GraphModule, matmul: Node) -> 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.
-    """
+    """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")
@@ -155,11 +142,7 @@ def pass_canonicalize(gm: GraphModule, real_inputs):
 
 
 def apply_autotp(gm: GraphModule, real_inputs, passes=None):
-    """Apply the AutoTP transformation passes to the graph.
-
-    The inserted collectives are shape-preserving (the shape-changing gather is emitted by the
-    layer forwards during capture), so unlike AutoSP this needs no shape re-propagation.
-    """
+    """Apply the AutoTP transformation passes to the graph."""
     for opt_pass in passes or AUTOTP_PASSES:
         opt_pass(gm, real_inputs)
     return gm
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 95d22fb5b2d9..6bde981b0b5c 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -305,9 +305,6 @@ 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:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 20dc63ff32b2..70b311a3a09b 100755
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1221,12 +1221,7 @@ def compile_autotp(self):
         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.
-        """
+        """Determines if the compiled graph comes from a parallelization pass rather than ZeRO."""
         return self.compile_autosp() or self.compile_autotp()
 
     def mics_shard_size(self):
@@ -2883,8 +2878,6 @@ 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"
 
-        # 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())
 
@@ -5500,8 +5493,6 @@ def get_autotp_backend(self, compile_kwargs):
                         "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
@@ -5509,8 +5500,6 @@ def get_autotp_backend(self, compile_kwargs):
                            "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)
 
diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md
index 16e8c4c0326c..2e53709d378f 100755
--- a/docs/_pages/config-json.md
+++ b/docs/_pages/config-json.md
@@ -2163,12 +2163,6 @@ DeepSpeed provides compiler-based optimization passes through the `compile` conf
 }
 ```
 
-<i>**passes**</i>: [array of strings]
-
-| Description                                                                       | Default |
-| ----------------------------------------------------------------------------------- | ------- |
-| List of compiler passes to apply. Currently supported: `["autosp", "autotp"]`.    | `[]`    |
-
 ### AutoTP options
 
 The `autotp` pass emits AutoTP's tensor-parallel collectives into the compiled graph instead of
@@ -2187,9 +2181,13 @@ and the pass reuses the same tensor-parallel group.
 }
 ```
 
-The pass compiles the module with `fullgraph=True`, because a graph break would leave part of the
-model without the collectives it suppressed at the module level. It cannot yet be combined with the
-ZeRO passes (`z1`, `z3`) or with `autosp`, and it does not support `tensor_parallel.tp_overlap_comm`.
+<i>**passes**</i>: [array of strings]
+
+| Description                                                                       | Default |
+| ----------------------------------------------------------------------------------- | ------- |
+| List of compiler passes to apply. Currently supported: `["autosp", "autotp"]`.    | `[]`    |
+
+
 
 ### Data Type options
 

From 40d92d6cf969d291d32645c9682677f3e75ef45f Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Tue, 11 Aug 2026 18:43:51 -0400
Subject: [PATCH 06/11] fixed for partial compiler flags, now raises error

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/passes/tp_compile.py        | 10 +-
 deepspeed/module_inject/layers.py             |  5 +-
 deepspeed/module_inject/tp_plan_converter.py  | 27 ++----
 tests/unit/compile/test_tp_compile.py         | 93 +++++++++++++++----
 .../module_inject/test_tp_plan_converter.py   | 26 +++---
 5 files changed, 104 insertions(+), 57 deletions(-)

diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
index 3152f59afddc..62cc00b85735 100644
--- a/deepspeed/compile/passes/tp_compile.py
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -49,11 +49,8 @@ def _in_family(layer_type, family) -> bool:
 
 
 def defer_collectives_to_compiler(model) -> None:
-    """Suppress the module-level TP collectives on layers this pass will handle in the graph.
-
-    Any tensor-parallel layer the pass cannot rewrite is rejected rather than left on the
-    module-level path.
-    """
+    """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
@@ -69,6 +66,9 @@ def defer_collectives_to_compiler(model) -> None:
             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
 
 
diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 6bde981b0b5c..3271de8b805c 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -1305,7 +1305,7 @@ def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
         self._mark_uc_metadata()
 
     def forward(self, input):
-        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:
@@ -1425,7 +1425,8 @@ def __init__(self, module, mp_group, shape, partition_dim=1, **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
diff --git a/deepspeed/module_inject/tp_plan_converter.py b/deepspeed/module_inject/tp_plan_converter.py
index 761f43ce8e82..17579b6788a2 100644
--- a/deepspeed/module_inject/tp_plan_converter.py
+++ b/deepspeed/module_inject/tp_plan_converter.py
@@ -3,12 +3,9 @@
 
 # DeepSpeed Team
 
-import logging
 from typing import List, Dict, Optional
 from .autotp_config import TPLayerSpec, PartitionType
 
-logger = logging.getLogger(__name__)
-
 SUPPORTED_STYLES = {"colwise", "colwise_rep", "colwise_gather_output", "rowwise", "replicated_with_grad_allreduce"}
 # `colwise_rep` was renamed to `colwise_gather_output` in huggingface/transformers#42809.
 
@@ -30,15 +27,17 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
         Returns None only when no entry is convertible, so the caller can fall back to the
         existing AutoTP path for models whose plan gives us nothing to work with.
         """
+        if not hf_tp_plan:
+            return None
+
         unsupported = {style for style in hf_tp_plan.values() if style.lower() not in SUPPORTED_STYLES}
         if unsupported:
-            logger.warning(
-                "HuggingFace tp_plan contains unsupported partition style(s): %s. "
-                "Layers with these styles are left unpartitioned; the remaining entries are still applied.",
-                sorted(unsupported))
+            raise ValueError(f"HuggingFace tp_plan contains unsupported partition style(s): {sorted(unsupported)}. "
+                             "Applying only the supported entries could shard one half of a column/row pair, so the "
+                             "plan is rejected as a whole. Provide an explicit 'tensor_parallel.partition_config' "
+                             "for this model instead.")
 
         layer_specs = []
-        convertible_entries = 0
 
         for pattern, partition in hf_tp_plan.items():
             regex_pattern = TPPlanConverter._wildcard_to_regex(pattern)
@@ -49,17 +48,12 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
             if partition_style in ("colwise", "colwise_rep", "colwise_gather_output"):
                 partition_type = PartitionType.COLUMN
                 gather_output = partition_style != "colwise"
-                convertible_entries += 1
             elif partition_style == "rowwise":
                 partition_type = PartitionType.ROW
-                convertible_entries += 1
-            elif partition_style == "replicated_with_grad_allreduce":
+            else:  # replicated_with_grad_allreduce, the only other supported style
                 # The parameter stays whole; only its gradient needs summing across the group.
                 partition_type = PartitionType.SKIP
                 grad_allreduce = True
-                convertible_entries += 1
-            else:
-                partition_type = PartitionType.SKIP
 
             # Only add .weight suffix if not already present
             if not regex_pattern.endswith(r"\.weight"):
@@ -75,11 +69,6 @@ def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
                     grad_allreduce=grad_allreduce,
                 ))
 
-        if convertible_entries == 0:
-            logger.warning("HuggingFace tp_plan has no convertible entries; styles=%s.",
-                           sorted(set(hf_tp_plan.values())))
-            return None
-
         return layer_specs
 
     @staticmethod
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index 0da686b798a8..bbbe623cd333 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -152,10 +152,12 @@ def test_matches_module_injection(self, gather_output_head):
 
 
 class FusedQKVBlock(torch.nn.Module):
-    """Attention-shaped block whose column-parallel projection is a shaped sub-param layer.
+    """Attention-shaped block whose projections are both shaped sub-param layers.
 
     AutoTP injects a fused QKV projection as SubParamLinearLayer rather than LinearLayer, which is
-    the case an exact-type check in the pass would miss.
+    the case an exact-type check in the pass would miss. The output projection is likewise given a
+    shape so it is injected as SubParamLinearAllreduce, whose forward must give up its module-level
+    all-reduce once the pass emits one into the graph — keeping both would reduce twice.
     """
 
     def __init__(self):
@@ -194,15 +196,22 @@ def build_fused_qkv_engine(tp_size, use_compile_pass):
             "partition_config": {
                 "use_default_specs":
                 False,
-                "layer_specs": [{
-                    "patterns": [".*\\.qkv_proj\\.weight$"],
-                    "partition_type": "column",
-                    "shape": [3, -1],
-                    "partition_dim": 0,
-                }, {
-                    "patterns": [".*\\.o_proj\\.weight$"],
-                    "partition_type": "row",
-                }],
+                "layer_specs": [
+                    {
+                        "patterns": [".*\\.qkv_proj\\.weight$"],
+                        "partition_type": "column",
+                        "shape": [3, -1],
+                        "partition_dim": 0,
+                    },
+                    {
+                        # A single sub-param spanning the input dim shards exactly like the plain row
+                        # split; the shape's only effect is injecting SubParamLinearAllreduce.
+                        "patterns": [".*\\.o_proj\\.weight$"],
+                        "partition_type": "row",
+                        "shape": [-1, 1],
+                        "partition_dim": 1,
+                    }
+                ],
             },
         },
         "zero_optimization": {
@@ -242,10 +251,13 @@ def test_fused_qkv_matches_module_injection(self):
         reference_engine = build_fused_qkv_engine(self.world_size, use_compile_pass=False)
         compiled_engine = build_fused_qkv_engine(self.world_size, use_compile_pass=True)
 
-        from deepspeed.module_inject.layers import SubParamLinearLayer
+        from deepspeed.module_inject.layers import SubParamLinearAllreduce, SubParamLinearLayer
         qkv = compiled_engine.module.layers[0].qkv_proj
         assert isinstance(qkv, SubParamLinearLayer), f"expected a shaped sub-param layer, got {type(qkv).__name__}"
         assert qkv.defer_collectives_to_compiler, "the fused QKV layer was left on the module-level path"
+        o_proj = compiled_engine.module.layers[0].o_proj
+        assert isinstance(o_proj, SubParamLinearAllreduce), f"expected a shaped row layer, got {type(o_proj).__name__}"
+        assert o_proj.defer_collectives_to_compiler, "the shaped row layer was left on the module-level path"
 
         torch.manual_seed(1234)
         x = torch.randn(1, 8, HIDDEN_DIM, device=device, dtype=torch.float32, requires_grad=True)
@@ -270,10 +282,11 @@ def test_fused_qkv_matches_module_injection(self):
 class TestAutoTPCompileMoE(DistributedTest):
     """A mixture-of-experts model must survive the pass.
 
-    MoE models are the reason the tp_plan converter has to skip unknown styles per entry rather
-    than discard the plan: their expert and router entries use styles AutoTP does not implement,
-    while their attention entries are ordinary colwise/rowwise. The experts themselves stay
-    replicated, so the pass only has to leave them alone and rewrite the attention projections.
+    Mixtral's expert and router entries use tp_plan styles AutoTP does not implement, and the
+    converter rejects a plan containing any unsupported style rather than applying it partially,
+    so the attention sharding is spelled out as an explicit partition config. The experts, router
+    and lm_head stay replicated, so the pass only has to leave them alone and rewrite the
+    attention projections.
     """
 
     world_size = 2
@@ -314,7 +327,23 @@ def build(use_compile_pass):
                     }
                 },
                 "tensor_parallel": {
-                    "autotp_size": self.world_size
+                    "autotp_size": self.world_size,
+                    "partition_config": {
+                        "use_default_specs":
+                        False,
+                        "layer_specs": [{
+                            "patterns": [
+                                ".*\\.self_attn\\.q_proj\\.weight$",
+                                ".*\\.self_attn\\.k_proj\\.weight$",
+                                ".*\\.self_attn\\.v_proj\\.weight$",
+                            ],
+                            "partition_type":
+                            "column",
+                        }, {
+                            "patterns": [".*\\.self_attn\\.o_proj\\.weight$"],
+                            "partition_type": "row",
+                        }],
+                    },
                 },
                 "zero_optimization": {
                     "stage": 0
@@ -394,6 +423,36 @@ def test_gradients_are_reduced_across_dp_group(self):
                 f"Gradient of {name} was not reduced across the data-parallel group"
 
 
+def _make_tp_layer(cls):
+    # Only the flag logic is under test, so the layer is built without weights or a process group.
+    layer = cls.__new__(cls)
+    torch.nn.Module.__init__(layer)
+    layer.mp_group = object()
+    layer.defer_collectives_to_compiler = False
+    return layer
+
+
+def test_defer_collectives_is_all_or_nothing():
+    """A rejected layer must leave every other layer's collectives untouched.
+
+    If the rejection is raised after some flags are already set and a caller catches it to fall
+    back to eager execution, the flagged layers would silently skip their collectives.
+    """
+    from deepspeed.compile.passes.tp_compile import defer_collectives_to_compiler
+    from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer, LmHeadLinearAllreduce
+
+    model = torch.nn.Module()
+    model.column = _make_tp_layer(LinearLayer)
+    model.row = _make_tp_layer(LinearAllreduce)
+    model.head = _make_tp_layer(LmHeadLinearAllreduce)
+
+    with pytest.raises(NotImplementedError, match="cannot rewrite"):
+        defer_collectives_to_compiler(model)
+
+    assert not model.column.defer_collectives_to_compiler
+    assert not model.row.defer_collectives_to_compiler
+
+
 class TestAutoTPCompileRejectsUnsupportedCombinations(DistributedTest):
 
     world_size = 1
diff --git a/tests/unit/module_inject/test_tp_plan_converter.py b/tests/unit/module_inject/test_tp_plan_converter.py
index c2b4ea2a161d..f8446d60980d 100644
--- a/tests/unit/module_inject/test_tp_plan_converter.py
+++ b/tests/unit/module_inject/test_tp_plan_converter.py
@@ -3,6 +3,8 @@
 
 # DeepSpeed Team
 
+import pytest
+
 from deepspeed.module_inject.tp_plan_converter import TPPlanConverter
 from deepspeed.module_inject.autotp_config import AutoTPConfig, PartitionType
 
@@ -108,25 +110,21 @@ def test_pattern_matches_param_name(self):
 
         assert re.match(down_pattern.patterns[0], "model.layers.5.mlp.down_proj.weight")
 
-    def test_unsupported_style_is_skipped_not_fatal(self):
-        """One unknown style must not discard the entries next to it.
+    def test_unsupported_style_rejects_whole_plan(self):
+        """Any unknown style must reject the plan outright, not skip the entry.
 
-        Dropping the whole plan sends the model down the heuristic path, which shards by name
-        alone and has no notion of the modules the plan deliberately excluded. On Llama4 that
-        wraps the MoE router, whose forward returns a tuple, and breaks the model outright.
+        The converter has no notion of which entries pair with which, so applying the supported
+        subset can shard one half of a column/row pair — here o_proj would be row-sharded while
+        the q_proj feeding it stays whole — and break the model. Failing loudly also keeps the
+        model off the heuristic path, which breaks models like Llama4 by wrapping its MoE router.
         """
         hf_plan = {"layers.*.q_proj": "local_colwise", "layers.*.o_proj": "rowwise"}
-        specs = TPPlanConverter.convert(hf_plan)
-
-        assert specs is not None
-        by_type = {spec.partition_type for spec in specs}
-        assert PartitionType.ROW in by_type, "the supported entry should still be applied"
-        assert PartitionType.SKIP in by_type, "the unsupported entry should be left unpartitioned"
+        with pytest.raises(ValueError, match="local_colwise"):
+            TPPlanConverter.convert(hf_plan)
 
-    def test_all_styles_unsupported_returns_none(self):
-        """With nothing convertible there is no plan to apply, so fall back."""
         hf_plan = {"layers.*.q_proj": "local_colwise", "layers.*.experts": "moe_tp_experts"}
-        assert TPPlanConverter.convert(hf_plan) is None
+        with pytest.raises(ValueError, match="unsupported partition style"):
+            TPPlanConverter.convert(hf_plan)
 
     def test_alternate_prefixes(self):
         """Test tp_plan with non-layers prefix"""

From a816df04a95d4efd436d7cc14181b786a9fdf205 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Tue, 11 Aug 2026 19:39:49 -0400
Subject: [PATCH 07/11] skip pytest if transformers not 5.x

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 tests/unit/compile/test_tp_compile.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index bbbe623cd333..86d28f1e9743 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -299,6 +299,10 @@ def test_mixtral_matches_module_injection(self):
         transformers = pytest.importorskip("transformers")
         if not hasattr(transformers, "MixtralForCausalLM"):
             pytest.skip("transformers build has no Mixtral")
+        # transformers 5.x wraps model forwards in a decorator that inspects __code__.co_varnames,
+        # which dynamo traces under fullgraph in torch 2.8
+        if not required_torch_version(min_version=2.8):
+            pytest.skip("tracing the transformers input-check wrapper requires torch >= 2.8")
 
         def build(use_compile_pass):
             config = transformers.MixtralConfig(vocab_size=256,

From 84ec6caaf21b79709abb6f78fdf4491929c35cf4 Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Tue, 11 Aug 2026 19:48:38 -0400
Subject: [PATCH 08/11] precommit

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/custom_ops/tp_collectives.py | 4 ++--
 deepspeed/compile/passes/tp_compile.py         | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/deepspeed/compile/custom_ops/tp_collectives.py b/deepspeed/compile/custom_ops/tp_collectives.py
index a71509759bec..7c64f4f899b5 100644
--- a/deepspeed/compile/custom_ops/tp_collectives.py
+++ b/deepspeed/compile/custom_ops/tp_collectives.py
@@ -32,7 +32,7 @@ def copy_to_tp_region_fake(input: torch.Tensor):
 @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.
     """
@@ -74,7 +74,7 @@ def gather_from_tp_region_fake(input: torch.Tensor):
 
 
 def _copy_to_tp_region_backward(ctx, grad):
-    # copy_to_tp_region and reduce_from_tp_region are duals, 
+    # 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())
 
diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py
index 62cc00b85735..a17ac02e256a 100644
--- a/deepspeed/compile/passes/tp_compile.py
+++ b/deepspeed/compile/passes/tp_compile.py
@@ -18,7 +18,7 @@
 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. 
+# already records the partitioning decision the pass needs.
 
 COLUMN_PARALLEL_LAYERS = (LinearLayer, SubParamLinearLayer)
 ROW_PARALLEL_LAYERS = (LinearAllreduce, SubParamLinearAllreduce)

From 218cf7b1d242a25da6efa1d7e7e8af544b1ba7ca Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Wed, 12 Aug 2026 07:01:13 -0400
Subject: [PATCH 09/11] added guards for test_tp_compile

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 tests/unit/compile/test_tp_compile.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index 86d28f1e9743..bbba9ff80fcf 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -299,10 +299,15 @@ def test_mixtral_matches_module_injection(self):
         transformers = pytest.importorskip("transformers")
         if not hasattr(transformers, "MixtralForCausalLM"):
             pytest.skip("transformers build has no Mixtral")
+        # Only transformers 5.x offers the static-shape batched_mm experts implementation; the 4.x
+        # eager experts route tokens through .nonzero(), a dynamic-shape op no full graph can
+        # capture (verified: 4.57 fails, 5.15 passes).
+        if int(transformers.__version__.split(".")[0]) < 5:
+            pytest.skip("the static-shape experts implementation requires transformers >= 5")
         # transformers 5.x wraps model forwards in a decorator that inspects __code__.co_varnames,
-        # which dynamo traces under fullgraph in torch 2.8
-        if not required_torch_version(min_version=2.8):
-            pytest.skip("tracing the transformers input-check wrapper requires torch >= 2.8")
+        # which dynamo only traces under fullgraph from torch 2.7 (verified: 2.6 fails, 2.7 passes).
+        if not required_torch_version(min_version=2.7):
+            pytest.skip("tracing the transformers input-check wrapper requires torch >= 2.7")
 
         def build(use_compile_pass):
             config = transformers.MixtralConfig(vocab_size=256,

From 0f71bad6f8e332330b9d4b30ecc64cc1369e8b3a Mon Sep 17 00:00:00 2001
From: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Date: Wed, 12 Aug 2026 18:38:45 -0400
Subject: [PATCH 10/11] added warning for transformers bug

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
---
 deepspeed/compile/init_tp.py          | 29 ++++++++++++++++--
 tests/unit/compile/test_tp_compile.py | 42 ++++++++++++++++++++++-----
 2 files changed, 61 insertions(+), 10 deletions(-)

diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py
index 9da3be3f9174..7ed3fe776f11 100644
--- a/deepspeed/compile/init_tp.py
+++ b/deepspeed/compile/init_tp.py
@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import torch
+from packaging.version import Version
 from torch.fx import GraphModule
 
 from deepspeed.utils.torch import required_torch_version
@@ -11,12 +12,36 @@
 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():
+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")
+    ]
+    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):
@@ -25,7 +50,7 @@ def init_autotp(model):
     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()
+    _check_autotp_compatibility(model)
     defer_collectives_to_compiler(model)
 
     def backend_fn(gm: GraphModule, real_inputs):
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index bbba9ff80fcf..c6510fd5fc95 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -5,11 +5,12 @@
 
 import pytest
 import torch
+from packaging.version import Version
 
 import deepspeed
 import deepspeed.comm as dist
 from deepspeed.accelerator import get_accelerator
-from deepspeed.compile.init_tp import AUTOTP_MIN_TORCH_VERSION
+from deepspeed.compile.init_tp import AUTOTP_MIN_TORCH_VERSION, BROKEN_TRANSFORMERS_MOE_VERSIONS
 from deepspeed.utils import groups
 from deepspeed.utils.torch import required_torch_version
 
@@ -299,13 +300,15 @@ def test_mixtral_matches_module_injection(self):
         transformers = pytest.importorskip("transformers")
         if not hasattr(transformers, "MixtralForCausalLM"):
             pytest.skip("transformers build has no Mixtral")
-        # Only transformers 5.x offers the static-shape batched_mm experts implementation; the 4.x
-        # eager experts route tokens through .nonzero(), a dynamic-shape op no full graph can
-        # capture (verified: 4.57 fails, 5.15 passes).
-        if int(transformers.__version__.split(".")[0]) < 5:
-            pytest.skip("the static-shape experts implementation requires transformers >= 5")
-        # transformers 5.x wraps model forwards in a decorator that inspects __code__.co_varnames,
-        # which dynamo only traces under fullgraph from torch 2.7 (verified: 2.6 fails, 2.7 passes).
+        tf_version = Version(transformers.__version__)
+        # The 4.x eager experts route tokens through .nonzero(), a dynamic-shape op no full graph
+        # can capture; the static-shape batched_mm implementation exists from 5.0.
+        if tf_version.major < 5:
+            pytest.skip("the static-shape batched_mm experts implementation requires transformers >= 5")
+        first_broken, first_fixed = BROKEN_TRANSFORMERS_MOE_VERSIONS
+        if Version(first_broken) <= tf_version < Version(first_fixed):
+            pytest.skip(f"transformers {first_broken}..{first_fixed} mutate the MoE routing tensor in place "
+                        "(huggingface/transformers#45621, fixed by #45634)")
         if not required_torch_version(min_version=2.7):
             pytest.skip("tracing the transformers input-check wrapper requires torch >= 2.7")
 
@@ -462,6 +465,29 @@ def test_defer_collectives_is_all_or_nothing():
     assert not model.row.defer_collectives_to_compiler
 
 
+def test_broken_transformers_moe_raises(monkeypatch):
+    """On an affected transformers release the pass must reject MoE models up front, with a
+    pointer to the fix, while leaving dense models untouched."""
+    transformers = pytest.importorskip("transformers")
+    from deepspeed.compile.init_tp import _check_broken_transformers_moe
+
+    first_broken, first_fixed = BROKEN_TRANSFORMERS_MOE_VERSIONS
+    monkeypatch.setattr(transformers, "__version__", first_broken)
+
+    moe_model = torch.nn.Module()
+    moe_model.experts = torch.nn.Module()
+    # The markers transformers' experts-implementation decorator leaves on a wrapped module.
+    moe_model.experts._apply_gate = lambda x: x
+    moe_model.experts.is_concatenated = True
+    with pytest.raises(RuntimeError, match="45634"):
+        _check_broken_transformers_moe(moe_model)
+
+    _check_broken_transformers_moe(torch.nn.Module())
+
+    monkeypatch.setattr(transformers, "__version__", first_fixed)
+    _check_broken_transformers_moe(moe_model)
+
+
 class TestAutoTPCompileRejectsUnsupportedCombinations(DistributedTest):
 
     world_size = 1

From 6cb92ef7b75fa0a7abf4a1f9a793491f1ded95e5 Mon Sep 17 00:00:00 2001
From: Masahiro Tanaka <mtanaka@anyscale.com>
Date: Wed, 12 Aug 2026 16:19:31 -0700
Subject: [PATCH 11/11] Fix AutoTP MoE compatibility guard

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
---
 deepspeed/compile/init_tp.py          |  1 +
 tests/unit/compile/test_tp_compile.py | 15 +++++++++++----
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py
index 7ed3fe776f11..c36402338699 100644
--- a/deepspeed/compile/init_tp.py
+++ b/deepspeed/compile/init_tp.py
@@ -35,6 +35,7 @@ def _check_broken_transformers_moe(model):
     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(
diff --git a/tests/unit/compile/test_tp_compile.py b/tests/unit/compile/test_tp_compile.py
index c6510fd5fc95..b17379981c99 100644
--- a/tests/unit/compile/test_tp_compile.py
+++ b/tests/unit/compile/test_tp_compile.py
@@ -3,6 +3,8 @@
 
 # DeepSpeed Team
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 from packaging.version import Version
@@ -465,9 +467,10 @@ def test_defer_collectives_is_all_or_nothing():
     assert not model.row.defer_collectives_to_compiler
 
 
-def test_broken_transformers_moe_raises(monkeypatch):
-    """On an affected transformers release the pass must reject MoE models up front, with a
-    pointer to the fix, while leaving dense models untouched."""
+@pytest.mark.parametrize("experts_implementation, should_raise", [("batched_mm", True), ("eager", False),
+                                                                  ("grouped_mm", False)])
+def test_broken_transformers_moe_raises(monkeypatch, experts_implementation, should_raise):
+    """Affected releases must reject only MoE models selecting the broken batched_mm path."""
     transformers = pytest.importorskip("transformers")
     from deepspeed.compile.init_tp import _check_broken_transformers_moe
 
@@ -479,7 +482,11 @@ def test_broken_transformers_moe_raises(monkeypatch):
     # The markers transformers' experts-implementation decorator leaves on a wrapped module.
     moe_model.experts._apply_gate = lambda x: x
     moe_model.experts.is_concatenated = True
-    with pytest.raises(RuntimeError, match="45634"):
+    moe_model.experts.config = SimpleNamespace(_experts_implementation=experts_implementation)
+    if should_raise:
+        with pytest.raises(RuntimeError, match="45634"):
+            _check_broken_transformers_moe(moe_model)
+    else:
         _check_broken_transformers_moe(moe_model)
 
     _check_broken_transformers_moe(torch.nn.Module())