Repository navigation
[RFC] Batch replica selection for affine transfer plans #8935
Description
Motivation
plan_transfer (#8623) takes an optional replica_selector: Callable[[ReplicaSelectionRequest], int]. When the target region of one piece meets a source region that several locations hold, the planner calls the selector once for that intersection and uses the returned location. Without a selector it takes the first holder, which is the lowest location.
This per-request form has two limits, which @jinyouzhi pointed out on #8885:
- No load balancing. The selector sees one request at a time, so it cannot spread bytes across holders, and its result depends on call order.
- No splitting. A region can come from only one holder, even when several hold it and the region is large.
The choice matters. On #8885, jinyouzhi measured repeated ZeRO-2 syncs: the lowest-rank default was the slowest policy, and "local first, then fewest assigned bytes" roughly halved the sync time (0.49 s vs 1.05–1.35 s).
On #8885, @jinyouzhi, @Achyuthan-S and I agreed on a batch form:
- the planner applies a pluggable policy, but hands it every request at once;
- the policy may split regions;
- the planner checks the result exactly.
@Achyuthan-S takes the policy itself from mid-November, as the decision layer above the plan in #8252. This issue is meant to agree on the interface before then.
A naming point first: locations vs processes
The source_rank / target_rank fields of a plan are location keys of the two maps (pieces_by_rank), not process ranks. Which process holds each location is known only to the executor's TransferEndpoints, which keeps separate source and target mappings. When a trainer syncs into a separate engine, source location 0 and target location 0 are different processes; in the #8809 external-group test they map to group ranks 0 and 1. So "local first" cannot be decided from the maps alone, and the interface below says "location" where it means a map key and "process" where it means a group rank.
Proposal (strawman)
A multi-parameter planning entry
Balancing only works across the whole sync, not one parameter at a time, so the batch has to span parameters:
plans = plan_transfers(
pairs, # Mapping[str, Tuple[ParamAffineMap, ParamAffineMap]]: name -> (target_map, source_map)
limits=None, # PlanBudget, per parameter as today
*,
endpoints=None, # Optional[TransferEndpoints]: lets the policy see processes
element_size=None, # Optional[Mapping[str, int]]: bytes per element per parameter
replica_policy=None, # Callable[[Sequence[ReplicaRequest]], Sequence[ReplicaAnswer]]; None -> default below
) # -> Dict[str, TransferPlan]plan_transfer(target_map, source_map, ...) stays as the one-parameter case. The executor already takes a batch (execute_transfers in #8809), so this mirrors it.
What the policy sees
One request per intersection of a target piece with a source region. Requests with a single holder are included too, so the policy knows the load it cannot move:
@dataclass(frozen=True)
class ReplicaRequest:
parameter: str
target_location: int # as TransferSegment.target_rank
target_process: Optional[int] # endpoints.target[target_location], if endpoints were given
logical_origin: Tuple[int, ...]
shape: Tuple[int, ...]
target_offset: int
target_strides: Tuple[int, ...]
candidates: Tuple[ReplicaCandidate, ...] # validated holders; same type as today (source_rank = location)
candidate_processes: Tuple[Optional[int], ...] # endpoints.source[...] for each candidate, if given
bytes: Optional[int] # elements x element_size[parameter], if givenRequests arrive in a fixed order, sorted by parameter, then target location, then origin.
What the policy returns
One answer per request, in the same order:
int: the whole region from that source location, as today.Sequence[ReplicaPart]: a split into parts.
@dataclass(frozen=True)
class ReplicaPart:
logical_origin: Tuple[int, ...]
shape: Tuple[int, ...]
source_location: intWhat the planner checks
Everything the policy returns is checked before any segment is emitted. A failure raises InvalidReplicaPolicyError, as an invalid rank does today:
- exactly one answer per request;
- every location is one of that request's validated candidates;
- every part lies inside its request's box, the parts are pairwise disjoint, and their volumes sum to the request's volume, so they tile it exactly;
- the extra segments from splits stay within
PlanBudget(max_segments,max_plan_bytes).
Each part becomes a segment. The planner computes its addresses from the chosen holder, as it does now, so a policy never handles offsets.
Default policy
The default needs no hardware information and does not split:
- With endpoints: if a candidate runs in the same process as the target, use it (a local copy). Otherwise pick the candidate whose process has the fewest bytes assigned so far in this batch, counting single-holder requests. Ties go to the lowest location.
- Without endpoints: fewest bytes assigned per source location, ties to the lowest location. There is no local preference, because the planner cannot tell what is local.
The default counts elements when element_size is not given, which is the same as bytes only if the whole sync uses one dtype.
It replaces "first holder" as the default, and the old behaviour stays available as a one-line policy.
Determinism
execute_transfers is collective: every member of the transfer group must run the same plan. If every rank plans for itself, the policy must be deterministic on the same input. If one process plans and ships the plan, it need not be. The interface doesn't choose between these, but the docs should say so.
Compatibility
An existing per-request replica_selector keeps working: the planner wraps it as a batch policy that answers each request on its own. We can deprecate it later or keep it as a convenience.
Out of scope
- The cost model itself: links, contention, topology. That is the policy layer, owned by @Achyuthan-S under Define affine IR for checkpointing #8252. [checkpoint] AffineIR transfer-plan analysis and link calibration (preview) jinyouzhi/DeepSpeed#6 (plan analysis and link calibration) could supply its inputs.
- Split heuristics beyond what the default needs.
- Executor changes. Segments from splits are ordinary segments.
Open questions
- Batch scope. Is the whole sync (
plan_transfers) the right unit? A policy that balances across several syncs would need state outside the planner. - Split shape. Arbitrary sub-boxes, or only cuts along axis 0? Arbitrary boxes are no harder to check, but they can produce more strided segments.
- Where processes come from. Is
TransferEndpointsthe right input, or should the planner take a plain location-to-process mapping so it does not depend on the executor module? - Placement of the default. Inside
affine_transfer.py, or a smallaffine_replica_policy.pynext to it, so policies don't grow the planner? - The old selector. Keep the adapter, or deprecate it once the batch form lands?
Plan
- Planner side: I can take
plan_transfers, the request and answer types, the checks above, the default policy, and the per-request adapter.- Tests: randomized replicated maps where every policy answer must give identical copied data.
- The existing adversarial planner cases must still be refused.
- Policies that return a wrong location, a gap, an overlap or a part outside the box must be rejected.
- Policy layer (@Achyuthan-S, from mid-November): cost-aware policies on top of this hook.
- Baseline (@jinyouzhi): the ZeRO-2 measurement above becomes the reference for the default.
Feedback on the interface is welcome before anyone starts on it. cc @delock, since this builds on #8252 and #8230.
Activity
Achyuthan-S commentedon Oct 28, 2026
Thanks @0z5a, this matches what we agreed on #8885, and I'm happy with the split. Thoughts on the open questions:
TransferEndpointscan produce that mapping.affine_replica_policy.py. The cost-aware policies will grow, and keeping them there leaves the planner small.One link to the correctness work: checks 2 and 3 are exactly what makes any answer safe. If every part uses a validated holder and the parts tile the request, each target address gets one holder of the same logical element, and when the replicas agree the copied data is identical whichever holder is chosen. That's Theorem 3 in the proof draft I'm sharing for #8252, and your randomized "every policy answer gives identical data" test is the matching check.
0z5a commentedon Oct 28, 2026
Thanks @Achyuthan-S, agreed on all five:
TransferEndpointscan produce them.affine_replica_policy.py.Thanks for the link to Theorem 3. I'll cite it in the randomized test's docstring.
Some transport measurements from my side bear on what a cost-aware policy needs to know. Each item names its hardware.
Splitting has a floor.
plan_transferstakes amin_split_bytes, defaulting to something tied to the staging budget, and the planner refuses smaller splits the same way it refuses an invalid location.On PCIe, concurrency sets the ceiling, not the link.
Across nodes the sign may flip.
For the interface I think this means two things:
process_nodemapping rather than a full topology object.cc @jinyouzhi @delock
Achyuthan-S commentedon Oct 28, 2026
Thanks @0z5a, these measurements are exactly the inputs the policy needs, and thanks for citing Theorem 3.
+1 to both additions:
process_node: location→process plus process→node gives each candidate's path (same process, same node, or network), which is what the first cost-aware policy needs. A full topology object can wait until something needs it.min_split_bytes: your numbers suggest deriving it rather than fixing it. A part pays off only if moving its bytes saves more time than the round it adds, so roughlybytes > α·βfor that path: 0.3 ms × 48 GB/s ≈ 14 MB on the GH200 link, in line with your "few tens of MB". A default tied to the staging budget is a good start, and a policy can pass a calibrated value.For mid-November I'm planning the policy as follows:
That last part is where measurements matter most, so I'd like to take you up on the GH200 offer for it.