[UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter

[0z5a]

Summary

Compose source and target affine maps into shard-to-shard copy descriptors without rebuilding the logical parameter. This PR remains metadata only and based on master.

Resolved copies now fold when both physical shard address mappings continue along an axis. Their descriptor no longer claims a single logical_origin: a copy can span two distinct logical regions. The planner keeps every validated replica candidate until a target intersection is resolved. It uses the lowest holder by default and accepts a caller-supplied replica_selector for other choices. Invalid selections fail explicitly.

The spec distinguishes homogeneous input map pieces from resolved copies. The regression suite is committed in tests/unit/checkpoint/test_affine_transfer_plan.py.

Review result

BigCode TP2→TP4, 48×8	Previous plan	Updated plan	Target elements written	Result
Lowest holder	8 segments	7 segments	768	Exact shard contents and guarded storage offsets
Holder with the target's private Q rows	Unavailable	6 segments	768	Exact shard contents and guarded storage offsets
Highest-ranked KV holder	Unavailable	7 segments	768	Exact shard contents and guarded storage offsets

The 7/6/7 counts apply to this layout and these three holder policies. For targets 0 and 1, the private Q rows are on source 0; for targets 2 and 3, they are on source 1. Selecting that same source for the replicated KV block makes the query and KV reads adjacent, allowing the sixth segment. Selecting the highest-ranked KV holder for every target does not. The previous 8-segment plan was valid; this change compresses its copy schedule. The selector does not change the input map's locations or relax its homogeneity rule.

On the 8× RTX 5090 host, Python 3.12.13 and torch 2.13.0+cu130, the B1 CPU suite passed: 98 passed (PYTHONPATH=$PWD CUDA_VISIBLE_DEVICES="" DS_ACCELERATOR=cpu python -m pytest tests/unit/checkpoint/test_affine_transfer_plan.py -q). Tests cover exact address oracles, nonzero storage offsets and guard regions, replica offsets, invalid selectors, serialization, budgets, determinism and the no-materialization contract.

The stacked B2 executor PR has two-GPU copy and model E2E results, including the direct-versus-rebuild speed table. B1 itself does not move data, so it makes no transport speed claim.

Scope

Scale 1 and matching logical shapes only. No checkpoint format change, topology search, CUDA graph, online weight update or optimizer coordinate transform. Refs #8230 and #8252.

[Achyuthan-S]

@0z5a Both numbers are reachable — §8.2's 7 is not conditional on choosing the replica per target.

I recomposed it from segmented_map at TP2 and TP4 and executed the plans against a real
tensor, checked element-for-element against extract: lowest holder gives 7, per-target
choice gives 6, both exact with nothing uncovered. Under lowest holder the fold is on target
rank 1 rather than 3 — source rank 0's query block ends exactly where its kv block begins,
and target rank 1's query rows are the last 8 of it, so src0[8:32] -> tgt1[0:24] is one
contiguous copy on both sides.

If you are getting 8, I would guess the merge stops at the query/kv boundary. That rule is
§5 consequence 4, and it is about map pieces: a piece carries locations, so it cannot span
a region that is rank-private in one half and replicated in the other. A plan segment carries
no locations — the holder is resolved before any copy is emitted — so the boundary is not
load-bearing there. Worth stating in the spec, since §5 reads like it forbids the fold
everywhere.

On the policy: agreed it should be a parameter. Lowest holder is not topology-neutral either
— it points every replica read at rank 0, which is the case a real transport least wants.

[0z5a]

@0z5a Both numbers are reachable — §8.2's 7 is not conditional on choosing the replica per target.

I recomposed it from segmented_map at TP2 and TP4 and executed the plans against a real

tensor, checked element-for-element against extract: lowest holder gives 7, per-target

choice gives 6, both exact with nothing uncovered. Under lowest holder the fold is on target

rank 1 rather than 3 — source rank 0's query block ends exactly where its kv block begins,

and target rank 1's query rows are the last 8 of it, so src0[8:32] -> tgt1[0:24] is one

contiguous copy on both sides.

If you are getting 8, I would guess the merge stops at the query/kv boundary. That rule is

§5 consequence 4, and it is about map pieces: a piece carries locations, so it cannot span

a region that is rank-private in one half and replicated in the other. A plan segment carries

no locations — the holder is resolved before any copy is emitted — so the boundary is not

load-bearing there. Worth stating in the spec, since §5 reads like it forbids the fold

everywhere.

On the policy: agreed it should be a parameter. Lowest holder is not topology-neutral either

— it points every replica read at rank 0, which is the case a real transport least wants.

Thanks, that clears up the distinction.

The current merge is indeed stopping at the source-piece boundary, so the TP2 -> TP4 BigCode case stays at 8. That is carrying the map-level homogeneity constraint from §5 into the resolved transfer plan, where locations is no longer part of the descriptor.

The update will separate those two rules:

• map pieces remain subject to the §5/P5 replication-boundary constraint;
• resolved transfer segments may fold across that boundary when the selected source rank is the same and both source and destination shard addresses are affine-contiguous.

With the existing lowest-holder policy, BigCode TP2 -> TP4 should therefore regress to 7 segments. I’ll also make replica-holder selection an explicit planner policy rather than baking lowest-holder into the geometry. A per-target policy can then exercise the 6-segment case separately, without treating that as the canonical/default plan.

This also seems worth clarifying next to the composer rules in §8.2: P5 constrains the IR pieces, not already-resolved transfer segments.

[Achyuthan-S]

Verified against the new head. I rebuilt the maps with segmented_map at TP2 and TP4 rather than using your fixture, planned with plan_transfer, and executed the segments element for element against extract: lowest holder gives 7, a per-target selector gives 6, both write 768 with nothing uncovered. Your committed suite also passes here — 97, CPU.

The spec text is right, and separating "resolved copies are not another homogeneous map" from P5 is the part I'd have got wrong if I'd written it alone. A folded copy having no single logical origin is the consequence I hadn't followed through.

One thing to tighten: "a specified per-target holder choice yields 6" depends on which per-target policy. A selector that takes the highest holder still gives 7 — I tried it. Only a selector that picks the rank holding the target's own private rows reaches 6, because that is what puts the query read adjacent to the kv read in the source shard. Worth saying, since
"per-target" alone does not imply the fold.

The selector hook is the right shape. It is the seam a cost model plugs into, which is the part I took on #8252.

[0z5a]

@Achyuthan-S Thanks, good catch. §8.2 and both draft PR descriptions now specify that the 6-segment fold requires choosing the holder of each target's private Q rows; choosing the highest KV holder still gives 7. I added planner and two-GPU execution regressions for 7/6/7. B1 passes 98 CPU tests, and B2 passes 12 cases per rank.