From cdb6ec6d825a69e0fee58f0da831e347e93b8056 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Sat, 31 Oct 2026 15:57:05 +0000
Subject: [PATCH 1/7] Check affine holders and source regions without pairwise
 comparisons

_validate_holders compared every piece with every other piece, and
_sources compared every source region with every other region, so a map
with 1024 pieces already exceeded the default max_pair_checks and
planning time grew quadratically.

When no rank repeats a box and distinct boxes do not overlap, a rank
covers a box exactly when it holds that very box. The holders of each box
can then be read from a table. Add boxes_disjoint to affine_geometry: it
sweeps along the most varied axis and compares a box only with the boxes
still open where it starts. Maps that fail the fast checks take the
original pairwise walk, which names the offending piece, so error
messages are unchanged.

A randomized test checks that plans and errors match the pairwise walk,
including replicas split at different points, overlapping pieces and
wrong holder lists.

Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 deepspeed/checkpoint/affine_geometry.py       | 23 ++++++
 deepspeed/checkpoint/affine_transfer.py       | 37 ++++++++--
 .../checkpoint/test_affine_transfer_plan.py   | 73 +++++++++++++++++++
 3 files changed, 125 insertions(+), 8 deletions(-)

diff --git a/deepspeed/checkpoint/affine_geometry.py b/deepspeed/checkpoint/affine_geometry.py
index 4e8a9b092952..c58d899713cb 100644
--- a/deepspeed/checkpoint/affine_geometry.py
+++ b/deepspeed/checkpoint/affine_geometry.py
@@ -119,3 +119,26 @@ def contains(self, other: 'LogicalBox') -> bool:
 
     def __repr__(self) -> str:
         return f'Box({self.origin}+{self.extent})'
+
+
+def boxes_disjoint(boxes: Sequence[LogicalBox]) -> bool:
+    """Whether no two boxes share an element.
+
+    Sweeps along the axis whose box origins are most varied, so a box is compared only with the boxes
+    still open where it starts. For the grid and run layouts the builders emit that is a handful of
+    boxes, not all of them.
+    """
+    boxes = [box for box in boxes if not box.is_empty]
+    if len(boxes) < 2:
+        return True
+    if not boxes[0].origin:
+        return False
+    axis = max(range(len(boxes[0].origin)), key=lambda a: len({box.origin[a] for box in boxes}))
+    open_boxes: List[LogicalBox] = []
+    for box in sorted(boxes, key=lambda b: b.origin[axis]):
+        start = box.origin[axis]
+        open_boxes = [other for other in open_boxes if other.origin[axis] + other.extent[axis] > start]
+        if any(not box.intersection(other).is_empty for other in open_boxes):
+            return False
+        open_boxes.append(box)
+    return True
diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py
index 4a127757d191..c89d44cd0925 100644
--- a/deepspeed/checkpoint/affine_transfer.py
+++ b/deepspeed/checkpoint/affine_transfer.py
@@ -28,11 +28,11 @@
 import collections.abc
 import json
 from dataclasses import dataclass
-from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Sequence, Set, Tuple, Union
 
 from .affine import AffinePiece, ParamAffineMap
-from .affine_geometry import (LogicalBox, coordinate as _coordinate, product as _product, row_major_strides as
-                              _row_major_strides)
+from .affine_geometry import (LogicalBox, boxes_disjoint, coordinate as _coordinate, product as _product,
+                              row_major_strides as _row_major_strides)
 
 __all__ = [
     'AffineTransferError', 'InvalidMapError', 'InvalidReplicaPolicyError', 'UnsupportedLayoutError', 'PlanBudget',
@@ -360,7 +360,26 @@ def _shards_tile(frames: List[_Frame], side: str, limits: PlanBudget) -> None:
 
 
 def _validate_holders(frames: List[_Frame], side: str, limits: PlanBudget) -> None:
-    """A piece's declared holders must cover exactly its logical box on every named rank."""
+    """A piece's declared holders must cover exactly its logical box on every named rank.
+
+    Most maps pass without comparing pieces pairwise: when no two distinct boxes overlap, a rank covers
+    a box exactly when it holds that very box, so the holders of each box can be read off a table.
+    Anything else, such as replicas split at different places, takes the pairwise walk, which also
+    names the piece at fault.
+    """
+    holders: Dict[LogicalBox, Set[int]] = {}
+    repeated = False
+    for frame in frames:
+        ranks = holders.setdefault(frame.box, set())
+        repeated = repeated or frame.rank in ranks
+        ranks.add(frame.rank)
+    if not repeated and boxes_disjoint(list(holders)):
+        if all(holders[frame.box] == frame.locations for frame in frames):
+            return
+    _validate_holders_pairwise(frames, side, limits)
+
+
+def _validate_holders_pairwise(frames: List[_Frame], side: str, limits: PlanBudget) -> None:
     if len(frames)**2 > limits.max_pair_checks:
         raise AffineTransferError(f'{side}: holder checks exceed max_pair_checks {limits.max_pair_checks}')
     by_rank: Dict[int, List[_Frame]] = {}
@@ -403,18 +422,20 @@ def _sources(frames: List[_Frame], limits: PlanBudget) -> List[Tuple[LogicalBox,
             by_box.setdefault(frame.box, []).append(frame)
 
     regions = sorted(by_box, key=lambda box: (box.origin, box.extent))
+    if not boxes_disjoint(regions):
+        _name_overlapping_regions(regions, limits)
+    return [(box, tuple(sorted(by_box[box], key=lambda frame: frame.rank))) for box in regions]
+
+
+def _name_overlapping_regions(regions: List[LogicalBox], limits: PlanBudget) -> None:
     if len(regions)**2 > limits.max_pair_checks:
         raise AffineTransferError(f'{len(regions)} source regions would take {len(regions)**2} checks, over '
                                   f'max_pair_checks {limits.max_pair_checks}')
-    grouped = []
     for index, box in enumerate(regions):
-        holders = by_box[box]
         for other in regions[index + 1:]:
             if not box.intersection(other).is_empty:
                 raise UnsupportedLayoutError(f'source: replicas of {box!r} and {other!r} have different '
                                              'piece boundaries')
-        grouped.append((box, tuple(sorted(holders, key=lambda frame: frame.rank))))
-    return grouped
 
 
 def _covers_exactly(region: LogicalBox, parts: Sequence[LogicalBox]) -> bool:
diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py
index b00fe855f1b4..fc14bfa14383 100644
--- a/tests/unit/checkpoint/test_affine_transfer_plan.py
+++ b/tests/unit/checkpoint/test_affine_transfer_plan.py
@@ -31,6 +31,7 @@
 
 from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, block_gather_map, contiguous_split_map,
                                          replicated_map, segmented_map, sub_param_map)
+from deepspeed.checkpoint.affine_geometry import LogicalBox, boxes_disjoint
 from deepspeed.checkpoint.affine_transfer import (AffineTransferError, InvalidMapError, PlanBudget, TransferSegment,
                                                   UnsupportedLayoutError, plan_transfer)
 
@@ -666,6 +667,78 @@ def test_merge_comparisons_have_their_own_budget():
         plan_transfer(target, source, PlanBudget(max_merge_checks=10))
 
 
+def test_many_consistent_pieces_need_no_pairwise_holder_checks():
+    rows, ranks = 1024, 8
+    source = block_gather_map((rows, 2), {rank: list(range(rank, rows, ranks)) for rank in range(ranks)}, 1, 0)
+    target = contiguous_split_map((rows, 2), [rows // ranks] * ranks, 0)
+    assert rows**2 > PlanBudget().max_pair_checks
+    full = torch.arange(rows * 2, dtype=torch.float64).reshape(rows, 2)
+    copied = _apply_plan(plan_transfer(target, source), target, _shards(source, lambda m, r: m.extract(full, r)))
+    for rank, shard in copied.items():
+        torch.testing.assert_close(shard, target.extract(full, rank))
+
+
+def _random_holder_map(rng, rows, cols, ranks):
+    """Row blocks on every rank, sometimes split at a replica's own points, with holders sometimes wrong."""
+    shared = sorted(rng.sample(range(1, rows), rng.randint(0, rows - 1)))
+    blocks_of = {}
+    for rank in range(ranks):
+        cuts = shared if rng.random() < 0.7 else sorted(rng.sample(range(1, rows), rng.randint(0, rows - 1)))
+        bounds = list(zip([0] + cuts, cuts + [rows]))
+        held = [block for block in bounds if rank == 0 or rng.random() < 0.6]
+        if rank and rng.random() < 0.1 and held:
+            low, high = held[0]
+            held.append((low, min(rows, high + 1)))
+        blocks_of[rank] = held
+    pieces, shapes = {}, {}
+    for rank, blocks in blocks_of.items():
+        offset, pieces[rank] = 0, []
+        for low, high in blocks:
+            holders = [other for other, theirs in blocks_of.items() if any(a <= low and high <= b for a, b in theirs)]
+            if rng.random() < 0.15:
+                holders = sorted(set(holders) ^ {rng.randrange(ranks)}) or [rank]
+            shape = (high - low, cols)
+            pieces[rank].append(AffinePiece(shape, low * cols, (cols, 1), offset, (cols, 1), holders))
+            offset += _product(shape)
+        shapes[rank] = (offset, )
+    return ParamAffineMap((rows, cols), shapes, pieces)
+
+
+def _outcome(target, source):
+    try:
+        plan = plan_transfer(target, source, PlanBudget(max_pair_checks=1 << 30))
+    except AffineTransferError as error:
+        return type(error), str(error)
+    return [segment.to_dict() for segment in plan.segments]
+
+
+@pytest.mark.parametrize('seed', range(200))
+def test_holder_check_agrees_with_pairwise_walk(seed, monkeypatch):
+    """The table lookup is a shortcut, so the exhaustive comparison it replaces is the oracle."""
+    from deepspeed.checkpoint import affine_transfer as transfer_module
+
+    rng = random.Random(seed)
+    rows, cols, ranks = rng.randint(1, 6), rng.randint(1, 3), rng.randint(1, 4)
+    mapped = _random_holder_map(rng, rows, cols, ranks)
+    other = replicated_map((rows, cols), 1)
+    target, source = (mapped, other) if seed % 2 else (other, mapped)
+    fast = _outcome(target, source)
+    monkeypatch.setattr(transfer_module, '_validate_holders', transfer_module._validate_holders_pairwise)
+    assert fast == _outcome(target, source)
+
+
+@pytest.mark.parametrize('seed', range(100))
+def test_boxes_disjoint_agrees_with_all_pairs(seed):
+    rng = random.Random(seed)
+    ndim = rng.randint(1, 3)
+    boxes = []
+    for _ in range(rng.randint(0, 8)):
+        origin = tuple(rng.randint(0, 4) for _ in range(ndim))
+        boxes.append(LogicalBox(origin, tuple(rng.randint(0, 3) for _ in range(ndim))))
+    overlapping = any(not a.intersection(b).is_empty for a, b in itertools.combinations(boxes, 2))
+    assert boxes_disjoint(boxes) == (not overlapping)
+
+
 # ---------------------------------------------------------------- never materializing
 
 

From 52faa3f93a72e71280355222afd05d993c2f0700 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Mon, 9 Nov 2026 12:49:43 +0000
Subject: [PATCH 2/7] Check each rank's shard addresses without pairwise
 comparisons

A rank whose pieces sweep as disjoint in its shard and fill it passes
_shards_tile with boxes_disjoint; any other rank takes the pairwise walk,
so error messages are unchanged. Maps that pass no longer hit the
per-rank len(held)**2 budget.

Suggested-by: 0z5a
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 deepspeed/checkpoint/affine_transfer.py            |  6 ++++++
 tests/unit/checkpoint/test_affine_transfer_plan.py | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py
index c89d44cd0925..4578cf6297f8 100644
--- a/deepspeed/checkpoint/affine_transfer.py
+++ b/deepspeed/checkpoint/affine_transfer.py
@@ -338,12 +338,18 @@ def _shards_tile(frames: List[_Frame], side: str, limits: PlanBudget) -> None:
     of them land on the same addresses of the shard, and their numels still sum to the shard size. A
     shard that two pieces claim cannot be read or written by address, so a plan built on such a map is
     wrong in whichever direction it is used.
+
+    A rank whose pieces sweep as disjoint and fill the shard passes without pairwise comparisons; any
+    other rank takes the pairwise walk, which names the pieces at fault.
     """
     by_rank: Dict[int, List[_Frame]] = {}
     for frame in frames:
         by_rank.setdefault(frame.rank, []).append(frame)
 
     for rank, held in sorted(by_rank.items()):
+        if boxes_disjoint([frame.local for frame in held]) and sum(frame.local.volume
+                                                                   for frame in held) == held[0].shard_numel:
+            continue
         if len(held)**2 > limits.max_pair_checks:
             raise AffineTransferError(f'{side}: rank {rank} holds {len(held)} pieces, over the budget for '
                                       'address checks')
diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py
index fc14bfa14383..53dd46b02d5a 100644
--- a/tests/unit/checkpoint/test_affine_transfer_plan.py
+++ b/tests/unit/checkpoint/test_affine_transfer_plan.py
@@ -678,6 +678,19 @@ def test_many_consistent_pieces_need_no_pairwise_holder_checks():
         torch.testing.assert_close(shard, target.extract(full, rank))
 
 
+def test_many_pieces_on_one_rank_need_no_pairwise_address_checks():
+    rows, ranks = 2048, 2
+    source = block_gather_map((rows, 2), {rank: list(range(rank, rows, ranks)) for rank in range(ranks)}, 1, 0)
+    target = contiguous_split_map((rows, 2), [rows // ranks] * ranks, 0)
+    assert (rows // ranks)**2 > PlanBudget().max_pair_checks
+    full = torch.arange(rows * 2, dtype=torch.float64).reshape(rows, 2)
+    # Merging this many segments is a separate quadratic step with its own budget.
+    plan = plan_transfer(target, source, PlanBudget(max_merge_checks=1 << 30))
+    copied = _apply_plan(plan, target, _shards(source, lambda m, r: m.extract(full, r)))
+    for rank, shard in copied.items():
+        torch.testing.assert_close(shard, target.extract(full, rank))
+
+
 def _random_holder_map(rng, rows, cols, ranks):
     """Row blocks on every rank, sometimes split at a replica's own points, with holders sometimes wrong."""
     shared = sorted(rng.sample(range(1, rows), rng.randint(0, rows - 1)))

From 82b30c79e91c324f54ae5ab2ee167947f64b223c Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Thu, 12 Nov 2026 14:48:00 +0000
Subject: [PATCH 3/7] Count the disjointness sweep against max_pair_checks

A disjoint layout can keep every box open for the whole sweep, so
boxes_disjoint was still quadratic without a budget of its own. Give it an
optional max_checks: once exceeded it reports the boxes as not shown disjoint,
and since the sweep compares at most n*(n-1)/2 pairs, the caller's pairwise
fallback under the same budget refuses at once. The planner passes
max_pair_checks at every call.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 deepspeed/checkpoint/affine_geometry.py       | 13 +++++++++--
 deepspeed/checkpoint/affine_transfer.py       |  8 +++----
 .../checkpoint/test_affine_transfer_plan.py   | 22 +++++++++++++++++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/deepspeed/checkpoint/affine_geometry.py b/deepspeed/checkpoint/affine_geometry.py
index c58d899713cb..1ef764f8e64b 100644
--- a/deepspeed/checkpoint/affine_geometry.py
+++ b/deepspeed/checkpoint/affine_geometry.py
@@ -121,12 +121,17 @@ def __repr__(self) -> str:
         return f'Box({self.origin}+{self.extent})'
 
 
-def boxes_disjoint(boxes: Sequence[LogicalBox]) -> bool:
+def boxes_disjoint(boxes: Sequence[LogicalBox], max_checks: Optional[int] = None) -> bool:
     """Whether no two boxes share an element.
 
     Sweeps along the axis whose box origins are most varied, so a box is compared only with the boxes
     still open where it starts. For the grid and run layouts the builders emit that is a handful of
-    boxes, not all of them.
+    boxes, not all of them, but a layout that keeps many boxes open still costs a comparison per pair.
+
+    With ``max_checks``, the sweep gives up and returns False once it has compared more pairs than
+    that, so False means "not shown disjoint". The sweep compares at most n*(n-1)/2 pairs, so giving
+    up implies n**2 exceeds ``max_checks``, and a caller's pairwise fallback under the same budget
+    refuses at once.
     """
     boxes = [box for box in boxes if not box.is_empty]
     if len(boxes) < 2:
@@ -135,7 +140,11 @@ def boxes_disjoint(boxes: Sequence[LogicalBox]) -> bool:
         return False
     axis = max(range(len(boxes[0].origin)), key=lambda a: len({box.origin[a] for box in boxes}))
     open_boxes: List[LogicalBox] = []
+    checks = 0
     for box in sorted(boxes, key=lambda b: b.origin[axis]):
+        checks += len(open_boxes)
+        if max_checks is not None and checks > max_checks:
+            return False
         start = box.origin[axis]
         open_boxes = [other for other in open_boxes if other.origin[axis] + other.extent[axis] > start]
         if any(not box.intersection(other).is_empty for other in open_boxes):
diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py
index 4578cf6297f8..fe4f5542ec7c 100644
--- a/deepspeed/checkpoint/affine_transfer.py
+++ b/deepspeed/checkpoint/affine_transfer.py
@@ -347,8 +347,8 @@ def _shards_tile(frames: List[_Frame], side: str, limits: PlanBudget) -> None:
         by_rank.setdefault(frame.rank, []).append(frame)
 
     for rank, held in sorted(by_rank.items()):
-        if boxes_disjoint([frame.local for frame in held]) and sum(frame.local.volume
-                                                                   for frame in held) == held[0].shard_numel:
+        if boxes_disjoint([frame.local for frame in held], limits.max_pair_checks) and sum(
+                frame.local.volume for frame in held) == held[0].shard_numel:
             continue
         if len(held)**2 > limits.max_pair_checks:
             raise AffineTransferError(f'{side}: rank {rank} holds {len(held)} pieces, over the budget for '
@@ -379,7 +379,7 @@ def _validate_holders(frames: List[_Frame], side: str, limits: PlanBudget) -> No
         ranks = holders.setdefault(frame.box, set())
         repeated = repeated or frame.rank in ranks
         ranks.add(frame.rank)
-    if not repeated and boxes_disjoint(list(holders)):
+    if not repeated and boxes_disjoint(list(holders), limits.max_pair_checks):
         if all(holders[frame.box] == frame.locations for frame in frames):
             return
     _validate_holders_pairwise(frames, side, limits)
@@ -428,7 +428,7 @@ def _sources(frames: List[_Frame], limits: PlanBudget) -> List[Tuple[LogicalBox,
             by_box.setdefault(frame.box, []).append(frame)
 
     regions = sorted(by_box, key=lambda box: (box.origin, box.extent))
-    if not boxes_disjoint(regions):
+    if not boxes_disjoint(regions, limits.max_pair_checks):
         _name_overlapping_regions(regions, limits)
     return [(box, tuple(sorted(by_box[box], key=lambda frame: frame.rank))) for box in regions]
 
diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py
index 53dd46b02d5a..df12f8dd1a44 100644
--- a/tests/unit/checkpoint/test_affine_transfer_plan.py
+++ b/tests/unit/checkpoint/test_affine_transfer_plan.py
@@ -691,6 +691,28 @@ def test_many_pieces_on_one_rank_need_no_pairwise_address_checks():
         torch.testing.assert_close(shard, target.extract(full, rank))
 
 
+def test_sweep_comparisons_count_against_the_pair_budget():
+    """A disjoint layout can keep every box open for the whole sweep, so the sweep is budgeted too.
+
+    h full-height one-column boxes plus 2h one-element boxes beside them tile a (2h, h+1) parameter.
+    Sweeping rows, each one-element box meets all h columns: 2h**2 comparisons for 3h pieces.
+    """
+    h = 32
+    rows, cols = 2 * h, h + 1
+    pieces = [AffinePiece((rows, 1), column, (cols, 1), column * rows, (1, 1), [0]) for column in range(h)]
+    pieces += [AffinePiece((1, 1), row * cols + h, (cols, 1), h * rows + row, (1, 1), [0]) for row in range(rows)]
+    source = ParamAffineMap((rows, cols), {0: (rows * cols, )}, {0: pieces})
+    target = replicated_map((rows, cols), 1)
+    assert boxes_disjoint([LogicalBox((0, column), (rows, 1))
+                           for column in range(h)] + [LogicalBox((row, h), (1, 1)) for row in range(rows)])
+
+    full = torch.arange(rows * cols, dtype=torch.float64).reshape(rows, cols)
+    copied = _apply_plan(plan_transfer(target, source), target, _shards(source, lambda m, r: m.extract(full, r)))
+    torch.testing.assert_close(copied[0], target.extract(full, 0))
+    with pytest.raises(AffineTransferError, match='max_pair_checks'):
+        plan_transfer(target, source, PlanBudget(max_pair_checks=2 * h * h - 1))
+
+
 def _random_holder_map(rng, rows, cols, ranks):
     """Row blocks on every rank, sometimes split at a replica's own points, with holders sometimes wrong."""
     shared = sorted(rng.sample(range(1, rows), rng.randint(0, rows - 1)))

From 733dc9135bbd870495e6b64e84016bf249bad7ee Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Fri, 13 Nov 2026 14:13:27 +0000
Subject: [PATCH 4/7] Check holders against an element-level oracle instead of
 the private pairwise walk

Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 .../checkpoint/test_affine_transfer_plan.py   | 57 ++++++++++++++-----
 1 file changed, 43 insertions(+), 14 deletions(-)

diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py
index df12f8dd1a44..f9c7196fc06b 100644
--- a/tests/unit/checkpoint/test_affine_transfer_plan.py
+++ b/tests/unit/checkpoint/test_affine_transfer_plan.py
@@ -729,7 +729,11 @@ def _random_holder_map(rng, rows, cols, ranks):
     for rank, blocks in blocks_of.items():
         offset, pieces[rank] = 0, []
         for low, high in blocks:
-            holders = [other for other, theirs in blocks_of.items() if any(a <= low and high <= b for a, b in theirs)]
+            # A replica may hold this block as several smaller pieces, so cover it by union.
+            holders = [
+                other for other, theirs in blocks_of.items() if all(
+                    any(a <= row < b for a, b in theirs) for row in range(low, high))
+            ]
             if rng.random() < 0.15:
                 holders = sorted(set(holders) ^ {rng.randrange(ranks)}) or [rank]
             shape = (high - low, cols)
@@ -739,27 +743,52 @@ def _random_holder_map(rng, rows, cols, ranks):
     return ParamAffineMap((rows, cols), shapes, pieces)
 
 
-def _outcome(target, source):
-    try:
-        plan = plan_transfer(target, source, PlanBudget(max_pair_checks=1 << 30))
-    except AffineTransferError as error:
-        return type(error), str(error)
-    return [segment.to_dict() for segment in plan.segments]
+def _holders_are_exact(map_):
+    """Element-by-element reading of the holder contract, independent of the planner.
 
+    No rank may hold an element twice, and every element of a piece must be held by
+    exactly the ranks the piece names.
+    """
+    full = torch.arange(_product(map_.logical_shape))
+    elements_of = {}
+    held_by = {}
+    for rank, pieces in map_.pieces_by_rank.items():
+        for piece in pieces:
+            elements = full.as_strided(piece.shape, piece.source_strides, piece.source_offset).flatten().tolist()
+            elements_of[(rank, id(piece))] = (elements, set(piece.locations))
+            for element in elements:
+                ranks = held_by.setdefault(element, [])
+                if rank in ranks:
+                    return False
+                ranks.append(rank)
+    return all(
+        set(held_by[element]) == locations for elements, locations in elements_of.values() for element in elements)
 
-@pytest.mark.parametrize('seed', range(200))
-def test_holder_check_agrees_with_pairwise_walk(seed, monkeypatch):
-    """The table lookup is a shortcut, so the exhaustive comparison it replaces is the oracle."""
-    from deepspeed.checkpoint import affine_transfer as transfer_module
 
+@pytest.mark.parametrize('seed', range(200))
+def test_holder_check_matches_an_element_level_oracle(seed):
+    """Replicas split at different points must still plan; a wrong holder set must be refused."""
     rng = random.Random(seed)
     rows, cols, ranks = rng.randint(1, 6), rng.randint(1, 3), rng.randint(1, 4)
     mapped = _random_holder_map(rng, rows, cols, ranks)
     other = replicated_map((rows, cols), 1)
     target, source = (mapped, other) if seed % 2 else (other, mapped)
-    fast = _outcome(target, source)
-    monkeypatch.setattr(transfer_module, '_validate_holders', transfer_module._validate_holders_pairwise)
-    assert fast == _outcome(target, source)
+    budget = PlanBudget(max_pair_checks=1 << 30)
+    if not _holders_are_exact(mapped):
+        with pytest.raises(InvalidMapError):
+            plan_transfer(target, source, budget)
+        return
+    try:
+        plan = plan_transfer(target, source, budget)
+    except UnsupportedLayoutError:
+        # A valid source whose replicas are cut at different points is a documented planner
+        # limitation, not a holder error; a target with such replicas must still plan.
+        assert source is mapped
+        return
+    full = torch.arange(rows * cols, dtype=torch.float64).reshape(rows, cols)
+    copied = _apply_plan(plan, target, _shards(source, lambda m, r: m.extract(full, r)))
+    for rank in target.shard_shapes:
+        torch.testing.assert_close(copied[rank], target.extract(full, rank))
 
 
 @pytest.mark.parametrize('seed', range(100))

From 885f8915bee1d54ff17778d4a75b1da82bea31e7 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Mon, 9 Nov 2026 13:51:03 +0000
Subject: [PATCH 5/7] Check strided shard destinations without pairwise
 comparisons

ParamAffineMap._validate_destinations compared every box-shaped piece of a
rank with every other one, and validate() runs whenever a ZeRO, AutoEP or
block-gather map is built. Sweep the boxes with boxes_disjoint first; only a
rank that fails the sweep takes the pairwise walk, which names the shared
address, so error messages are unchanged. Validating 2048 interleaved
column pieces per rank drops from 6.7 s to 68 ms.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 deepspeed/checkpoint/affine.py                | 12 ++++-
 .../unit/checkpoint/test_affine_shard_map.py  | 44 +++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py
index f2e98f353295..59395b3fe4d5 100644
--- a/deepspeed/checkpoint/affine.py
+++ b/deepspeed/checkpoint/affine.py
@@ -20,6 +20,7 @@
 
 import torch
 
+from .affine_geometry import LogicalBox, boxes_disjoint
 from .affine_geometry import (dot as _dot, product as _product, row_major_range_boxes as _row_major_range_boxes,
                               row_major_strides as _row_major_strides, strided_address as _strided_address)
 
@@ -292,8 +293,8 @@ def _validate_destinations(self, rank, pieces, shard_shape):
             cursor = stop
 
         # A column split's pieces have no interval, so the check above skips them. Check them as
-        # boxes against every other box-shaped piece of the rank instead. Ranks whose pieces are
-        # all dense are fully covered above and skip this pairwise pass.
+        # boxes of the shard grid instead. Ranks whose pieces are all dense are fully covered above
+        # and skip this pass.
         if not has_strided_piece:
             return
         boxes = []
@@ -313,6 +314,13 @@ def _validate_destinations(self, rank, pieces, shard_shape):
                 # not a box of the shard's grid, so the walk in validate_coverage() checks it instead.
                 continue
             boxes.append(box)
+        # A sweep proves most layouts disjoint without comparing every pair; the pairwise walk
+        # only runs to name the address that two boxes share.
+        if boxes_disjoint([
+                LogicalBox(tuple(starts), tuple(stop - start for start, stop in zip(starts, stops)))
+                for starts, stops in boxes
+        ]):
+            return
         for index, box in enumerate(boxes):
             for other_box in boxes[index + 1:]:
                 shared = _first_shared_address(box, other_box, shard_shape)
diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py
index 42ff675630e9..af6cc9c80e22 100644
--- a/tests/unit/checkpoint/test_affine_shard_map.py
+++ b/tests/unit/checkpoint/test_affine_shard_map.py
@@ -18,6 +18,8 @@
 process group or an accelerator.
 """
 
+import random
+
 import pytest
 import torch
 
@@ -589,6 +591,48 @@ def refuse_walk(self, base, strides):
     affine_map.validate()
 
 
+def _random_column_block_map(rng):
+    """Column blocks of row bands tiling one shard, with one block sometimes shifted onto another."""
+    rows, cols = rng.randint(1, 6), rng.randint(2, 8)
+    col_cuts = sorted(rng.sample(range(1, cols), rng.randint(1, cols - 1)))
+    blocks = []
+    for low, high in zip([0] + col_cuts, col_cuts + [cols]):
+        row_cuts = sorted(rng.sample(range(1, rows), rng.randint(0, rows - 1))) if rows > 1 else []
+        blocks += [[top, low, bottom - top, high - low] for top, bottom in zip([0] + row_cuts, row_cuts + [rows])]
+    if rng.random() < 0.5:
+        block = rng.choice(blocks)
+        block[0] = rng.randint(0, rows - block[2])
+        block[1] = rng.randint(0, cols - block[3])
+    rng.shuffle(blocks)
+    pieces = [
+        AffinePiece((height, width), top * cols + left, (cols, 1), top * cols + left, (cols, 1), [0])
+        for top, left, height, width in blocks
+    ]
+    return ParamAffineMap((rows, cols), {0: (rows, cols)}, {0: pieces})
+
+
+def _validation_outcome(affine_map):
+    try:
+        affine_map.validate()
+    except ValueError as error:
+        return str(error)
+    return None
+
+
+@pytest.mark.parametrize('seed', range(200))
+def test_strided_destination_sweep_agrees_with_pairwise_walk(seed, monkeypatch):
+    """The sweep accepts exactly the layouts that write no address twice, and errors read as before."""
+    from deepspeed.checkpoint import affine as affine_module
+
+    affine_map = _random_column_block_map(random.Random(seed))
+    written = [(piece.dest_offset + row * piece.dest_strides[0] + column) for piece in affine_map.pieces_by_rank[0]
+               for row in range(piece.shape[0]) for column in range(piece.shape[1])]
+    fast = _validation_outcome(affine_map)
+    assert (fast is None) == (len(set(written)) == len(written))
+    monkeypatch.setattr(affine_module, 'boxes_disjoint', lambda boxes: False)
+    assert fast == _validation_outcome(affine_map)
+
+
 # Guards added in review. Each one covers a way a map could produce a plausible but wrong
 # parameter rather than failing, which is the failure mode that matters for a checkpoint.
 

From ca89f2e28de909fdaab4a54b1b73f89fc6baff56 Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Wed, 11 Nov 2026 13:41:24 +0000
Subject: [PATCH 6/7] Use affine_geometry for destination boxes and address
 ranges

_dest_box() now returns a LogicalBox built with coordinate(), so
_validate_destinations() hands boxes straight to boxes_disjoint() and
names a shared address from box.intersection() and dot(). The
hand-written _first_shared_address() and divmod loop are gone.

address_range() and overlaps_itself() move to affine_geometry. The
transfer executor reuses address_range() for its buffer-reach check and
dot() for chunk offsets. Error messages and behavior are unchanged.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 deepspeed/checkpoint/affine.py                | 90 +++++--------------
 deepspeed/checkpoint/affine_geometry.py       | 32 ++++++-
 .../checkpoint/affine_transfer_executor.py    | 20 ++---
 3 files changed, 61 insertions(+), 81 deletions(-)

diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py
index 59395b3fe4d5..0d424de577c1 100644
--- a/deepspeed/checkpoint/affine.py
+++ b/deepspeed/checkpoint/affine.py
@@ -20,7 +20,7 @@
 
 import torch
 
-from .affine_geometry import LogicalBox, boxes_disjoint
+from .affine_geometry import LogicalBox, address_range, boxes_disjoint, coordinate, overlaps_itself
 from .affine_geometry import (dot as _dot, product as _product, row_major_range_boxes as _row_major_range_boxes,
                               row_major_strides as _row_major_strides, strided_address as _strided_address)
 
@@ -166,49 +166,26 @@ def _dest_interval(piece):
 
 
 def _dest_box(piece, shard_shape):
-    """Return a piece's destination as per-axis ``(starts, stops)`` in the shard, or None.
+    """Return a piece's destination as a box of the shard's grid, or None.
 
     A column split writes each piece into a few columns of every shard row, so its
     destination is not one interval. It is still a box in the shard's own grid, because its
     strides are the shard's row-major strides, and a box can be checked for bounds and
-    overlap from its corners without walking any elements.
+    overlap from its corners without walking any elements. A view that wraps past the end of
+    a shard row is not a box of that grid; the walk in validate_coverage() checks it instead.
     """
     if piece.numel == 0 or len(piece.shape) != len(shard_shape) or len(piece.dest_strides) != len(shard_shape):
         return None
     shard_strides = _row_major_strides(shard_shape)
     if any(size > 1 and got != want for size, got, want in zip(piece.shape, piece.dest_strides, shard_strides)):
         return None
-
-    starts = []
-    remainder = piece.dest_offset
-    for stride in shard_strides:
-        start, remainder = divmod(remainder, stride)
-        starts.append(start)
-    stops = [start + size for start, size in zip(starts, piece.shape)]
-    return starts, stops
-
-
-def _dest_address_range(piece):
-    """Return the lowest and highest shard address a piece writes."""
-    low = piece.dest_offset
-    high = piece.dest_offset
-    for size, stride in zip(piece.shape, piece.dest_strides):
-        reach = (size - 1) * stride
-        if reach < 0:
-            low += reach
-        else:
-            high += reach
-    return low, high
-
-
-def _first_shared_address(box, other_box, shard_shape):
-    """Return the first shard address two boxes both write, or None when they are disjoint."""
-    corner = []
-    for start, stop, other_start, other_stop in zip(box[0], box[1], other_box[0], other_box[1]):
-        if max(start, other_start) >= min(stop, other_stop):
-            return None
-        corner.append(max(start, other_start))
-    return sum(position * stride for position, stride in zip(corner, _row_major_strides(shard_shape)))
+    origin = coordinate(piece.dest_offset, shard_strides, shard_shape)
+    if origin is None:
+        return None
+    box = LogicalBox(origin, piece.shape)
+    if any(high > dim for high, dim in zip(box.upper(), shard_shape)):
+        return None
+    return box
 
 
 class ParamAffineMap:
@@ -262,7 +239,7 @@ def _validate_rank(self, rank):
                              f'account for {held}.')
         for piece in pieces:
             for side, strides in (('source', piece.source_strides), ('destination', piece.dest_strides)):
-                if _overlaps_itself(piece.shape, strides):
+                if overlaps_itself(piece.shape, strides):
                     raise ValueError(f'Rank {rank} piece {piece} reaches one {side} address twice within itself, '
                                      'so one of its elements would silently overwrite another.')
         self._validate_destinations(rank, pieces, self.shard_shapes[rank])
@@ -300,33 +277,24 @@ def _validate_destinations(self, rank, pieces, shard_shape):
         boxes = []
         for piece in pieces:
             if piece.numel and _dest_interval(piece) is None:
-                low, high = _dest_address_range(piece)
+                low, high = address_range(piece.dest_offset, piece.shape, piece.dest_strides)
                 if low < 0 or high >= shard_numel:
                     raise ValueError(f'Rank {rank} piece {piece} writes shard addresses [{low}, {high}], '
                                      f'outside the shard of {shard_numel} elements.')
             box = _dest_box(piece, shard_shape)
-            if box is None:
-                continue
-            starts, stops = box
-            fits = all(start >= 0 and stop <= dim for start, stop, dim in zip(starts, stops, shard_shape))
-            if not fits:
-                # A view may wrap past the end of a shard row while staying inside the shard. It is
-                # not a box of the shard's grid, so the walk in validate_coverage() checks it instead.
-                continue
-            boxes.append(box)
+            if box is not None:
+                boxes.append(box)
         # A sweep proves most layouts disjoint without comparing every pair; the pairwise walk
         # only runs to name the address that two boxes share.
-        if boxes_disjoint([
-                LogicalBox(tuple(starts), tuple(stop - start for start, stop in zip(starts, stops)))
-                for starts, stops in boxes
-        ]):
+        if boxes_disjoint(boxes):
             return
+        shard_strides = _row_major_strides(shard_shape)
         for index, box in enumerate(boxes):
             for other_box in boxes[index + 1:]:
-                shared = _first_shared_address(box, other_box, shard_shape)
-                if shared is not None:
-                    raise ValueError(f'Rank {rank} writes shard address {shared} twice, so one piece would '
-                                     'silently overwrite another.')
+                shared = box.intersection(other_box)
+                if not shared.is_empty:
+                    raise ValueError(f'Rank {rank} writes shard address {_dot(shared.origin, shard_strides)} '
+                                     'twice, so one piece would silently overwrite another.')
 
     def validate_coverage(self):
         """Full check: each shard address is written once, the shards cover the parameter,
@@ -520,22 +488,6 @@ def _scaled(tensor, scale, scale_power):
 _UNWRITTEN, _SINGLE_HOLDER, _REPLICA = 0, 1, 2
 
 
-def _overlaps_itself(shape, strides):
-    """Whether two positions of a view with this shape and these strides share an address.
-
-    Taking the axes that move in order of stride, the view is free of overlap when each stride
-    clears every address the axes before it reach. That holds for every layout the builders here
-    produce. A view that interleaves its axes some other way is refused even where it would not
-    collide, which keeps this to a few comparisons per piece instead of a walk over its elements.
-    """
-    reach = 0
-    for stride, size in sorted((stride, size) for size, stride in zip(shape, strides) if size > 1):
-        if stride <= reach:
-            return True
-        reach += stride * (size - 1)
-    return False
-
-
 def _place(placed, piece, rank):
     """Mark the shard addresses ``piece`` occupies in ``placed``, refusing one already taken.
 
diff --git a/deepspeed/checkpoint/affine_geometry.py b/deepspeed/checkpoint/affine_geometry.py
index 1ef764f8e64b..d02562ab1bd7 100644
--- a/deepspeed/checkpoint/affine_geometry.py
+++ b/deepspeed/checkpoint/affine_geometry.py
@@ -10,7 +10,8 @@
 from typing import List, Optional, Sequence, Tuple
 
 __all__ = [
-    'LogicalBox', 'coordinate', 'dot', 'product', 'row_major_range_boxes', 'row_major_strides', 'strided_address'
+    'LogicalBox', 'address_range', 'coordinate', 'dot', 'overlaps_itself', 'product', 'row_major_range_boxes',
+    'row_major_strides', 'strided_address'
 ]
 
 
@@ -46,6 +47,35 @@ def dot(index, strides):
     return sum(position * stride for position, stride in zip(index, strides))
 
 
+def address_range(offset: int, shape: Sequence[int], strides: Sequence[int]) -> Tuple[int, int]:
+    """The lowest and highest address a strided view reaches, both inclusive."""
+    low = offset
+    high = offset
+    for size, stride in zip(shape, strides):
+        reach = (size - 1) * stride
+        if reach < 0:
+            low += reach
+        else:
+            high += reach
+    return low, high
+
+
+def overlaps_itself(shape: Sequence[int], strides: Sequence[int]) -> bool:
+    """Whether two positions of a view with this shape and these strides share an address.
+
+    Taking the axes that move in order of stride, the view is free of overlap when each stride
+    clears every address the axes before it reach. That holds for every layout the builders
+    produce. A view that interleaves its axes some other way is refused even where it would not
+    collide, which keeps this to a few comparisons per view instead of a walk over its elements.
+    """
+    reach = 0
+    for stride, size in sorted((stride, size) for size, stride in zip(shape, strides) if size > 1):
+        if stride <= reach:
+            return True
+        reach += stride * (size - 1)
+    return False
+
+
 def strided_address(shape, offset, strides, index):
     address = offset
     for size, stride in reversed(list(zip(shape, strides))):
diff --git a/deepspeed/checkpoint/affine_transfer_executor.py b/deepspeed/checkpoint/affine_transfer_executor.py
index a9d769aae731..2ec10ad02124 100644
--- a/deepspeed/checkpoint/affine_transfer_executor.py
+++ b/deepspeed/checkpoint/affine_transfer_executor.py
@@ -51,6 +51,7 @@
 import torch
 from deepspeed.accelerator import get_accelerator
 
+from .affine_geometry import address_range, dot
 from .affine_transfer import TransferPlan, TransferSegment
 
 __all__ = [
@@ -137,16 +138,13 @@ def _chunks(segment: TransferSegment, max_elements: int):
     origins = [range(0, extent, step) for extent, step in zip(segment.shape, tile)]
     for origin in product(*origins):
         shape = tuple(min(step, extent - low) for step, extent, low in zip(tile, segment.shape, origin))
-        yield TransferSegment(
-            source_rank=segment.source_rank,
-            source_offset=segment.source_offset + sum(low * stride
-                                                      for low, stride in zip(origin, segment.source_strides)),
-            source_strides=segment.source_strides,
-            target_rank=segment.target_rank,
-            target_offset=segment.target_offset + sum(low * stride
-                                                      for low, stride in zip(origin, segment.target_strides)),
-            target_strides=segment.target_strides,
-            shape=shape)
+        yield TransferSegment(source_rank=segment.source_rank,
+                              source_offset=segment.source_offset + dot(origin, segment.source_strides),
+                              source_strides=segment.source_strides,
+                              target_rank=segment.target_rank,
+                              target_offset=segment.target_offset + dot(origin, segment.target_strides),
+                              target_strides=segment.target_strides,
+                              shape=shape)
 
 
 def build_operations(plan: TransferPlan, endpoints: TransferEndpoints, budget: StagingBudget,
@@ -236,7 +234,7 @@ def _validate_buffers(plan: TransferPlan, source_buffers: Mapping[int, torch.Ten
                 continue
             if location not in buffers:
                 raise AffineTransferExecutionError(f'{which} location {location} is owned here but holds no buffer')
-            reach = offset + sum((size - 1) * stride for size, stride in zip(segment.shape, strides))
+            _, reach = address_range(offset, segment.shape, strides)
             if reach >= buffers[location].numel():
                 raise AffineTransferExecutionError(
                     f'{which} of {segment!r} reaches element {reach} of a {buffers[location].numel()}-element buffer')

From b2ddefe93b0b64bb3e6fad4b9184d5fe7242f27e Mon Sep 17 00:00:00 2001
From: "Jin, Youzhi" <youzhi.jin@intel.com>
Date: Thu, 12 Nov 2026 14:23:59 +0000
Subject: [PATCH 7/7] Assert the refused address instead of forcing the
 pairwise walk

The randomized destination test monkeypatched affine.boxes_disjoint to
compare the sweep with the pairwise walk, which ties it to that helper.
Check the contract instead: a refusal names an address the map writes twice.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
---
 .../unit/checkpoint/test_affine_shard_map.py  | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py
index af6cc9c80e22..d8d1b3803465 100644
--- a/tests/unit/checkpoint/test_affine_shard_map.py
+++ b/tests/unit/checkpoint/test_affine_shard_map.py
@@ -19,6 +19,7 @@
 """
 
 import random
+import re
 
 import pytest
 import torch
@@ -620,17 +621,20 @@ def _validation_outcome(affine_map):
 
 
 @pytest.mark.parametrize('seed', range(200))
-def test_strided_destination_sweep_agrees_with_pairwise_walk(seed, monkeypatch):
-    """The sweep accepts exactly the layouts that write no address twice, and errors read as before."""
-    from deepspeed.checkpoint import affine as affine_module
-
+def test_strided_destination_check_refuses_exactly_double_writes(seed):
+    """validate() accepts exactly the layouts that write no address twice, and a refusal names
+    an address that is written twice."""
     affine_map = _random_column_block_map(random.Random(seed))
     written = [(piece.dest_offset + row * piece.dest_strides[0] + column) for piece in affine_map.pieces_by_rank[0]
                for row in range(piece.shape[0]) for column in range(piece.shape[1])]
-    fast = _validation_outcome(affine_map)
-    assert (fast is None) == (len(set(written)) == len(written))
-    monkeypatch.setattr(affine_module, 'boxes_disjoint', lambda boxes: False)
-    assert fast == _validation_outcome(affine_map)
+    outcome = _validation_outcome(affine_map)
+    assert (outcome is None) == (len(set(written)) == len(written))
+    if outcome is not None:
+        match = re.fullmatch(
+            r'Rank 0 writes shard address (\d+) twice, so one piece would silently overwrite '
+            r'another\.', outcome)
+        assert match is not None, outcome
+        assert written.count(int(match.group(1))) > 1
 
 
 # Guards added in review. Each one covers a way a map could produce a plausible but wrong