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..65e7a3070110 --- /dev/null +++ b/deepspeed/compile/custom_ops/tp_collectives.py @@ -0,0 +1,72 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +import deepspeed.comm as dist +from deepspeed.utils import groups + + +def get_tp_group(): + """Return the tensor-parallel group created by the existing AutoTP setup. + + The AutoTP pass reuses the groups that ``TpTrainingManager`` already builds, so the compiled + collectives always communicate over the same group as the module-level ones they replace. + """ + return groups.get_tensor_model_parallel_group() + + +@torch.library.custom_op("autotp::copy_to_tp_region", mutates_args=()) +def copy_to_tp_region(input: torch.Tensor) -> torch.Tensor: + """Identity in the forward pass, all-reduce in the backward pass. + + This is Megatron's ``f``. It is inserted before a column-parallel matmul: the activation is + already replicated across the tensor-parallel group, so nothing has to happen in the forward + pass, while each rank contributes a partial gradient that must be summed in the backward pass. + """ + return input.clone() + + +@torch.library.register_fake("autotp::copy_to_tp_region") +def copy_to_tp_region_fake(input: torch.Tensor): + return torch.empty_like(input) + + +@torch.library.custom_op("autotp::reduce_from_tp_region", mutates_args=()) +def reduce_from_tp_region(input: torch.Tensor) -> torch.Tensor: + """All-reduce in the forward pass, identity in the backward pass. + + This is Megatron's ``g``. It is inserted after a row-parallel matmul, whose output is only a + partial sum because each rank holds a slice of the input dimension. + """ + output = input.contiguous().clone() + dist.all_reduce(output, group=get_tp_group()) + return output + + +@torch.library.register_fake("autotp::reduce_from_tp_region") +def reduce_from_tp_region_fake(input: torch.Tensor): + return torch.empty_like(input) + + +def _copy_to_tp_region_backward(ctx, grad): + # f and g are duals, so f's backward is simply g. + return reduce_from_tp_region(grad.contiguous()) + + +def _reduce_from_tp_region_backward(ctx, grad): + return grad + + +def _setup_context_without_saved_tensors(ctx, inputs, output): + # Both collectives are shape-preserving and stateless, so their backwards need nothing saved. + pass + + +torch.library.register_autograd("autotp::copy_to_tp_region", + _copy_to_tp_region_backward, + setup_context=_setup_context_without_saved_tensors) +torch.library.register_autograd("autotp::reduce_from_tp_region", + _reduce_from_tp_region_backward, + setup_context=_setup_context_without_saved_tensors) diff --git a/deepspeed/compile/init_tp.py b/deepspeed/compile/init_tp.py new file mode 100644 index 000000000000..9da3be3f9174 --- /dev/null +++ b/deepspeed/compile/init_tp.py @@ -0,0 +1,35 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import torch +from torch.fx import GraphModule + +from deepspeed.utils.torch import required_torch_version + +from .passes.tp_compile import apply_autotp, defer_collectives_to_compiler + +AUTOTP_MIN_TORCH_VERSION = 2.6 + + +def _check_autotp_compatibility(): + if not required_torch_version(min_version=AUTOTP_MIN_TORCH_VERSION): + raise RuntimeError(f"The AutoTP compile pass requires PyTorch >= {AUTOTP_MIN_TORCH_VERSION}, found " + f"{torch.__version__}.") + + +def init_autotp(model): + """Hand the tensor-parallel collectives of an AutoTP-partitioned model over to the compiler. + + The model is expected to have been partitioned already by the regular AutoTP path, so this only + suppresses the module-level collectives and returns a backend that emits them as graph nodes. + """ + _check_autotp_compatibility() + defer_collectives_to_compiler(model) + + def backend_fn(gm: GraphModule, real_inputs): + apply_autotp(gm, real_inputs) + return torch._inductor.compile(gm, real_inputs) + + return backend_fn diff --git a/deepspeed/compile/passes/tp_compile.py b/deepspeed/compile/passes/tp_compile.py new file mode 100644 index 000000000000..db82f5ac4c64 --- /dev/null +++ b/deepspeed/compile/passes/tp_compile.py @@ -0,0 +1,138 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from typing import Dict, List + +import torch +from torch.fx import GraphModule, Node + +from deepspeed.module_inject.layers import LinearAllreduce, LinearLayer + +from ..custom_ops import tp_collectives # noqa: F401 + +COLUMN_PARALLEL_OP = torch.ops.autotp.copy_to_tp_region.default +ROW_PARALLEL_OP = torch.ops.autotp.reduce_from_tp_region.default + +# AutoTP replaces nn.Linear with these layers and shards their weights, so the injected layer type +# already records the partitioning decision the pass needs. Reading it back is more robust than +# re-deriving column/row from parameter-name patterns. +COLUMN_PARALLEL_LAYER = LinearLayer +ROW_PARALLEL_LAYER = LinearAllreduce + +# The injected layers compute their matmul with torch.matmul; the plain nn.Linear spelling is +# accepted too so the pass keeps working if a layer is lowered differently. +_MATMUL_TARGETS = { + torch.matmul, + torch.ops.aten.matmul.default, + torch.ops.aten.linear.default, + torch._C._nn.linear, +} + + +def defer_collectives_to_compiler(model) -> int: + """Suppress the module-level TP collectives on layers this pass will handle in the graph. + + Returns the number of layers handed over to the pass. Layers the pass does not rewrite (the + fused sub-param variants, conv and embedding layers) keep their module-level collectives and + stay correct as-is. + """ + deferred = 0 + for name, module in model.named_modules(): + is_row_parallel = type(module) is ROW_PARALLEL_LAYER + is_column_parallel = type(module) is COLUMN_PARALLEL_LAYER + if not (is_row_parallel or is_column_parallel): + continue + if module.mp_group is None: + continue + if type(module).tp_overlap_comm: + raise NotImplementedError("AutoTP compile pass does not support tp_overlap_comm. Set " + "'tp_overlap_comm': false to emit the collectives into the graph.") + # GatherFromTensorParallelRegion reads the gathered shard sizes back into Python, which the + # full graph this pass needs cannot capture. Leaving such a layer on the module-level path + # is not an option either: the pass identifies column-parallel layers by type, so it would + # add a second collective on top of the module's own and reduce the input gradient twice. + if is_column_parallel and module.gather_output: + raise NotImplementedError( + f"AutoTP compile pass does not support gather_output layers, but '{name}' is one. Partition it " + "without gather_output, or drop 'autotp' from the DeepCompile passes for this model.") + module.defer_collectives_to_compiler = True + deferred += 1 + return deferred + + +def _originating_layer_type(node: Node): + """Return the innermost nn.Module type a node was traced from, or None.""" + module_stack = node.meta.get("nn_module_stack") + if not module_stack: + return None + _, module_type = list(module_stack.values())[-1] + return module_type + + +def _insert_row_collective(gm: GraphModule, matmul: Node) -> Node: + """Insert g after a row-parallel matmul. + + Every consumer has to read the reduced value, which is also what the module-level + RowParallel.apply this replaces produces. + """ + with gm.graph.inserting_after(matmul): + collective_node = gm.graph.call_function(ROW_PARALLEL_OP, args=(matmul, )) + collective_node.meta["val"] = matmul.meta.get("val") + matmul.replace_all_uses_with(collective_node) + collective_node.update_arg(0, matmul) + return collective_node + + +def _insert_column_collective(gm: GraphModule, activation: Node, consumers: List[Node]) -> Node: + """ + Insert f in front of the column-parallel matmuls that share activation. + """ + with gm.graph.inserting_before(consumers[0]): + collective_node = gm.graph.call_function(COLUMN_PARALLEL_OP, args=(activation, )) + collective_node.meta["val"] = activation.meta.get("val") + for consumer in consumers: + consumer.replace_input_with(activation, collective_node) + return collective_node + + +def pass_insert_tp_collectives(gm: GraphModule, real_inputs): + """Insert the tensor-parallel collectives around the matmuls of the injected AutoTP layers.""" + column_consumers: Dict[Node, List[Node]] = {} + + for node in list(gm.graph.nodes): + if node.op != "call_function" or node.target not in _MATMUL_TARGETS: + continue + + layer_type = _originating_layer_type(node) + if layer_type is ROW_PARALLEL_LAYER: + _insert_row_collective(gm, node) + elif layer_type is COLUMN_PARALLEL_LAYER: + activation = node.args[0] + column_consumers.setdefault(activation, []).append(node) + + for activation, consumers in column_consumers.items(): + _insert_column_collective(gm, activation, consumers) + + +def pass_canonicalize(gm: GraphModule, real_inputs): + gm.graph.eliminate_dead_code() + gm.graph.lint() + gm.recompile() + + +AUTOTP_PASSES = [ + pass_insert_tp_collectives, + pass_canonicalize, +] + + +def apply_autotp(gm: GraphModule, real_inputs, passes=None): + """Apply the AutoTP transformation passes to the graph. + + The collectives are shape-preserving, so unlike AutoSP this needs no shape re-propagation. + """ + for opt_pass in passes or AUTOTP_PASSES: + opt_pass(gm, real_inputs) + return gm 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..b42e65b17cb4 --- /dev/null +++ b/tests/unit/compile/test_tp_compile.py @@ -0,0 +1,216 @@ +# Copyright (c) DeepSpeed Team. +# 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.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=AUTOTP_MIN_TORCH_VERSION), + reason=f"The AutoTP compile pass requires PyTorch >= {AUTOTP_MIN_TORCH_VERSION}") + +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): + # 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): + + 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 self.head(x) + + +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": { + "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", + }, head_spec], + }, + }, + "zero_optimization": { + "stage": 0, + }, + } + if use_compile_pass: + config["compile"] = {"deepcompile": True, "passes": ["autotp"]} + return config + + +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, gather_output_head)) + 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, 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), \ + "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}" + + 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. + + 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 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 + + 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()