diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index f2e98f353295..0d424de577c1 100644 --- a/deepspeed/checkpoint/affine.py +++ b/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. diff --git a/deepspeed/checkpoint/affine_geometry.py b/deepspeed/checkpoint/affine_geometry.py index 4e8a9b092952..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))): @@ -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 diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py index 4a127757d191..fe4f5542ec7c 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', @@ -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: 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') diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 42ff675630e9..d8d1b3803465 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -18,6 +18,9 @@ process group or an accelerator. """ +import random +import re + import pytest import torch @@ -589,6 +592,51 @@ 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_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])] + 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 # parameter rather than failing, which is the failure mode that matters for a checkpoint. diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py index b00fe855f1b4..f9c7196fc06b 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,142 @@ 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 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 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))) + 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: + # 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) + 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 _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_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) + 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)) +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