Create Reduce-Scatter Primitive Benchmark

Summary:
This diff introduces a new benchmark for evaluating the performance of the reduce-scatter collective in the `PooledEmbeddingsReduceScatter` module from `dist_data.py`. 
The benchmark is integrated with the existing `benchmark_primitive.py` file. 
The goal is to measure the reduce-scatter performance and analyze its impact on the overall efficiency of the `output_dist` process.

### Key Changes:

*   Updated `benchmark_primitive.py` to include the reduce-scatter benchmark.
*   Added comments to explain the newly added benchmark in the context of `PooledEmbeddingsReduceScatter` and its usage in `output_dist`.

Differential Revision: D113960618

Author: Siqi Huang

torchrec/distributed/benchmark/benchmark_primitive.py

@@ -24,7 +24,10 @@
 ``dist_data.py``). The second is ``kt_a2a``, the All-to-All performance of
 ``PooledEmbeddingsAllToAll`` -- the dense pooled-embedding (``KeyedTensor``) collective
 ``output_dist`` uses to redistribute real embedding outputs. Unlike ``kjt_a2a`` it
-exchanges float tensors rather than sparse indices.
+exchanges float tensors rather than sparse indices. The third is ``reduce_scatter``, the
+reduce-scatter performance of ``PooledEmbeddingsReduceScatter`` (from ``dist_data.py``) --
+the collective ``output_dist`` uses instead of the A2A for row-wise / table-row-wise
+sharding, summing each rank's partial pooled embeddings and scattering the batch dimension.
 
 A follow-up launcher binary will call ``runner`` explicitly with options to run on
 MAST or locally.
@@ -36,7 +39,11 @@
 
 import torch
 from torchrec.distributed.benchmark.base import benchmark_func, BenchmarkResult
-from torchrec.distributed.dist_data import KJTAllToAll, PooledEmbeddingsAllToAll
+from torchrec.distributed.dist_data import (
+    KJTAllToAll,
+    PooledEmbeddingsAllToAll,
+    PooledEmbeddingsReduceScatter,
+)
 from torchrec.distributed.test_utils.process_runner import SingleProcessContext
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
@@ -358,12 +365,150 @@ def _benchmark_kt_a2a(
     return result
 
 
+def _make_reduce_scatter_input(
+    batch_size: int,
+    dim: int,
+    values_dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    """Build this rank's input for ``PooledEmbeddingsReduceScatter``.
+
+    Returns a ``[batch_size, dim]`` tensor standing in for this rank's *partial* pooled
+    embeddings over the full global batch. reduce-scatter sums these partials across ranks
+    and scatters the batch dimension, so each rank ends up with
+    ``[batch_size // world_size, dim]`` -- its slice of the reduced result. Content is
+    random -- only the transport size matters for this benchmark, not the values.
+    """
+    return torch.rand((batch_size, dim), dtype=values_dtype, device=device)
+
+
+def _run_reduce_scatter(
+    _batch_inputs: List[Any],
+    *,
+    rs: PooledEmbeddingsReduceScatter,
+    local_embs: torch.Tensor,
+) -> None:
+    """One measured iteration: full reduce-scatter, then touch the output.
+
+    Rank alignment against the straggler effect is handled by ``PerfWrapper`` (it barriers
+    before each iteration, outside the timing window); this function only runs the
+    collective. ``PooledEmbeddingsReduceScatter`` returns a single-stage awaitable -- one
+    ``wait()`` sums each rank's ``local_embs`` across the group and scatters the batch
+    dimension, yielding this rank's ``[batch_size // world_size, dim]`` slice. ``numel()``
+    only reads the output's shape metadata: no data read, no kernel, and -- like ``wait()``
+    on CUDA -- no host sync (the collective runs async on the stream). The input tensor is
+    reused across iterations.
+
+    So we ``torch.cuda.synchronize()`` at the end to actually block the host on the
+    collective inside the measured region; that is what makes the wall-clock timer reflect
+    end-to-end collective latency (GPU-event timing is unaffected either way).
+    """
+    out = rs(local_embs).wait()
+    out.numel()
+    if local_embs.is_cuda:
+        torch.cuda.synchronize(local_embs.device)
+
+
+def _benchmark_reduce_scatter(
+    ctx: SingleProcessContext,
+    rank: int,
+    world_size: int,
+    **kwargs: Any,
+) -> BenchmarkResult:
+    """``PooledEmbeddingsReduceScatter`` (dense pooled-embedding reduce-scatter) benchmark.
+
+    Builds a dense pooled-embedding input tensor of a configurable size, then measures the
+    latency of reducing-and-scattering it through ``PooledEmbeddingsReduceScatter`` (the
+    ``dist_data.py`` module, the same one row-wise / table-row-wise / grid sharding
+    instantiate for ``output_dist``) over ``ctx.pg`` -- each rank holds partial pooled sums
+    for the global batch that must be summed across ranks and scattered back to each rank's
+    local batch slice. Correctness of the reduced output is intentionally not verified.
+
+    Args:
+        ctx: live single-process context (device + process group) injected by the
+            process runner; use ``ctx.device`` / ``ctx.pg`` directly.
+        rank: this process' global rank.
+        world_size: total number of ranks.
+        **kwargs: benchmark options:
+            batch_size (int): global batch size (rows of the input) -- must be divisible by
+                ``world_size`` (reduce-scatter splits the batch evenly across ranks).
+                Default 32 * 1024.
+            dim (int): embedding width. The headline transport size is ``batch_size * dim``
+                (with the defaults, ``32768 * 3072 * 4B ~= 400 MB`` of float32 per rank,
+                comparable to ``kt_a2a``). Default 3072.
+            values_dtype (torch.dtype): dtype of the embedding tensor; must be a floating
+                dtype. Default float32.
+            num_benchmarks (int): number of measured iterations. Default 20.
+            num_profiles (int): number of profiled iterations (requires profile_dir).
+                Default 0.
+            profile_dir (str): directory for chrome traces; empty disables profiling.
+            name (str): human-readable benchmark name. Default "reduce_scatter".
+
+    Returns:
+        This rank's ``BenchmarkResult``.
+    """
+    batch_size: int = int(kwargs.get("batch_size", 32 * 1024))
+    dim: int = int(kwargs.get("dim", 3072))
+    values_dtype: torch.dtype = kwargs.get("values_dtype", torch.float32)
+    num_benchmarks: int = int(kwargs.get("num_benchmarks", 20))
+    num_profiles: int = int(kwargs.get("num_profiles", 0))
+    profile_dir: str = str(kwargs.get("profile_dir", ""))
+    name: str = str(kwargs.get("name", "reduce_scatter"))
+
+    pg: Optional[torch.distributed.ProcessGroup] = ctx.pg
+    assert pg is not None, "ctx.pg must be initialized by the process runner"
+    assert batch_size % world_size == 0, (
+        f"batch_size ({batch_size}) must be divisible by world_size ({world_size}): "
+        "PooledEmbeddingsReduceScatter scatters the global batch evenly across ranks."
+    )
+
+    local_embs = _make_reduce_scatter_input(
+        batch_size=batch_size,
+        dim=dim,
+        values_dtype=values_dtype,
+        device=ctx.device,
+    )
+    rs = PooledEmbeddingsReduceScatter(pg)
+
+    logger.info(
+        "rank=%d local_rank=%d host=%s running reduce-scatter benchmark: batch_size=%d "
+        "dim=%d device=%s",
+        rank,
+        ctx.local_rank,
+        socket.gethostname(),
+        batch_size,
+        dim,
+        ctx.device,
+    )
+
+    result = benchmark_func(
+        name=name,
+        rank=rank,
+        world_size=world_size,
+        func_to_benchmark=_run_reduce_scatter,
+        bench_inputs=[],
+        prof_inputs=[],
+        benchmark_func_kwargs={"rs": rs, "local_embs": local_embs},
+        num_profiles=num_profiles,
+        num_benchmarks=num_benchmarks,
+        profile_dir=profile_dir,
+        device_type=ctx.device.type,
+        pg=pg,
+    )
+
+    if rank == 0:
+        logger.info("reduce-scatter benchmark result:\n%s", result)
+
+    return result
+
+
 # Registry of available primitive benchmarks, keyed by the ``primitive`` flag.
 # Add new primitive benchmarks here -- each is called as
 # ``fn(ctx, rank, world_size, **kwargs)`` and returns a per-rank ``BenchmarkResult``.
 _BENCHMARKS: Dict[str, Callable[..., BenchmarkResult]] = {
     "kjt_a2a": _benchmark_kjt_a2a,
     "kt_a2a": _benchmark_kt_a2a,
+    "reduce_scatter": _benchmark_reduce_scatter,
 }