From eeb58590975d3df2b34cd98003a5ab5464ff4a7a Mon Sep 17 00:00:00 2001 From: Siqi Huang Date: Thu, 6 Aug 2026 09:59:30 -0700 Subject: [PATCH] Create Reduce-Scatter Primitive Benchmark (#4483) 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`. Reviewed By: TroyGarden Differential Revision: D113960618 --- .../benchmark/benchmark_primitive.py | 149 +++++++++++++++++- 1 file changed, 147 insertions(+), 2 deletions(-) diff --git a/torchrec/distributed/benchmark/benchmark_primitive.py b/torchrec/distributed/benchmark/benchmark_primitive.py index 8972718a4..f841ae136 100644 --- a/torchrec/distributed/benchmark/benchmark_primitive.py +++ b/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, }