From 4a4bee9ed55e2e0afb0fe30c85ccc750d72bda9d Mon Sep 17 00:00:00 2001 From: 0z5a Date: Mon, 12 Oct 2026 15:40:35 +0200 Subject: [PATCH 1/3] [checkpoint] Execute affine transfers with bounded peer staging Signed-off-by: 0z5a --- .../checkpoint/affine_transfer_executor.py | 503 ++++++++++++++++++ docs/code-docs/source/autoep.rst | 26 + .../test_affine_transfer_executor.py | 443 +++++++++++++++ 3 files changed, 972 insertions(+) create mode 100644 deepspeed/checkpoint/affine_transfer_executor.py create mode 100644 tests/unit/checkpoint/test_affine_transfer_executor.py diff --git a/deepspeed/checkpoint/affine_transfer_executor.py b/deepspeed/checkpoint/affine_transfer_executor.py new file mode 100644 index 000000000000..5bedf3f2f9e4 --- /dev/null +++ b/deepspeed/checkpoint/affine_transfer_executor.py @@ -0,0 +1,503 @@ +# 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. + +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.""" + 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() + 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') + 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 device.type == 'cuda' and callable(block_stream): + # Order arena reuse on the caller stream without a host polling delay per round. + 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 device.type != 'cpu': + 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/autoep.rst b/docs/code-docs/source/autoep.rst index 2dff03ed742a..92ef52faba9e 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -638,3 +638,29 @@ or ``None`` for automatic sizing, even when the selected backend does not use row blocking. Invalid settings are rejected before a backend runs. A model with a non-writable stock ``loss_function`` raises an actionable error rather than being silently left unmodified. + +Executing affine transfers +------------------------- + +A frozen shard can be transferred into an existing packed target without +reconstructing the full logical tensor. 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``. 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..4e02c6ea9e5e --- /dev/null +++ b/tests/unit/checkpoint/test_affine_transfer_executor.py @@ -0,0 +1,443 @@ +# 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 + + +@pytest.mark.skipif(get_accelerator().device_name() != 'cuda' or get_accelerator().device_count() < 2, + reason='requires two visible CUDA devices before distributed worker launch') +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 get_accelerator().device_name() == 'cuda': + 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) + + +@pytest.mark.skipif(get_accelerator().device_name() != 'cuda' or get_accelerator().device_count() < 8, + reason='requires eight visible CUDA devices before distributed worker launch') +class TestB2DisjointEndpoints(DistributedTest): + world_size = 8 + + def test_source_and_target_ranks_do_not_overlap(self): + _join_transfer_group() + source, target = _row_to_column(4, tp_size=4) + dtype = torch.float32 + full = _full(32, 32, dtype) + endpoints = TransferEndpoints(source={location: location + for location in source.shard_shapes}, + target={location: location + 4 + 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) + + assert stats.local_copy_bytes == 0, 'no process holds both ends, so nothing may be copied locally' + assert stats.received_bytes > 0 + _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])) From ba9e8b65219efdde2833abda3d66ceb03f907d6d Mon Sep 17 00:00:00 2001 From: 0z5a Date: Tue, 20 Oct 2026 10:58:36 +0200 Subject: [PATCH 2/3] [checkpoint] Gate affine transfer completion on device capability Address review on the affine transfer executor: - Decide stream-ordered completion with is_synchronized_device() instead of the device type, for both the per-round handle and the final stream synchronization. - Run the multi-rank executor tests on any accelerator. The disjoint-endpoint case now uses two or four ranks, so it fits the two-device PR runners and the four-rank CPU nightly. - Move the executor documentation from the AutoEP page to a new Checkpointing API page, and record that the staging arena and StagingBudget belong to the current torch point-to-point transport. Co-Authored-By: Claude Opus 5.5 Signed-off-by: 0z5a --- .../checkpoint/affine_transfer_executor.py | 22 +++++++-- docs/code-docs/source/affine-transfer.rst | 48 +++++++++++++++++++ docs/code-docs/source/autoep.rst | 26 ---------- docs/code-docs/source/index.rst | 1 + .../test_affine_transfer_executor.py | 31 +++++++----- 5 files changed, 85 insertions(+), 43 deletions(-) create mode 100644 docs/code-docs/source/affine-transfer.rst diff --git a/deepspeed/checkpoint/affine_transfer_executor.py b/deepspeed/checkpoint/affine_transfer_executor.py index 5bedf3f2f9e4..ebcd8ad2acf8 100644 --- a/deepspeed/checkpoint/affine_transfer_executor.py +++ b/deepspeed/checkpoint/affine_transfer_executor.py @@ -28,6 +28,15 @@ 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. @@ -64,7 +73,10 @@ class TransferEndpoints: @dataclass(frozen=True) class StagingBudget: - """Bytes a rank may stage per peer, in each direction, at one time.""" + """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 @@ -307,6 +319,8 @@ def execute_transfers(requests: Sequence[TransferRequest], 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 = {}, {} @@ -436,8 +450,8 @@ def execute_transfers(requests: Sequence[TransferRequest], if operations: for handle in dist.batch_isend_irecv(operations): block_stream = getattr(handle, 'block_current_stream', None) - if device.type == 'cuda' and callable(block_stream): - # Order arena reuse on the caller stream without a host polling delay per round. + 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: @@ -467,7 +481,7 @@ def execute_transfers(requests: Sequence[TransferRequest], _region(source, operation.source_offset, operation.chunk_shape, operation.segment.source_strides)) copied += operation.payload_bytes - if device.type != 'cpu': + 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) 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/autoep.rst b/docs/code-docs/source/autoep.rst index 92ef52faba9e..2dff03ed742a 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -638,29 +638,3 @@ or ``None`` for automatic sizing, even when the selected backend does not use row blocking. Invalid settings are rejected before a backend runs. A model with a non-writable stock ``loss_function`` raises an actionable error rather than being silently left unmodified. - -Executing affine transfers -------------------------- - -A frozen shard can be transferred into an existing packed target without -reconstructing the full logical tensor. 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``. 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 index 4e02c6ea9e5e..632988d5ed06 100644 --- a/tests/unit/checkpoint/test_affine_transfer_executor.py +++ b/tests/unit/checkpoint/test_affine_transfer_executor.py @@ -63,8 +63,6 @@ def _row_to_column(world_size, tp_size=2): return source, target -@pytest.mark.skipif(get_accelerator().device_name() != 'cuda' or get_accelerator().device_count() < 2, - reason='requires two visible CUDA devices before distributed worker launch') class TestB2Transfer(DistributedTest): world_size = 2 @@ -198,7 +196,7 @@ def test_mixed_packed_and_strided_regions_preserve_storage_windows(self): 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 get_accelerator().device_name() == 'cuda': + if not get_accelerator().is_synchronized_device(): stream = get_accelerator().Stream() stream.wait_stream(get_accelerator().current_stream()) with get_accelerator().stream(stream): @@ -229,19 +227,24 @@ def test_batch_parameter_order_must_match(self): execute_transfers(requests, endpoints, torch.float32) -@pytest.mark.skipif(get_accelerator().device_name() != 'cuda' or get_accelerator().device_count() < 8, - reason='requires eight visible CUDA devices before distributed worker launch') class TestB2DisjointEndpoints(DistributedTest): - world_size = 8 + """Every process is either a reader or a writer, never both. - def test_source_and_target_ranks_do_not_overlap(self): + 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(4, tp_size=4) + source, target = _row_to_column(2) dtype = torch.float32 - full = _full(32, 32, dtype) - endpoints = TransferEndpoints(source={location: location + 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: location + 4 + target={location: half + location % half for location in target.shard_shapes}) source_buffers = { location: source.extract(full, location).contiguous() @@ -252,8 +255,10 @@ def test_source_and_target_ranks_do_not_overlap(self): stats = execute_transfer(plan_transfer(target, source, PlanBudget()), source_buffers, target_buffers, endpoints, dtype) - assert stats.local_copy_bytes == 0, 'no process holds both ends, so nothing may be copied locally' - assert stats.received_bytes > 0 + # 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) From 99aef488416ecf3648be418ae133de7358263370 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Mon, 26 Oct 2026 18:13:02 +0100 Subject: [PATCH 3/3] [checkpoint] Take the transport rank from the group it names RL frameworks build their trainer <-> engine weight-update group outside the default world: each process keeps its own default group (world size 1 for a TP=1 trainer or engine) and joins a separate group created with _new_process_group_helper on its own store, as slime, SGLang and vLLM do. torch's rank bookkeeping for such a group maps every member's default-world rank 0 to group rank 0, so dist.get_rank(group) returned 0 on both ends and the preflight refused the whole batch ("target location 0 is not owned by this process"). When a group is passed, read the rank and size from the group itself. Every other path is unchanged. Add a CPU reproduction: two processes, each with a world-size-1 default group, joined by an externally created gloo group. It fails before this change and passes after it. The new test file builds that group with torch.distributed directly, so it joins the existing check-torchdist exclusion next to tests/unit/comm/test_dist.py. Co-Authored-By: Claude Opus 5.5 Signed-off-by: 0z5a --- .pre-commit-config.yaml | 2 +- .../checkpoint/affine_transfer_executor.py | 9 ++- ...affine_transfer_executor_external_group.py | 64 +++++++++++++++++++ 3 files changed, 72 insertions(+), 3 deletions(-) create mode 100644 tests/unit/checkpoint/test_affine_transfer_executor_external_group.py 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 index ebcd8ad2acf8..a9d769aae731 100644 --- a/deepspeed/checkpoint/affine_transfer_executor.py +++ b/deepspeed/checkpoint/affine_transfer_executor.py @@ -316,8 +316,13 @@ def execute_transfers(requests: Sequence[TransferRequest], import deepspeed.comm as dist initialized = dist.is_initialized() - rank = dist.get_rank(group=group) if initialized else 0 - world_size = dist.get_world_size(group=group) if initialized else 1 + 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() 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)