diff --git a/deepspeed/checkpoint/affine_replica_policy.py b/deepspeed/checkpoint/affine_replica_policy.py new file mode 100644 index 000000000000..f381400aa9e2 --- /dev/null +++ b/deepspeed/checkpoint/affine_replica_policy.py @@ -0,0 +1,88 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Replica policies for ``affine_transfer.plan_transfers``. + +A policy takes every ``ReplicaRequest`` of one sync and returns one answer per request: the +location of a validated holder, or ``ReplicaPart``\\ s that cut the request along axis 0. The planner +checks every answer before it emits a segment, so a policy decides only where data is read from; +what arrives is fixed by the maps. Policies live here, not in the planner, so cost-aware ones can +grow without the planner growing with them. + +A policy that keeps state across syncs can be any callable object; the planner holds none. +""" + +from typing import Callable, Dict, Hashable, List, Sequence + +from .affine_transfer import ReplicaAnswer, ReplicaPolicy, ReplicaRequest, ReplicaSelectionRequest + +__all__ = ['first_holder', 'local_first_least_loaded', 'per_request'] + +ReplicaSelector = Callable[[ReplicaSelectionRequest], int] + + +def first_holder(requests: Sequence[ReplicaRequest]) -> List[ReplicaAnswer]: + """The lowest holder location for every request, which is what ``plan_transfer`` does by default.""" + return [request.candidates[0].source_rank for request in requests] + + +def local_first_least_loaded(requests: Sequence[ReplicaRequest]) -> List[ReplicaAnswer]: + """Read locally when the target's own process holds the region, otherwise from the least loaded holder. + + Load is what a sending process has been assigned so far in this batch: bytes when the requests + carry them, elements otherwise. Requests with a single holder count, since that load cannot + move, and local copies do not, since they use no link. Without process mappings a local copy + cannot be recognised, so each holder location counts as its own sender. Ties go to the lowest + location and requests arrive in the planner's fixed order, so the answer is deterministic. Nothing + is split. + """ + load: Dict[Hashable, int] = {} + answers: List[ReplicaAnswer] = [] + for request in requests: + local = [ + candidate.source_rank for candidate, process in zip(request.candidates, request.candidate_processes) + if process is not None and process == request.target_process + ] + if local: + answers.append(min(local)) + continue + + def sender(index: int) -> Hashable: + process = request.candidate_processes[index] + if process is not None: + return ('process', process) + return ('location', request.candidates[index].source_rank) + + chosen = min(range(len(request.candidates)), + key=lambda index: (load.get(sender(index), 0), request.candidates[index].source_rank)) + size = request.bytes if request.bytes is not None else request.numel + load[sender(chosen)] = load.get(sender(chosen), 0) + size + answers.append(request.candidates[chosen].source_rank) + return answers + + +def per_request(selector: ReplicaSelector) -> ReplicaPolicy: + """Run a ``plan_transfer``-style ``replica_selector`` as a batch policy. + + As in ``plan_transfer``, the selector is consulted only where a region has several holders, and + it sees each request on its own. + """ + + def policy(requests: Sequence[ReplicaRequest]) -> List[ReplicaAnswer]: + answers: List[ReplicaAnswer] = [] + for request in requests: + if len(request.candidates) == 1: + answers.append(request.candidates[0].source_rank) + continue + answers.append( + selector( + ReplicaSelectionRequest(target_rank=request.target_location, + logical_origin=request.logical_origin, + shape=request.shape, + target_offset=request.target_offset, + target_strides=request.target_strides, + candidates=request.candidates))) + return answers + + return policy diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py index 87452bd6ec20..727f65f0f3f9 100644 --- a/deepspeed/checkpoint/affine_transfer.py +++ b/deepspeed/checkpoint/affine_transfer.py @@ -25,15 +25,17 @@ and ``_covers_exactly`` states what that argument does and does not buy. """ +import collections.abc import json from dataclasses import dataclass -from typing import Callable, Dict, FrozenSet, List, Optional, Sequence, Tuple +from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Sequence, Tuple, Union from .affine import AffinePiece, ParamAffineMap __all__ = [ 'AffineTransferError', 'InvalidMapError', 'InvalidReplicaPolicyError', 'UnsupportedLayoutError', 'PlanBudget', - 'PlanStatistics', 'ReplicaCandidate', 'ReplicaSelectionRequest', 'TransferSegment', 'TransferPlan', 'plan_transfer' + 'PlanStatistics', 'ReplicaCandidate', 'ReplicaSelectionRequest', 'ReplicaRequest', 'ReplicaPart', + 'TransferSegment', 'TransferPlan', 'plan_transfer', 'plan_transfers' ] @@ -154,6 +156,43 @@ class ReplicaSelectionRequest: candidates: Tuple[ReplicaCandidate, ...] +@dataclass(frozen=True) +class ReplicaRequest: + """One region a target piece needs, with every validated holder that can supply it. + + ``target_location`` and each candidate's ``source_rank`` are keys of the target and source maps, + not processes: the same number names different processes on the two sides. ``target_process`` + and ``candidate_processes`` are processes, and are ``None`` unless the caller passed the + location-to-process mappings. ``bytes`` is ``None`` unless the caller passed element sizes. + """ + parameter: str + target_location: int + target_process: Optional[int] + logical_origin: Tuple[int, ...] + shape: Tuple[int, ...] + target_offset: int + target_strides: Tuple[int, ...] + candidates: Tuple[ReplicaCandidate, ...] + candidate_processes: Tuple[Optional[int], ...] + bytes: Optional[int] + + @property + def numel(self) -> int: + return _product(self.shape) + + +@dataclass(frozen=True) +class ReplicaPart: + """A cut of a request's region along axis 0, and the holder location it is read from.""" + logical_origin: Tuple[int, ...] + shape: Tuple[int, ...] + source_location: int + + +ReplicaAnswer = Union[int, Sequence[ReplicaPart]] +ReplicaPolicy = Callable[[Sequence[ReplicaRequest]], Sequence[ReplicaAnswer]] + + @dataclass(frozen=True) class TransferSegment: """One directly executable strided copy, from a source shard into a target shard. @@ -475,6 +514,24 @@ def _covers_union(region: LogicalBox, boxes: Sequence[LogicalBox], limits: PlanB return not uncovered +@dataclass(frozen=True) +class _Intersection: + """One target piece's share of one source region, before a holder is chosen for it.""" + target: _Frame + shared: LogicalBox + holders: Tuple[_Frame, ...] + + +@dataclass(frozen=True) +class _Prepared: + """A validated pair of maps, cut into intersections that only lack a choice of holder.""" + logical_shape: Tuple[int, ...] + intersections: Tuple[_Intersection, ...] + source_pieces: int + target_pieces: int + pair_checks: int + + def plan_transfer(target_map: ParamAffineMap, source_map: ParamAffineMap, limits: Optional[PlanBudget] = None, @@ -487,6 +544,209 @@ def plan_transfer(target_map: ParamAffineMap, have matched by byte count before -- so this checks only that the two geometries compose. """ budget = limits or PlanBudget() + prepared = _prepare(target_map, source_map, budget) + + choices: List[Tuple[Tuple[LogicalBox, _Frame], ...]] = [] + for item in prepared.intersections: + target, shared, holders = item.target, item.shared, item.holders + source = holders[0] + if replica_selector is not None and len(holders) > 1: + request = ReplicaSelectionRequest(target_rank=target.rank, + logical_origin=shared.origin, + shape=shared.extent, + target_offset=target.address(shared), + target_strides=target.strides, + candidates=tuple( + ReplicaCandidate(holder.rank, holder.address(shared), holder.strides) + for holder in holders)) + rank = replica_selector(request) + if type(rank) is not int or rank not in {holder.rank for holder in holders}: + raise InvalidReplicaPolicyError(f'selector returned invalid source rank {rank!r} ' + f'for target rank {target.rank}') + source = next(holder for holder in holders if holder.rank == rank) + choices.append(((shared, source), )) + return _emit(prepared, choices, budget) + + +def plan_transfers(pairs: Mapping[str, Tuple[ParamAffineMap, ParamAffineMap]], + limits: Optional[PlanBudget] = None, + *, + source_processes: Optional[Mapping[int, int]] = None, + target_processes: Optional[Mapping[int, int]] = None, + element_size: Optional[Mapping[str, int]] = None, + min_split_bytes: Optional[int] = None, + replica_policy: Optional[ReplicaPolicy] = None) -> Dict[str, TransferPlan]: + """Plan every parameter of one sync, with one replica policy that sees every request at once. + + ``pairs`` maps a parameter name to its ``(target_map, source_map)``. Each pair is validated + exactly as ``plan_transfer`` validates it, and ``limits`` applies to each parameter on its own. + + ``replica_policy`` is then called once, with a ``ReplicaRequest`` for every intersection of a + target piece with a source region, in a fixed order: parameter, target location, origin. Regions + with a single holder are included, because their load is load the policy cannot move. The policy + returns one answer per request: the location of one validated holder for the whole region, or + ``ReplicaPart``\\ s that cut the region along axis 0 and name a validated holder for each cut. + + Every answer is checked before any segment is emitted, and a failure raises + ``InvalidReplicaPolicyError``. There must be one answer per request, every location must be a + validated holder of that request, and the parts must tile the request exactly along axis 0. When + ``min_split_bytes`` is given, each part of a split must also be at least that large. So a policy + decides where data is read from but not what arrives: when the replicas agree, every accepted + answer copies the same data. + + Plan locations are keys of the two maps, not processes. ``source_processes`` and + ``target_processes`` say which process holds each location. When given, the requests carry them, + so a policy can tell a local copy from a transfer. ``element_size`` gives bytes per element for + each parameter, which the maps do not know, and ``min_split_bytes`` requires it. Without a + policy, ``affine_replica_policy.local_first_least_loaded`` is used. + """ + budget = limits or PlanBudget() + names = sorted(pairs) + for which, mapping in (('source_processes', source_processes), ('target_processes', target_processes)): + if mapping is not None and any( + type(location) is not int or type(process) is not int for location, process in mapping.items()): + raise AffineTransferError(f'{which} must map int locations to int processes') + if element_size is not None: + for name in names: + size = element_size.get(name) + if type(size) is not int or size <= 0: + raise AffineTransferError(f'element_size for {name!r} is {size!r}, not a positive int') + if min_split_bytes is not None: + if type(min_split_bytes) is not int or min_split_bytes <= 0: + raise AffineTransferError(f'min_split_bytes is {min_split_bytes!r}, not a positive int') + if element_size is None: + raise AffineTransferError('min_split_bytes needs element_size: the maps count elements, not bytes') + + prepared = {} + for name in names: + target_map, source_map = pairs[name] + prepared[name] = _prepare(target_map, source_map, budget) + + entries: List[Tuple[Tuple[object, ...], str, int, _Intersection]] = [] + for name in names: + for index, item in enumerate(prepared[name].intersections): + entries.append(((name, item.target.rank, item.shared.origin, item.shared.extent), name, index, item)) + entries.sort(key=lambda entry: entry[0]) + + requests = [] + for _, name, _, item in entries: + requests.append( + ReplicaRequest(parameter=name, + target_location=item.target.rank, + target_process=_process(target_processes, item.target.rank, 'target'), + logical_origin=item.shared.origin, + shape=item.shared.extent, + target_offset=item.target.address(item.shared), + target_strides=item.target.strides, + candidates=tuple( + ReplicaCandidate(holder.rank, holder.address(item.shared), holder.strides) + for holder in item.holders), + candidate_processes=tuple( + _process(source_processes, holder.rank, 'source') for holder in item.holders), + bytes=None if element_size is None else item.shared.volume * element_size[name])) + requests = tuple(requests) + + if replica_policy is None: + from .affine_replica_policy import local_first_least_loaded + replica_policy = local_first_least_loaded + answers = replica_policy(requests) + if not _is_sequence(answers) or len(answers) != len(requests): + count = len(answers) if _is_sequence(answers) else type(answers).__name__ + raise InvalidReplicaPolicyError(f'the policy must return one answer per request: {len(requests)} ' + f'requests, got {count}') + + choices: Dict[str, List[Tuple[Tuple[LogicalBox, _Frame], ...]]] = { + name: [()] * len(prepared[name].intersections) + for name in names + } + parts_per_name = {name: 0 for name in names} + for (_, name, index, item), request, answer in zip(entries, requests, answers): + parts = _resolve(request, item, answer, min_split_bytes, element_size) + choices[name][index] = parts + parts_per_name[name] += len(parts) + if parts_per_name[name] > budget.max_pair_checks: + raise AffineTransferError(f'{name!r}: the policy split regions into more than max_pair_checks ' + f'{budget.max_pair_checks} parts') + return {name: _emit(prepared[name], choices[name], budget) for name in names} + + +def _is_sequence(value: object) -> bool: + """A real sequence: indexable and sized, which a generator is not and a string should not count as.""" + return isinstance(value, collections.abc.Sequence) and not isinstance(value, (str, bytes)) + + +def _process(mapping: Optional[Mapping[int, int]], location: int, side: str) -> Optional[int]: + if mapping is None: + return None + if location not in mapping: + raise AffineTransferError(f'{side} location {location} has no process in {side}_processes') + return mapping[location] + + +def _resolve(request: ReplicaRequest, item: _Intersection, answer: object, min_split_bytes: Optional[int], + element_size: Optional[Mapping[str, int]]) -> Tuple[Tuple[LogicalBox, _Frame], ...]: + """Turn one policy answer into (region, holder) pairs, or refuse it. + + The parts are checked against the request, never against each other's claims: each is inside the + request's box by construction (it spans every axis but the first in full), and walking them in + axis-0 order from the request's first row to its last proves they are disjoint and leave no gap. + """ + where = (f'answer for {request.parameter!r} at {request.logical_origin} on target location ' + f'{request.target_location}') + holders = {holder.rank: holder for holder in item.holders} + box = item.shared + if type(answer) is int: + if answer not in holders: + raise InvalidReplicaPolicyError(f'{where}: location {answer} is not a validated holder ' + f'{sorted(holders)}') + return ((box, holders[answer]), ) + if not _is_sequence(answer) or not answer: + raise InvalidReplicaPolicyError(f'{where}: expected a holder location or a non-empty sequence of ' + f'ReplicaPart, got {answer!r}') + + ndim = len(box.extent) + parts = [] + for part in answer: + if not isinstance(part, ReplicaPart): + raise InvalidReplicaPolicyError(f'{where}: {part!r} is not a ReplicaPart') + if type(part.source_location) is not int or part.source_location not in holders: + raise InvalidReplicaPolicyError(f'{where}: location {part.source_location!r} is not a validated ' + f'holder {sorted(holders)}') + origin, shape = tuple(part.logical_origin), tuple(part.shape) + if len(origin) != ndim or len(shape) != ndim or any(type(value) is not int for value in origin + shape): + raise InvalidReplicaPolicyError(f'{where}: part {origin}+{shape} is not a {ndim}-axis box of ints') + if origin[1:] != box.origin[1:] or shape[1:] != box.extent[1:]: + raise InvalidReplicaPolicyError(f'{where}: part {origin}+{shape} does not span every axis but the ' + f'first of {box!r}; only cuts along axis 0 are supported') + if any(size <= 0 for size in shape): + raise InvalidReplicaPolicyError(f'{where}: part {origin}+{shape} is empty') + parts.append((LogicalBox(origin, shape), holders[part.source_location])) + + if ndim == 0: + if len(parts) != 1: + raise InvalidReplicaPolicyError(f'{where}: a scalar cannot be split') + return tuple(parts) + parts.sort(key=lambda entry: entry[0].origin[0]) + row = box.origin[0] + for part_box, _ in parts: + if part_box.origin[0] != row: + kind = 'overlap' if part_box.origin[0] < row else 'gap' + raise InvalidReplicaPolicyError(f'{where}: parts leave a {kind} at row {min(row, part_box.origin[0])} ' + f'of {box!r}') + row += part_box.extent[0] + if row != box.origin[0] + box.extent[0]: + raise InvalidReplicaPolicyError(f'{where}: parts stop at row {row}, inside or past {box!r}') + if len(parts) > 1 and min_split_bytes is not None: + for part_box, _ in parts: + size = part_box.volume * element_size[request.parameter] + if size < min_split_bytes: + raise InvalidReplicaPolicyError(f'{where}: part {part_box!r} is {size} bytes, under ' + f'min_split_bytes {min_split_bytes}') + return tuple(parts) + + +def _prepare(target_map: ParamAffineMap, source_map: ParamAffineMap, budget: PlanBudget) -> _Prepared: + """Validate both maps and cut the target into intersections with the source regions.""" if tuple(source_map.logical_shape) != tuple(target_map.logical_shape): raise InvalidMapError(f'source describes {tuple(source_map.logical_shape)} and target describes ' f'{tuple(target_map.logical_shape)}') @@ -506,7 +766,7 @@ def plan_transfer(target_map: ParamAffineMap, raise AffineTransferError(f'{pair_checks} candidate pairs exceeds max_pair_checks ' f'{budget.max_pair_checks}') - pending: List[Tuple[_Frame, LogicalBox, Tuple[_Frame, ...]]] = [] + intersections: List[_Intersection] = [] for frame in sorted(targets, key=lambda frame: (frame.rank, frame.box.origin, frame.box.extent)): filled: List[LogicalBox] = [] for box, holders in sources: @@ -514,51 +774,47 @@ def plan_transfer(target_map: ParamAffineMap, if shared.is_empty: continue filled.append(shared) - pending.append((frame, shared, holders)) + intersections.append(_Intersection(frame, shared, holders)) if not _covers_exactly(frame.box, filled): raise InvalidMapError(f'target rank {frame.rank} region {frame.box!r} would receive ' f'{sum(part.volume for part in filled)} of its {frame.box.volume} elements') + return _Prepared(logical_shape=logical_shape, + intersections=tuple(intersections), + source_pieces=len(sources), + target_pieces=len(targets), + pair_checks=pair_checks) + +def _emit(prepared: _Prepared, choices: Sequence[Tuple[Tuple[LogicalBox, _Frame], ...]], + budget: PlanBudget) -> TransferPlan: + """Turn chosen (region, holder) pairs into merged segments and check the plan's budget.""" candidates: List[TransferSegment] = [] - for target, shared, holders in pending: - source = holders[0] - if replica_selector is not None and len(holders) > 1: - request = ReplicaSelectionRequest(target_rank=target.rank, - logical_origin=shared.origin, - shape=shared.extent, - target_offset=target.address(shared), - target_strides=target.strides, - candidates=tuple( - ReplicaCandidate(holder.rank, holder.address(shared), holder.strides) - for holder in holders)) - rank = replica_selector(request) - if type(rank) is not int or rank not in {holder.rank for holder in holders}: - raise InvalidReplicaPolicyError(f'selector returned invalid source rank {rank!r} ' - f'for target rank {target.rank}') - source = next(holder for holder in holders if holder.rank == rank) - candidates.append( - TransferSegment(source_rank=source.rank, - source_offset=source.address(shared), - source_strides=source.strides, - target_rank=target.rank, - target_offset=target.address(shared), - target_strides=target.strides, - shape=shared.extent)) + for item, parts in zip(prepared.intersections, choices): + target = item.target + for region, source in parts: + candidates.append( + TransferSegment(source_rank=source.rank, + source_offset=source.address(region), + source_strides=source.strides, + target_rank=target.rank, + target_offset=target.address(region), + target_strides=target.strides, + shape=region.extent)) segments = _coalesce(candidates, budget) if len(segments) > budget.max_segments: raise AffineTransferError(f'{len(segments)} segments exceeds max_segments {budget.max_segments}') - plan_bytes = _plan_bytes(segments, logical_shape) + plan_bytes = _plan_bytes(segments, prepared.logical_shape) if plan_bytes > budget.max_plan_bytes: raise AffineTransferError(f'plan metadata is {plan_bytes} bytes, over max_plan_bytes ' f'{budget.max_plan_bytes}') return TransferPlan(segments=segments, - logical_shape=tuple(target_map.logical_shape), + logical_shape=prepared.logical_shape, moved_elements=sum(segment.numel for segment in segments), - statistics=PlanStatistics(source_pieces=len(sources), - target_pieces=len(targets), - pair_checks=pair_checks, + statistics=PlanStatistics(source_pieces=prepared.source_pieces, + target_pieces=prepared.target_pieces, + pair_checks=prepared.pair_checks, segments_before_merge=len(candidates), segments=len(segments), plan_bytes=plan_bytes)) diff --git a/tests/unit/checkpoint/test_affine_replica_policy.py b/tests/unit/checkpoint/test_affine_replica_policy.py new file mode 100644 index 000000000000..85bcc3e8c9fa --- /dev/null +++ b/tests/unit/checkpoint/test_affine_replica_policy.py @@ -0,0 +1,352 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Check ``plan_transfers``: a replica policy may choose where data is read from, never what arrives. + +The planner validates every policy answer before it emits a segment: one answer per request, only +validated holders, and parts that tile each request exactly along axis 0. If that check is right, +then whenever the replicas agree, every accepted answer moves the same data -- which is what the +correctness argument for the affine IR (#8252, Theorem 3 of the proof draft) needs from a policy. +The randomized test below is that statement as a check: many random policies, including random +splits, all producing the same target shards. The rejection tests make sure an answer that would +break it is refused rather than planned. +""" + +import dataclasses +import random + +import pytest +import torch + +from deepspeed.checkpoint.affine import contiguous_split_map, replicated_map, segmented_map +from deepspeed.checkpoint.affine_replica_policy import first_holder, local_first_least_loaded, per_request +from deepspeed.checkpoint.affine_transfer import (AffineTransferError, InvalidReplicaPolicyError, PlanBudget, + ReplicaCandidate, ReplicaPart, ReplicaRequest, plan_transfer, + plan_transfers) + +SENTINEL = -12345.0 + + +def _product(shape): + count = 1 + for size in shape: + count *= size + return count + + +def _apply_plan(plan, target_map, source_shards): + """Execute a plan over plain buffers, counting writes. Test scaffolding, never a production path.""" + targets = {} + for rank, shape in target_map.shard_shapes.items(): + values = torch.full((_product(shape), ), SENTINEL, dtype=torch.float64) + writes = torch.zeros(_product(shape), dtype=torch.int64) + targets[rank] = (values, writes) + for segment in plan.segments: + source = torch.as_strided(source_shards[segment.source_rank].reshape(-1), segment.shape, + segment.source_strides, segment.source_offset) + values, writes = targets[segment.target_rank] + torch.as_strided(values, segment.shape, segment.target_strides, segment.target_offset).copy_(source) + torch.as_strided(writes, segment.shape, segment.target_strides, segment.target_offset).add_(1) + return targets + + +def _assert_moves_exactly(plan, source_map, target_map, full): + sources = {rank: source_map.extract(full, rank).contiguous() for rank in source_map.shard_shapes} + for rank, (values, writes) in _apply_plan(plan, target_map, sources).items(): + assert torch.equal(values.view(target_map.shard_shapes[rank]), target_map.extract(full, rank)), rank + assert bool(torch.all(writes == 1)), rank + + +def _pairs(): + """Parameters with replicated sources: whole-parameter replicas, a partly replicated fused block, + a vector and a scalar.""" + bigcode_source = segmented_map((48, 8), [(32, False), (16, True)], 0, 2, split_widths=[[16, 16]]) + bigcode_target = segmented_map((48, 8), [(32, False), (16, True)], 0, 4, split_widths=[[8, 8, 8, 8]]) + return { + 'dense': (contiguous_split_map((12, 6), [2, 4], 1), replicated_map((12, 6), 3)), + 'rows': (contiguous_split_map((10, 4), [3, 7], 0), replicated_map((10, 4), 2)), + 'bigcode': (bigcode_target, bigcode_source), + 'vector': (contiguous_split_map((9, ), [4, 5], 0), replicated_map((9, ), 3)), + 'scalar': (replicated_map((), 2), replicated_map((), 3)), + } + + +def _random_policy(rng): + + def policy(requests): + answers = [] + for request in requests: + locations = [candidate.source_rank for candidate in request.candidates] + rows = request.shape[0] if request.shape else 1 + if not request.shape or rows == 1 or rng.random() < 0.4: + answers.append(rng.choice(locations)) + continue + cuts = sorted(rng.sample(range(1, rows), rng.randint(1, min(3, rows - 1)))) + starts = [0] + cuts + ends = cuts + [rows] + parts = [ + ReplicaPart((request.logical_origin[0] + start, ) + request.logical_origin[1:], + (end - start, ) + request.shape[1:], rng.choice(locations)) + for start, end in zip(starts, ends) + ] + rng.shuffle(parts) + answers.append(parts) + return answers + + return policy + + +@pytest.mark.parametrize('seed', range(24)) +def test_every_accepted_answer_moves_the_same_data(seed): + pairs = _pairs() + plans = plan_transfers(pairs, replica_policy=_random_policy(random.Random(seed))) + assert sorted(plans) == sorted(pairs) + for name, (target_map, source_map) in pairs.items(): + torch.manual_seed(seed) + full = torch.randn(tuple(target_map.logical_shape), dtype=torch.float64) + _assert_moves_exactly(plans[name], source_map, target_map, full) + + +def test_first_holder_and_per_request_reproduce_plan_transfer(): + pairs = _pairs() + + def select(request): + return max(candidate.source_rank for candidate in request.candidates) + + for policy, selector in ((first_holder, None), (per_request(select), select)): + plans = plan_transfers(pairs, replica_policy=policy) + for name, (target_map, source_map) in pairs.items(): + assert plans[name] == plan_transfer(target_map, source_map, replica_selector=selector), name + + +def _recorded(policy): + seen = [] + + def wrapper(requests): + seen.extend(requests) + return policy(requests) + + return wrapper, seen + + +def test_requests_arrive_in_a_fixed_order_with_processes_and_bytes(): + pairs = { + 'b': (contiguous_split_map((4, 2), [2, 2], 0), replicated_map((4, 2), 2)), + 'a': (contiguous_split_map((6, ), [6], 0), contiguous_split_map((6, ), [6], 0)), + } + policy, seen = _recorded(first_holder) + plan_transfers(pairs, + source_processes={ + 0: 7, + 1: 8 + }, + target_processes={ + 0: 8, + 1: 9 + }, + element_size={ + 'a': 4, + 'b': 2 + }, + replica_policy=policy) + + assert [(request.parameter, request.target_location, request.logical_origin) for request in seen] == [ + ('a', 0, (0, )), + ('b', 0, (0, 0)), + ('b', 1, (2, 0)), + ] + single = seen[0] + assert [candidate.source_rank for candidate in single.candidates] == [0] + assert single.candidate_processes == (7, ) + assert single.target_process == 8 + assert single.bytes == 6 * 4 and single.numel == 6 + replicated = seen[2] + assert [candidate.source_rank for candidate in replicated.candidates] == [0, 1] + assert replicated.candidate_processes == (7, 8) + assert replicated.target_process == 9 + assert replicated.bytes == 4 * 2 + + +def test_without_mappings_requests_carry_no_processes_or_bytes(): + policy, seen = _recorded(first_holder) + plan_transfers({'w': (contiguous_split_map((4, ), [4], 0), replicated_map((4, ), 2))}, replica_policy=policy) + assert seen[0].target_process is None + assert seen[0].candidate_processes == (None, None) + assert seen[0].bytes is None + + +def _sources(plans): + return {name: sorted({segment.source_rank for segment in plan.segments}) for name, plan in plans.items()} + + +def test_default_policy_copies_locally_when_the_target_process_holds_the_region(): + pairs = {'w': (contiguous_split_map((8, 4), [4, 4], 0), replicated_map((8, 4), 3))} + plans = plan_transfers(pairs, source_processes={0: 10, 1: 11, 2: 12}, target_processes={0: 12, 1: 11}) + by_target = {segment.target_rank: segment.source_rank for segment in plans['w'].segments} + assert by_target == {0: 2, 1: 1} + + +def test_default_policy_balances_and_counts_load_it_cannot_move(): + # 'pinned' has one holder (location 0); its 64 bytes make location 0 the busier sender, so the + # replicated 'w' goes to location 1 even though 0 is the lower location. + pairs = { + 'pinned': (contiguous_split_map((16, ), [16], 0), contiguous_split_map((16, ), [16], 0)), + 'w': (contiguous_split_map((4, ), [4], 0), replicated_map((4, ), 2)), + } + plans = plan_transfers(pairs, + source_processes={ + 0: 0, + 1: 1 + }, + target_processes={0: 5}, + element_size={ + 'pinned': 4, + 'w': 4 + }) + assert _sources(plans) == {'pinned': [0], 'w': [1]} + + +def test_default_policy_spreads_equal_regions_over_holders(): + pairs = {'w': (contiguous_split_map((8, 2), [2, 2, 2, 2], 0), replicated_map((8, 2), 2))} + plans = plan_transfers(pairs) + by_target = {segment.target_rank: segment.source_rank for segment in plans['w'].segments} + assert by_target == {0: 0, 1: 1, 2: 0, 3: 1} + + +def test_default_policy_without_processes_never_assumes_a_local_copy(): + # The same number names different processes on the two sides. Without mappings, target + # location 1 must not prefer source location 1 as if it were local: the tie goes to location 0. + request = ReplicaRequest(parameter='w', + target_location=1, + target_process=None, + logical_origin=(0, ), + shape=(4, ), + target_offset=0, + target_strides=(1, ), + candidates=(ReplicaCandidate(0, 0, (1, )), ReplicaCandidate(1, 0, (1, ))), + candidate_processes=(None, None), + bytes=None) + assert local_first_least_loaded([request]) == [0] + with_processes = dataclasses.replace(request, target_process=11, candidate_processes=(10, 11)) + assert local_first_least_loaded([with_processes]) == [1] + + +def test_a_split_reads_each_cut_from_its_holder(): + target_map, source_map = contiguous_split_map((6, 3), [6], 0), replicated_map((6, 3), 2) + + def split(requests): + (request, ) = requests + assert request.shape == (6, 3) + return [[ReplicaPart((4, 0), (2, 3), 0), ReplicaPart((0, 0), (4, 3), 1)]] + + plan = plan_transfers({'w': (target_map, source_map)}, replica_policy=split)['w'] + assert sorted((segment.source_rank, segment.target_offset, segment.shape) for segment in plan.segments) == [ + (0, 12, (2, 3)), + (1, 0, (4, 3)), + ] + _assert_moves_exactly(plan, source_map, target_map, torch.randn(6, 3, dtype=torch.float64)) + + +def _one_request_pair(): + return {'w': (contiguous_split_map((6, 4), [6], 0), replicated_map((6, 4), 2))} + + +def _answering(answer): + return lambda requests: [answer for _ in requests] + + +@pytest.mark.parametrize( + 'answer', + [ + 2, # not a holder + True, # a bool is not a location + 1.0, + None, + '0', # a string is not a sequence of parts + [], + [3], # not a ReplicaPart + [ReplicaPart((0, 0), (6, 4), 5)], # not a holder + [ReplicaPart((0, 0), (6, 4), True)], + [ReplicaPart((0, 0), (2, 4), 0), ReplicaPart((3, 0), (3, 4), 1)], # gap at row 2 + [ReplicaPart((0, 0), (4, 4), 0), ReplicaPart((3, 0), (3, 4), 1)], # overlap at row 3 + [ReplicaPart((0, 0), (4, 4), 0)], # stops early + [ReplicaPart((0, 0), (7, 4), 0)], # runs past the region + [ReplicaPart((0, 0), (6, 2), 0), ReplicaPart((0, 2), (6, 2), 1)], # a cut along axis 1 + [ReplicaPart((0, 0), (0, 4), 0), ReplicaPart((0, 0), (6, 4), 1)], # an empty part + [ReplicaPart((0, ), (6, ), 0)], # wrong number of axes + [ReplicaPart((0, 0), (6.0, 4), 0)], + ]) +def test_answers_that_could_change_what_arrives_are_refused(answer): + with pytest.raises(InvalidReplicaPolicyError): + plan_transfers(_one_request_pair(), replica_policy=_answering(answer)) + + +def test_one_answer_per_request_is_required(): + pairs = {'w': (contiguous_split_map((4, ), [2, 2], 0), replicated_map((4, ), 2))} + with pytest.raises(InvalidReplicaPolicyError): + plan_transfers(pairs, replica_policy=lambda requests: [0]) + with pytest.raises(InvalidReplicaPolicyError): + plan_transfers(pairs, replica_policy=lambda requests: (0 for _ in requests)) + # Any real sequence will do, for the answers and for the parts of one answer. + plan_transfers(pairs, replica_policy=lambda requests: tuple(0 for _ in requests)) + plan_transfers(pairs, + replica_policy=lambda requests: [(ReplicaPart(request.logical_origin, request.shape, 1), ) + for request in requests]) + + +def test_a_scalar_cannot_be_split(): + pairs = {'s': (replicated_map((), 1), replicated_map((), 2))} + plan_transfers(pairs, replica_policy=_answering([ReplicaPart((), (), 1)])) + with pytest.raises(InvalidReplicaPolicyError): + plan_transfers(pairs, replica_policy=_answering([ReplicaPart((), (), 0), ReplicaPart((), (), 1)])) + + +def test_min_split_bytes_bounds_every_part_of_a_split(): + split = _answering([ReplicaPart((0, 0), (2, 4), 0), ReplicaPart((2, 0), (4, 4), 1)]) + # The smaller part is 2 x 4 elements x 4 bytes = 32 bytes. + plan_transfers(_one_request_pair(), element_size={'w': 4}, min_split_bytes=32, replica_policy=split) + with pytest.raises(InvalidReplicaPolicyError): + plan_transfers(_one_request_pair(), element_size={'w': 4}, min_split_bytes=33, replica_policy=split) + # An unsplit answer is never too small. + plan_transfers(_one_request_pair(), element_size={'w': 4}, min_split_bytes=10**9, replica_policy=_answering(0)) + + +@pytest.mark.parametrize( + 'kwargs', + [ + dict(min_split_bytes=8), # needs element_size + dict(element_size={'w': 4}, min_split_bytes=0), + dict(element_size={}), + dict(element_size={'w': True}), + dict(source_processes={0: 0}), # location 1 has no process + dict(target_processes={}), + dict(source_processes={ + 0: 0, + 1: 1.0 + }), + ]) +def test_inconsistent_arguments_are_refused(kwargs): + with pytest.raises(AffineTransferError): + plan_transfers(_one_request_pair(), **kwargs) + + +def test_splits_count_against_the_budget(): + pairs = {'w': (contiguous_split_map((64, 1), [64], 0), replicated_map((64, 1), 2))} + rows = _answering([ReplicaPart((row, 0), (1, 1), row % 2) for row in range(64)]) + plan_transfers(pairs, PlanBudget(max_pair_checks=64), replica_policy=rows) + with pytest.raises(AffineTransferError): + plan_transfers(pairs, PlanBudget(max_pair_checks=63), replica_policy=rows) + + +def test_maps_are_validated_before_the_policy_runs(): + called = [] + + def policy(requests): + called.append(True) + return first_holder(requests) + + pairs = {'w': (contiguous_split_map((4, 2), [4], 0), replicated_map((4, 3), 2))} + with pytest.raises(AffineTransferError): + plan_transfers(pairs, replica_policy=policy) + assert not called