Link-balance the sharded relay schedules (up to 3.2x vs NCCL)

Summary:
The sharded relay collectives circumvent the MI3XX single-XGMI-link limit by
recruiting the idle GPUs on a node as relay helpers. This retunes their
schedules against the actual per-link cost model, which roughly doubles the
2-active speedups and turns the two 4-active collectives that were *slower* than
NCCL into wins.

Cost model used throughout: on MI350X every GPU pair has exactly one XGMI link,
so each GPU has 7 links at ~56 GB/s measured. A schedule's runtime is
`max over (link, direction) of bytes carried`, summed over serialized
`ncclGroup` boundaries.

**1. A=2: fold the direct exchange into the two relay groups (`numChunks = H+2`)**
All four A=2 paths ran three serialized groups -- scatter, forward, then a
separate active<->active direct exchange -- with `numChunks = H+1`. That costs
`3*count/7 = 0.43*count` and, worse, leaves the active<->active link completely
idle during the two relay groups: 1 of 7 links wasted for 2/3 of the runtime.
Balancing a rank's egress (`2*count - d` over 7 links) against the direct link's
own bound (`d`) puts the optimum at `d = count/4`. Realizing it needs just two
groups with `numChunks = H+2`, one direct chunk riding along with each relay
group, so every link carries exactly one chunk per direction per group:
`0.43*count -> 0.25*count`, i.e. a 2.33x ceiling becomes 4.0x.

**2. A=2 allreduce: reduce at the helper instead of forwarding both slots**
Both active ranks send the *same* logical chunk index to a helper, so their sum
is already the final allreduced value. The helper now sums its two slots and
returns one reduced chunk to each active rank. Link cost is identical (the
helper still sends one chunk per active rank), but this drops the active rank's
relay scratch and its fused add+scale over 6/8 of the buffer, and spreads the
reduction across every helper GPU instead of piling it on the two actives.
Deliberately NOT applied to reduce-scatter: there slot 0 is a0's contribution to
a1's *output* and slot 1 is a1's contribution to a0's output -- different
outputs, not summable -- so helpers stay passthrough there.

**3. A=2 allreduce small messages: one full exchange instead of RS+AG**
The small-message pure-direct path did a reduce-scatter swap plus an all-gather
swap. Both move `count` per link direction, but RS+AG needs two group
boundaries, so a single full exchange is strictly better in the latency-bound
regime.

**4. A=4 allreduce: offload fraction 780 -> 500 permille**
Per group the intra links carry `pD/A` and the cross links `pO/H`, so the
two-group critical path is `2*max(pD, pO)/A` -- minimized when the direct and
offload regions are EQUAL. 780 skewed everything onto the cross links for a
1.28x ceiling, which is exactly the ~1.03x that was measured. 500 gives a 2.0x
ceiling. Also restored a 2 MB pure-direct floor so small messages skip the
2-hop hop entirely.

**5. A=4 all-gather: drop the 16-stage pipeline for a balanced 2-group schedule**
The pipeline existed to "overlap the helper-forward against the next
active-send", but on the A=4 / 2-group topology a rank's helpers ARE the active
ranks of the other group, so scatter and forward are egress on the *same cross
link in the same direction*. They add rather than overlap, making the 17 group
boundaries x ~38 p2p ops per superstep pure launch overhead. Replaced with two
groups; since group 2's cross links carry `(A-1)x` group 1's, the direct region
is split 1:(A-1) across the groups to keep both balanced.

**6. A=4 reduce-scatter: replace recursive-halving with flat reduce-at-helper**
Each helper now owns one position slice of every block, collects that slice from
the A-1 non-owner sources, sums them, and forwards a single reduced chunk to the
owner -- woven with a direct all-to-all reduce-scatter over the intra links.
Reducing at the helper is what keeps the return hop cheap: A-1 chunks in, one
out. The scratch mirrors the output layout so the whole reduction collapses to
two fused multi-input passes.

Because that helper reduces rather than forwards, it needs one chunk per
(owner, source) pair -- `A*(A-1)*chunk`, i.e. 1.5x recvCount on an 8-GPU node --
so the A>2 reduce-scatter helper-buffer contract grows from the two-slot
passthrough size to `2 * recvCount`. `sharded_relay_utils.py` and the benchmark
are updated to match; the C++ tests already allocated `A * recvCount`.

The torchrec unit test that pins that contract changes in this commit too, so the
expectation never lags the production sizing: `test_helper_buffers_passthrough_sized_4active` becomes `test_helper_buffers_sized_to_2x_recv_count_4active` and
asserts `2 * recv[g]` rather than `_passthrough_helper_size(...)`.

**7. Crossover retuning, measured separately for fused and parallel**
Every pure-direct/offload threshold was re-measured now that the relay is ~1.7x
faster. Notably reduce-scatter A=2 fused dropped 8 MB -> 2 MB (fixing a 4.5 MB
dip), the A=4 reduce-scatter offload only pays past 48 MB, and A=4 all-gather
past 12 MB fused / 8 MB parallel.

**Tried and reverted: helper offload for A=4 all-to-all.** The link model
promised 1.67x, but a permutation gives the helper `A*(A-1) = 12` distinct
(dest, source) chunks per group with no reduction to amortize the op count. It
measured 0.75-0.97x against pure-direct from 13.5 MB to 135 MB and only
1.03-1.07x at 256 MB-1 GB, so pure-direct was kept. The open lead (coalescing
the helper's sends per dest, which needs gather/scatter kernels because both ends
are strided by segmentCount) is recorded in the `shardedRelayAllToAllFlat`
docblock.

Differential Revision: D115998361

Author: srinathb

torchrec/distributed/sharded_relay_utils.py

@@ -478,6 +478,26 @@ def _get_helper_flat_buf(
     return buf if buf.numel() == total else buf.narrow(0, 0, total)
 
 
+def _relay_helper_size(
+    total_g: int,
+    sparse_group_size: int,
+) -> int:
+    """Helper scratch size (in elements) for one group when the relay reduces at
+    the helper.
+
+    The A>2 relay paths give each helper one position slice of every block and
+    have it sum the contributing sources before forwarding, which needs one chunk
+    per (owner, source) pair:
+
+        A * (A - 1) * chunk,  chunk = total_g // (A + H)
+
+    On the 8-GPU node (A == H == 4) that is 1.5 * total_g. Rather than tracking H
+    here, allocate 2 * total_g, which covers every supported (A, H) split and
+    matches the allreduce A>2 contract.
+    """
+    return 2 * total_g
+
+
 def _passthrough_helper_size(
     total_g: int,
     sparse_group_size: int,
@@ -718,9 +738,16 @@ def reduce_scatter_tensors_with_sharded_relay(
                         unpack_flat = out_flat
                 else:
                     recv_count_g = per_group_recv_counts[g]
-                    helper_size_g = _passthrough_helper_size(
-                        recv_count_g, sparse_group_size, num_chunks
-                    )
+                    if sparse_group_size > 2:
+                        # A>2 reduces at the helper, which needs one chunk per
+                        # (owner, source) pair rather than two passthrough slots.
+                        helper_size_g = _relay_helper_size(
+                            recv_count_g, sparse_group_size
+                        )
+                    else:
+                        helper_size_g = _passthrough_helper_size(
+                            recv_count_g, sparse_group_size, num_chunks
+                        )
                     helper_buf = _get_helper_flat_buf(
                         state, g, helper_size_g, dtype, device
                     )

torchrec/distributed/tests/bench_sharded_relay_perf.py

@@ -1698,6 +1698,9 @@ def _relay_helper_size(
         return _passthrough_helper_size(elements, active_ranks, num_chunks)
     if collective == "reduce_scatter":
         recv = elements // active_ranks
+        if active_ranks > 2:
+            # A>2 reduces at the helper: one chunk per (owner, source) pair.
+            return 2 * recv
         return _passthrough_helper_size(recv, active_ranks, num_chunks)
     if collective == "all_to_all":
         if active_ranks > 2:

torchrec/distributed/tests/test_sharded_relay_utils.py

@@ -2005,7 +2005,7 @@ def test_single_call_two_groups(self) -> None:
             100,
         )
 
-    def test_helper_buffers_passthrough_sized_4active(self) -> None:
+    def test_helper_buffers_sized_to_2x_recv_count_4active(self) -> None:
         state = _make_state(rank=0, sparse_group_size=4, local_size=8)
         # input total 800 -> recv_count 200
         reduce_scatter_tensors_with_sharded_relay(
@@ -2018,18 +2018,16 @@ def test_helper_buffers_passthrough_sized_4active(self) -> None:
         in_tensors = kwargs["input_tensors"]
         out_tensors = kwargs["output_tensors"]
         recv = kwargs["per_group_recv_counts"]
-        # num_chunks = local_size - sparse_group_size + 1 = 8 - 4 + 1 = 5
-        num_chunks = (state.local_size - state.sparse_group_size) + 1
 
         helper_ptrs = set()
         for g in range(state.num_sparse_groups):
             if g == state.my_sparse_group:
                 self.assertEqual(in_tensors[g].numel(), 800)
                 self.assertEqual(out_tensors[g].numel(), 200)
                 continue
-            expected = _passthrough_helper_size(
-                recv[g], state.sparse_group_size, num_chunks
-            )
+            # A>2 reduce-scatter reduces at the helper and retains the
+            # 2 * recvCount helper-buffer contract.
+            expected = 2 * recv[g]
             self.assertEqual(in_tensors[g].numel(), expected)
             # Helper uses one scratch buffer for both send and recv.
             self.assertEqual(in_tensors[g].data_ptr(), out_tensors[g].data_ptr())