diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0f940a703887..f6d9831788d3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,7 +39,7 @@ repos: name: check-torchdist entry: ./scripts/check-torchdist.py language: python - exclude: ^(.agents/skills/pr-checklist/SKILL.md|deepspeed/comm/|docs/|benchmarks/|scripts/check-torchdist.py|deepspeed/moe/sharded_moe.py|deepspeed/runtime/comm/coalesced_collectives.py|deepspeed/launcher/launch.py|tests/unit/comm/test_dist.py) + exclude: ^(.agents/skills/pr-checklist/SKILL.md|deepspeed/comm/|docs/|benchmarks/|scripts/check-torchdist.py|deepspeed/moe/sharded_moe.py|deepspeed/runtime/comm/coalesced_collectives.py|deepspeed/launcher/launch.py|tests/unit/comm/test_dist.py|tests/unit/checkpoint/test_affine_transfer_executor_external_group.py) # Specific deepspeed/ files are excluded for now until we wrap ProcessGroup in deepspeed.comm - repo: local diff --git a/deepspeed/checkpoint/affine_transfer_executor.py b/deepspeed/checkpoint/affine_transfer_executor.py new file mode 100644 index 000000000000..a9d769aae731 --- /dev/null +++ b/deepspeed/checkpoint/affine_transfer_executor.py @@ -0,0 +1,522 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Move a parameter between two shardings by executing a plan from ``affine_transfer``. + +The plan names, for each copy, a region of one source shard and the region of one target shard that +should receive it. This module turns those names into data movement: pack out of the source buffer, +across a transport group, through a staging buffer, and into the target buffer. The full logical +parameter is never formed anywhere, which was the point of planning instead of converting. + +Three things are deliberately not decided here. + +*Endpoints.* Which process actually holds a plan location is supplied by the caller. The plan's two +rank namespaces are not one world size: a source rank 0 has no claim to be the process that is target +rank 0, and folding the two together is how a transfer reads one rank's shard as another's. + +*Budget.* How many bytes may be in flight is the caller's, because what fits depends on what else is +resident on the device. A budget that is exhausted is an error, never a silently larger transfer. + +*Dtype.* Passed in rather than read from the buffers, because a process that holds no buffer at all +still has to reach the same conclusion as the ones that do -- see ``plan_digest``. + +The protocol is: agree on a digest of the schedule before any communication, then work in rounds. In +each round every rank arms all of its receives first and only then posts its sends, so two ranks that +each have something for the other cannot both be waiting on the other's send. That ordering is the +entire deadlock argument, and it only holds because the digest made every rank compute the same +rounds. Returns only once every payload it posted has landed, so a caller can read the target buffers +as soon as it is given control. + +Staging is an artifact of the current transport, not of affine transfer. Torch point-to-point +operations move one contiguous tensor per peer per message, so a round whose regions are strided or come +from several segments is packed through a per-peer arena, and ``StagingBudget`` bounds that arena; a +round that is one region already contiguous in its buffer is sent or received in place. +A zero-copy transport such as GPUDirect RDMA would instead take each round's ``(offset, length)`` pairs +as a scatter/gather list over registered memory: that removes the arena copy, but needs a registration +lifetime this API has no place for. Expect the wire contract -- one packed message per peer per round -- +to change when such a path lands. + +One device per process, with explicit transport-group rank bindings, synchronous on return. No CUDA graphs, no overlap with compute, +no custom kernels: those all change the schedule, which is the last thing a first version of a data +mover should be perturbing. +""" + +import sys +import zlib +from dataclasses import dataclass +from itertools import product +from typing import Dict, List, Mapping, Optional, Sequence, Tuple + +import torch +from deepspeed.accelerator import get_accelerator + +from .affine_transfer import TransferPlan, TransferSegment + +__all__ = [ + 'AffineTransferExecutionError', 'TransferEndpoints', 'StagingBudget', 'TransferStats', 'TransferRequest', + 'plan_digest', 'execute_transfer', 'execute_transfers' +] + + +class AffineTransferExecutionError(ValueError): + """Wrong enough that no data should have moved.""" + + +@dataclass(frozen=True) +class TransferEndpoints: + """Which process holds each location named by the plan. ``source`` covers source ranks, ``target`` + covers target ranks, and the two are separate mappings for the reason given above.""" + source: Mapping[int, int] + target: Mapping[int, int] + + +@dataclass(frozen=True) +class StagingBudget: + """Bytes a rank may stage per peer, in each direction, at one time. + + A property of the current torch point-to-point transport rather than of the plan; see the module + docstring.""" + bytes_per_peer: int = 8 << 20 + + +@dataclass(frozen=True) +class TransferStats: + segments: int + chunks: int + rounds: int + local_copy_bytes: int + sent_bytes: int + received_bytes: int + + +@dataclass(frozen=True) +class TransferRequest: + """One parameter; ``key`` must identify it identically on every process.""" + key: str + plan: Optional[TransferPlan] + source_buffers: Mapping[int, torch.Tensor] + target_buffers: Mapping[int, torch.Tensor] + preflight_error: Optional[str] = None + + +@dataclass(frozen=True) +class _Operation: + """One chunk of one segment, resolved to a pair of processes.""" + segment: TransferSegment + chunk: int + chunk_shape: Tuple[int, ...] + source_offset: int + target_offset: int + from_process: int + to_process: int + elements: int + + @property + def key(self) -> Tuple: + return (self.segment.target_rank, self.segment.target_offset, self.segment.source_rank, + self.segment.source_offset, self.chunk) + + @property + def payload_bytes(self) -> int: + return self.elements * self._element_size + + _element_size: int = 0 + + +def _chunks(segment: TransferSegment, max_elements: int): + """Tile a strided rectangle by its packed payload, retaining both address spaces.""" + if segment.numel <= max_elements: + yield segment + return + tile = [1] * len(segment.shape) + remaining = max_elements + for axis in range(len(tile) - 1, -1, -1): + tile[axis] = min(segment.shape[axis], remaining) + remaining //= tile[axis] + 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) + + +def build_operations(plan: TransferPlan, endpoints: TransferEndpoints, budget: StagingBudget, + element_size: int) -> List[_Operation]: + """The complete schedule, as a list every rank derives identically.""" + if element_size <= 0 or budget.bytes_per_peer < element_size: + raise AffineTransferExecutionError('the staging budget must hold at least one element') + operations: List[_Operation] = [] + for segment in plan.segments: + _validate_segment(segment, len(plan.logical_shape)) + for index, chunk in enumerate(_chunks(segment, budget.bytes_per_peer // element_size)): + operation = _Operation(segment=segment, + chunk=index, + chunk_shape=chunk.shape, + source_offset=chunk.source_offset, + target_offset=chunk.target_offset, + from_process=endpoints.source[segment.source_rank], + to_process=endpoints.target[segment.target_rank], + elements=chunk.numel, + _element_size=element_size) + operations.append(operation) + operations.sort(key=lambda operation: operation.key) + return operations + + +def plan_digest(plan: TransferPlan, endpoints: TransferEndpoints, budget: StagingBudget, element_size: int) -> int: + """A checksum over the schedule: geometry, addressing, endpoints and budget. + + Deliberately not over buffer contents -- those differ by design. This is the thing every rank must + be able to compute before it owns any data, which is why dtype is an argument and not an + inspection. + """ + payload = [(segment.source_rank, segment.target_rank, segment.source_offset, segment.target_offset, segment.shape, + segment.source_strides, segment.target_strides) for segment in plan.segments] + framed = { + 'schema': 3, + 'segments': payload, + 'logical_shape': plan.logical_shape, + 'source_endpoints': dict(sorted(endpoints.source.items())), + 'target_endpoints': dict(sorted(endpoints.target.items())), + 'bytes_per_peer': budget.bytes_per_peer, + 'element_size': element_size, + } + return zlib.crc32(repr(sorted(framed.items(), key=lambda item: item[0])).encode()) + + +def _validate_segment(segment: TransferSegment, ndim: int) -> None: + if len(segment.shape) != ndim: + raise AffineTransferExecutionError('segment dimensionality does not match its logical shape') + for size in segment.shape: + if not isinstance(size, int) or isinstance(size, bool) or size <= 0: + raise AffineTransferExecutionError('segment extents must be positive integers') + for offset, strides in ((segment.source_offset, segment.source_strides), (segment.target_offset, + segment.target_strides)): + if not isinstance(offset, int) or isinstance(offset, bool) or offset < 0: + raise AffineTransferExecutionError('segment offsets must be non-negative integers') + if len(strides) != ndim: + raise AffineTransferExecutionError('segment stride dimensionality does not match its shape') + for size, stride in zip(segment.shape, strides): + if not isinstance(stride, int) or isinstance(stride, bool) or stride < 0 or (size > 1 and stride == 0): + raise AffineTransferExecutionError('segment strides must be positive on non-singleton axes') + + +def _validate_buffers(plan: TransferPlan, source_buffers: Mapping[int, torch.Tensor], + target_buffers: Mapping[int, torch.Tensor], dtype: torch.dtype, device: torch.device, + endpoints: TransferEndpoints, rank: int) -> None: + for which, buffers, bindings in (('source', source_buffers, endpoints.source), ('target', target_buffers, + endpoints.target)): + for location, buffer in buffers.items(): + if bindings.get(location) != rank: + raise AffineTransferExecutionError(f'{which} location {location} is not owned by this process') + _flat(location, buffers, dtype) + if buffer.device != device: + raise AffineTransferExecutionError( + f'{which} location {location}: buffer is on {buffer.device}, transport is on {device}') + for segment in plan.segments: + _validate_segment(segment, len(plan.logical_shape)) + for which, location, offset, strides, buffers, bindings in (('source', segment.source_rank, + segment.source_offset, segment.source_strides, + source_buffers, endpoints.source), + ('target', segment.target_rank, + segment.target_offset, segment.target_strides, + target_buffers, endpoints.target)): + if location not in bindings: + raise AffineTransferExecutionError(f'{which} location {location} has no endpoint binding') + if bindings[location] != rank: + 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)) + if reach >= buffers[location].numel(): + raise AffineTransferExecutionError( + f'{which} of {segment!r} reaches element {reach} of a {buffers[location].numel()}-element buffer') + + +def _validate_aliases(requests: Sequence[TransferRequest]) -> None: + sources = [ + _data_range(buffer) for request in requests for buffer in request.source_buffers.values() if buffer.numel() + ] + targets = [ + _data_range(buffer) for request in requests for buffer in request.target_buffers.values() if buffer.numel() + ] + for low, high in targets: + if any(low < source_high and source_low < high for source_low, source_high in sources): + raise AffineTransferExecutionError('a source and a target overlap within the transfer batch') + for index, (low, high) in enumerate(targets): + if any(low < other_high and other_low < high for other_low, other_high in targets[:index]): + raise AffineTransferExecutionError('two targets overlap within the transfer batch') + + +def execute_transfer(plan: TransferPlan, + source_buffers: Mapping[int, torch.Tensor], + target_buffers: Mapping[int, torch.Tensor], + endpoints: TransferEndpoints, + dtype: torch.dtype, + budget: StagingBudget = StagingBudget(), + group: Optional[object] = None) -> TransferStats: + """Execute one plan with the same preflight and completion contract as a batch. + + Bindings use transport-group ranks. Buffers must be packed views on the owning process; + their storage offsets are retained. Rejection on any participant precedes all target writes. + Transport failures after payload begins may leave partially written targets. + """ + request = TransferRequest('parameter', plan, source_buffers, target_buffers) + return execute_transfers([request], endpoints, dtype, budget, group) + + +def _packed_rounds(indexed_operations: Sequence[Tuple[int, _Operation]], max_elements: int): + rounds = [] + for item, operation in indexed_operations: + if not rounds or rounds[-1][1] + operation.elements > max_elements: + rounds.append(([], 0)) + entries, used = rounds[-1] + entries.append((item, operation, used)) + rounds[-1] = (entries, used + operation.elements) + return rounds + + +def _direct_payload(entries, requests, dtype, source): + if len(entries) != 1: + return None + item, operation, _ = entries[0] + request, segment = requests[item], operation.segment + if source: + flat = _flat(segment.source_rank, request.source_buffers, dtype) + offset, strides = operation.source_offset, segment.source_strides + else: + flat = _flat(segment.target_rank, request.target_buffers, dtype) + offset, strides = operation.target_offset, segment.target_strides + region = _region(flat, offset, operation.chunk_shape, strides) + # A view can replace staging only when the transport sees the same packed element order. + return region.view(-1) if region.is_contiguous() else None + + +@torch.no_grad() +def execute_transfers(requests: Sequence[TransferRequest], + endpoints: TransferEndpoints, + dtype: torch.dtype, + budget: StagingBudget = StagingBudget(), + group: Optional[object] = None) -> TransferStats: + """Move a frozen parameter batch with collective preflight and bounded peer messages. + + Each direction stages at most ``bytes_per_peer`` per active peer. All requests and staging + allocations are checked before payload, including read/write aliases between parameters. + The caller stream's final unpack and local copies complete before this function returns. + """ + import deepspeed.comm as dist + + initialized = dist.is_initialized() + if initialized and group is not None and callable(getattr(group, 'rank', None)): + # Ask the group itself: a weight-update group built outside the default world (as trainer <-> engine + # groups are) maps every member's default-world rank 0 to group rank 0 in torch's bookkeeping. + rank, world_size = group.rank(), group.size() + else: + rank = dist.get_rank(group=group) if initialized else 0 + world_size = dist.get_world_size(group=group) if initialized else 1 + device = torch.device(get_accelerator().current_device_name()) if initialized else torch.device('cpu') + # Without a process group every buffer is on the host, whatever the accelerator is. + stream_ordered = initialized and not get_accelerator().is_synchronized_device() + indexed, local_copies = [], [] + recv_rounds, send_rounds, recv_arenas, send_arenas = {}, {}, {}, {} + recv_direct, send_direct = {}, {} + digest, rejected = 0, 0 + error = None + try: + element_size = torch.empty((), dtype=dtype).element_size() + if (not isinstance(budget.bytes_per_peer, int) or isinstance(budget.bytes_per_peer, bool) + or budget.bytes_per_peer < element_size): + raise AffineTransferExecutionError('the staging budget must hold at least one element') + for bindings in (endpoints.source, endpoints.target): + for location, process in bindings.items(): + if (not isinstance(location, int) or isinstance(location, bool) or location < 0 + or not isinstance(process, int) or isinstance(process, bool) or not 0 <= process < world_size): + raise AffineTransferExecutionError('endpoint bindings must name valid transport-group ranks') + keys = [request.key for request in requests] + if any(not isinstance(key, str) or not key for key in keys) or len(set(keys)) != len(keys): + raise AffineTransferExecutionError('transfer batch keys must be distinct non-empty strings') + schedules = {} + for item, request in enumerate(requests): + if request.preflight_error is not None: + raise AffineTransferExecutionError(request.preflight_error) + if request.plan is None: + raise AffineTransferExecutionError('a transfer request has no plan') + _validate_buffers(request.plan, request.source_buffers, request.target_buffers, dtype, device, endpoints, + rank) + # Many parameters share the same plan. Reuse only geometry within this call; + # every request still validates its own buffers and participates in the digest. + identity = id(request.plan) + if identity not in schedules: + schedules[identity] = (build_operations(request.plan, endpoints, budget, element_size), + plan_digest(request.plan, endpoints, budget, element_size)) + indexed.extend((item, operation) for operation in schedules[identity][0]) + _validate_aliases(requests) + fingerprint = (3, str(dtype), budget.bytes_per_peer, tuple(sorted(endpoints.source.items())), + tuple(sorted(endpoints.target.items())), + tuple((request.key, schedules[id(request.plan)][1]) for request in requests)) + digest = zlib.crc32(repr(fingerprint).encode()) + incoming: Dict[int, List[Tuple[int, _Operation]]] = {} + outgoing: Dict[int, List[Tuple[int, _Operation]]] = {} + for item, operation in indexed: + if operation.from_process == operation.to_process == rank: + local_copies.append((item, operation)) + elif operation.to_process == rank: + incoming.setdefault(operation.from_process, []).append((item, operation)) + elif operation.from_process == rank: + outgoing.setdefault(operation.to_process, []).append((item, operation)) + max_elements = budget.bytes_per_peer // element_size + recv_rounds = {peer: _packed_rounds(ops, max_elements) for peer, ops in incoming.items()} + send_rounds = {peer: _packed_rounds(ops, max_elements) for peer, ops in outgoing.items()} + recv_direct = { + peer: [_direct_payload(entries, requests, dtype, False) for entries, _ in rounds] + for peer, rounds in recv_rounds.items() + } + send_direct = { + peer: [_direct_payload(entries, requests, dtype, True) for entries, _ in rounds] + for peer, rounds in send_rounds.items() + } + recv_arenas = { + peer: + torch.empty(max(count for index, (_, count) in enumerate(rounds) if recv_direct[peer][index] is None), + dtype=dtype, + device=device) + for peer, rounds in recv_rounds.items() if any(payload is None for payload in recv_direct[peer]) + } + send_arenas = { + peer: + torch.empty(max(count for index, (_, count) in enumerate(rounds) if send_direct[peer][index] is None), + dtype=dtype, + device=device) + for peer, rounds in send_rounds.items() if any(payload is None for payload in send_direct[peer]) + } + except (ValueError, TypeError, KeyError, AttributeError, RuntimeError, OverflowError, MemoryError) as caught: + error = caught + rejected = 1 + recv_arenas, send_arenas = {}, {} + print(f'[transfer] rank {rank} rejected the batch: {caught}', file=sys.stderr, flush=True) + + # Invalid metadata must not be inspected again while constructing the agreement payload. + if world_size > 1: + local = torch.tensor([digest, rejected], dtype=torch.int64, device=device) + seen = torch.empty((world_size, 2), dtype=torch.int64, device=device) + dist.all_gather(list(seen.unbind(0)), local, group=group) + # Read the complete decision once instead of synchronizing each rank's scalar separately. + decisions = seen.cpu().tolist() + if any(pair[1] for pair in decisions): + raise AffineTransferExecutionError( + 'a rank rejected the transfer batch before it began; no data moved') from error + if len({pair[0] for pair in decisions}) != 1: + raise AffineTransferExecutionError('ranks disagree about the transfer schedule (batch schedule)') + elif rejected: + raise AffineTransferExecutionError(f'the transfer batch was rejected before data moved: {error}') from error + + rounds = max((len(value) for value in list(recv_rounds.values()) + list(send_rounds.values())), default=0) + sent = received = copied = 0 + pending = [] + + for round_index in range(rounds): + p2p = [] + for peer, peer_rounds in sorted(recv_rounds.items()): + if round_index < len(peer_rounds): + payload = recv_direct[peer][round_index] + if payload is None: + payload = recv_arenas[peer][:peer_rounds[round_index][1]] + p2p.append(('recv', payload, peer)) + for peer, peer_rounds in sorted(send_rounds.items()): + if round_index >= len(peer_rounds): + continue + entries, count = peer_rounds[round_index] + payload = send_direct[peer][round_index] + if payload is None: + arena = send_arenas[peer] + for item, operation, offset in entries: + source = _flat(operation.segment.source_rank, requests[item].source_buffers, dtype) + arena[offset:offset + operation.elements].view(operation.chunk_shape).copy_( + _region(source, operation.source_offset, operation.chunk_shape, + operation.segment.source_strides)) + payload = arena[:count] + sent += count * element_size + p2p.append(('send', payload, peer)) + operations = [ + dist.P2POp(dist.irecv if direction == 'recv' else dist.isend, + buffer, + dist.get_global_rank(group, peer) if group is not None else peer, + group=group) for direction, buffer, peer in p2p + ] + if operations: + for handle in dist.batch_isend_irecv(operations): + block_stream = getattr(handle, 'block_current_stream', None) + if stream_ordered and callable(block_stream): + # Order arena reuse on the caller stream; transport errors surface at the bounded drain below. + block_stream() + pending.append(handle) + else: + handle.wait() + # Stream ordering already protects arena reuse. Bound retained handles while avoiding + # a host polling delay after every small group of rounds; all errors are observed on return. + if (round_index + 1) % 32 == 0 or round_index + 1 == rounds: + for handle in pending: + handle.wait() + pending.clear() + for peer, peer_rounds in recv_rounds.items(): + if round_index >= len(peer_rounds): + continue + received += peer_rounds[round_index][1] * element_size + if recv_direct[peer][round_index] is not None: + continue + for item, operation, offset in peer_rounds[round_index][0]: + target = _flat(operation.segment.target_rank, requests[item].target_buffers, dtype) + payload = recv_arenas[peer][offset:offset + operation.elements].view(operation.chunk_shape) + _region(target, operation.target_offset, operation.chunk_shape, + operation.segment.target_strides).copy_(payload) + + for item, operation in local_copies: + source = _flat(operation.segment.source_rank, requests[item].source_buffers, dtype) + target = _flat(operation.segment.target_rank, requests[item].target_buffers, dtype) + _region(target, operation.target_offset, operation.chunk_shape, operation.segment.target_strides).copy_( + _region(source, operation.source_offset, operation.chunk_shape, operation.segment.source_strides)) + copied += operation.payload_bytes + + if stream_ordered: + get_accelerator().current_stream().synchronize() + return TransferStats(sum(len(request.plan.segments) for request in requests), len(indexed), rounds, copied, sent, + received) + + +def _data_range(buffer: torch.Tensor) -> Tuple[int, int]: + """Byte extent a tensor's elements occupy in storage.""" + start = buffer.data_ptr() + span = buffer.numel() * buffer.element_size() + return start, start + span + + +def _flat(location: int, buffers: Mapping[int, torch.Tensor], dtype: torch.dtype) -> torch.Tensor: + """The buffer for ``location`` as its own elements, without copying it.""" + try: + buffer = buffers[location] + except KeyError: + raise AffineTransferExecutionError(f'location {location} is bound to this process but holds no buffer') + if buffer.layout != torch.strided or not buffer.is_contiguous(): + raise AffineTransferExecutionError(f'location {location}: shard is not contiguous in its own storage') + if buffer.dtype != dtype: + raise AffineTransferExecutionError(f'location {location}: {buffer.dtype} is not {dtype}') + return buffer.view(-1) + + +def _region(flat: torch.Tensor, offset: int, shape: Tuple[int, ...], strides: Tuple[int, ...]) -> torch.Tensor: + """``flat`` viewed as one planned copy region. + + Plan offsets count elements from the first element of the shard; ``as_strided`` counts from the + start of the storage. Adding the buffer's own storage offset is the explicit form of that, and the + alternative -- making the buffer contiguous -- would copy the whole shard to fix arithmetic. + """ + return torch.as_strided(flat, size=shape, stride=strides, storage_offset=flat.storage_offset() + offset) diff --git a/docs/code-docs/source/affine-transfer.rst b/docs/code-docs/source/affine-transfer.rst new file mode 100644 index 000000000000..6f27352ae90a --- /dev/null +++ b/docs/code-docs/source/affine-transfer.rst @@ -0,0 +1,48 @@ +Affine Shard Transfers +====================== + +AutoTP and AutoEP both describe the shards they create with affine maps from +``deepspeed.checkpoint.affine``. ``deepspeed.checkpoint.affine_transfer`` +plans how to move a parameter from one such sharding to another, and +``deepspeed.checkpoint.affine_transfer_executor`` executes that plan directly +into existing packed target shards without reconstructing the full logical +tensor. Neither depends on which parallelism produced the maps. + +Executing affine transfers +-------------------------- + +Endpoint mappings use transport-group ranks and distinguish source locations +from target locations:: + + from deepspeed.checkpoint.affine_transfer import plan_transfer + from deepspeed.checkpoint.affine_transfer_executor import ( + StagingBudget, TransferEndpoints, execute_transfer) + + plan = plan_transfer(target_map, source_map) + execute_transfer(plan, local_sources, local_targets, + TransferEndpoints(source_bindings, target_bindings), + dtype, StagingBudget(bytes_per_peer=8 << 20)) + +Use ``execute_transfers`` with uniquely keyed ``TransferRequest`` objects to pack +multiple parameters into peer messages. Each process supplies only its own +contiguous buffers; strided plan regions and buffer storage offsets are retained. +All participants must agree on requests, endpoints, dtype and staging budget. +Preflight refusal leaves targets untouched. Return guarantees target readiness; +a transport failure after payload begins can leave partially written targets. +The budget applies to each peer direction, not total device memory. Weights must +remain frozen during transfer; optimizer migration and live AutoEP swaps are outside +this API. Maintenance contact: ``@0z5a``. + +Transport and staging +--------------------- + +The staging arena and ``StagingBudget`` belong to the current transport, not to +affine transfer itself. Torch point-to-point operations move one contiguous +tensor per peer per message, so a round whose regions are strided or +come from several segments is packed through a per-peer arena; a round that is +one region already contiguous in its buffer is sent or received in place. A +zero-copy transport such as GPUDirect RDMA would instead take each round's +``(offset, length)`` pairs as a scatter/gather list over registered memory. That +removes the arena copy but needs a registration lifetime this API has no place +for, so the wire contract (one packed message per peer per round) is expected to +change when such a path lands. diff --git a/docs/code-docs/source/index.rst b/docs/code-docs/source/index.rst index 5ddf605557cc..4639081120b0 100644 --- a/docs/code-docs/source/index.rst +++ b/docs/code-docs/source/index.rst @@ -32,6 +32,7 @@ Checkpointing API :maxdepth: 2 model-checkpointing + affine-transfer activation-checkpointing diff --git a/tests/unit/checkpoint/test_affine_transfer_executor.py b/tests/unit/checkpoint/test_affine_transfer_executor.py new file mode 100644 index 000000000000..632988d5ed06 --- /dev/null +++ b/tests/unit/checkpoint/test_affine_transfer_executor.py @@ -0,0 +1,448 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Check transfer payloads against extraction from the logical tensor.""" + +import math +import os +from dataclasses import replace + +import pytest +import torch +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.affine import contiguous_split_map, replicated_map +from deepspeed.checkpoint.affine_transfer import PlanBudget, plan_transfer +from deepspeed.checkpoint.affine_transfer_executor import (AffineTransferExecutionError, StagingBudget, + TransferEndpoints, TransferRequest, execute_transfer, + execute_transfers) +from unit.common import DistributedTest + +UNWRITTEN = float('nan') +DTYPES = [torch.float32, torch.float16, torch.bfloat16] + + +def _join_transfer_group(): + if not dist.is_initialized(): + dist.init_distributed(dist_backend=get_accelerator().communication_backend_name(), init_method='env://') + + +def _device(): + local_rank = int(os.environ.get('LOCAL_RANK', dist.get_rank())) + get_accelerator().set_device(local_rank) + return torch.device(get_accelerator().current_device_name()) + + +def _full(rows, cols, dtype): + """Identical markers that stay distinct after fp16 or bf16 conversion.""" + values = torch.arange(rows * cols, dtype=torch.float32, device=_device()).reshape(rows, cols).to(dtype) + flat = values.reshape(-1) + assert torch.unique(flat).numel() == flat.numel(), f'{dtype} markers collided; this case cannot detect a swap' + return values + + +def _targets(map_, endpoints, dtype): + return { + location: torch.full(shape, UNWRITTEN, dtype=dtype, device=_device()) + for location, shape in map_.shard_shapes.items() if endpoints.target[location] == dist.get_rank() + } + + +def _assert_matches(plan_targets, target_map, full): + for location, buffer in plan_targets.items(): + want = target_map.extract(full, location) + assert torch.equal(buffer, want), (f'target location {location}: {int((~torch.eq(buffer, want)).sum())} ' + f'of {want.numel()} elements differ from the extracted shard') + + +def _row_to_column(world_size, tp_size=2): + rows, cols = 8 * tp_size, 8 * tp_size + source = contiguous_split_map((rows, cols), [rows // tp_size] * tp_size, 0) + target = contiguous_split_map((rows, cols), [cols // world_size] * world_size, 1) + return source, target + + +class TestB2Transfer(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize('dtype', DTYPES, ids=['fp32', 'fp16', 'bf16']) + def test_pure_copy_lands_every_element(self, dtype): + _join_transfer_group() + source, target = _row_to_column(self.world_size) + full = _full(16, 16, dtype) + endpoints = TransferEndpoints(source={loc: loc + for loc in source.shard_shapes}, + target={loc: loc + for loc in target.shard_shapes}) + source_buffers = { + location: source.extract(full, location).contiguous() + for location, process in endpoints.source.items() if process == dist.get_rank() + } + target_buffers = _targets(target, endpoints, dtype) + + plan = plan_transfer(target, source, PlanBudget()) + stats = execute_transfer(plan, source_buffers, target_buffers, endpoints, dtype) + + _assert_matches(target_buffers, target, full) + assert stats.chunks > 0 + assert sum(segment.numel for segment in plan.segments) == math.prod(target.logical_shape), \ + 'the plan must cover the parameter exactly once' + + @pytest.mark.parametrize('batched', [False, True]) + def test_scalar_plan_crosses_processes(self, batched): + _join_transfer_group() + rank = dist.get_rank() + map_ = replicated_map((), 1) + plan, endpoints = plan_transfer(map_, map_), TransferEndpoints({0: 0}, {0: 1}) + sources = {0: torch.tensor(3.0, device=_device())} if rank == 0 else {} + storage = torch.full((3, ), -9.0, device=_device()) + targets = {0: storage[1].view(())} if rank == 1 else {} + if batched: + execute_transfers([TransferRequest('scalar', plan, sources, targets)], endpoints, torch.float32, + StagingBudget(4)) + else: + execute_transfer(plan, sources, targets, endpoints, torch.float32, StagingBudget(4)) + expected = torch.tensor([-9.0, 3.0 if rank == 1 else -9.0, -9.0], device=_device()) + assert torch.equal(storage, expected) + + def test_ranks_disagreeing_about_the_schedule_refuse_everywhere(self): + _join_transfer_group() + """A lone refusal would strand its peers, so disagreement has to become a collective failure. + + One rank sizes its staging differently, which changes the chunk schedule and therefore the + digest. Nobody may post communication after that, so this asserts that all ranks raise rather + than that some rank does. + """ + source, target = _row_to_column(self.world_size) + dtype = torch.float32 + full = _full(16, 16, dtype) + endpoints = TransferEndpoints(source={loc: loc + for loc in source.shard_shapes}, + target={loc: loc + for loc in target.shard_shapes}) + source_buffers = { + location: source.extract(full, location).contiguous() + for location, process in endpoints.source.items() if process == dist.get_rank() + } + target_buffers = _targets(target, endpoints, dtype) + budget = StagingBudget(bytes_per_peer=4 << 20 if dist.get_rank() == 0 else 1 << 20) + + with pytest.raises(AffineTransferExecutionError, match='transfer schedule'): + execute_transfer(plan_transfer(target, source, PlanBudget()), source_buffers, target_buffers, endpoints, + dtype, budget) + + def test_shared_plan_batch_keeps_distinct_payloads_and_revalidates(self): + _join_transfer_group() + rank, world = dist.get_rank(), dist.get_world_size() + map_ = contiguous_split_map((4 * world, 8), [4] * world, 0) + plan = plan_transfer(map_, map_) + endpoints = TransferEndpoints({location: location + for location in range(world)}, + {location: (location + world // 2) % world + for location in range(world)}) + target_location = (rank + world // 2) % world + requests, checks = [], [] + for index in range(6): + full = _full(4 * world, 8, torch.float32) + index * 32 * world + source = map_.extract(full, rank).contiguous() + target = torch.full((4, 8), UNWRITTEN, device=_device()) + requests.append(TransferRequest(str(index), plan, {rank: source}, {target_location: target})) + checks.append((target, map_.extract(full, target_location))) + previous_rounds = 0 + # The smallest budget spans 48 rounds, crossing the bounded CUDA work retention interval. + for budget in (1024, 32, 16): + stats = execute_transfers(requests, endpoints, torch.float32, StagingBudget(budget)) + assert stats.rounds > previous_rounds + assert stats.sent_bytes == stats.received_bytes == 6 * 4 * 8 * 4 + previous_rounds = stats.rounds + for request, (target, expected) in zip(requests, checks): + assert torch.equal(target, expected) + request.source_buffers[rank].add_(1) + expected.add_(1) + target.fill_(UNWRITTEN) + if rank == 0: + requests[1].target_buffers[target_location] = checks[1][0][:1] + with pytest.raises(AffineTransferExecutionError, match='rejected the transfer batch'): + execute_transfers(requests, endpoints, torch.float32, StagingBudget(1024)) + assert all(torch.isnan(target).all().item() for target, _ in checks) + + def test_mixed_packed_and_strided_regions_preserve_storage_windows(self): + _join_transfer_group() + world, rank = dist.get_world_size(), dist.get_rank() + width = 4 * world + endpoints = TransferEndpoints({location: location + for location in range(world)}, + {location: (location + world // 2) % world + for location in range(world)}) + requests, checks, guards = [], [], [] + # Flattening a noncontiguous region or ignoring a shard's storage offset corrupts this batch. + for index, (source_axis, target_axis) in enumerate(((0, 0), (0, 1), (1, 0))): + full = _full(width, width, torch.float32) + index * width * width + source_map = contiguous_split_map((width, width), [4] * world, source_axis) + target_map = contiguous_split_map((width, width), [4] * world, target_axis) + sources, targets = {}, {} + for map_, bindings, buffers, is_source in ((source_map, endpoints.source, sources, True), + (target_map, endpoints.target, targets, False)): + for location, shape in map_.shard_shapes.items(): + if bindings[location] != rank: + continue + size = math.prod(shape) + storage = torch.full((size + 12, ), -17.0, device=_device()) + buffer = storage[5:5 + size].view(shape) + if is_source: + buffer.copy_(map_.extract(full, location)) + buffers[location] = buffer + guards.append((storage, size)) + requests.append(TransferRequest(f'mixed-{index}', plan_transfer(target_map, source_map), sources, targets)) + checks.append((targets, target_map, full)) + if not get_accelerator().is_synchronized_device(): + stream = get_accelerator().Stream() + stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(stream): + execute_transfers(requests, endpoints, torch.float32, StagingBudget(256)) + else: + execute_transfers(requests, endpoints, torch.float32, StagingBudget(256)) + # The return contract makes targets ready even after leaving the caller's stream context. + for targets, target_map, full in checks: + _assert_matches(targets, target_map, full) + for storage, size in guards: + assert torch.all(storage[:5] == -17).item() + assert torch.all(storage[5 + size:] == -17).item() + + def test_batch_parameter_order_must_match(self): + _join_transfer_group() + source, target = _row_to_column(self.world_size) + full = _full(16, 16, torch.float32) + endpoints = TransferEndpoints(source={0: 0, 1: 1}, target={0: 0, 1: 1}) + sources = { + location: source.extract(full, location).contiguous() + for location in source.shard_shapes if endpoints.source[location] == dist.get_rank() + } + plan = plan_transfer(target, source) + first = TransferRequest('first', plan, sources, _targets(target, endpoints, torch.float32)) + second = TransferRequest('second', plan, sources, _targets(target, endpoints, torch.float32)) + requests = [first, second] if dist.get_rank() == 0 else [second, first] + with pytest.raises(AffineTransferExecutionError, match='batch schedule'): + execute_transfers(requests, endpoints, torch.float32) + + +class TestB2DisjointEndpoints(DistributedTest): + """Every process is either a reader or a writer, never both. + + A transfer that only appears to work by falling back to local copies fails here. Two ranks fit the + two-device pull-request runners; four fit the CPU multi-rank nightly and give each writer two readers. + """ + + @pytest.mark.parametrize('world_size', [2, 4]) + def test_source_and_target_ranks_do_not_overlap(self, world_size): + _join_transfer_group() + source, target = _row_to_column(2) + dtype = torch.float32 + full = _full(16, 16, dtype) + half = world_size // 2 + # Readers are the lower half of the ranks and writers the upper half. + endpoints = TransferEndpoints(source={location: location % half + for location in source.shard_shapes}, + target={location: half + location % half + for location in target.shard_shapes}) + source_buffers = { + location: source.extract(full, location).contiguous() + for location, process in endpoints.source.items() if process == dist.get_rank() + } + target_buffers = _targets(target, endpoints, dtype) + + stats = execute_transfer(plan_transfer(target, source, PlanBudget()), source_buffers, target_buffers, + endpoints, dtype) + + # With no process holding both ends, each element crosses the transport exactly once. + assert stats.local_copy_bytes == 0 + assert stats.sent_bytes == sum(buffer.numel() for buffer in source_buffers.values()) * 4 + assert stats.received_bytes == sum(buffer.numel() for buffer in target_buffers.values()) * 4 + _assert_matches(target_buffers, target, full) + + +class TestB2TransferCPU: + + def test_shared_plan_uses_each_buffer_and_rechecks_each_call(self): + map_ = replicated_map((3, 4), 1) + plan = plan_transfer(map_, map_) + endpoints = TransferEndpoints({0: 0}, {0: 0}) + sources = [torch.arange(12, dtype=torch.float32).view(3, 4) + i * 100 for i in range(2)] + targets = [torch.full((3, 4), -1.0) for _ in sources] + requests = [ + TransferRequest(str(i), plan, {0: source}, {0: target}) + for i, (source, target) in enumerate(zip(sources, targets)) + ] + for budget, chunks in ((48, 2), (8, 12)): + for source, target in zip(sources, targets): + source.add_(1) + target.fill_(-1) + stats = execute_transfers(requests, endpoints, torch.float32, StagingBudget(budget)) + assert stats.chunks == chunks + assert stats.local_copy_bytes == 96 + for source, target in zip(sources, targets): + assert torch.equal(target, source) + for target in targets: + target.fill_(-1) + # Reusing geometry must not reuse a different request's buffer validation. + requests[1].target_buffers[0] = targets[1][:1] + with pytest.raises(AffineTransferExecutionError): + execute_transfers(requests, endpoints, torch.float32, StagingBudget(48)) + for target in targets: + assert torch.equal(target, torch.full((3, 4), -1.0)) + + @pytest.mark.parametrize('shape', [(3, 7, 11), (13, 17, 1), (2, 5, 7, 3)]) + def test_nd_tiles_fit_a_small_packed_budget_and_preserve_guards(self, shape): + # Last-axis-only splitting rejects these valid shapes even when one element fits. + numel = math.prod(shape) + full = torch.arange(numel, dtype=torch.float32).reshape(shape) + source_storage = torch.full((numel + 10, ), -7.0) + target_storage = torch.full((numel + 12, ), -9.0) + source = source_storage[5:5 + numel].view(shape) + target = target_storage[6:6 + numel].view(shape) + source.copy_(full) + map_ = replicated_map(shape, 1) + request = TransferRequest('expert', plan_transfer(map_, map_), {0: source}, {0: target}) + stats = execute_transfers([request], TransferEndpoints({0: 0}, {0: 0}), torch.float32, + StagingBudget(bytes_per_peer=28)) + assert torch.equal(target, full) + assert torch.equal(target_storage[:6], torch.full((6, ), -9.0)) + assert torch.equal(target_storage[-6:], torch.full((6, ), -9.0)) + assert stats.local_copy_bytes == numel * 4 + assert stats.chunks > 1 + + @pytest.mark.parametrize('shape', [(4, 7, 11), (4, 5, 7, 3)]) + def test_nd_strided_pieces_land_in_the_target_map(self, shape): + source_map = contiguous_split_map(shape, [2, 2], 0) + last = shape[-1] + target_map = contiguous_split_map(shape, [last // 2, last - last // 2], len(shape) - 1) + full = torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) + sources = {rank: source_map.extract(full, rank).contiguous() for rank in source_map.shard_shapes} + targets = {rank: torch.full(local_shape, -1.0) for rank, local_shape in target_map.shard_shapes.items()} + endpoints = TransferEndpoints({0: 0, 1: 0}, {0: 0, 1: 0}) + execute_transfer(plan_transfer(target_map, source_map), sources, targets, endpoints, torch.float32, + StagingBudget(bytes_per_peer=20)) + for rank, target in targets.items(): + assert torch.equal(target, target_map.extract(full, rank)) + + @pytest.mark.parametrize('invalid', [ + 'source_noncontiguous', + 'target_noncontiguous', + 'missing_source', + 'missing_target', + 'negative_offset', + 'short_stride', + 'zero_stride', + 'wrong_dtype', + 'zero_budget', + 'sub_element_budget', + 'unowned_source', + 'unowned_target', + 'invalid_endpoint', + ]) + @pytest.mark.parametrize('batched', [False, True]) + def test_invalid_buffers_and_metadata_are_rejected_before_any_write(self, invalid, batched): + map_ = replicated_map((2, 3), 1) + plan = plan_transfer(map_, map_) + source = torch.arange(6, dtype=torch.float32).reshape(2, 3) + target = torch.full((2, 3), -11.0) + targets = {0: target} + sources = {0: source} + endpoints = TransferEndpoints({0: 0}, {0: 0}) + budget = StagingBudget(16) + if invalid == 'source_noncontiguous': + sources[0] = source.t() + elif invalid == 'target_noncontiguous': + targets[0] = target.t() + elif invalid == 'missing_source': + sources = {} + elif invalid == 'missing_target': + targets = {} + elif invalid in ('negative_offset', 'short_stride', 'zero_stride'): + changes = { + 'negative_offset': { + 'source_offset': -1 + }, + 'short_stride': { + 'source_strides': (1, ) + }, + 'zero_stride': { + 'target_strides': (0, 1) + } + } + segment = replace(plan.segments[0], **changes[invalid]) + plan = replace(plan, segments=(segment, )) + elif invalid == 'wrong_dtype': + sources[0] = source.to(torch.bfloat16) + elif invalid == 'zero_budget': + budget = StagingBudget(0) + elif invalid == 'sub_element_budget': + budget = StagingBudget(3) + elif invalid == 'unowned_source': + sources = {1: source} + elif invalid == 'unowned_target': + targets = {1: target} + elif invalid == 'invalid_endpoint': + endpoints = TransferEndpoints({0: 1}, {0: 0}) + first_target = torch.full((2, 3), -13.0) + with pytest.raises(AffineTransferExecutionError): + if batched: + first = TransferRequest('first', plan_transfer(map_, map_), {0: source}, {0: first_target}) + second = TransferRequest('second', plan, sources, targets) + execute_transfers([first, second], endpoints, torch.float32, budget) + else: + execute_transfer(plan, sources, targets, endpoints, torch.float32, budget) + assert torch.equal(target, torch.full((2, 3), -11.0)) + assert torch.equal(first_target, torch.full((2, 3), -13.0)) + + @pytest.mark.parametrize('alias', ['read_write', 'write_write']) + def test_aliases_between_requests_are_rejected_before_any_write(self, alias): + map_ = replicated_map((6, ), 1) + plan = plan_transfer(map_, map_) + shared = torch.full((6, ), -11.0) + other = torch.full((6, ), -13.0) + first = TransferRequest('first', plan, {0: torch.arange(6, dtype=torch.float32)}, {0: shared}) + second = TransferRequest('second', plan, + {0: shared if alias == 'read_write' else torch.arange(6, dtype=torch.float32)}, + {0: other if alias == 'read_write' else shared}) + with pytest.raises(AffineTransferExecutionError): + execute_transfers([first, second], TransferEndpoints({0: 0}, {0: 0}), torch.float32) + assert torch.equal(shared, torch.full((6, ), -11.0)) + assert torch.equal(other, torch.full((6, ), -13.0)) + + def test_disjoint_windows_of_one_storage_are_valid(self): + map_ = replicated_map((6, ), 1) + storage = torch.arange(12, dtype=torch.float32) + source, target = storage[:6], storage[6:] + expected = source.clone() + execute_transfer(plan_transfer(map_, map_), {0: source}, {0: target}, TransferEndpoints({0: 0}, {0: 0}), + torch.float32, StagingBudget(8)) + assert torch.equal(target, expected) + assert torch.equal(source, expected) + + def test_frozen_parameters_keep_their_identity_and_storage(self): + map_ = replicated_map((2, 3), 1) + source = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(2, 3)) + target = torch.nn.Parameter(torch.full((2, 3), -1.0)) + identity, pointer = id(target), target.data_ptr() + execute_transfer(plan_transfer(map_, map_), {0: source}, {0: target}, TransferEndpoints({0: 0}, {0: 0}), + torch.float32, StagingBudget(8)) + assert id(target) == identity and target.data_ptr() == pointer + assert target.requires_grad and target.grad is None + assert torch.equal(target, source) + + @pytest.mark.parametrize('batched', [False, True]) + def test_scalar_plan_preserves_storage_offsets_and_guards(self, batched): + source_storage = torch.tensor([-7.0, 3.0, -7.0]) + target_storage = torch.full((3, ), -9.0) + source, target = source_storage[1].view(()), target_storage[1].view(()) + map_ = replicated_map((), 1) + plan, endpoints = plan_transfer(map_, map_), TransferEndpoints({0: 0}, {0: 0}) + if batched: + execute_transfers([TransferRequest('scalar', plan, {0: source}, {0: target})], endpoints, torch.float32, + StagingBudget(4)) + else: + execute_transfer(plan, {0: source}, {0: target}, endpoints, torch.float32, StagingBudget(4)) + assert torch.equal(target_storage, torch.tensor([-9.0, 3.0, -9.0])) + assert torch.equal(source_storage, torch.tensor([-7.0, 3.0, -7.0])) diff --git a/tests/unit/checkpoint/test_affine_transfer_executor_external_group.py b/tests/unit/checkpoint/test_affine_transfer_executor_external_group.py new file mode 100644 index 000000000000..1edf30666dec --- /dev/null +++ b/tests/unit/checkpoint/test_affine_transfer_executor_external_group.py @@ -0,0 +1,64 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""execute_transfer over a weight-update group built outside the default world. + +Trainer and rollout processes each own a default process group (here of world size 1, as a TP=1 engine or trainer +has) and join a separate two-process group created with ``_new_process_group_helper`` on its own store -- the way +RL frameworks build their trainer <-> engine weight-update group. In that group torch's rank bookkeeping maps every +member's default-world rank 0 to group rank 0, so the executor must take the rank from the group itself. +""" +import datetime +import os +import tempfile + +import pytest +import torch +import torch.multiprocessing as mp + +from deepspeed.accelerator import get_accelerator + + +def _worker(rank, port, tmp): + import torch.distributed as tdist + from torch.distributed.distributed_c10d import PrefixStore, _new_process_group_helper, _world + + tdist.init_process_group('gloo', init_method=f'file://{tmp}/default-{rank}', rank=0, world_size=1) + import deepspeed + deepspeed.init_distributed(dist_backend='gloo') + store = tdist.TCPStore('127.0.0.1', port, 2, rank == 0, timeout=datetime.timedelta(seconds=60)) + group, _ = _new_process_group_helper(2, + rank, [], + tdist.Backend('gloo'), + PrefixStore('weight-update', store), + group_name='weight-update', + backend_options=None, + timeout=datetime.timedelta(seconds=60)) + _world.pg_group_ranks[group] = {0: 0, 1: 1} + + from deepspeed.checkpoint.affine import contiguous_split_map + from deepspeed.checkpoint.affine_transfer import plan_transfer + from deepspeed.checkpoint.affine_transfer_executor import TransferEndpoints, execute_transfer + + full = torch.arange(96, dtype=torch.float32).reshape(8, 12) + source = contiguous_split_map((8, 12), [4, 4], 0) + target = contiguous_split_map((8, 12), [6, 6], 1) + plan = plan_transfer(target, source) + # Group rank 0 reads every source shard, group rank 1 writes every target shard. + endpoints = TransferEndpoints(source={loc: 0 + for loc in source.shard_shapes}, + target={loc: 1 + for loc in target.shard_shapes}) + sources = {loc: source.extract(full, loc).contiguous() for loc in source.shard_shapes} if rank == 0 else {} + targets = {loc: torch.full(shape, -1.0) for loc, shape in target.shard_shapes.items()} if rank == 1 else {} + execute_transfer(plan, sources, targets, endpoints, torch.float32, group=group) + for loc, shard in targets.items(): + assert torch.equal(shard, target.extract(full, loc)), loc + + +@pytest.mark.skipif(get_accelerator().device_name() != 'cpu', + reason='gloo groups over host buffers; CPU accelerator only') +def test_transfer_over_a_weight_update_group_outside_the_default_world(): + with tempfile.TemporaryDirectory() as tmp: + mp.spawn(_worker, args=(29000 + os.getpid() % 1000, tmp), nprocs=2, join=True)