[AffineIR] Check strided shard destinations without pairwise comparisons

deepspeed/checkpoint/affine.py

@@ -20,6 +20,7 @@
 
 import torch
 
+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)
 
@@ -165,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:
@@ -261,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])
@@ -292,33 +270,31 @@ 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 = []
         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(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,
@@ -512,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.
 

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))):
@@ -119,3 +149,35 @@ def contains(self, other: 'LogicalBox') -> bool:
 
     def __repr__(self) -> str:
         return f'Box({self.origin}+{self.extent})'
+
+
+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, 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:
+        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] = []
+    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):
+            return False
+        open_boxes.append(box)
+    return True

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',
@@ -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], 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 '
                                       'address checks')
@@ -360,7 +366,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), limits.max_pair_checks):
+        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 +428,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, 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]
+
+
+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:

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')