From 8abe7ed428e31378132894da764f969ef0ca90b7 Mon Sep 17 00:00:00 2001 From: Naveenraj Kamalakannan Date: Sat, 1 Aug 2026 22:56:12 +0000 Subject: [PATCH 1/3] added support for autotp Signed-off-by: Naveenraj Kamalakannan --- 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 Date: Wed, 5 Aug 2026 19:50:06 -0400 Subject: [PATCH 2/3] fix for 2x residuals Signed-off-by: Naveenraj Kamalakannan --- .../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 Date: Wed, 5 Aug 2026 19:53:32 -0400 Subject: [PATCH 3/3] restore AutoTP torch version check Signed-off-by: Naveenraj Kamalakannan --- 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):