Skip to content

TL/UCP: Relax Allreduce Ring Constraints - #1336

Open
wfaderhold21 wants to merge 2 commits into
openucx:masterfrom
wfaderhold21:topic/relax-ar-ring
Open

TL/UCP: Relax Allreduce Ring Constraints#1336
wfaderhold21 wants to merge 2 commits into
openucx:masterfrom
wfaderhold21:topic/relax-ar-ring

Conversation

@wfaderhold21

Copy link
Copy Markdown
Collaborator

What

Addresses #1293

wfaderhold21 and others added 2 commits July 30, 2026 13:54
Ring reduce-scatter only implements post-op AVG and rejects UCC_OP_AVG
when REDUCE_AVG_PRE_OP=1 (the default). With the default selector now
routing >=4MB allreduce to ring, a large default AVG allreduce is first
offered to ring; the score-map fallback must transparently complete it
on a pre-op-capable algorithm (knomial/sra).

Add two gtests to the test_allreduce_avg_order fixture:
 - avg_ring_pre_op_fallback: force @Ring for all sizes under default
   pre-op and verify the collective still completes with correctly
   averaged results (deterministic fallback proof).
 - avg_large_default_selection: default config with a >4MB message so
   the real default selector routes to ring; regression guard for the
   4m-inf ring routing.
@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR allows UCP ring allreduce to process counts that are not divisible by the team size and selects ring by default for allreduce messages of at least 4 MiB.

  • Uses remainder-aware block sizes and offsets in ring allgather communication and initial source copying.
  • Removes the allreduce ring divisibility restriction.
  • Adds non-divisible-count and AVG fallback coverage across supported host and GPU memory types.
  • Restores test-process environment variables to their original set or unset state after context creation.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete changed-code failure identified.

The uneven ring paths consistently calculate matching block sizes and offsets, score selection retains a working generic fallback for unsupported ring operations, and the test harness restores environment state after its configuration consumers run.

Important Files Changed

Filename Overview
src/components/tl/ucp/allgather/allgather_ring.c Replaces equal-block arithmetic with consistent remainder-aware counts and offsets for sends, receives, and the initial copy.
src/components/tl/ucp/allreduce/allreduce.h Changes default allreduce selection to knomial below 4 KiB, SRA knomial through 4 MiB, and ring for larger messages.
src/components/tl/ucp/allreduce/allreduce_ring.c Removes the divisibility guard so the composed reduce-scatter and allgather ring can use uneven block partitioning.
test/gtest/coll/test_allreduce.cc Adds broad correctness coverage for non-divisible ring counts and AVG fallback under forced and default selection.
test/gtest/common/test_ucc.cc Preserves each unique environment variable's original presence and value, including unsetting variables that were initially absent.

Reviews (1): Last reviewed commit: "TEST: add allreduce pre-op fall back tes..." | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant