diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 2c4dd31a485a..6ff033ea2b43 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -429,7 +429,12 @@ def calculate_num_chunks(self, all_rank_num_tokens: List[int]) -> int: if self.use_dp and self.comm is not None: num_rows = self._dp_padded_num_rows(all_rank_num_tokens) else: - num_rows = sum(all_rank_num_tokens) + # non-DP: no cross-rank dispatch. The scheduler fills all_rank_num_tokens + # from [x.shape[0]] before calling here, so it must be a single-element list. + assert len(all_rank_num_tokens) == 1, ( + f"non-DP path expects a single-element list, got {len(all_rank_num_tokens)}" + ) + num_rows = all_rank_num_tokens[0] return (num_rows + self.moe_max_num_tokens - 1) // self.moe_max_num_tokens def split_chunk(self, split_token_num: int, split_num_chunks: int) -> List[int]: diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index ca227bc4c5ff..fe0153386763 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -160,6 +160,34 @@ def forward( ): raise_moe_lora_multichunk_unsupported(num_chunks) + # ========== 0-token rank deadlock fix ========== + # When some ranks have 0 tokens in single-chunk forward with collective comm, + # those ranks hang in CUDA kernels (e.g. NVFP4 quantize_input with 0-row tensor) + # before reaching moe.comm.dispatch(), causing NCCL AllGather deadlock on + # non-zero ranks. Fix: activate DP padding uniformly across all ranks so every + # rank uses sizes=None (uniform allgather) and pads x/router_logits to max_tokens. + # Mirrors the empty-chunk substitution in _forward_multiple_chunks (line ~597-620). + # Existing truncation at Step 4 discards dummy-token outputs automatically. + # NOTE: kept after the multi-chunk rejection above so unit tests that stub `moe` + # with a minimal namespace (no `.comm`/`.use_dp`) still exercise that path; only + # relevant to single-chunk collective comm anyway. + if ( + moe.comm is not None + and moe.use_dp + and all_rank_max_num_tokens > 0 + and not use_dp_padding + and any(t == 0 for t in all_rank_num_tokens_padded) + ): + use_dp_padding = True + all_rank_num_tokens_padded = [all_rank_max_num_tokens] * len(all_rank_num_tokens) + local_n = x.shape[0] + if local_n < all_rank_max_num_tokens: + pad = all_rank_max_num_tokens - local_n + x = torch.cat([x, x.new_zeros((pad, x.shape[1]))], dim=0) + router_logits = torch.cat( + [router_logits, router_logits.new_zeros((pad, router_logits.shape[1]))], dim=0 + ) + # May fall back AllToAll -> AllGather; this is the only sanctioned # mutation of ``moe.comm`` from a scheduler. moe.determine_communication_method(all_rank_num_tokens_padded, num_chunks) diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index bb3829703bbc..5e7fae73d90d 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/tests/microbenchmarks/bench_moe/case_runner.py @@ -344,6 +344,7 @@ def _select_routing_inputs( routing_plan: RoutingPlan, rank: int, moe_ep_size: int, + enable_attention_dp: bool, base_router_logits: torch.Tensor, device: torch.device, act_dtype: torch.dtype, @@ -411,6 +412,21 @@ def _select_routing_inputs( except Exception as exc: return None, _RoutingSkip(f"native logits projection error: {type(exc).__name__}: {exc}") + # In attention-DP + MoE-TP layouts (DTP / CUSTOM-DP), _project_router_logits + # returns logits shaped [agg_tokens, E] covering all DP shards aggregated + # onto ep_axis_rank. The MoE internally allgathers each rank's local + # router_logits before routing, so each rank must supply only its local + # slice [offset_r : offset_r + n_r] of the full projected tensor. + world_size_inferred = len(routing_plan.per_rank_num_tokens) + if enable_attention_dp and int(moe_ep_size) < world_size_inferred: + offset = sum( + routing_plan.per_rank_num_tokens[s] + for s in range(world_size_inferred) + if s % int(moe_ep_size) == ep_axis_rank and s < rank + ) + local_n = routing_plan.per_rank_num_tokens[rank] + new_logits = new_logits[offset : offset + local_n] + if projection_status != "exact" and rc_spec.projection_policy == "reject": return None, _RoutingSkip( skip_reason=( @@ -536,21 +552,6 @@ def _resolve_layout_and_plan( except ValueError as exc: return _short_circuit(result, "skipped", str(exc)) - # Routing-control's dispatch_matrix axis is ``moe_ep_size`` while - # ``per_rank_num_tokens`` follows the world (DP source) axis. When the two - # disagree (DTP/TTP/CUSTOM with ``moe_ep_size != world_size``) the plan - # either crashes inside ``_build_routing_plan`` or silently drops the - # tokens of world ranks beyond ``moe_ep_size``. Skip cleanly. - if rc_active and int(moe_ep_size) != int(world_size): - return _short_circuit( - result, - "skipped", - f"routing-control requires moe_ep_size == world_size " - f"(got moe_ep_size={moe_ep_size}, world_size={world_size}); " - "the dispatch_matrix axis would not align with the per-rank token " - "distribution. Use parallel_mode in {DEP, TEP} or drop routing-control.", - ) - routing_plan: Optional[RoutingPlan] = None if rc_active: try: @@ -561,6 +562,7 @@ def _resolve_layout_and_plan( top_k=int(model.top_k), num_experts=int(model.num_experts), moe_ep_size=int(moe_ep_size), + enable_dp=bool(_enable_dp), ) except Exception as exc: reason = f"routing plan error: {type(exc).__name__}: {exc}" @@ -568,7 +570,7 @@ def _resolve_layout_and_plan( return _short_circuit(result, "skipped", reason) per_rank = list(routing_plan.per_rank_num_tokens) else: - per_rank = _per_rank_tokens(workload, world_size) + per_rank = _per_rank_tokens(workload, world_size, enable_dp=bool(_enable_dp)) return int(moe_ep_size), per_rank, routing_plan @@ -663,6 +665,11 @@ def _run_one_candidate( result.moe_tp_size = int(mapping.moe_tp_size) result.enable_attention_dp = bool(mapping.enable_attention_dp) + # TEP/TTP (no attention DP): no cross-rank dispatch; the scheduler fills + # all_rank_num_tokens from x.shape[0]. Pass None to follow that path. + if not mapping.enable_attention_dp: + all_rank_num_tokens = None + AutoTuner.get().setup_distributed_state(mapping) AutoTuner.get().clear_cache() @@ -683,7 +690,11 @@ def _run_one_candidate( mapping=mapping, moe_backend=config.backend, use_cuda_graph=bool(config.cuda_graph), - max_num_tokens=max(int(local_num_tokens), 1), + # Symmetric-memory comm backends (e.g. NVLINK_ONE_SIDED) size their + # workspace from max_num_tokens and require every rank to allocate the + # same size, so use the global per-rank maximum rather than this rank's + # local token count (which differs under uneven attention-DP shards). + max_num_tokens=max(int(max(per_rank)) if per_rank else 0, 1), use_low_precision_moe_combine=bool(config.use_low_precision_moe_combine), enable_perfect_router=enable_perfect_router, dtype=act_dtype, @@ -735,6 +746,7 @@ def _run_one_candidate( routing_plan=routing_plan, rank=rank, moe_ep_size=int(moe_ep_size), + enable_attention_dp=bool(result.enable_attention_dp), base_router_logits=router_logits, device=device, act_dtype=act_dtype, diff --git a/tests/microbenchmarks/bench_moe/cli.py b/tests/microbenchmarks/bench_moe/cli.py index 8cefcfe4f902..cc27bd49fa49 100644 --- a/tests/microbenchmarks/bench_moe/cli.py +++ b/tests/microbenchmarks/bench_moe/cli.py @@ -29,6 +29,7 @@ from tensorrt_llm.models.modeling_utils import QuantAlgo from .backend import MoeBackendType +from .mapping import _resolve_mapping_layout from .routing import _per_rank_tokens from .search import ( _coerce_str_tuple, @@ -91,7 +92,14 @@ def _build_worker_header(ctx: _BenchmarkContext, launcher: str, world_size: int) "world_size": world_size, "analysis": list(ctx.analysis) or ["summary"], "workloads": [ - w.to_dict(per_rank_num_tokens=_per_rank_tokens(w, world_size)) for w in ctx.workloads + w.to_dict( + per_rank_num_tokens=_per_rank_tokens( + w, + world_size, + enable_dp=bool(_resolve_mapping_layout(ctx.base_config, world_size)[2]), + ) + ) + for w in ctx.workloads ], "base_config": ctx.base_config.to_dict(), } @@ -214,8 +222,9 @@ def parse_args() -> argparse.Namespace: nargs="+", required=False, help=( - "Global token counts to sweep. Each value is balanced across ranks " - "with any remainder on rank 0. Example: --balanced_total_num_tokens 64 256 1024." + "Global token counts to sweep. Each value is balanced across ranks, " + "spreading any remainder one token per leading rank (e.g. world_size=4, " + "tokens=2 -> [1, 1, 0, 0]). Example: --balanced_total_num_tokens 64 256 1024." ), ) diff --git a/tests/microbenchmarks/bench_moe/mapping.py b/tests/microbenchmarks/bench_moe/mapping.py index 5074fe0da437..545b9109edcd 100644 --- a/tests/microbenchmarks/bench_moe/mapping.py +++ b/tests/microbenchmarks/bench_moe/mapping.py @@ -90,12 +90,17 @@ def _resolve_mapping_layout(config: ConfigSpec, world_size: int) -> Tuple[int, i def _build_mapping_from_config(config: ConfigSpec, world_size: int) -> Mapping: """Build ``Mapping`` from a ``ConfigSpec`` + world size; sets ``rank=mpi_rank()``.""" moe_ep, moe_tp, enable_dp = _resolve_mapping_layout(config, world_size) + # gpus_per_node must match actual visible GPUs per node so that + # mapping.local_rank (= rank % gpus_per_node) gives the correct device index. + # The Mapping default (8) is wrong for multi-node runs with fewer GPUs per node. + gpus_per_node = torch.cuda.device_count() mapping = Mapping( world_size=world_size, tp_size=world_size, moe_ep_size=moe_ep, moe_tp_size=moe_tp, enable_attention_dp=enable_dp, + gpus_per_node=gpus_per_node, ) mapping.rank = mpi_rank() return mapping diff --git a/tests/microbenchmarks/bench_moe/results.py b/tests/microbenchmarks/bench_moe/results.py index e5637edefe18..43870a8e53ef 100644 --- a/tests/microbenchmarks/bench_moe/results.py +++ b/tests/microbenchmarks/bench_moe/results.py @@ -21,6 +21,7 @@ from tensorrt_llm._utils import mpi_allgather +from .mapping import _resolve_mapping_layout from .routing import _per_rank_tokens from .specs import ConfigSpec, ModelSpec, RunResult, WorkloadSpec from .utils import _compute_stats @@ -407,7 +408,8 @@ def _make_skipped_run_result( r = RunResult(model=model, workload=workload, config=config) r.status = "skipped" r.skip_reason = reason - r.per_rank_num_tokens = _per_rank_tokens(workload, world_size) + _, _, _enable_dp = _resolve_mapping_layout(config, world_size) + r.per_rank_num_tokens = _per_rank_tokens(workload, world_size, enable_dp=bool(_enable_dp)) r.status_per_rank = {f"rank{i}": "skipped" for i in range(world_size)} r.instrumentation = { "level": ",".join(sorted(analysis)) if analysis else "summary", diff --git a/tests/microbenchmarks/bench_moe/routing/builders.py b/tests/microbenchmarks/bench_moe/routing/builders.py index af3d792d756a..3f7ec660c9c4 100644 --- a/tests/microbenchmarks/bench_moe/routing/builders.py +++ b/tests/microbenchmarks/bench_moe/routing/builders.py @@ -115,43 +115,88 @@ def _build_per_rank_num_tokens( spec: RoutingControlSpec, num_tokens: int, world_size: int, + enable_dp: bool, ) -> List[int]: """Resolve ``per_rank_num_tokens`` for a workload. - Explicit ``spec.per_rank_num_tokens`` wins; otherwise tokens are split - evenly across ranks with any remainder on rank 0. + Explicit ``spec.per_rank_num_tokens`` wins; otherwise the token count per + rank depends on the attention-DP setting: + + * ``enable_dp=True`` (DEP / DTP): tokens are DP-sharded across ranks, so + each rank holds ``num_tokens / world_size``. + * ``enable_dp=False`` (TEP / TTP): attention is tensor-parallel, so every + rank sees the complete batch and holds ``num_tokens``. + + When an explicit list is provided its sum is validated against the expected + total (``num_tokens`` for DP modes, ``num_tokens * world_size`` for non-DP). """ if spec.per_rank_num_tokens is None: + if not enable_dp: + return [int(num_tokens)] * world_size return _distribute_tokens(int(num_tokens), world_size) + expected_total = int(num_tokens) * (1 if enable_dp else world_size) return _validate_per_rank_token_list( - spec.per_rank_num_tokens, world_size=world_size, expected_total=int(num_tokens) + spec.per_rank_num_tokens, world_size=world_size, expected_total=expected_total ) -def _per_rank_tokens(workload: WorkloadSpec, world_size: int) -> List[int]: +def _per_rank_tokens(workload: WorkloadSpec, world_size: int, enable_dp: bool) -> List[int]: """Materialize the ``per_rank_num_tokens`` list for a workload + world size.""" return _build_per_rank_num_tokens( - workload.routing_control, int(workload.num_tokens), world_size + workload.routing_control, int(workload.num_tokens), world_size, enable_dp ) +def _aggregate_dispatch_source_tokens( + per_rank_num_tokens: List[int], + ep_size: int, + enable_dp: bool, +) -> List[int]: + """Project world-rank token counts onto EP-source rows. + + TRT-LLM Mapping orders MoE ranks with ``moe_ep_rank = tp_rank % moe_ep_size``. + In attention-DP modes each world rank owns a distinct token shard, so TP + shards targeting the same EP row are summed. In non-DP MoE-TP modes those TP + shards carry the same logical tokens, so only the first TP shard contributes + to the logical dispatch plan. + """ + if ep_size <= 0: + return [] + if len(per_rank_num_tokens) == ep_size: + return [int(v) for v in per_rank_num_tokens] + + source_tokens = [0] * ep_size + if not enable_dp: + for ep_rank in range(ep_size): + if ep_rank < len(per_rank_num_tokens): + source_tokens[ep_rank] = int(per_rank_num_tokens[ep_rank]) + else: + for rank, num_tokens in enumerate(per_rank_num_tokens): + source_tokens[rank % ep_size] += int(num_tokens) + + return source_tokens + + def _build_dispatch_matrix( comm_pattern: str, per_rank_num_tokens: List[int], top_k: int, ep_size: int, + enable_dp: bool, seed: int = 0, ) -> List[List[int]]: """Build the canonical slot ``dispatch_matrix`` for ``comm_pattern``. - Row sums always equal ``per_rank_num_tokens[src] * top_k``. The matrix is - a pure planning artefact: it does not enforce per-token uniqueness yet. - That constraint is checked at materialisation time. + Row sums equal the EP-source token counts projected from + ``per_rank_num_tokens`` times ``top_k``. When world ranks outnumber EP + ranks (DTP / TTP / CUSTOM MoE-TP layouts), multiple world-rank rows are + aggregated onto the same EP-source row. """ name, kwargs = _parse_comm_pattern(comm_pattern) + source_tokens = _aggregate_dispatch_source_tokens(per_rank_num_tokens, ep_size, enable_dp) matrix: List[List[int]] = [[0] * ep_size for _ in range(ep_size)] for src in range(ep_size): - row_total = int(per_rank_num_tokens[src]) * int(top_k) + row_total = int(source_tokens[src]) * int(top_k) if row_total == 0: continue if name == "file": @@ -309,6 +354,7 @@ def _build_routing_plan( top_k: int, num_experts: int, moe_ep_size: int, + enable_dp: bool, ) -> RoutingPlan: """Translate a ``RoutingControlSpec`` into a canonical normalised plan.""" if moe_ep_size <= 0 or num_experts % moe_ep_size != 0: @@ -318,11 +364,10 @@ def _build_routing_plan( experts_per_rank = num_experts // moe_ep_size if top_k > num_experts: raise ValueError(f"top_k ({top_k}) must be <= num_experts ({num_experts})") - per_rank = _build_per_rank_num_tokens(spec, num_tokens, world_size) - # The dispatch matrix is indexed by EP rank on both axes. The current - # worker only calls routing-control planning when ``moe_ep_size`` equals - # ``world_size`` so that this EP-axis matrix also matches the user-visible - # per-rank token list. + per_rank = _build_per_rank_num_tokens(spec, num_tokens, world_size, enable_dp) + # The dispatch matrix stays on EP axes. When MoE-TP makes multiple world + # ranks share one EP rank, the world-rank token counts are aggregated onto + # the corresponding EP-source row before building the matrix. if spec.routing_pattern_file: default_patterns = {("balanced_alltoall", "balanced"), ("random", "random")} if (spec.comm_pattern, spec.expert_pattern) not in default_patterns: @@ -335,26 +380,32 @@ def _build_routing_plan( ) else: dispatch_matrix = _build_dispatch_matrix( - spec.comm_pattern, per_rank, top_k, moe_ep_size, seed=spec.seed + spec.comm_pattern, + per_rank, + top_k, + moe_ep_size, + enable_dp=enable_dp, + seed=spec.seed, ) expert_histogram = _build_expert_histogram( spec.expert_pattern, dispatch_matrix, experts_per_rank, moe_ep_size, seed=spec.seed ) # Per-row sums are an invariant; emit a clearer error than the materialiser would. + source_tokens = _aggregate_dispatch_source_tokens(per_rank, moe_ep_size, enable_dp) for src in range(moe_ep_size): - expected = int(per_rank[src]) * int(top_k) if src < len(per_rank) else 0 + expected = int(source_tokens[src]) * int(top_k) if src < len(source_tokens) else 0 actual = sum(dispatch_matrix[src]) if actual != expected: raise ValueError( - f"dispatch_matrix row {src} sums to {actual}, expected per_rank_num_tokens[{src}] * top_k = {expected}" + f"dispatch_matrix row {src} sums to {actual}, expected aggregate source tokens * top_k = {expected}" ) # Global expert histogram total must match total slots. - total_slots = sum(int(t) for t in per_rank) * int(top_k) + total_slots = sum(int(t) for t in source_tokens) * int(top_k) hist_total = sum(sum(row) for row in expert_histogram) if hist_total != total_slots: raise ValueError( - f"expert_histogram sum={hist_total} must equal sum(per_rank_num_tokens) * top_k = {total_slots}" + f"expert_histogram sum={hist_total} must equal aggregate source tokens * top_k = {total_slots}" ) return RoutingPlan( diff --git a/tests/microbenchmarks/bench_moe/routing/materialize.py b/tests/microbenchmarks/bench_moe/routing/materialize.py index fe43eedf9336..c3d1fc8a85e5 100644 --- a/tests/microbenchmarks/bench_moe/routing/materialize.py +++ b/tests/microbenchmarks/bench_moe/routing/materialize.py @@ -55,14 +55,21 @@ def _flatten_plan_slots_for_rank( experts_per_rank: int, moe_ep_size: int, ) -> List[int]: - """Flatten one plan row into expert ids while preserving slot counts.""" - local_num_tokens = int(plan.per_rank_num_tokens[src_rank]) + """Flatten one plan row into expert ids while preserving slot counts. + + ``local_num_tokens`` is derived from the dispatch-matrix row sum rather + than from ``per_rank_num_tokens[src_rank]``. In MoE-TP + attention-DP + layouts (DTP / CUSTOM-DP) the dispatch matrix is EP-axis indexed while + ``per_rank_num_tokens`` is world-rank indexed; the row sum is always the + correct EP-axis aggregate (``source_tokens[src_rank] * top_k``). + """ row = list(plan.dispatch_matrix[src_rank]) - if sum(row) != local_num_tokens * top_k: + row_sum = sum(row) + if top_k > 0 and row_sum % top_k != 0: raise ValueError( - f"dispatch_matrix row sum ({sum(row)}) must equal local_num_tokens*top_k " - f"({local_num_tokens * top_k}) for rank {src_rank}" + f"dispatch_matrix row {src_rank} sum ({row_sum}) is not divisible by top_k ({top_k})" ) + local_num_tokens = row_sum // top_k if top_k > 0 else 0 flat: List[int] = [] for dst in range(moe_ep_size): @@ -174,7 +181,13 @@ def _materialize_selected_experts_for_rank( 4. Run a small repair pass that swaps duplicated expert ids between rows until each token has ``top_k`` distinct experts. """ - local_num_tokens = int(plan.per_rank_num_tokens[src_rank]) + # Derive the effective token count from the dispatch-matrix row sum so that + # MoE-TP + attention-DP layouts (DTP / CUSTOM-DP) are handled correctly. + # In those layouts the row sum equals the aggregated source tokens for the + # EP rank, while per_rank_num_tokens[src_rank] would only reflect one DP + # shard's contribution. + row_sum = sum(plan.dispatch_matrix[src_rank]) + local_num_tokens = row_sum // max(top_k, 1) if local_num_tokens == 0: ids = torch.zeros((0, top_k), dtype=torch.int32, device=device) scales = torch.zeros((0, top_k), dtype=scale_dtype, device=device) diff --git a/tests/microbenchmarks/bench_moe/routing/native_logits.py b/tests/microbenchmarks/bench_moe/routing/native_logits.py index 449c38dd71ac..b28c0a65bd5a 100644 --- a/tests/microbenchmarks/bench_moe/routing/native_logits.py +++ b/tests/microbenchmarks/bench_moe/routing/native_logits.py @@ -123,7 +123,13 @@ def _project_router_logits_for_plan( Returns ``(router_logits, status, reason)`` where ``status`` is one of ``"exact"``, ``"projected"``, or ``"rejected"``. """ - local_num_tokens = int(plan.per_rank_num_tokens[src_rank]) + # Derive the effective token count from the dispatch-matrix row sum. + # In MoE-TP + attention-DP layouts (DTP / CUSTOM-DP) the row sum equals + # the aggregated source tokens for the EP rank (which is what the router + # sees after the in-MoE allgather), while per_rank_num_tokens[src_rank] + # would only cover one DP shard. + row_sum = sum(plan.dispatch_matrix[src_rank]) + local_num_tokens = row_sum // max(top_k, 1) if row_sum > 0 else 0 if local_num_tokens == 0: return ( torch.empty((0, num_experts), dtype=dtype, device=device), diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index 49f539f212e7..f692911e7ff2 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -24,15 +24,34 @@ import torch +from tensorrt_llm._utils import local_mpi_size from tensorrt_llm.models.modeling_utils import QuantAlgo from .backend import MoeBackendType, get_backend_class -from .mapping import _resolve_mapping_layout +from .mapping import _PARALLEL_MODE_LAYOUTS, _resolve_mapping_layout from .specs import _ALL_BACKENDS, _FORCED_COMM_ENV_VALUES, ConfigSpec, ModelSpec, SearchSpec _FUSED_COMM_BACKENDS = frozenset({"MEGAMOE_DEEPGEMM"}) +def _is_deepep_feasible(num_ranks: int) -> bool: + """Return True if DeepEP supports the given EP rank count on this node topology. + + Intranode: num_ranks in {2, 4, 8} and num_ranks == local_mpi_size(). + Internode: exactly 8 ranks per node, with 2/4/8/16 RDMA nodes. + Mirrors the feasibility check in fused_moe_wide_ep.py::select_alltoall_method_type. + """ + _INTRANODE_RANKS = {2, 4, 8} + _REQUIRED_LOCAL_SIZE = 8 + _INTERNODE_RDMA_NODES = {2, 4, 8, 16} + mpi_size = local_mpi_size() + if num_ranks == mpi_size and num_ranks in _INTRANODE_RANKS: + return True + if mpi_size != _REQUIRED_LOCAL_SIZE: + return False + return (num_ranks // mpi_size) in _INTERNODE_RDMA_NODES + + def _check_backend_can_implement( backend_str: str, quant_algo: Optional[QuantAlgo], @@ -65,6 +84,23 @@ def _comm_axis_for_backend(backend: Any, comm_methods: Tuple[Any, ...]) -> Tuple return comm_methods +def _comm_axis_for_parallel_mode(pmode: str, comm_methods: Tuple[Any, ...]) -> Tuple[Any, ...]: + """Collapse comm axis to AUTO for parallel modes without attention DP. + + Non-AUTO forced comm methods require enable_attention_dp=True (see + is_candidate_valid). TEP and TTP have enable_dp=False, so only AUTO + is ever valid for them. Generating forced-comm candidates for these + modes only produces prune rows — handle it at generation time instead. + CUSTOM mode is passed through unchanged (validated separately). + """ + layout = _PARALLEL_MODE_LAYOUTS.get(str(pmode).upper()) + if layout is None: + return comm_methods # CUSTOM: unknown layout, keep as-is + if not layout["enable_attention_dp"]: + return ("AUTO",) + return comm_methods + + def expand_search( base_config: ConfigSpec, search: SearchSpec, @@ -88,7 +124,13 @@ def expand_search( for backend, pmode, cgraph, combine in itertools.product( backends, parallel_modes, cuda_graph_options, combine_options ): - for comm in _comm_axis_for_backend(backend, comm_methods): + effective_comm = _comm_axis_for_backend(backend, comm_methods) + # For non-fused backends apply parallel-mode comm constraint at + # generation time so TEP/TTP always get comm=AUTO instead of + # generating forced-comm candidates that are immediately pruned. + if effective_comm != ("NONE",): + effective_comm = _comm_axis_for_parallel_mode(pmode, effective_comm) + for comm in effective_comm: candidate = replace( base_config, backend=str(backend).upper(), @@ -121,6 +163,68 @@ def is_candidate_valid( except ValueError as exc: return False, str(exc) + # DenseGEMM only supports TP; any EP configuration (TEP, DEP, custom ep>1) is unsupported. + if config.backend.upper() == "DENSEGEMM" and moe_ep > 1: + return False, ( + f"DENSEGEMM does not support EP (ep_size={moe_ep}); " + "use TEP/DEP only with other backends" + ) + + # MegaMoEDeepGemm is EP-only (asserts moe_tp_size == 1 in __init__); DTP/TTP are invalid. + if config.backend.upper() == "MEGAMOE_DEEPGEMM" and moe_tp > 1: + return False, ( + f"MEGAMOE_DEEPGEMM does not support MoE-TP (moe_tp_size={moe_tp}); " + "use DEP/TEP modes only" + ) + + # DENSEGEMM DTP: FC2 kernel requires (intermediate_size / moe_tp_size) % 256 == 0. + # DENSEGEMM __init__ only checks the full intermediate_size, so a model like + # DeepSeek V3 (intermediate_size=2048, 2048%256=0) passes __init__ but fails + # at runtime with moe_tp_size=16 (2048/16=128, 128%256!=0). + if config.backend.upper() == "DENSEGEMM" and moe_ep == 1 and moe_tp > 1: + if model.intermediate_size % moe_tp != 0: + return False, ( + f"DENSEGEMM DTP: intermediate_size={model.intermediate_size} " + f"not divisible by moe_tp_size={moe_tp}" + ) + per_tp_k = model.intermediate_size // moe_tp + _DENSEGEMM_MMA_TILE_K = 256 + if per_tp_k % _DENSEGEMM_MMA_TILE_K != 0: + return False, ( + f"DENSEGEMM DTP moe_tp_size={moe_tp}: intermediate_size/tp={per_tp_k} " + f"not aligned to FC2 MMA tile-K={_DENSEGEMM_MMA_TILE_K}" + ) + + # NVFP4 on CuteDSL / TRTLLM-Gen requires the per-partition intermediate size + # (intermediate_size / moe_tp_size) to be a multiple of the NVFP4 weight + # alignment (128). Unlike CUTLASS (which pads intermediate_size_per_partition + # up to 128), these backends use the unpadded logical size when laying out the + # block-scale tensor and fail during weight load: CUTEDSL raises a reshape + # RuntimeError (e.g. "shape '[-1, 192, 448]' is invalid for input of size + # 114688" — 192 padded to 256) and TRTLLM-Gen hits `assert intermediate_size % + # weight_alignment == 0`. Prune the unsupported combo with a clear reason + # instead of letting it crash mid-sweep. Example: DeepSeek-V4-Pro + # (intermediate_size=3072) at moe_tp_size=32 -> 3072/32=96, 96%128!=0. + if ( + config.backend.upper() in ("CUTEDSL", "TRTLLM") + and model.quant_algo_enum == QuantAlgo.NVFP4 + and moe_tp > 1 + ): + _NVFP4_WEIGHT_ALIGNMENT = 128 + if model.intermediate_size % moe_tp != 0: + return False, ( + f"{config.backend.upper()} NVFP4: intermediate_size=" + f"{model.intermediate_size} not divisible by moe_tp_size={moe_tp}" + ) + per_tp_k = model.intermediate_size // moe_tp + if per_tp_k % _NVFP4_WEIGHT_ALIGNMENT != 0: + return False, ( + f"{config.backend.upper()} NVFP4 moe_tp_size={moe_tp}: " + f"intermediate_size/tp={per_tp_k} not aligned to NVFP4 weight " + f"alignment={_NVFP4_WEIGHT_ALIGNMENT} (CUTLASS pads to 128, " + f"CUTEDSL/TRTLLM do not)" + ) + # Forced communication on non-DP / MoE-TP paths. forced = config.comm_method.upper() if forced not in ("AUTO", "NONE"): @@ -130,6 +234,12 @@ def is_candidate_valid( return False, f"comm_method={forced} requires moe_tp_size=1 (got {moe_tp})" if world_size == 1: return False, f"comm_method={forced} has no effect at world_size=1" + if forced == "DEEPEP" and not _is_deepep_feasible(moe_ep): + return False, ( + f"comm_method={forced}: moe_ep_size={moe_ep} not supported by DeepEP topology " + f"(local_mpi_size={local_mpi_size()}; supported: intranode {{2,4,8}}, " + f"internode 8-ranks/node x {{2,4,8,16}} nodes)" + ) return True, None diff --git a/tests/microbenchmarks/bench_moe/utils.py b/tests/microbenchmarks/bench_moe/utils.py index 2000cbbc8e62..d4874f5d91e1 100644 --- a/tests/microbenchmarks/bench_moe/utils.py +++ b/tests/microbenchmarks/bench_moe/utils.py @@ -106,15 +106,18 @@ def _compute_stats(values: List[float]) -> Dict[str, float]: def _distribute_tokens(total: int, world_size: int) -> List[int]: - """Distribute ``total`` global tokens evenly across ``world_size`` ranks.""" + """Distribute ``total`` global tokens evenly across ranks. + + Remainder tokens are spread one-per-rank over the leading ranks (instead of + piling the entire remainder on rank 0), so e.g. (total=2, world_size=4) -> + [1, 1, 0, 0]. An even, non-degenerate split keeps every rank's per-rank token + count within 1 of each other, which the downstream symmetric-memory workspace + sizing relies on. + """ if world_size <= 0 or total < 0: raise ValueError(f"invalid args: total={total}, world_size={world_size}") - if world_size == 1: - return [total] - base = total // world_size - out = [base] * world_size - out[0] += total - base * world_size - return out + base, rem = divmod(total, world_size) + return [base + (1 if i < rem else 0) for i in range(world_size)] def _validate_per_rank_token_list( diff --git a/tests/microbenchmarks/bench_moe/worker.py b/tests/microbenchmarks/bench_moe/worker.py index 1019b6bfc9a5..e4c8e44804d2 100644 --- a/tests/microbenchmarks/bench_moe/worker.py +++ b/tests/microbenchmarks/bench_moe/worker.py @@ -28,7 +28,7 @@ import time import traceback from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import torch from mpi4py import MPI @@ -70,6 +70,18 @@ def _try_import(module_path: str, attr: Optional[str] = None, default: Any = Non POISON_HERE_PREFIX = "cuda_context_poisoned_after_success" POISON_UPSTREAM_PREFIX = "cuda_context_poisoned_upstream" WATCHDOG_UPSTREAM_PREFIX = "watchdog_timeout_upstream" +# Terminal (status="failed") marker for a candidate the watchdog killed for +# exceeding its wall-clock budget. NOT suffixed "_upstream": is_completed_for_resume +# treats status="failed" as terminal, so --resume_from SKIPS it (does not re-attempt +# and re-hang) while still surfacing the hang as a result row with a clear reason. +WATCHDOG_TIMEOUT_PREFIX = "watchdog_timeout" +# Terminal (status="failed") placeholder pre-written for the in-flight candidate +# BEFORE it runs. If the process dies mid-candidate in a way nothing else can +# record -- a CUDA device-side assert that aborts the MPI step, OOM-kill, +# SIGSEGV, node loss -- this persisted row makes the candidate terminal so +# --resume_from skips it and advances. Replaced with the real result on normal +# completion. Like WATCHDOG_TIMEOUT_PREFIX, NOT suffixed "_upstream". +INCOMPLETE_PREFIX = "incomplete" BENCH_MOE_POISON_EXIT_CODE = 75 @@ -182,11 +194,29 @@ def allreduce_poison_reason(local_reason: Optional[str]) -> Optional[str]: class CandidateWatchdog: - """Hard wall-clock guard around one candidate; SIGKILLs the process on timeout.""" + """Hard wall-clock guard around one candidate; SIGKILLs the process on timeout. + + On timeout the guard first invokes ``on_timeout`` (used to record the hung + candidate as a terminal ``failed`` result + checkpoint so it is not silently + lost and is skipped on ``--resume_from`` rather than re-attempted), then + SIGKILLs to break the wedged CUDA/NCCL state. A genuine hang cannot be + recovered in-process, so the kill is unavoidable; ``on_timeout`` makes it a + recorded outcome instead of a vanished one. + """ - def __init__(self, budget_s: float, label: str): + def __init__( + self, + budget_s: float, + label: str, + on_timeout: Optional[Callable[[], None]] = None, + rank0_persist_grace_s: float = 8.0, + ): self._budget_s = float(budget_s) self._label = label + self._on_timeout = on_timeout + # Non-rank-0 ranks wait this long before SIGKILL so rank 0 can persist the + # checkpoint before the first task exit tears down the whole srun step. + self._rank0_persist_grace_s = float(rank0_persist_grace_s) self._cancelled = threading.Event() self._thread: Optional[threading.Thread] = None @@ -211,16 +241,35 @@ def __exit__(self, exc_type, exc, tb) -> bool: def _guard(self) -> None: if self._cancelled.wait(self._budget_s): return + rank = mpi_rank() try: sys.stderr.write( f"[bench_moe watchdog] candidate '{self._label}' exceeded " f"{self._budget_s:.1f}s budget on pid={os.getpid()} " - f"rank={mpi_rank()}; sending SIGKILL to break suspected " - f"NCCL deadlock or CUDA hang.\n" + f"rank={rank}; recording it as a failed (timeout) result, then " + f"sending SIGKILL to break suspected NCCL deadlock or CUDA hang.\n" ) sys.stderr.flush() except Exception: pass + # Record the hung candidate as a terminal failed row + checkpoint. Rank 0 + # writes; other ranks no-op (see _emit_checkpoint_report). The main thread + # is blocked in a GIL-releasing CUDA call, so this watchdog thread can run. + if self._on_timeout is not None: + try: + self._on_timeout() + except Exception as exc: # never let bookkeeping block the kill + try: + sys.stderr.write( + f"[bench_moe watchdog] on_timeout callback failed " + f"({type(exc).__name__}: {exc}); killing anyway.\n" + ) + sys.stderr.flush() + except Exception: + pass + # Let rank 0 flush its checkpoint before the first SIGKILL aborts the step. + if rank != 0 and self._rank0_persist_grace_s > 0: + time.sleep(self._rank0_persist_grace_s) os.kill(os.getpid(), signal.SIGKILL) @@ -518,8 +567,52 @@ def _run_benchmark_worker_under_current_mpi(args: argparse.Namespace, launcher: ) _maybe_print_rank0(f"[bench_moe] running {case_label}") + # Pre-write a terminal "failed" placeholder for THIS candidate and + # checkpoint it BEFORE running. If the process then dies mid-candidate in + # a way nothing else can catch -- a CUDA device-side assert that aborts + # the MPI step, OOM-kill, SIGSEGV, node loss -- this persisted row keeps + # the candidate terminal, so --resume_from skips it and advances to the + # next one instead of re-attempting (and re-crashing on) the same + # candidate forever. On normal completion it is replaced with the real + # result below. Only the in-flight candidate gets a placeholder; + # not-yet-run candidates have no row and are still attempted on resume. + placeholder = _make_skipped_run_result( + model=ctx.model, + workload=workload, + config=cand, + world_size=world_size, + analysis=ctx.analysis, + reason=( + f"{INCOMPLETE_PREFIX}: process died before this candidate " + f"finished (crash/abort/OOM/kill) ({case_label})" + ), + ) + placeholder.status = "failed" + placeholder.status_per_rank = {f"rank{i}": "incomplete" for i in range(world_size)} + accumulated_rows.append(_runresult_to_row(placeholder)) + _emit_checkpoint_report(args=args, ctx=ctx, rows=accumulated_rows, world_size=world_size) + + # If the watchdog fires (suspected hang), overwrite the placeholder's + # reason with the precise timeout text and re-checkpoint before SIGKILL, + # so the hang is surfaced as a clear result (the row is already terminal). + def _record_watchdog_timeout( + _label: str = case_label, + _budget_s: float = watchdog_budget_s, + ) -> None: + if accumulated_rows: + accumulated_rows[-1]["skip_reason"] = ( + f"{WATCHDOG_TIMEOUT_PREFIX}: exceeded {_budget_s:.0f}s; " + f"suspected NCCL/CUDA hang ({_label})" + ) + accumulated_rows[-1]["status_per_rank"] = { + f"rank{i}": "timeout" for i in range(world_size) + } + _emit_checkpoint_report( + args=args, ctx=ctx, rows=accumulated_rows, world_size=world_size + ) + # Hard wall-clock guard around the actual candidate execution. - with CandidateWatchdog(watchdog_budget_s, case_label): + with CandidateWatchdog(watchdog_budget_s, case_label, on_timeout=_record_watchdog_timeout): with torch.device(device): r = _run_one_candidate( model=ctx.model, @@ -539,8 +632,10 @@ def _run_benchmark_worker_under_current_mpi(args: argparse.Namespace, launcher: input_cache=input_cache, enable_perfect_router_requested=bool(args.enable_perfect_router), ) + # Candidate finished normally: replace the pre-written placeholder (the + # last row) with the real result. row = _runresult_to_row(r) - accumulated_rows.append(row) + accumulated_rows[-1] = row if rank == 0: print(json.dumps(row, indent=2), flush=True) diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/modules/moe/quantize_utils.py index fcf1e3abcf5d..2633e839284e 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/modules/moe/quantize_utils.py @@ -2011,21 +2011,21 @@ def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]: w1_weight, None, scaling_vector_size, True ) w1_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w1_sf_block.cpu().view(intermediate_size, -1) + w1_sf_block.view(intermediate_size, -1) ) w2_weight_mxfp4, w2_sf_block = torch.ops.trtllm.fp4_quantize( w2_weight, None, scaling_vector_size, True ) w2_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w2_sf_block.cpu().view(hidden_size_out, -1) + w2_sf_block.view(hidden_size_out, -1) ) w3_weight_mxfp4, w3_sf_block = torch.ops.trtllm.fp4_quantize( w3_weight, None, scaling_vector_size, True ) w3_sf_block_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - w3_sf_block.cpu().view(intermediate_size, -1) + w3_sf_block.view(intermediate_size, -1) ) weights[f"{expert_id}.w1.weight"] = w1_weight_mxfp4 diff --git a/tests/unittest/_torch/modules/moe/test_moe_module.py b/tests/unittest/_torch/modules/moe/test_moe_module.py index 77277f705f46..f9fcaec58a9c 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_module.py +++ b/tests/unittest/_torch/modules/moe/test_moe_module.py @@ -527,7 +527,9 @@ def _test_moe_worker_impl( # Setup mapping mapping = mapping or Mapping() mapping.rank = mpi_rank() - all_rank_num_tokens = [seq_len] * mapping.world_size + # DP modes need per-rank token counts for cross-rank dispatch; non-DP modes + # (TTP/TEP) have no dispatch and the scheduler derives [x.shape[0]] from None. + all_rank_num_tokens = [seq_len] * mapping.world_size if mapping.enable_attention_dp else None torch.cuda.set_device(mapping.rank) _ensure_dist_for_megamoe(moe_backend, mapping.rank, mapping.world_size)