remove some code

Author: Binyang2014

python/mscclpp_benchmark/bench_collective.py

@@ -69,7 +69,6 @@ class DTypeSpec:
     cupy_dtype: Any
     mscclpp_dtype: Any
     accum_dtype: Any | None = None
-    supports_reduction_correctness: bool = True
     fp8_format: str | None = None
 
 
@@ -79,7 +78,6 @@ class CandidateSpec:
     min_message_size: int | None = None
     max_message_size: int | None = None
     max_nblocks: int | None = None
-    min_nthreads: int | None = None
     supported_skus: tuple[str, ...] | None = None
     requires_nvls: bool = False
     requires_symmetric_memory: bool = False
@@ -93,7 +91,6 @@ class BenchmarkCase:
     input: cp.ndarray
     output: cp.ndarray
     dtype_spec: DTypeSpec
-    allgather_mode: str
     symmetric_memory: bool = False
 
 
@@ -170,7 +167,6 @@ def _with_accum_type(dtype_spec: DTypeSpec, accum_type: str | None) -> DTypeSpec
         cupy_dtype=dtype_spec.cupy_dtype,
         mscclpp_dtype=dtype_spec.mscclpp_dtype,
         accum_dtype=accum_dtype,
-        supports_reduction_correctness=dtype_spec.supports_reduction_correctness,
         fp8_format=dtype_spec.fp8_format,
     )
 
@@ -193,36 +189,40 @@ def _parse_int_list(raw: str | None, default: tuple[int, ...]) -> tuple[int, ...
     return values
 
 
-def _candidate_specs(
-    collective: str, message_size: int, *, symmetric_memory: bool = False
-) -> tuple[CandidateSpec, ...]:
+def _candidate_specs(collective: str, *, symmetric_memory: bool = False) -> tuple[CandidateSpec, ...]:
     if collective == _ALLGATHER:
         return (CandidateSpec("default_allgather_fullmesh2", max_nblocks=64, supported_skus=("MI300X",)),)
     if collective != _ALLREDUCE:
         raise ValueError(f"Unsupported collective: {collective}")
-    if message_size <= 512 * 1024:
-        candidates = (
-            CandidateSpec(
-                "default_allreduce_nvls_packet",
-                max_nblocks=16,
-                supported_skus=("H100", "GB300"),
-                requires_nvls=True,
-            ),
-            CandidateSpec("default_allreduce_packet", max_nblocks=56),
-            CandidateSpec("default_allreduce_allpair_packet", max_nblocks=56),
-        )
-    elif message_size <= 4 * 1024 * 1024:
-        candidates = (
-            CandidateSpec("default_allreduce_packet", max_nblocks=56),
-            CandidateSpec("default_allreduce_allpair_packet", max_nblocks=56),
-            CandidateSpec("default_allreduce_rsag_zero_copy"),
-            CandidateSpec("default_allreduce_fullmesh", max_nblocks=64, supported_skus=("MI300X",)),
-        )
-    else:
-        candidates = (
-            CandidateSpec("default_allreduce_rsag_zero_copy"),
-            CandidateSpec("default_allreduce_fullmesh", max_nblocks=64, supported_skus=("MI300X",)),
-        )
+    candidates = (
+        CandidateSpec(
+            "default_allreduce_nvls_packet",
+            max_message_size=512 * 1024,
+            max_nblocks=16,
+            supported_skus=("H100", "GB300"),
+            requires_nvls=True,
+        ),
+        CandidateSpec(
+            "default_allreduce_packet",
+            max_message_size=4 * 1024 * 1024,
+            max_nblocks=56,
+        ),
+        CandidateSpec(
+            "default_allreduce_allpair_packet",
+            max_message_size=4 * 1024 * 1024,
+            max_nblocks=56,
+        ),
+        CandidateSpec(
+            "default_allreduce_rsag_zero_copy",
+            min_message_size=512 * 1024 + 1,
+        ),
+        CandidateSpec(
+            "default_allreduce_fullmesh",
+            min_message_size=512 * 1024 + 1,
+            max_nblocks=64,
+            supported_skus=("MI300X",),
+        ),
+    )
     if symmetric_memory:
         return (
             CandidateSpec(
@@ -244,20 +244,17 @@ def _candidate_algorithms(comm: Comm, case: BenchmarkCase) -> list[tuple[Any, Ca
     symmetric_memory = case.symmetric_memory
     profile = getattr(comm, "hardware_profile", None)
     filtered_out = False
-    for candidate in _candidate_specs(case.collective, case.message_size, symmetric_memory=symmetric_memory):
+    for candidate in _candidate_specs(case.collective, symmetric_memory=symmetric_memory):
         if not _candidate_supports_profile(candidate, profile):
             filtered_out = True
             continue
-        if candidate.requires_nvls and not _mscclpp().is_nvls_supported():
-            filtered_out = True
-            continue
-        if candidate.requires_symmetric_memory and not symmetric_memory:
+        if not _candidate_supports_message_size(candidate, case.message_size):
             filtered_out = True
             continue
-        if candidate.min_message_size is not None and case.message_size < candidate.min_message_size:
+        if candidate.requires_nvls and not _mscclpp().is_nvls_supported():
             filtered_out = True
             continue
-        if candidate.max_message_size is not None and case.message_size > candidate.max_message_size:
+        if candidate.requires_symmetric_memory and not symmetric_memory:
             filtered_out = True
             continue
         algorithm = available.get(candidate.algorithm)
@@ -281,6 +278,14 @@ def _candidate_supports_profile(candidate: CandidateSpec, profile: HardwareProfi
     return sku in candidate.supported_skus
 
 
+def _candidate_supports_message_size(candidate: CandidateSpec, message_size: int) -> bool:
+    if candidate.min_message_size is not None and message_size < candidate.min_message_size:
+        return False
+    if candidate.max_message_size is not None and message_size > candidate.max_message_size:
+        return False
+    return True
+
+
 def _make_case(
     *,
     collective: str,
@@ -299,7 +304,6 @@ def _make_case(
             input=memory,
             output=memory,
             dtype_spec=dtype_spec,
-            allgather_mode=allgather_mode,
             symmetric_memory=symmetric_memory,
         )
 
@@ -323,7 +327,6 @@ def _make_case(
         input=input_buffer,
         output=output,
         dtype_spec=dtype_spec,
-        allgather_mode=allgather_mode,
         symmetric_memory=symmetric_memory,
     )
 
@@ -522,7 +525,6 @@ def main(argv: list[str] | None = None) -> None:
         candidate_algorithms=_candidate_algorithms,
         check_correctness=_check_correctness,
         measure=_try_measure_case,
-        symmetric_memory=args.symmetric_memory,
     )
 
     rows: list[list[str]] = []

python/mscclpp_benchmark/comm.py

@@ -4,8 +4,7 @@
 from __future__ import annotations
 
 import logging
-from contextlib import contextmanager
-from typing import Any, Iterator
+from typing import Any
 
 logger = logging.getLogger(__name__)
 _ALLREDUCE_COLLECTIVE = "allreduce"
@@ -60,41 +59,6 @@ def data_ptr(self) -> int:
         return int(self.buffer.data())
 
 
-@contextmanager
-def init_mscclpp_comm_group_for_ranks(ranks: list[int], *, name: str) -> Iterator[Any]:
-    del name
-    from mpi4py import MPI
-
-    world_comm = MPI.COMM_WORLD
-    unique_ranks = [int(rank) for rank in ranks]
-    if len(unique_ranks) != len(set(unique_ranks)):
-        raise ValueError(f"Duplicate ranks are not allowed: {unique_ranks}")
-    if world_comm.Get_rank() not in unique_ranks:
-        raise ValueError(f"Rank {world_comm.Get_rank()} is not a member of subgroup {unique_ranks}")
-
-    sub_comm = None
-    if unique_ranks == list(range(world_comm.Get_size())):
-        mpi_comm = world_comm
-    else:
-        subgroup = world_comm.group.Incl(unique_ranks)
-        sub_comm = world_comm.Create_group(subgroup)
-        mpi_comm = sub_comm
-
-    _ensure_device()
-    comm_group = _mscclpp().CommGroup(mpi_comm)
-    setattr(comm_group, "_mpi_comm", mpi_comm)
-    try:
-        yield comm_group
-    finally:
-        destroy = getattr(comm_group, "destroy", None)
-        if callable(destroy):
-            destroy()
-        if sub_comm is not None:
-            free = getattr(sub_comm, "Free", None)
-            if callable(free):
-                free()
-
-
 class _AllReduceOp:
     def __init__(self, comm: "Comm", x: Any, *, symmetric_memory: bool = False) -> None:
         self._comm = comm

python/mscclpp_benchmark/correctness.py

@@ -55,12 +55,6 @@ def check_correctness(
     *,
     niter: int = 1,
 ) -> CorrectnessStats:
-    if case.collective == "allreduce" and not case.dtype_spec.supports_reduction_correctness:
-        raise ValueError(
-            f"Correctness checking for {case.collective} with {case.dtype_spec.name} is not implemented; "
-            "use --skip-correctness or a numeric dtype"
-        )
-
     all_ok = True
     local_max_abs_diff = 0.0
     local_sum_abs_diff = 0.0

python/mscclpp_benchmark/tuner.py

@@ -23,7 +23,6 @@ def __init__(
         candidate_algorithms: Callable[[Any, Any], list[tuple[Any, Any]]],
         check_correctness: Callable[..., bool],
         measure: Callable[..., float | None],
-        symmetric_memory: bool,
     ) -> None:
         self.comm = comm
         self.candidate_nblocks = tuple(candidate_nblocks)
@@ -34,12 +33,11 @@ def __init__(
         self._candidate_algorithms = candidate_algorithms
         self._check_correctness = check_correctness
         self._measure = measure
-        self._symmetric_memory = symmetric_memory
 
     def tune(self, case: Any) -> TunedConfig | None:
         best_config: TunedConfig | None = None
         best_time_us = float("inf")
-        symmetric_memory = bool(getattr(case, "symmetric_memory", self._symmetric_memory))
+        symmetric_memory = bool(getattr(case, "symmetric_memory", False))
         candidates = self._candidate_algorithms(self.comm, case)
         if not candidates:
             if self.comm.rank == 0:
@@ -54,8 +52,6 @@ def tune(self, case: Any) -> TunedConfig | None:
                 if candidate_spec.max_nblocks is not None and nblocks > candidate_spec.max_nblocks:
                     continue
                 for nthreads in self.candidate_nthreads:
-                    if candidate_spec.min_nthreads is not None and nthreads < candidate_spec.min_nthreads:
-                        continue
                     config = TunedConfig(
                         algorithm=algorithm.name,
                         nblocks=nblocks,