diff --git a/deepspeed/module_inject/auto_ep_config.py b/deepspeed/module_inject/auto_ep_config.py index 0129a2cb79e1..953bb05f6ace 100644 --- a/deepspeed/module_inject/auto_ep_config.py +++ b/deepspeed/module_inject/auto_ep_config.py @@ -122,6 +122,22 @@ def validate_autoep_config( if not config.enabled: return + # Reject configurations that would bypass the requested fused reduction. + if config.combine_impl == "fused_weighted_sum": + if tp_size > 1: + raise ValueError('combine_impl="fused_weighted_sum" does not support folded tensor parallelism ' + f"(tensor_parallel.autotp_size={tp_size}), which restores combined tokens from " + "assignment metadata instead of the weighted reduction it implements. Set " + 'tensor_parallel.autotp_size to 1, or leave combine_impl unset.') + if config.expert_tensor_parallel_size > 1: + raise ValueError('combine_impl="fused_weighted_sum" requires expert_tensor_parallel_size=1, but got ' + f"{config.expert_tensor_parallel_size}. Set expert_tensor_parallel_size to 1, or leave " + "combine_impl unset.") + if config.comm_backend == "deepep" and config.autoep_size > 1: + raise ValueError('combine_impl="fused_weighted_sum" cannot be used with comm_backend="deepep" because ' + "DeepEP already restores and reduces the routed rows. Set comm_backend to " + '"comm", or leave combine_impl unset.') + folding_spec = build_folding_spec( world_size=world_size, pp_size=pp_size, @@ -156,7 +172,7 @@ def validate_autoep_config( f"got '{config.score_apply}'") # Validate combine_impl - valid_combine_impl = ("auto", "weighted_sum", "legacy_bmm") + valid_combine_impl = ("auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm") if config.combine_impl not in valid_combine_impl: raise ValueError(f"combine_impl must be one of {valid_combine_impl}, " f"got '{config.combine_impl}'") @@ -297,6 +313,15 @@ def validate_autoep_post_detection( return for spec in specs: + # The fused reduction folds the routing weight into the top-k reduction, + # which only exists when scores are applied after the experts. + if config.combine_impl == "fused_weighted_sum": + resolved_score_apply = config.score_apply if config.score_apply != "auto" else spec.score_apply + if resolved_score_apply != "post": + raise ValueError(f'combine_impl="fused_weighted_sum" requires score_apply="post", but layer ' + f"'{spec.moe_module_name}' resolved score_apply=\"{resolved_score_apply}\". " + "Leave combine_impl unset.") + # ep_size must not exceed num_experts if config.autoep_size > spec.num_experts: valid_divisors = _divisors(spec.num_experts) diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 69215e07fa46..2167ed73ae84 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -21,6 +21,7 @@ import deepspeed.comm as dist from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter +from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_token_ops from deepspeed.utils import logger from deepspeed.module_inject.auto_ep_comm import (DEEPEP_BACKEND, DeepEPExchange, assert_dtype_supported, deepep_combine, deepep_dispatch) @@ -63,7 +64,8 @@ def resolve_score_apply_mode( def resolve_combine_impl( - config_override: Literal["auto", "weighted_sum", "legacy_bmm"], ) -> Literal["weighted_sum", "legacy_bmm"]: + config_override: Literal["auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm"], +) -> Literal["weighted_sum", "fused_weighted_sum", "legacy_bmm"]: """Resolve combine implementation from config override or default.""" if config_override != "auto": return config_override @@ -379,6 +381,7 @@ def __init__( self.top_k = spec.top_k self.score_apply = resolve_score_apply_mode(spec, config.score_apply) self.combine_impl = resolve_combine_impl(config.combine_impl) + self._fused_combine_checked = False route_norm = spec.route_norm if config.route_norm is None else config.route_norm self.ep_size = ep_size self.ep_rank = ep_rank @@ -559,6 +562,11 @@ def set_deepspeed_parallelism( if folding_group_handles is not None: self.folding_group_handles = folding_group_handles + if self.combine_impl == "fused_weighted_sum" and folding_group_handles.spec.tp_size > 1: + # Folded TP restores tokens through a different path. + raise ValueError('combine_impl="fused_weighted_sum" does not support folded tensor parallelism ' + f"(tensor_parallel.autotp_size={folding_group_handles.spec.tp_size}). Set " + 'tensor_parallel.autotp_size to 1, or leave combine_impl unset.') if self.comm_backend == DEEPEP_BACKEND and folding_group_handles.spec.tp_size > 1: # DeepEP's combine returns token-major rows, which folded TP's # assignment-metadata restore can't consume. Refuse rather than @@ -672,6 +680,11 @@ def forward( bsz, seqlen, hdim = hidden_states.shape x = hidden_states.reshape(-1, hdim) # [T, H] + # Fail all ranks before any collective can stall. + if self.combine_impl == "fused_weighted_sum" and not self._fused_combine_checked: + fused_token_ops.assert_supported(x, score_apply=self.score_apply) + self._fused_combine_checked = True + # Router ro: RouterOutput = RouterOutput(*self.router(x, self.expert_bias)) @@ -790,6 +803,14 @@ def forward( # backend being selected: with ep_size == 1 the local path runs # instead and still has one row per assignment to reduce. output = expert_output.reshape(bsz, seqlen, hdim) + elif self.combine_impl == "fused_weighted_sum": + output = fused_token_ops.fused_weighted_restore( + expert_output, + top_scores=ro.top_scores, + token_indices_sorted=token_indices_sorted, + top_k=self.top_k, + shape=(bsz, seqlen, hdim), + ) else: output = combine_from_routed( expert_output, diff --git a/deepspeed/module_inject/auto_ep_presets/base.py b/deepspeed/module_inject/auto_ep_presets/base.py index aa8985871643..14bb739538a2 100644 --- a/deepspeed/module_inject/auto_ep_presets/base.py +++ b/deepspeed/module_inject/auto_ep_presets/base.py @@ -109,7 +109,7 @@ class AutoEPConfig: route_norm: bool | None = None route_scale: float = 1.0 score_apply: Literal["auto", "pre", "post"] = "auto" - combine_impl: Literal["auto", "weighted_sum", "legacy_bmm"] = "auto" + combine_impl: Literal["auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm"] = "auto" comm_backend: Literal["comm", "deepep"] = "comm" comm_num_sm: int = 12 comm_qp_margin: int = 4 diff --git a/deepspeed/ops/triton_ops/autoep_fused_token_ops.py b/deepspeed/ops/triton_ops/autoep_fused_token_ops.py new file mode 100644 index 000000000000..1036eb500608 --- /dev/null +++ b/deepspeed/ops/triton_ops/autoep_fused_token_ops.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Fused AutoEP token restore without the eager scatter and FP32 intermediate. + +The kernel reduces each token's top-k rows in FP32. Communication, routing, +expert reorder, and grouped GEMM remain unchanged. +""" + +from __future__ import annotations + +import torch + +from deepspeed.ops.triton_ops._triton import _TRITON_AVAILABLE, triton, tl + +_IS_ROCM_PYTORCH = getattr(torch.version, "hip", None) is not None + +SUPPORTED_ROW_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + +_MAX_BLOCK_HIDDEN = 512 +_INVERT_INDEX_BLOCK = 256 +# The kernels hold a [slots, BLOCK_H] FP32 block live, so the hidden tile shrinks +# as top-k grows to keep that block in registers rather than spilling. +_MAX_BLOCK_ELEMENTS = 2048 + +if _TRITON_AVAILABLE: + + @triton.jit + def _invert_index_kernel( + index_ptr, + inverse_ptr, + num_indices, + num_inverse_rows, + BLOCK: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + in_range = offsets < num_indices + + targets = tl.load(index_ptr + offsets, mask=in_range, other=-1).to(tl.int64) + writable = in_range & (targets >= 0) & (targets < num_inverse_rows) + tl.store(inverse_ptr + tl.where(writable, targets, 0), offsets.to(tl.int32), mask=writable) + + @triton.jit + def _weighted_restore_forward_kernel( + rows_ptr, + inverse_ptr, + scores_ptr, + out_ptr, + hidden, + rows_stride, + scores_stride, + out_stride, + TOP_K: tl.constexpr, + K_PADDED: tl.constexpr, + BLOCK_H: tl.constexpr, + ): + token = tl.program_id(0) + hidden_offsets = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + + slots = tl.arange(0, K_PADDED) + slot_mask = slots < TOP_K + + source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64) + row_valid = slot_mask & (source_rows >= 0) + safe_rows = tl.where(row_valid, source_rows, 0) + + scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32) + + block_mask = row_valid[:, None] & hidden_mask[None, :] + values = tl.load( + rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :], + mask=block_mask, + other=0.0, + ).to(tl.float32) + + # Match the eager path's FP32 product and accumulation. + weighted = tl.sum(values * scores[:, None], axis=0) + tl.store( + out_ptr + token * out_stride + hidden_offsets, + weighted.to(out_ptr.dtype.element_ty), + mask=hidden_mask, + ) + + @triton.jit + def _weighted_restore_backward_kernel( + grad_out_ptr, + rows_ptr, + inverse_ptr, + scores_ptr, + grad_rows_ptr, + grad_scores_ptr, + hidden, + grad_out_stride, + rows_stride, + scores_stride, + grad_rows_stride, + grad_scores_stride, + TOP_K: tl.constexpr, + K_PADDED: tl.constexpr, + BLOCK_H: tl.constexpr, + ): + token = tl.program_id(0) + + slots = tl.arange(0, K_PADDED) + slot_mask = slots < TOP_K + + source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64) + row_valid = slot_mask & (source_rows >= 0) + safe_rows = tl.where(row_valid, source_rows, 0) + scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32) + + grad_rows_dtype = grad_rows_ptr.dtype.element_ty + # Keeping one token per program avoids a second reduction pass for scores. + score_partials = tl.zeros([K_PADDED, BLOCK_H], dtype=tl.float32) + + for hidden_start in range(0, hidden, BLOCK_H): + hidden_offsets = hidden_start + tl.arange(0, BLOCK_H) + hidden_mask = hidden_offsets < hidden + block_mask = row_valid[:, None] & hidden_mask[None, :] + + upstream = tl.load( + grad_out_ptr + token * grad_out_stride + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + values = tl.load( + rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :], + mask=block_mask, + other=0.0, + ).to(tl.float32) + score_partials += values * upstream[None, :] + + tl.store( + grad_rows_ptr + safe_rows[:, None] * grad_rows_stride + hidden_offsets[None, :], + (upstream[None, :] * scores[:, None]).to(grad_rows_dtype), + mask=block_mask, + ) + + grad_scores = tl.sum(score_partials, axis=1) + tl.store( + grad_scores_ptr + token * grad_scores_stride + slots, + grad_scores.to(grad_scores_ptr.dtype.element_ty), + mask=slot_mask, + ) + + +def is_available() -> bool: + """Whether this build can run the fused weighted restore at all.""" + return _TRITON_AVAILABLE and not _IS_ROCM_PYTORCH + + +def assert_supported(rows: torch.Tensor, *, score_apply: str) -> None: + """Reject unsupported configurations before collectives begin.""" + if not _TRITON_AVAILABLE: + raise RuntimeError('combine_impl="fused_weighted_sum" needs Triton, which is not installed in this ' + "environment. Install Triton, or leave combine_impl unset.") + if _IS_ROCM_PYTORCH: + raise RuntimeError('combine_impl="fused_weighted_sum" is not yet supported on ROCm. Leave combine_impl ' + "unset to run here.") + if rows.device.type != "cuda": + raise RuntimeError('combine_impl="fused_weighted_sum" runs CUDA kernels but this layer is on device ' + f'"{rows.device.type}". Leave combine_impl unset to run here.') + if rows.dtype not in SUPPORTED_ROW_DTYPES: + raise RuntimeError('combine_impl="fused_weighted_sum" supports bfloat16, float16, and float32 rows, got ' + f"{rows.dtype}. Leave combine_impl unset, or use a supported floating-point dtype.") + if score_apply != "post": + raise RuntimeError('combine_impl="fused_weighted_sum" folds the routing weight into the top-k reduction, ' + f'which only exists for score_apply="post", but this layer resolved ' + f'score_apply="{score_apply}". Leave combine_impl unset.') + + +def _block_hidden(hidden: int, slots: int) -> int: + """Choose a power-of-two tile within the FP32 register budget.""" + budget = max(16, _MAX_BLOCK_ELEMENTS // slots) + return min(_MAX_BLOCK_HIDDEN, budget, max(16, triton.next_power_of_2(hidden))) + + +def _padded_top_k(top_k: int) -> int: + """Round top-k up to a power of two, which ``tl.arange`` requires.""" + return max(2, triton.next_power_of_2(top_k)) + + +def _invert_index(index: torch.Tensor, num_inverse_rows: int) -> torch.Tensor: + """Invert the row permutation produced by sorting routed assignments.""" + inverse = torch.empty((num_inverse_rows, ), dtype=torch.int32, device=index.device) + num_indices = index.numel() + + grid = (triton.cdiv(num_indices, _INVERT_INDEX_BLOCK), ) + _invert_index_kernel[grid]( + index.contiguous(), + inverse, + num_indices, + num_inverse_rows, + BLOCK=_INVERT_INDEX_BLOCK, + ) + return inverse + + +class _FusedWeightedRestore(torch.autograd.Function): + """Weight rows by their routing score and reduce over top-k in one pass.""" + + @staticmethod + def forward(ctx, combined_rows, top_scores, inverse, top_k): + combined_rows = combined_rows.contiguous() + n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1] + output = torch.empty((n_tokens, hidden), dtype=combined_rows.dtype, device=combined_rows.device) + + ctx.save_for_backward(combined_rows, top_scores, inverse) + ctx.top_k = top_k + + k_padded = _padded_top_k(top_k) + block_hidden = _block_hidden(hidden, slots=k_padded) + grid = (n_tokens, triton.cdiv(hidden, block_hidden)) + _weighted_restore_forward_kernel[grid]( + combined_rows, + inverse, + top_scores, + output, + hidden, + combined_rows.stride(0), + top_scores.stride(0), + output.stride(0), + TOP_K=top_k, + K_PADDED=k_padded, + BLOCK_H=block_hidden, + ) + return output + + @staticmethod + def backward(ctx, grad_output): + combined_rows, top_scores, inverse = ctx.saved_tensors + grad_output = grad_output.contiguous() + + grad_rows = torch.empty_like(combined_rows) + grad_scores = torch.empty_like(top_scores) + + n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1] + + k_padded = _padded_top_k(ctx.top_k) + _weighted_restore_backward_kernel[(n_tokens, )]( + grad_output, + combined_rows, + inverse, + top_scores, + grad_rows, + grad_scores, + hidden, + grad_output.stride(0), + combined_rows.stride(0), + top_scores.stride(0), + grad_rows.stride(0), + grad_scores.stride(0), + TOP_K=ctx.top_k, + K_PADDED=k_padded, + BLOCK_H=_block_hidden(hidden, slots=k_padded), + ) + return grad_rows, grad_scores, None, None + + +def fused_weighted_restore( + combined_rows: torch.Tensor, + top_scores: torch.Tensor, + token_indices_sorted: torch.Tensor, + top_k: int, + shape: tuple[int, int, int], +) -> torch.Tensor: + """Restore ``[T * K, H]`` rows directly to weighted ``[B, S, H]`` output.""" + bsz, seqlen, hidden = shape + n_tokens = bsz * seqlen + expected_rows = n_tokens * top_k + inverse = _invert_index(token_indices_sorted, expected_rows) + output = _FusedWeightedRestore.apply(combined_rows, top_scores.contiguous(), inverse, top_k) + return output.reshape(bsz, seqlen, hidden) diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 4fda25f47167..cd1c2def9562 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -998,6 +998,12 @@ smoke coverage used for this AutoEP surface produced the following version gates | -------------------------------------------------------------------------------------------------------------- | -------- | | When to apply router scores: `"pre"` (before experts), `"post"` (during combine), or `"auto"` (from preset). | `"auto"` | +***combine_impl***: [string] + +| Description | Default | +| -------------------------------------------------------------------------------------------------------------- | -------- | +| How expert outputs are weighted by their router scores and reduced over top-k. `"auto"` resolves to `"weighted_sum"`. `"fused_weighted_sum"` is experimental and computes the same reduction in one Triton pass, without materializing the scattered assignment buffer or the `[tokens, top_k, hidden]` FP32 intermediate; it requires CUDA, Triton, bfloat16/float16 activations, `tensor_parallel.autotp_size=1`, `expert_tensor_parallel_size=1`, and a resolved `score_apply="post"`, and is rejected rather than silently ignored when any of those does not hold. `"legacy_bmm"` is a debug reduction retained for model-family verification. | `"auto"` | + ***route_norm***: [boolean] | Description | Default | diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index b772771dd164..38a4f510d260 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -144,6 +144,51 @@ Requirements and limits: - Not compatible with folded tensor parallelism (``expert_tensor_parallel_size > 1``), which is rejected at setup. +**Fused weighted restore (experimental):** + +After the combine all-to-all, AutoEP holds one row per routed assignment and has +to turn it back into one row per token. ``combine_impl`` selects how: + +.. code-block:: json + + { + "expert_parallel": { + "enabled": true, + "autoep_size": 16, + "preset_model": "qwen3_moe", + "combine_impl": "fused_weighted_sum" + } + } + +``"auto"`` (default) resolves to ``"weighted_sum"``, which scatters the rows into +a zero-filled ``[tokens * top_k, hidden]`` buffer, widens it to FP32 to apply the +routing weights, and reduces over top-k. ``"fused_weighted_sum"`` computes the +same result in a single pass: each program owns one token and one slice of the +hidden dimension, walks its top-k rows in registers and accumulates in FP32, so +neither the scattered buffer nor the FP32 intermediate is allocated. At the +canonical shape the FP32 intermediate alone is 64 MiB per layer. + +Routing weights are still accumulated in FP32 and cast once, so the result +matches the eager reduction to within the order of the top-k summation. Only the +reduction changes: the collectives, the router, the grouped GEMM and the +expert-major reorder are untouched. + +``"fused_weighted_sum"`` is rejected, rather than quietly ignored, when it would +have nothing to replace or would change semantics: + +- ``tensor_parallel.autotp_size`` greater than 1, which uses folded tensor + parallelism and restores combined tokens from assignment metadata instead; +- ``expert_tensor_parallel_size`` greater than 1; +- ``comm_backend="deepep"`` with expert parallelism, because DeepEP already + restores and reduces its routed rows; +- a resolved ``score_apply`` other than ``"post"``; +- activations that are not bfloat16, float16, or float32, a non-CUDA device, or + a build without Triton. + +Failing fast matters for measurement: a run that asked for the fused reduction +and silently got the eager one would report the difference between an +implementation and itself. + **Constraints:** - ``autoep_size`` must divide ``num_experts`` for all detected MoE layers. diff --git a/tests/unit/v1/moe/test_autoep_fused_parity.py b/tests/unit/v1/moe/test_autoep_fused_parity.py new file mode 100644 index 000000000000..00211912ae0e --- /dev/null +++ b/tests/unit/v1/moe/test_autoep_fused_parity.py @@ -0,0 +1,196 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""End-to-end parity for the eager and fused combine implementations.""" + +import functools + +import deepspeed +import pytest +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer +from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_ops +from deepspeed.utils import safe_get_full_grad +from unit.common import DistributedTest +from unit.v1.moe.autoep_test_utils import ( + MockMoETransformer, + engine_input_dtype, + mixed_precision_config, + seed_everything, +) + +HIDDEN_SIZE = 64 +SEQ_LEN = 16 +NUM_EXPERTS = 4 + +# The top-k accumulation order differs, so parity allows last-bit noise only. +PARITY_TOLERANCE = {"rtol": 1e-2, "atol": 1e-3} + + +def _fused_engine_available(): + accelerator = get_accelerator() + return (accelerator.is_available() and accelerator.device_name().startswith("cuda") and fused_ops.is_available()) + + +pytestmark = pytest.mark.skipif(not _fused_engine_available(), + reason="the fused weighted restore needs CUDA and Triton") + + +def _config(combine_impl, ep_size): + return { + **mixed_precision_config(), + "train_micro_batch_size_per_gpu": 1, + "gradient_clipping": 0.0, + "optimizer": { + "type": "AdamW", + "params": { + "lr": 1e-3, + "betas": [0.9, 0.999], + "eps": 1e-8, + }, + }, + "expert_parallel": { + "enabled": True, + "autoep_size": ep_size, + "preset_model": "mixtral", + "load_balance_coeff": None, + "combine_impl": combine_impl, + }, + } + + +def _build_engine(combine_impl, ep_size, reference_state, seed): + seed_everything(seed) + model = MockMoETransformer(num_layers=2, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size=2 * HIDDEN_SIZE) + model.load_state_dict(reference_state) + engine, _, _, _ = deepspeed.initialize(model=model, config=_config(combine_impl, ep_size)) + return engine + + +def _checkpoint_moe_layers(engine): + """Recompute each MoE block in backward, as the benchmarked runs do.""" + for module in engine.module.modules(): + if isinstance(module, AutoEPMoELayer): + module.forward = functools.partial(torch.utils.checkpoint.checkpoint, module.forward, use_reentrant=False) + + +def _named_gradients(engine): + gradients = {} + for name, param in engine.module.named_parameters(): + if not param.requires_grad: + continue + grad = safe_get_full_grad(param) + if grad is not None: + gradients[name] = grad.detach().float().cpu().clone() + return gradients + + +def _parameters(engine): + return { + name: param.detach().float().cpu().clone() + for name, param in engine.module.named_parameters() if param.requires_grad + } + + +def _take_one_step(engine, seed, *, checkpoint_activations): + if checkpoint_activations: + _checkpoint_moe_layers(engine) + + generator = torch.Generator().manual_seed(seed) + batch = torch.randn((1, SEQ_LEN, HIDDEN_SIZE), generator=generator, dtype=torch.float32) + batch = batch.to(engine.device, dtype=engine_input_dtype(engine)).requires_grad_(True) + + before = _parameters(engine) + output = engine(batch) + loss = output.float().pow(2).mean() + engine.backward(loss) + + gradients = _named_gradients(engine) + input_grad = batch.grad.detach().float().cpu().clone() + engine.step() + + delta = {name: _parameters(engine)[name] - value for name, value in before.items()} + return { + "loss": loss.detach().float().cpu().clone(), + "output": output.detach().float().cpu().clone(), + "input_grad": input_grad, + "gradients": gradients, + "delta": delta, + } + + +def _assert_step_matches(fused, eager): + torch.testing.assert_close(fused["loss"], eager["loss"], **PARITY_TOLERANCE) + torch.testing.assert_close(fused["output"], eager["output"], **PARITY_TOLERANCE) + torch.testing.assert_close(fused["input_grad"], eager["input_grad"], **PARITY_TOLERANCE) + + assert fused["gradients"], "no gradients were captured, so the comparison would be vacuous" + assert set(fused["gradients"]) == set(eager["gradients"]) + # Ensure both sides of the fused restore reached the comparison. + assert any(".router." in name for name in fused["gradients"]), "no router gradient was captured" + assert any(".experts.w" in name for name in fused["gradients"]), "no expert gradient was captured" + + for name in sorted(eager["gradients"]): + torch.testing.assert_close(fused["gradients"][name], + eager["gradients"][name], + msg=lambda formatted, name=name: f"gradient mismatch for {name}\n{formatted}", + **PARITY_TOLERANCE) + + for name in sorted(eager["delta"]): + torch.testing.assert_close(fused["delta"][name], + eager["delta"][name], + msg=lambda formatted, name=name: f"parameter update mismatch for {name}\n" + f"{formatted}", + **PARITY_TOLERANCE) + + assert any(value.abs().sum() > 0 for value in eager["delta"].values()), "the optimizer step changed nothing" + + +class TestAutoEPFusedParityExpertParallel(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize("checkpoint_activations", [True, False]) + def test_fused_matches_eager_through_a_full_step(self, checkpoint_activations): + seed = 4321 + seed_everything(seed) + reference_state = MockMoETransformer(num_layers=2, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size=2 * HIDDEN_SIZE).state_dict() + + eager_engine = _build_engine("weighted_sum", 2, reference_state, seed) + eager = _take_one_step(eager_engine, seed, checkpoint_activations=checkpoint_activations) + + fused_engine = _build_engine("fused_weighted_sum", 2, reference_state, seed) + assert all(module.combine_impl == "fused_weighted_sum" for module in fused_engine.module.modules() + if isinstance(module, AutoEPMoELayer)), "the fused reduction was not actually selected" + fused = _take_one_step(fused_engine, seed, checkpoint_activations=checkpoint_activations) + + _assert_step_matches(fused, eager) + + +class TestAutoEPFusedParityLocalExperts(DistributedTest): + world_size = 1 + + def test_fused_matches_eager_without_expert_parallelism(self): + seed = 8765 + seed_everything(seed) + reference_state = MockMoETransformer(num_layers=2, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size=2 * HIDDEN_SIZE).state_dict() + + eager = _take_one_step(_build_engine("weighted_sum", 1, reference_state, seed), + seed, + checkpoint_activations=False) + fused = _take_one_step(_build_engine("fused_weighted_sum", 1, reference_state, seed), + seed, + checkpoint_activations=False) + + _assert_step_matches(fused, eager) diff --git a/tests/unit/v1/moe/test_autoep_unit.py b/tests/unit/v1/moe/test_autoep_unit.py index d28eb96c8036..fb598eb55002 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -234,6 +234,56 @@ def test_validate_folding_routing_requires_boolean(self): tp_size=1, sp_size=1) + def test_combine_impl_rejects_unknown_value(self): + config = parse_autoep_config({"enabled": True, "combine_impl": "triton"}) + with pytest.raises(ValueError, match="combine_impl must be one of"): + validate_autoep_config(config, world_size=1, pp_size=1, tp_size=1, sp_size=1) + + def test_fused_combine_rejects_folded_tensor_parallelism(self): + config = parse_autoep_config({ + "enabled": True, + "autoep_size": 2, + "combine_impl": "fused_weighted_sum", + }) + with pytest.raises(ValueError, match=r"tensor_parallel\.autotp_size=2"): + validate_autoep_config(config, world_size=4, pp_size=1, tp_size=2, sp_size=1) + + def test_fused_combine_rejects_expert_tensor_parallelism(self): + config = parse_autoep_config({ + "enabled": True, + "autoep_size": 2, + "expert_tensor_parallel_size": 2, + "combine_impl": "fused_weighted_sum", + }) + with pytest.raises(ValueError, match="requires expert_tensor_parallel_size=1"): + validate_autoep_config(config, world_size=4, pp_size=1, tp_size=1, sp_size=1) + + def test_fused_combine_rejects_deepep(self): + config = parse_autoep_config({ + "enabled": True, + "autoep_size": 2, + "combine_impl": "fused_weighted_sum", + "comm_backend": "deepep", + "comm_max_tokens_per_rank": 4096, + }) + with pytest.raises(ValueError, match='cannot be used with comm_backend="deepep"'): + validate_autoep_config(config, world_size=2, pp_size=1, tp_size=1, sp_size=1) + + @pytest.mark.parametrize("score_apply, spec_score_apply", [("auto", "pre"), ("pre", "post")]) + def test_fused_combine_requires_post_score_apply(self, score_apply, spec_score_apply): + config = parse_autoep_config({ + "enabled": True, + "combine_impl": "fused_weighted_sum", + "score_apply": score_apply, + }) + with pytest.raises(ValueError, match='requires score_apply="post"'): + validate_autoep_post_detection(config, [_make_spec(score_apply=spec_score_apply)]) + + def test_fused_combine_accepts_the_standard_path(self): + config = parse_autoep_config({"enabled": True, "autoep_size": 2, "combine_impl": "fused_weighted_sum"}) + validate_autoep_config(config, world_size=2, pp_size=1, tp_size=1, sp_size=1) + validate_autoep_post_detection(config, [_make_spec(num_experts=4, score_apply="post")]) + @pytest.mark.parametrize("value", UNSUPPORTED_LOAD_BALANCE_VALUES) def test_load_balance_coeff_rejected_at_parse(self, value): with pytest.raises(ValueError) as exc_info: diff --git a/tests/unit/v1/ops/triton_ops/test_autoep_fused_token_ops.py b/tests/unit/v1/ops/triton_ops/test_autoep_fused_token_ops.py new file mode 100644 index 000000000000..2e065c4f7a08 --- /dev/null +++ b/tests/unit/v1/ops/triton_ops/test_autoep_fused_token_ops.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Compare the fused weighted restore with the eager reference.""" + +import pytest +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.module_inject.auto_ep_layer import combine_from_routed +from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_ops + + +def _fused_engine_available(): + accelerator = get_accelerator() + return (accelerator.is_available() and accelerator.device_name().startswith("cuda") and fused_ops.is_available()) + + +pytestmark = pytest.mark.skipif(not _fused_engine_available(), + reason="the fused weighted restore needs CUDA and Triton") + + +def _device(): + return get_accelerator().current_device_name() + + +@pytest.mark.parametrize("top_k", [2, 4, 6, 8]) +@pytest.mark.parametrize("hidden", [128, 130]) +@pytest.mark.parametrize("row_dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("score_dtype", [torch.float32, torch.bfloat16]) +def test_fused_weighted_restore_matches_eager_including_gradients(top_k, hidden, row_dtype, score_dtype): + device = _device() + num_tokens, num_experts = 24, 8 + generator = torch.Generator(device=device).manual_seed(20260824) + + selected_experts = torch.randint(0, num_experts, (num_tokens, top_k), device=device, generator=generator) + token_indices_sorted = torch.argsort(selected_experts.view(-1), stable=True) + assert not torch.equal(token_indices_sorted, torch.arange(num_tokens * top_k, device=device)) + + rows = torch.randn(num_tokens * top_k, hidden, device=device, dtype=row_dtype, generator=generator) + scores = torch.rand(num_tokens, top_k, device=device, dtype=score_dtype, generator=generator) + upstream = torch.randn(1, num_tokens, hidden, device=device, dtype=row_dtype, generator=generator) + + eager_rows = rows.clone().requires_grad_(True) + eager_scores = scores.clone().requires_grad_(True) + eager_output = combine_from_routed( + eager_rows, + top_scores=eager_scores, + token_indices_sorted=token_indices_sorted, + top_k=top_k, + score_apply="post", + combine_impl="weighted_sum", + shape=(1, num_tokens, hidden), + ) + + fused_rows = rows.clone().requires_grad_(True) + fused_scores = scores.clone().requires_grad_(True) + fused_output = fused_ops.fused_weighted_restore( + fused_rows, + top_scores=fused_scores, + token_indices_sorted=token_indices_sorted, + top_k=top_k, + shape=(1, num_tokens, hidden), + ) + + output_tolerance = {"rtol": 1e-5, "atol": 1e-6} if row_dtype == torch.float32 else {} + torch.testing.assert_close(fused_output, eager_output, **output_tolerance) + + eager_output.backward(upstream) + fused_output.backward(upstream) + + torch.testing.assert_close(fused_rows.grad, eager_rows.grad, **output_tolerance) + # Hidden reduction order only affects the last bits of FP32 score gradients. + score_tolerance = {"rtol": 1e-4, "atol": 1e-5} if score_dtype == torch.float32 else {} + torch.testing.assert_close(fused_scores.grad, eager_scores.grad, **score_tolerance) + + +def test_fused_engine_names_what_it_cannot_run(): + device = _device() + for dtype in fused_ops.SUPPORTED_ROW_DTYPES: + fused_ops.assert_supported(torch.randn(8, 16, device=device, dtype=dtype), score_apply="post") + + with pytest.raises(RuntimeError, match="bfloat16, float16, and float32"): + fused_ops.assert_supported(torch.randn(8, 16, device=device, dtype=torch.float64), score_apply="post") + + supported = torch.randn(8, 16, device=device, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match='resolved score_apply="pre"'): + fused_ops.assert_supported(supported, score_apply="pre") + + with pytest.raises(RuntimeError, match="CUDA kernels"): + fused_ops.assert_supported(torch.randn(8, 16, dtype=torch.bfloat16), score_apply="post")