Link-balance the sharded relay schedules (up to 3.2x vs NCCL) by srinathb-meta · Pull Request #4570 · meta-pytorch/torchrec

srinathb-meta · 2026-08-15T02:56:03Z

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

meta-codesync · 2026-08-15T02:56:13Z

@srinathb-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115998367.

…ailing (meta-pytorch#4570) Summary: X-link: meta-pytorch/torchcomms#3678 When a group's aligned relay chunk rounded down to zero the four sharded relay collectives returned `ncclInvalidArgument` and required the caller to retry on a plain collective. In a fused multi-group call that is worse than it sounds: a single small or awkwardly sized group forced the whole fused call to fail even when the other groups were large enough to relay, so heterogeneous group sizes were effectively unsupported. Give every collective a phase-symmetric per-group direct fallback instead. When `chunkSize` aligns down to zero the segment is covered by the two direct regions (`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link, the helper scatter/forward is skipped for that group, and the helper-side reduction is skipped as well. Groups that do have a workable chunk size keep the existing relay schedule unchanged, so each group in a fused call now picks its own schedule independently. - track `dirASizes` per group so the direct regions are sized from the selected schedule rather than assumed equal to `chunkSize` - guard the helper scatter, forward, and reduction on `chunkSize > 0` - all-gather additionally tracks `dirAOffsets`, which is no longer derivable from `relayTotals` once the fallback can move the direct region to offset 0 - correct the minimum helper geometry to H+2 chunks Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was `_passthrough_helper_size(...)` when production has always sized it through `_relay_helper_size()` (`2 * recvCount`); the expectation now matches the unchanged production contract. Behavior for buffers that already had a nonzero aligned chunk is unchanged. Differential Revision: D115998367

…ailing (meta-pytorch#3678) Summary: Pull Request resolved: meta-pytorch#3678 X-link: meta-pytorch/torchrec#4570 When a group's aligned relay chunk rounded down to zero the four sharded relay collectives returned `ncclInvalidArgument` and required the caller to retry on a plain collective. In a fused multi-group call that is worse than it sounds: a single small or awkwardly sized group forced the whole fused call to fail even when the other groups were large enough to relay, so heterogeneous group sizes were effectively unsupported. Give every collective a phase-symmetric per-group direct fallback instead. When `chunkSize` aligns down to zero the segment is covered by the two direct regions (`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link, the helper scatter/forward is skipped for that group, and the helper-side reduction is skipped as well. Groups that do have a workable chunk size keep the existing relay schedule unchanged, so each group in a fused call now picks its own schedule independently. - track `dirASizes` per group so the direct regions are sized from the selected schedule rather than assumed equal to `chunkSize` - guard the helper scatter, forward, and reduction on `chunkSize > 0` - all-gather additionally tracks `dirAOffsets`, which is no longer derivable from `relayTotals` once the fallback can move the direct region to offset 0 - correct the minimum helper geometry to H+2 chunks Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was `_passthrough_helper_size(...)` when production has always sized it through `_relay_helper_size()` (`2 * recvCount`); the expectation now matches the unchanged production contract. Behavior for buffers that already had a nonzero aligned chunk is unchanged. Differential Revision: D115998367

…ailing (meta-pytorch#3678) Summary: Pull Request resolved: meta-pytorch#3678 X-link: meta-pytorch/torchrec#4570 When a group's aligned relay chunk rounded down to zero the four sharded relay collectives returned `ncclInvalidArgument` and required the caller to retry on a plain collective. In a fused multi-group call that is worse than it sounds: a single small or awkwardly sized group forced the whole fused call to fail even when the other groups were large enough to relay, so heterogeneous group sizes were effectively unsupported. Give every collective a phase-symmetric per-group direct fallback instead. When `chunkSize` aligns down to zero the segment is covered by the two direct regions (`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link, the helper scatter/forward is skipped for that group, and the helper-side reduction is skipped as well. Groups that do have a workable chunk size keep the existing relay schedule unchanged, so each group in a fused call now picks its own schedule independently. - track `dirASizes` per group so the direct regions are sized from the selected schedule rather than assumed equal to `chunkSize` - guard the helper scatter, forward, and reduction on `chunkSize > 0` - all-gather additionally tracks `dirAOffsets`, which is no longer derivable from `relayTotals` once the fallback can move the direct region to offset 0 - correct the minimum helper geometry to H+2 chunks Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was `_passthrough_helper_size(...)` when production has always sized it through `_relay_helper_size()` (`2 * recvCount`); the expectation now matches the unchanged production contract. Behavior for buffers that already had a nonzero aligned chunk is unchanged. Reviewed By: JinghanHuang Differential Revision: D115998367

Summary: Add the core C++ implementation of 2-rank sharded relay reduce-scatter to rcclx, the reduce-scatter analogue of the sharded relay allreduce landed in D103262937. Like allreduce, the logical collective is a 2-rank reduce-scatter between the two active ranks per sparse group, accelerated by passthrough helpers that relay sharded chunks so XGMI traffic stays unidirectional under phase-synchronized execution. Algorithm: each active rank's sendBuff holds nActiveRanksPerGroup x recvCount elements (block[i] = the slice destined for active index i); each active recvBuff holds recvCount elements and receives the sum/avg of both ranks' block[myActiveIndex]. Each rank keeps its own block and relays block[otherActiveIndex] to the other rank, sharded across helpers, then reduces it into the output block. Same 5-phase passthrough structure as allreduce, restricted to one block: - Phase 1 active->helpers scatter of the sendBlock chunks - Phase 2 helpers->active batched passthrough forward - Phase 3 active fused add (+scale for AVG) into the seeded output block - Phase 4 active<->active direct-exchange of the last chunk - Phase 5 active final reduction of the direct chunk Supports both in-place (recvBuff == sendBuff + ownBlockOffset) and out-of-place; ncclSum and ncclAvg only. Reuses the dtype-generic kernels in sharded_relay_allreduce_kernels.{h,cu} (no new .cu); the ScratchBufferCache, rank-config builder, and DISPATCH host macros are re-declared file-locally (the allreduce copies are file-scoped and not linkable across TUs). This diff includes: - New relay algorithm in meta/relay/sharded_relay_reduce_scatter.{h,cc} - API surface additions in nccl.h and rccl.h - Collective dispatch integration in collectives.cc No def_build.bzl change is needed: meta/relay/*.{h,cc,cu} are auto-globbed. This diff exposes the C++ API only; Python bindings and torchrec utilities follow in later diffs of the stack. There is no model_parallel.py integration for reduce-scatter. As the bottom of the stack, this commit also sets up the two things every commit above it depends on. **Gate the sharded-relay suites off for the intermediate commits.** The 8-GPU suites are meaningless on a partially applied stack, and running them in per-diff CI would burn 8 GPUs to fail. Add a `GTEST_FILTER` env that matches no test to the pre-existing `sharded_relay_allreduce_test` target (the three new suites carry it from the commit that introduces them), and `labels = [tpx_labels.disabled]` to `test_sharded_relay_utils` and `bench_sharded_relay_perf`. Every target still builds at every commit, so compile breakage is caught per diff. All of it is removed at the stack tip, where the suites are re-enabled and run. **Bring the memory footprints inside a shared-host budget.** The busBW and benchmark defaults were sized for an exclusively owned 8-GPU host and OOM as soon as anything else is resident: - ShardedRelayAllReduceTest.cu: the two 4-group busBW tests drop from 24GB to 1GB per group (`Z_BusBW_4Groups_{InPlace,OutOfPlace}_24GB` -> `..._1GB`). A rank is active for one group and a helper for the other three, so it holds `dataBytes + 3 x (nActiveRanksPerGroup x dataBytes)` -- 24GB per group is ~168GB per rank. The suites added above this commit are sized at 1GB per group from the start. - bench_sharded_relay_perf.py: the production per-group totals (~22 GiB per group, ~69 GiB per rank with helper and scratch buffers) are kept as documented constants but capped to a 1 GiB-per-group default via `_default_totals()`. `BENCH_TABLE_SIZE` / `BENCH_KERNEL_SIZE_GB` still measure at production scale on an idle host; the module docstring documents both. Neither change affects the relay algorithms -- only test/benchmark sizes and which targets execute. Differential Revision: D115998376

Summary: Add the distributed layer for sharded relay reduce-scatter, mirroring the allreduce utilities (D103262936) but WITHOUT model_parallel integration (this diff exposes the utility wrappers only). This includes: - FusedShardedRelayMultiGroup.reduce_scatter_multi_group in caffe2/torch/distributed/fb/sharded_relay_process_group.py, a thin wrapper over the torchcomms sharded_relay_multi_group_reduce_scatter pybind taking separate input/output tensor lists and per_group_recv_counts - reduce_scatter_tensors_with_sharded_relay in torchrec/distributed/sharded_relay_utils.py, the flat-concat helper that packs the active group's input (nActiveRanks x recv_count, two blocks), makes ONE fused phase-synchronized call across all sparse groups, and unpacks the reduced output block into the caller's output tensors - _get_active_output_flat_buf + a _active_output_flat_cache field on ShardedRelayState for the out-of-place output buffer; helper buffers reuse the existing _passthrough_helper_size, whose first parameter is renamed total_g -> count_g with an expanded docstring: the value the kernel receives is the group total for allreduce/reduce-scatter but the per-segment count for all-to-all, and the old name read as if the group total were always correct - _pack_into_flat / _unpack_from_flat, the shared fused pack (torch.cat into a pre-allocated flat buffer) and unpack (torch._foreach_copy_ over split views) helpers. One kernel launch each instead of one per caller tensor. Every collective added by this stack uses them, and the pre-existing allreduce util is moved onto them later in the stack In-place vs out-of-place: the primitive takes separate input/output tensors so the caller selects the mode by the buffers passed (in-place = output aliases the input's local-contribution block at ownBlockOffset; out-of-place = distinct output). The util exposes an in_place flag that drives the corresponding kernel path internally; both produce identical results in output_tensors_dict. Reduce-scatter is inherently out-of-place at the tensor level (input is nActiveRanks x the output size). SUM and AVG only. No torchrec/distributed/model_parallel.py changes. Differential Revision: D115998378

Summary: Add tests and a benchmark path for the sharded relay reduce-scatter distributed utilities, mirroring the allreduce tests (D103262938): - FlatReduceScatterTest: CPU/mock unit tests for reduce_scatter_tensors_with_sharded_relay covering call counts, recv_count = input_total / nActiveRanks, active input/output buffer sizing, divisibility and output-mismatch errors, passthrough-sized + distinct helper buffers, in-place output aliasing the input's local block at ownBlockOffset, out-of-place distinct output, value write-back, multi-tensor unpack, and metadata-cache reuse. - FusedReduceScatterValidationTest: validation for FusedShardedRelayMultiGroup.reduce_scatter_multi_group (active input/output too-small raise ValueError; recv_count=0 groups skip validation). Both too-small tests pass `skip_validation=False`, since `reduce_scatter_multi_group` skips its shape checks by default (see the previous commit) and would otherwise reach the missing-native-RCCLX RuntimeError instead of the expected ValueError. - bench_reduce_scatter_flat + a [D] REDUCE-SCATTER timed path in the perf benchmark worker. It releases the Bench C kernel tensors first and uses recv_count = prod_total/2 so the input (2 x recv_count, two blocks) matches the allreduce active footprint and stays within GPU memory. De-tautologise the pre-existing _passthrough_helper_size tests in this file. Three of them recomputed the function's own min(total_g, A * align_down(total_g // num_chunks, 128)) expression and compared it to the function's output, so both sides moved together and the assertions could not fail: - test_passthrough_size_matches_2x_chunkSize_for_realistic_totals - test_python_meets_cpp_min_required_at_alignment_boundary (which did also pin a literal, so only its mirrored assertion was dead) - test_total_per_rank_helper_memory_is_6x_chunkSize They now assert constants worked out by hand from the contract -- 3_429_423_360 for the 12_002_982_488/A=2/num_chunks=7 production shape, 256_000 at the exact alignment boundary, and 10_288_270_080 for the per-rank total across 3 helper groups -- plus the property the alignment exists for, that each of the A slots is a whole number of 128-element chunks. The two genuinely independent tests in the class (the tiny-count fallback and the cap at total_g) are unchanged. Differential Revision: D115998394

Summary: Add the distributed layer for sharded relay all-to-all, mirroring the reduce-scatter utilities but WITHOUT model_parallel integration (this diff exposes the utility wrappers only). This includes: - FusedShardedRelayMultiGroup.all_to_all_multi_group in caffe2/torch/distributed/fb/sharded_relay_process_group.py, a thin wrapper over the torchcomms sharded_relay_multi_group_all_to_all pybind taking separate input/output tensor lists and per_group_segment_counts - all_to_all_tensors_with_sharded_relay in torchrec/distributed/sharded_relay_utils.py, the flat-concat helper that packs the active group's input (nActiveRanks x segment_count, two segments), makes ONE fused phase-synchronized call across all sparse groups, and unpacks the transposed output into the caller's output tensors All-to-all performs NO reduction, so there is no reduce op. It is OUT-OF-PLACE ONLY (matching native ncclAllToAll): the util always uses a separate output flat buffer, and the process-group method rejects an active group whose input and output tensors alias (data_ptr() equal) and validates that the active input and output each hold nActiveRanks x segment_count elements. No torchrec/distributed/model_parallel.py changes. Document why the helper-scratch sizing passes the per-segment count. _passthrough_helper_size's first parameter is named total_g, which reads as "the caller's group total" and makes this call site look like it under-sizes by a factor of nActiveRanks. It does not: the parameter is really the per-group count the KERNEL receives. All-to-all hands the kernel segmentCounts[g], so the kernel computes chunkSize = align_down(seg / numChunks, 128) from the segment count and each helper stages nActiveRanks slots of that -- exactly what the function returns. The allreduce path passes its group total for the same reason: that is the count its kernel sees. num_chunks here, (local_size - sparse_group_size) + 1, equals the kernel's numHelpers + 1. The min(...) clamp cannot under-size either: it only binds when nActiveRanks >= numChunks, and numChunks = (local_size - nActiveRanks) + 1 keeps nActiveRanks strictly smaller for every supported 8-rank topology (A=2 -> 7, A=4 -> 5). When the aligned chunk floors to zero the kernel does not read a helper buffer at all -- it rejects the call outright at this point in the stack, and once the zero-chunk direct fallback lands it runs direct-only with relayTotals[g] == 0 and every helper block gated on chunkSizes[g] > 0 -- so the value the fallback returns can only over-allocate. Differential Revision: D115998370

Summary: Add tests and a benchmark path for the sharded relay all-to-all distributed utilities, mirroring the reduce-scatter tests: - FlatAllToAllTest: CPU/mock unit tests for all_to_all_tensors_with_sharded_relay covering segment_count = input_total / nActiveRanks, active input/output buffer sizing (both = full size), divisibility and output-mismatch errors, passthrough-sized + distinct helper buffers, out-of-place output flat buffer (distinct from the input flat buffer), value write-back, multi-tensor unpack, and metadata-cache reuse. - FusedAllToAllValidationTest: validation for FusedShardedRelayMultiGroup.all_to_all_multi_group (active input/output too-small raise ValueError; in-place aliasing raises ValueError; segment_count=0 groups skip validation). All three raising tests pass `skip_validation=False`, since `all_to_all_multi_group` skips its shape checks by default and would otherwise reach the missing-native-RCCLX RuntimeError instead of the expected ValueError. - bench_all_to_all_flat + a [D] ALL-TO-ALL timed path and an NCCL all_to_all_single baseline in the perf benchmark, with a per-collective ALL-TO-ALL section and speedup. Sized segment = prod_total/4 and frees prior collectives' buffers to stay within GPU memory. All-to-all is out-of-place only, so there are no in-place correctness tests (instead the validation test asserts in-place is rejected). Differential Revision: D115998384

Summary: Add the distributed layer for sharded relay all-gather, mirroring the reduce-scatter utilities but WITHOUT model_parallel integration (this diff exposes the utility wrappers only). This includes: - FusedShardedRelayMultiGroup.all_gather_multi_group in caffe2/torch/distributed/fb/sharded_relay_process_group.py, a thin wrapper over the torchcomms sharded_relay_multi_group_all_gather pybind taking separate input/output tensor lists and per_group_send_counts - all_gather_tensors_with_sharded_relay in torchrec/distributed/sharded_relay_utils.py, the flat-concat helper that packs the active group's input (send_count, this rank's contribution), makes ONE fused phase-synchronized call across all sparse groups, and unpacks the gathered output (nActiveRanks x send_count) into the caller's output tensors All-gather performs NO reduction, so there is no reduce op. It is the dual of reduce-scatter and supports both in-place and out-of-place. The util's in_place flag packs the input into the active rank's own slot of the output flat buffer (kernel in-place path: sendbuff == recvbuff + myActiveIndex x send_count) or uses a separate input flat buffer (out-of-place). The process-group method validates that the active input holds >= send_count elements and the active output holds >= nActiveRanks x send_count elements. No torchrec/distributed/model_parallel.py changes. Differential Revision: D115998385

Summary: Add tests and a benchmark path for the sharded relay all-gather distributed utilities, mirroring the reduce-scatter tests: - FlatAllGatherTest: CPU/mock unit tests for all_gather_tensors_with_sharded_relay covering output total = nActiveRanks x send_count, active input/output buffer sizing, output-mismatch error, passthrough-sized + distinct helper buffers, in-place (input packed into the active rank's own slot of the output flat buffer) and out-of-place (separate input flat), value write-back, multi-tensor unpack, and metadata-cache reuse. - FusedAllGatherValidationTest: validation for FusedShardedRelayMultiGroup.all_gather_multi_group (active input/output too-small raise ValueError; send_count=0 groups skip validation). Both too-small tests pass `skip_validation=False`, since `all_gather_multi_group` skips its shape checks by default and would otherwise reach the missing-native-RCCLX RuntimeError instead of the expected ValueError. - bench_all_gather_flat + a [F] ALL-GATHER timed path and an NCCL all_gather_into_tensor baseline in the perf benchmark, with a per-collective ALL-GATHER section and speedup. Sized send_count = prod_total/4 and frees prior collectives' buffers to stay within GPU memory. All-gather supports both in-place and out-of-place (like reduce-scatter), so the in_place path is exercised in addition to out-of-place. Differential Revision: D115998375

Summary: Allow the sharded relay allreduce to run with 4 active ranks per sparse group (in addition to 2) from the TorchRec distributed layer. This plumbs the 4-rank support added in the rcclx kernel + torchcomms bindings up through setup_sharded_relay / model_parallel, so a 2D-parallel model configured with num_parallel_worlds=4 (e.g. omni-fmv3) uses the 4-active path automatically. - _validate_sharded_relay_preconditions: relaxed the hard sharding_group_size == 2 guard to accept 2 or 4 (power of two), and replaced the fixed "local_size < 4" check with the correct general "need at least one helper" check (local_size > sharding_group_size). - setup_sharded_relay docstring updated (supports 2 or 4). - allreduce_tensors_with_sharded_relay: for A>2 the flat helper-reduce-and- broadcast kernel needs a larger helper scratch than the 2-active passthrough, so the helper buffer is sized 2*total_g (covers (A+1)*oChunk for any offload fraction); the 2-active path keeps the passthrough size. model_parallel.py already converts replica_group_size = world_size // model_parallel_group_size and passes it to setup_sharded_relay, so num_parallel_worlds=4 -> replica_group_size=4 flows through; the fused binding and group layout derive the active-rank count generically. Behavior: - BM-FM is unchanged: it stays at 2 active ranks (NCCL_SHARDED_RELAY_MODE_ENABLE=1, num_parallel_worlds=2). - A model with num_parallel_worlds=4 now enables the 4-active sharded relay. Note: the 4 replicas of a group must map to 4 consecutive local ranks per node ({0,1,2,3},{4,5,6,7}), since all_active_ranks is built as consecutive ranges. Differential Revision: D115998382

Summary: Add CPU unit tests for the 4-active sharded relay allreduce distributed path and extend the perf benchmark to report results for BOTH 2-active and 4-active groups. test_sharded_relay_utils.py: - FlatAllreduce4ActiveTest: drives allreduce_tensors_with_sharded_relay with sparse_group_size=4 (8-rank node -> 2 groups of 4). Verifies one fused call, per_group_sizes, the active full-size buffer, and that the A>2 helper buffers are sized 2*total_g (the flat helper-reduce-and-broadcast scratch, not the 2-active passthrough size) and distinct, plus the consecutive active-rank layout [[0,1,2,3],[4,5,6,7]]. - FusedAllreduce4ActiveValidationTest: allreduce_multi_group accepts 4 active ranks (RuntimeError without a native comm, not ValueError). bench_sharded_relay_perf.py: - The benchmark sweeps BOTH 2-active and 4-active sharded relay groups and prints a full report for each; the header shows active-ranks/group. The workers take sharding_group_size + port_offset args (per-iteration ports avoid endpoint collisions). - 4-active allreduce helper buffers (flat_bufs, kernel_scratch, and the bench_fused_flat scratch) are sized 2*total_g to match the flat kernel. - Adds an optional BENCH_ONLY env (allreduce|reduce_scatter|all_to_all| all_gather) that restricts the run to a single collective so tuning one 4-active collective does not pay for the other three (default "all"); bw() is guarded against a zero (skipped) measurement. - Fix: kernel-direct (Benchmark C) built per-group sizes as num_sparse_groups entries (was list(prod_totals), which mismatched num_sparse_groups=2 at 4-active). The _passthrough_helper_size 4-active test asserts hand-computed constants rather than re-deriving them. As first written it recomputed min(total_g, A * align_down(total_g // num_chunks, 128)) -- the exact expression the function implements -- and compared that to the function's output, so any change to the formula moved both sides together and the assertion could never fail. It now pins literals worked out by hand from the contract (12_002_982_488 with A=4/num_chunks=5 -> 9_602_385_920, and the 2-active shape -> 3_429_423_360) and covers the two branches the original missed: the min() clamp when A >= num_chunks, and the align_down(...) == 0 fallback below one chunk. A separate case pins the property the 128-alignment exists for -- each of the A slots is a whole number of 128-element chunks whenever min() has not clamped. Differential Revision: D115998388

Summary: Extend the reduce-scatter distributed-layer tests and the perf benchmark to cover 4 active ranks per sparse group (in addition to 2), completing the 4-rank reduce-scatter work. The distributed plumbing was already generic over the active-rank count, so this diff is test/benchmark-only. test_sharded_relay_utils.py: - FlatReduceScatter4ActiveTest: drives reduce_scatter_tensors_with_sharded_relay with sparse_group_size=4 (8-rank node -> 2 groups of 4). Verifies one fused call, num_groups=2, recv_count = input_total // 4, active input (A x recv_count) / output (recv_count) sizing, passthrough-sized + distinct helper buffers (numChunks = local_size - 4 + 1 = 5) reused for send and recv, and the consecutive active-rank layout [[0,1,2,3],[4,5,6,7]]. - FusedReduceScatter4ActiveValidationTest: reduce_scatter_multi_group accepts 4 active ranks (RuntimeError without a native comm, not ValueError). bench_sharded_relay_perf.py: - The benchmark now sweeps BOTH 2-active and 4-active reduce-scatter (the existing enumerate((2, 4)) sweep). The reduce-scatter NCCL baseline, fused bench, and printout are pulled out of the `if sparse_group_size == 2:` guards so they run for both 2 and 4; all-to-all / all-gather remain 2-active only. Barrier balance is preserved because all ranks share the same env-driven sparse_group_size. - Generalized the reduce-scatter buffer math from the 2-active `// 2` to `// sparse_group_size` (input = A x recv_count). - Fix: rs_recv_counts is now built over range(num_sparse_groups) instead of `for t in prod_totals`. prod_totals always has 4 entries while num_sparse_groups = local_size // sparse_group_size (2 at 4-active), and the binding requires per_group_recv_counts to match input_tensors length; the old list comprehension produced 4 entries and raised "per_group_recv_counts size must match input_tensors size" at 4-active. The 2-active output is byte-identical (range(4) yields the same list). Differential Revision: D115998387

…rk sweep Summary: Cover 4 active ranks per sparse group for all-to-all in the distributed layer: the A>2 helper-buffer handling, CPU unit tests, and the perf benchmark sweep. sharded_relay_utils.py: - all_to_all_tensors_with_sharded_relay: A>2 all-to-all is PURE-DIRECT (no helper relay), so a helper rank does no work for that group and needs no helper buffer -- the util passes a tiny size-1 placeholder for helper groups. The 2-active path keeps its real passthrough-sized helper buffer. test_sharded_relay_utils.py: - FlatAllToAll4ActiveTest: drives all_to_all_tensors_with_sharded_relay with sparse_group_size=4 (8-rank node -> 2 groups of 4). Verifies one fused call, num_groups=2, segment_count = input_total // 4, OUT-OF-PLACE active in/out, the A>2 helper-group slots are size-1 placeholders (pure-direct, no helper work) and distinct, and the layout [[0,1,2,3],[4,5,6,7]]. - FusedAllToAll4ActiveValidationTest: all_to_all_multi_group accepts 4 active ranks (RuntimeError without a native comm, not ValueError) and still rejects in-place at 4 active ranks. `test_4active_in_place_rejected` passes `skip_validation=False`, since `all_to_all_multi_group` skips its shape checks by default and would otherwise reach the missing-native-RCCLX RuntimeError instead of the expected ValueError. bench_sharded_relay_perf.py: - Sweeps BOTH 2-active and 4-active all-to-all; the a2a A>2 helper slot is a size-1 placeholder (pure-direct); the a2a baseline/fused bench are gated by the BENCH_ONLY env. - Generalized the a2a segment math and built a2a_seg_counts over range(num_sparse_groups) (2-active output byte-identical). Differential Revision: D115998358

Summary: Enable and document 4-active sharded relay all-gather in the Python plumbing. This is NOT doc-only: it includes a functional helper-buffer sizing change in the torchrec distributed layer, plus docstring generalizations. Functional change: - torchrec/distributed/sharded_relay_utils.py: size the A>2 all-gather helper passthrough buffer as nActiveRanks * send_count (the flat scatter->forward relay needs each helper to hold A source chunks). The 2-active path keeps the original _passthrough_helper_size. Without this the helper buffer is too small for the 4-active flat relay. Docstring generalizations (no behavior change; the binding and the FusedShardedRelayMultiGroup.all_gather_multi_group process-group method were already generic over the active-rank count): - TorchCommRCCLXPy.cpp / _comms_rcclx.pyi: the all-gather binding docstring now says "a power-of-two number of active ranks (2 or 4)" instead of "exactly 2 active ranks ... a 2-rank all-gather", describes the A>2 flat scatter->forward relay, and describes the helper buffer as an "nActiveRanks-slot scratch" rather than "two-slot". - caffe2/.../sharded_relay_process_group.py (+ xplat mirror): same docstring generalization for all_gather_multi_group. The validation already derives n_active = len(all_active_ranks[g]) and checks send_count (input) / n_active x send_count (output), with in-place allowed, so a 4-wide all_active_ranks is accepted unchanged. Differential Revision: D115998397

Summary: Extend the all-gather distributed-layer tests and the perf benchmark to cover 4 active ranks per sparse group (in addition to 2), completing the 4-rank sharded relay project. With this diff all four collectives sweep at both 2 and 4 active ranks; the benchmark no longer has any `if sparse_group_size == 2:` guards. test_sharded_relay_utils.py: - FlatAllGather4ActiveTest: drives all_gather_tensors_with_sharded_relay with sparse_group_size=4 (8-rank node -> 2 groups of 4). Verifies one fused call, num_groups=2, output = nActiveRanks x send_count, the A>2 helper buffers sized sparse_group_size * send_count (matching the flat all-gather's A*cs broadcast scratch that the util already allocates -- corrected from the earlier passthrough-size assertion), in-place (active input aliases the output's own slot), and the layout [[0,1,2,3],[4,5,6,7]]. - FusedAllGather4ActiveValidationTest: all_gather_multi_group accepts 4 active ranks (RuntimeError without a native comm, not ValueError). bench_sharded_relay_perf.py: - Sweeps BOTH 2-active and 4-active all-gather; the all-gather baseline/fused bench are gated by the BENCH_ONLY env. Generalized the send-count math and built ag_send_counts over range(num_sparse_groups) (2-active byte-identical). Differential Revision: D115998383

Summary: `allreduce_tensors_with_sharded_relay` did its own packing and unpacking inline: a `torch.cat(out=)` into the active flat buffer before the per-group loop, and a `split()` + `torch._foreach_copy_` copy-out after the fused call. The reduce-scatter utils earlier in this stack introduced `_pack_into_flat` / `_unpack_from_flat`, which do exactly that (one fused kernel each) and are shared by every collective. Move allreduce onto them so all four collectives pack and unpack the same way. - torchrec/distributed/sharded_relay_utils.py: drop the inline cat/copy-out and call `_pack_into_flat` where the active group's tensor is built inside the per-group loop. Track the resulting `(flat buffer, destination list)` pair in `unpack_flat` / `unpack_dst` so step 5 no longer has to re-derive which buffer holds the active group's result. Keeping that pair explicit is what lets the out-of-place variant later in the stack retarget it without touching the unpack. - caffe2/torch/distributed/fb/sharded_relay_process_group.py (+ xplat mirror): note in the size-validation comment that each group is a single contiguous tensor, which is the invariant the loop below it relies on. No behavior change: the same buffers are packed, reduced in place, and unpacked into the same caller tensors. Differential Revision: D115998366

Summary: The RCCLX sharded-relay multi-group allreduce kernel is fully out-of-place capable, but the Python path only ever did in-place allreduce -- an asymmetry vs reduce_scatter and all_gather, which already expose both variants. This threads an optional output through every layer so Python callers can request out-of-place allreduce (reduced result written to a separate destination buffer, inputs preserved), matching the kernel's capability. When the new argument is omitted/None the behavior is identical to today (in-place). Changes (bottom-up): - comms/torchcomms/rcclx/TorchCommRCCLX.{hpp,cpp}: add optional output_tensors param to sharded_relay_multi_group_all_reduce. When provided, the active group's output points at the caller's separate contiguous output tensor and recvBuffs[segGroup] is retargeted; the active output numel() must match the input. Absent -> unchanged in-place aliasing. - comms/torchcomms/rcclx/TorchCommRCCLXPy.cpp: add output_tensors=None keyword pybind arg forwarding to the impl. - comms/torchcomms/rcclx/_comms_rcclx.pyi: reflect the new optional arg. - caffe2/torch/distributed/fb/sharded_relay_process_group.py (+ xplat mirror): add optional output_tensors to allreduce_multi_group, validate length, and pass it through. `skip_validation` also flips to default True here, matching the reduce_scatter / all_to_all / all_gather methods: the fused util validates its own shapes and does not want to pay for a second pass. - torchrec/distributed/sharded_relay_utils.py: add optional output_tensors_dict to allreduce_tensors_with_sharded_relay; per dtype build a parallel output group list (active group -> caller's output buffer, helper groups -> passthrough scratch) and forward it. None -> in-place (unchanged). That list is a plain list gated on `output_tensors_dict`, matching the sibling reduce_scatter / all_to_all / all_gather utils, and is passed as None on the in-place path so a single Optional drives the whole branch. Because of that default flip, a caller that wants the shape checks has to ask for them. `FusedShardedRelayValidationTest.test_raises_value_error_on_tensor_size_mismatch` now passes `skip_validation=False`; without it the call skips validation and reaches the "requires TorchCommRCCLX with native ... support" RuntimeError instead of raising the ValueError the test asserts. The out-of-place output is also dtype-checked. It is the buffer the kernel writes the active group's reduced result into, so it has to carry the same dtype as the input being reduced; a mismatch would write with the wrong element size. The in-place path is already covered by the per-group uniformity check earlier in the stack, and this extends the same invariant to the output introduced here. Differential Revision: D115998369

Summary: bench_sharded_relay_perf.py had a single fixed-shape benchmark (`test_benchmark`) and no way to see where the relay wins or loses across message sizes. Add a fused message-size sweep covering all four collectives at both 2 and 4 active ranks. - Adds `test_collectives_msg_size_sweep` to `BenchShardedRelayPerfTest`. It sweeps the fixed message sizes in the new `_MSG_SWEEP_SIZES` (4 KB .. 1 GB, bfloat16) for every collective (allreduce, reduce-scatter, all-to-all, all-gather) at both 2 and 4 active ranks, and prints one summary table per (collective, active-rank-count) — 8 tables total. - Size axis: the swept `nbytes` is the per-active-rank *input* tensor byte size (same meaning as the original allreduce sweep), so sizes stay comparable across collectives. - One scenario per size: FUSED — (NUM_GPUS // A) concurrent A-rank groups in one multi-group relay kernel call vs the matching NCCL baseline run in parallel on each rank's A-rank sub-group. The uncontended single-group scenario and the contended separate-job scenario are benchmarked by their own sweeps in the next commit, so each table here has one comparison column. - The sweep machinery is table-driven (parameterized by (collective, A)) rather than one hard-coded path per collective: shared per-collective helpers build the input/output/helper buffers and dispatch to the existing `bench_*` relay helpers and the NCCL baselines. - NCCL allreduce/reduce-scatter baselines keep SUM + manual divide (RCCL on MI350X lacks the AVG kernel); the relay path uses AVG in the kernel. all-to-all / all-gather do no reduction. Times are best-of-N (min), barrier-aligned + hipEvent-timed, reusing the existing `_measure_ms` helper. The change is purely additive; the existing `test_benchmark` is unchanged. Run the sweep selectively (the buck2 test runner imports the module, so the selector is the fully-qualified module.Class.method path): ``` buck2 run mode/opt-amd-gpu -m rocm70 -m rcclx_dev \ //torchrec/distributed/tests:bench_sharded_relay_perf -- \ torchrec.distributed.tests.bench_sharded_relay_perf.BenchShardedRelayPerfTest.test_collectives_msg_size_sweep ``` Both sweep workers now shut down the same way. `_msg_sweep_relay_worker` already called `dist.barrier()` before `dist.destroy_process_group()`; `_msg_sweep_nccl_worker` destroyed the group directly, so the pair had different shutdown semantics and a fast rank could tear down the process group while a peer was still finishing its collectives. The barrier is added to the NCCL worker to match. Both workers run an 8-rank NCCL default group with one rank per GPU, and the trailing barrier is straight-line code at the end of the worker that every rank reaches (the only `continue` in the loop is on `elements % active_ranks != 0`, which is identical on every rank), so it cannot deadlock. Differential Revision: D115998363

…sults Summary: Adds the two non-fused message-size sweeps to bench_sharded_relay_perf.py and moves all three sweeps onto file-based results emission. **Parallel independent-comm sweep.** Generalizes the parallel independent-comm sweep from allreduce-only at 2 active ranks to all four collectives at both 2 and 4 active ranks. This is the contended counterpart to test_collectives_msg_size_sweep's FUSED scenario: it measures single-group A-rank sharded relay performance under the XGMI link contention that arises when several independent single-group relays overlap on separate communicators. - Renames the test to `test_parallel_collectives_msg_size_sweep` on `BenchShardedRelayPerfTest`. It sweeps the fixed message sizes in `_MSG_SWEEP_SIZES` (4 KB .. 1 GB, bfloat16) for every collective (allreduce, reduce-scatter, all-to-all, all-gather) at both 2 and 4 active ranks, and prints one summary table per (collective, active-rank-count). - For each (collective, A) it models N = NUM_GPUS // A independent workloads, each with its OWN 8-rank RCCLX communicator and a disjoint active group (e.g. A=2: 0-1, 2-3, 4-5, 6-7; A=4: 0-3, 4-7). Each rank is active in exactly one comm and a relay helper in the others. All N single-group (num_groups=1) A-rank relays are launched concurrently via async_op (one per comm, each on its comm's dedicated stream) so they overlap. The N raw comms (world_size // min(A) = 4) are created once and reused across A. - NCCL baseline: N disjoint A-rank NCCL collectives run concurrently — identical to the FUSED sweep's baseline, so it reuses `_nccl_baseline_op` (allreduce and reduce-scatter use SUM + manual divide since RCCL on MI350X lacks the AVG kernel; the relay path uses AVG). Times are best-of-N (min), barrier-aligned + hipEvent-timed via `_measure_ms`. - Requires the collective wrappers to expose `async_op`; the reduce-scatter, all-to-all, and all-gather wrappers gained that in their respective distributed-utilities diffs (symmetric to allreduce) earlier in the stack. - Reuses the shared per-collective buffer/shape helpers (`_active_io`, `_relay_helper_size`, `_relay_counts`) and the report formatters (`_sweep_get_ms` / `_sweep_fmt_ms` / `_sweep_fmt_speedup`) from the FUSED sweep. Each phase picks a free TCPStore port at runtime (`_find_free_port`) to avoid EADDRINUSE from a socket left in TIME_WAIT. **Single-group best-case sweep.** Adds `test_single_group_collectives_msg_size_sweep`, which runs exactly ONE A-rank relay group on a full 8-rank communicator with no co-resident jobs, against an NCCL baseline of one A-rank collective with the remaining ranks idle. It is the uncontended upper bound for the same grid, and it isolates per-call overhead from the link contention the parallel sweep measures. It shares the parallel sweep's worker path and adds `_format_single_group_sweep_table` / `_print_single_group_msg_sweep_report`. **File-based results emission.** The sweep spawns up to N x NUM_GPUS workers, all writing glog / thrift / RCCLX init logging to the same stdout, which mangles tables printed line by line. `_emit_report` now assembles each report as one string and writes it atomically to both stdout and a dedicated file (`bench_sharded_relay_{fused,parallel,single_group}_sweep_results.txt` under tmpdir; `BENCH_RESULTS_FILE` overrides the path). `_print_sweep_table` is replaced by `_format_sweep_table` so the FUSED sweep from the previous commit goes through the same path, and `_measure_ms` is reworked to be barrier-aligned and hipEvent-timed so the concurrent-job timings are comparable across ranks. `test_benchmark` is otherwise unchanged. Small-message note: the A=2 reduce-scatter/all-to-all 2-active relay path requires elements >= 1792, which every swept size (>= 4 KB = 2048 elements) satisfies; A=4 uses the recursive path with no such floor. Run the sweep selectively (the buck2 test runner imports the module, so the selector is the fully-qualified module.Class.method path): ``` buck2 run mode/opt-amd-gpu -m rocm70 -m rcclx_dev \ //torchrec/distributed/tests:bench_sharded_relay_perf -- \ torchrec.distributed.tests.bench_sharded_relay_perf.BenchShardedRelayPerfTest.test_parallel_collectives_msg_size_sweep ``` Both parallel sweep workers now shut down the same way. `_parallel_msg_sweep_relay_worker` already called `dist.barrier()` before `dist.destroy_process_group()`; `_parallel_msg_sweep_nccl_worker` destroyed the group directly, so the pair had different shutdown semantics and a fast rank could tear down the process group while a peer was still finishing its collectives. The barrier is added to the NCCL worker to match. The default group here is gloo across all N * NUM_GPUS procs (the per-job NCCL device group is separate), so this is the same cheap global barrier the relay worker already performs, and every rank reaches it (the only `continue` is on `elements % active_ranks != 0`, identical on every rank, and there is no early return before the destroy). Differential Revision: D115998386

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

…ailing (meta-pytorch#3678) Summary: Pull Request resolved: meta-pytorch#3678 X-link: meta-pytorch/torchrec#4570 When a group's aligned relay chunk rounded down to zero the four sharded relay collectives returned `ncclInvalidArgument` and required the caller to retry on a plain collective. In a fused multi-group call that is worse than it sounds: a single small or awkwardly sized group forced the whole fused call to fail even when the other groups were large enough to relay, so heterogeneous group sizes were effectively unsupported. Give every collective a phase-symmetric per-group direct fallback instead. When `chunkSize` aligns down to zero the segment is covered by the two direct regions (`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link, the helper scatter/forward is skipped for that group, and the helper-side reduction is skipped as well. Groups that do have a workable chunk size keep the existing relay schedule unchanged, so each group in a fused call now picks its own schedule independently. - track `dirASizes` per group so the direct regions are sized from the selected schedule rather than assumed equal to `chunkSize` - guard the helper scatter, forward, and reduction on `chunkSize > 0` - all-gather additionally tracks `dirAOffsets`, which is no longer derivable from `relayTotals` once the fallback can move the direct region to offset 0 - correct the minimum helper geometry to H+2 chunks Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was `_passthrough_helper_size(...)` when production has always sized it through `_relay_helper_size()` (`2 * recvCount`); the expectation now matches the unchanged production contract. Behavior for buffers that already had a nonzero aligned chunk is unchanged. Reviewed By: JinghanHuang Differential Revision: D115998367

…ailing (#3678) Summary: Pull Request resolved: #3678 X-link: meta-pytorch/torchrec#4570 When a group's aligned relay chunk rounded down to zero the four sharded relay collectives returned `ncclInvalidArgument` and required the caller to retry on a plain collective. In a fused multi-group call that is worse than it sounds: a single small or awkwardly sized group forced the whole fused call to fail even when the other groups were large enough to relay, so heterogeneous group sizes were effectively unsupported. Give every collective a phase-symmetric per-group direct fallback instead. When `chunkSize` aligns down to zero the segment is covered by the two direct regions (`dirA` takes half, `dirB` absorbs the remainder) over the active-to-active link, the helper scatter/forward is skipped for that group, and the helper-side reduction is skipped as well. Groups that do have a workable chunk size keep the existing relay schedule unchanged, so each group in a fused call now picks its own schedule independently. - track `dirASizes` per group so the direct regions are sized from the selected schedule rather than assumed equal to `chunkSize` - guard the helper scatter, forward, and reduction on `chunkSize > 0` - all-gather additionally tracks `dirAOffsets`, which is no longer derivable from `relayTotals` once the fallback can move the direct region to offset 0 - correct the minimum helper geometry to H+2 chunks Also fixes a torchrec test that asserted the A>2 reduce-scatter helper buffer was `_passthrough_helper_size(...)` when production has always sized it through `_relay_helper_size()` (`2 * recvCount`); the expectation now matches the unchanged production contract. Behavior for buffers that already had a nonzero aligned chunk is unchanged. Reviewed By: JinghanHuang Differential Revision: D115998367 fbshipit-source-id: 1664292a7a3acd72c7a4b78a890e41e1388db822

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 15, 2026

meta-codesync Bot added the meta-exported label Aug 15, 2026

srinathb-meta mentioned this pull request Aug 15, 2026

Serve zero-chunk relay geometries with a direct fallback instead of failing (#3678) meta-pytorch/torchcomms#3678

Closed

meta-codesync Bot changed the title ~~Serve zero-chunk relay geometries with a direct fallback instead of failing~~ Serve zero-chunk relay geometries with a direct fallback instead of failing (#4570) Aug 15, 2026

srinathb-meta force-pushed the export-D115998367 branch from 6ea6196 to ccd475d Compare August 15, 2026 03:20

srinathb-meta force-pushed the export-D115998367 branch from ccd475d to de6df81 Compare August 15, 2026 06:04

meta-codesync Bot changed the title ~~Serve zero-chunk relay geometries with a direct fallback instead of failing (#4570)~~ Link-balance the sharded relay schedules (up to 3.2x vs NCCL) Aug 15, 2026

srinathb-meta force-pushed the export-D115998367 branch from de6df81 to d3522a6 Compare August 16, 2026 04:46

srinathb-meta force-pushed the export-D115998367 branch from d3522a6 to d13aae9 Compare August 17, 2026 06:08

srinathb added 15 commits August 17, 2026 09:16

srinathb added 3 commits August 17, 2026 09:16

srinathb-meta force-pushed the export-D115998367 branch 2 times, most recently from bafb162 to c2ca39a Compare August 17, 2026 16:28

Navigation Menu

Uh oh!

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

Link-balance the sharded relay schedules (up to 3.2x vs NCCL)#4570
srinathb-meta wants to merge 18 commits into
meta-pytorch:mainfrom
srinathb-meta:export-D115998367

srinathb-meta commented Aug 15, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

meta-codesync Bot commented Aug 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

srinathb-meta commented Aug 15, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Aug 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

srinathb-meta commented Aug 15, 2026 •

edited by meta-codesync Bot

Loading