From 6450ca1d3aacc4cb509e7b0004c1b8f26954706d Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Mon, 2 Nov 2026 10:17:27 +0000 Subject: [PATCH 1/2] [FEAT] Extend affine live-state transfer to CPU-offloaded optimizer storage (#8905) --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 13e79bacce4a..29df5b2e6562 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,4 @@ tests/unit/saved_checkpoint/ # virtual env directory for format venv +.venv/ From 9c6e548d6eefb4b354de281a0aeb172bba975bd8 Mon Sep 17 00:00:00 2001 From: Arjun Pakhan Date: Wed, 4 Nov 2026 12:21:00 +0000 Subject: [PATCH 2/2] [FEAT] Extend affine live-state transfer to CPU-offloaded optimizer storage (#8905) --- .gitignore | 1 + deepspeed/checkpoint/affine.py | 221 +++++- deepspeed/checkpoint/affine_ir_spec.md | 73 ++ deepspeed/checkpoint/affine_transfer.py | 16 +- .../checkpoint/affine_transfer_executor.py | 589 ++++++++++++++++ deepspeed/checkpoint/zero_affine.py | 340 +++++++++ deepspeed/runtime/base_optimizer.py | 51 ++ deepspeed/runtime/zero/stage3.py | 11 + docs/code-docs/source/affine-transfer.rst | 54 ++ docs/code-docs/source/index.rst | 1 + .../test_affine_transfer_executor.py | 666 ++++++++++++++++++ .../checkpoint/test_affine_transfer_plan.py | 37 +- .../checkpoint/test_convert_checkpoint.py | 4 +- .../checkpoint/test_universal_checkpoint.py | 194 ++++- tests/unit/checkpoint/test_zero_affine.py | 598 ++++++++++++++++ 15 files changed, 2844 insertions(+), 12 deletions(-) create mode 100644 deepspeed/checkpoint/affine_transfer_executor.py create mode 100644 deepspeed/checkpoint/zero_affine.py create mode 100644 docs/code-docs/source/affine-transfer.rst create mode 100644 tests/unit/checkpoint/test_affine_transfer_executor.py create mode 100644 tests/unit/checkpoint/test_zero_affine.py diff --git a/.gitignore b/.gitignore index 29df5b2e6562..49d34c817f76 100644 --- a/.gitignore +++ b/.gitignore @@ -67,3 +67,4 @@ tests/unit/saved_checkpoint/ # virtual env directory for format venv .venv/ +.venv/ diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 312076fb8e93..6d886dca7b3e 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -22,7 +22,7 @@ __all__ = [ 'AffinePiece', 'ParamAffineMap', 'AFFINE_MAP_FORMAT_VERSION', 'SCALE_POWER_BY_STATE', 'replicated_map', - 'contiguous_split_map', 'sub_param_map', 'segmented_map', 'block_gather_map' + 'contiguous_split_map', 'sub_param_map', 'segmented_map', 'block_gather_map', 'nest_affine_maps' ] # Encoding version of the stored map, independent of the universal checkpoint version so @@ -634,3 +634,222 @@ def block_gather_map(shape, block_ids_by_rank, block_size, partition_dim): pieces_by_rank[rank] = pieces return ParamAffineMap(logical_shape=shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) + + +def nest_affine_maps(outer, inner_by_outer_rank): + """Compose a layout map with a description of how each of its shards is stored. + + ``outer`` describes the parameter's layout: the shard each layout rank (a tensor- or + expert-parallel rank) holds. ``inner_by_outer_rank[r]`` describes how layout shard ``r``, + flattened to ``(numel,)``, is spread over the processes that store it, as + ``deepspeed.checkpoint.zero_affine.zero_storage_affine_map`` does for ZeRO replicas and + partitions. Its pieces must be contiguous intervals. + + The result describes the parameter in its own shape, keyed by the inner ranks, and each shard is + what a process stores: a flat run of its layout shard. A stored run is generally not a box of the + parameter, since a ZeRO partition may begin or end inside a row, so it is described by a few + boxes, each contiguous in storage: a partial row at either end and the whole rows between. A + block that is not contiguous in its layout shard, such as one split along a trailing axis inside + the shard, is cut further until each box is. An element listed on several layout ranks is held by + the union of their holders, so pieces are cut wherever any of those holders' storage begins or + ends. + """ + if set(inner_by_outer_rank) != set(outer.shard_shapes): + raise ValueError(f'Storage is described for layout ranks {sorted(inner_by_outer_rank)}, but the layout has ' + f'ranks {sorted(outer.shard_shapes)}.') + storage_shapes = {} + for outer_rank, inner in inner_by_outer_rank.items(): + shard_numel = _product(outer.shard_shapes[outer_rank]) + if inner.logical_shape != (shard_numel, ): + raise ValueError(f'Layout rank {outer_rank} holds {shard_numel} elements, but its storage map describes ' + f'shape {inner.logical_shape}.') + for storage_rank, shape in inner.shard_shapes.items(): + if storage_rank in storage_shapes: + # One shard per process: a process storing two layout shards would need two buffers. + raise ValueError(f'Rank {storage_rank} stores more than one layout shard.') + storage_shapes[storage_rank] = shape + intervals_by_outer_rank = {rank: _storage_intervals(inner) for rank, inner in inner_by_outer_rank.items()} + + # Replicated layout pieces describe the same block on each of their holders. + replicas = {} + for outer_rank, pieces in outer.pieces_by_rank.items(): + for piece in pieces: + if piece.numel == 0: + continue + key = (piece.shape, piece.source_offset, piece.source_strides) + replicas.setdefault(key, []).append((outer_rank, piece)) + + pieces_by_rank = {rank: [] for rank in storage_shapes} + for copies in replicas.values(): + listed = {outer_rank for outer_rank, _ in copies} + for _, piece in copies: + if piece.locations != listed: + raise ValueError(f'Layout piece {piece} lists holders {sorted(piece.locations)}, but layout ranks ' + f'{sorted(listed)} hold it.') + _nest_replicated_block(copies, intervals_by_outer_rank, pieces_by_rank) + + affine_map = ParamAffineMap(logical_shape=outer.logical_shape, + shard_shapes=storage_shapes, + pieces_by_rank=pieces_by_rank) + affine_map.validate() + return affine_map + + +def _storage_intervals(inner): + # (start, stop, storage rank, offset in that rank's shard, scale) for every stored interval. + intervals = [] + for storage_rank, pieces in inner.pieces_by_rank.items(): + for piece in pieces: + if piece.numel == 0: + continue + if piece.source_strides != (1, ) or piece.dest_strides != (1, ): + raise ValueError(f'Storage piece {piece} of rank {storage_rank} is not a contiguous interval.') + intervals.append( + (piece.source_offset, piece.source_offset + piece.numel, storage_rank, piece.dest_offset, piece.scale)) + return sorted(intervals) + + +def _nest_replicated_block(copies, intervals_by_outer_rank, pieces_by_rank): + # Cuts are row-major element indices of the block. Between two cuts, every copy lies inside one + # stored interval of every holder; each range between cuts is then split into boxes that are + # contiguous in every copy's shard. + block = copies[0][1] + cuts = {0, block.numel} + for outer_rank, piece in copies: + _check_monotone_dest(piece) + for start, stop, _, _, _ in intervals_by_outer_rank[outer_rank]: + for boundary in (start, stop): + cuts.add(_first_index_at_or_after(piece, boundary)) + cuts = sorted(cuts) + + for low, high in zip(cuts, cuts[1:]): + if low == high: + continue + stored = [] + for outer_rank, piece in copies: + address = _strided_address(piece.shape, piece.dest_offset, piece.dest_strides, low) + covering = [ + interval for interval in intervals_by_outer_rank[outer_rank] if interval[0] <= address < interval[1] + ] + if not covering: + raise ValueError(f'Layout rank {outer_rank} holds shard element {address}, but its storage map ' + 'stores it on no rank.') + stored.extend((piece, address, interval) for interval in covering) + locations = {interval[2] for _, _, interval in stored} + # A box must be one run of addresses in every copy, so the most fragmented copy decides. + dense_rows = max(_dense_rows(piece) for piece, _, _ in stored) + for origin, extent in _row_major_range_boxes(block.shape, low, high, dense_rows): + source_offset = block.source_offset + _dot(origin, block.source_strides) + for piece, _, (start, _, storage_rank, storage_offset, storage_scale) in stored: + address = piece.dest_offset + _dot(origin, piece.dest_strides) + pieces_by_rank[storage_rank].append( + AffinePiece(shape=extent, + source_offset=source_offset, + source_strides=block.source_strides, + dest_offset=storage_offset + address - start, + dest_strides=_row_major_strides(extent), + locations=locations, + scale=piece.scale * storage_scale)) + + +def _dot(index, strides): + return sum(position * stride for position, stride in zip(index, strides)) + + +def _dense_rows(piece): + # The number of leading axes after which the piece is contiguous in its shard: a box of the piece + # that spans every later axis in full, and only those, is one run of shard addresses. + run = 1 + for axis in range(len(piece.shape) - 1, -1, -1): + size, stride = piece.shape[axis], piece.dest_strides[axis] + if size > 1 and stride != run: + return axis + 1 + run *= size + return 0 + + +def _row_major_range_boxes(shape, low, high, dense_rows=0): + """Split the row-major index range ``[low, high)`` of ``shape`` into ``(origin, extent)`` boxes. + + At most two boxes per axis: a partial run at either end and whole runs between. The first + ``dense_rows`` axes are never merged, so a box spans more than one position of such an axis only + where the storage is contiguous across it. + """ + if low >= high: + return [] + if not shape: + return [((), ())] + if len(shape) == 1: + return [((low, ), (high - low, ))] + inner = _product(shape[1:]) + first, low_rest = divmod(low, inner) + last, high_rest = divmod(high, inner) + + def within(row, start, stop): + return [((row, ) + origin, (1, ) + extent) + for origin, extent in _row_major_range_boxes(shape[1:], start, stop, dense_rows - 1)] + + if first == last: + return within(first, low_rest, high_rest) + boxes = [] + if low_rest: + boxes.extend(within(first, low_rest, inner)) + first += 1 + if dense_rows > 0: + for row in range(first, last): + boxes.extend(within(row, 0, inner)) + elif last > first: + boxes.append(((first, ) + (0, ) * (len(shape) - 1), (last - first, ) + tuple(shape[1:]))) + if high_rest: + boxes.extend(within(last, 0, high_rest)) + return boxes + + +def _check_monotone_dest(piece): + # Storage intervals are ranges of shard addresses. A piece whose addresses grow with its + # row-major element index meets such a range in one run of indices; any other piece could + # interleave with it. + later_span = 0 + for size, stride in reversed(list(zip(piece.shape, piece.dest_strides))): + if size > 1: + if stride <= later_span: + raise NotImplementedError(f'Layout piece {piece} does not store its elements in row-major order, so ' + 'its storage cannot be described.') + later_span += stride * (size - 1) + + +def _strided_address(shape, offset, strides, index): + address = offset + for size, stride in reversed(list(zip(shape, strides))): + index, position = divmod(index, size) + address += position * stride + return address + + +def _first_index_at_or_after(piece, address): + # The smallest row-major index whose shard address is at least ``address``; ``numel`` if none. + remaining = address - piece.dest_offset + if remaining <= 0: + return 0 + spans = [] + later_span = 0 + for size, stride in reversed(list(zip(piece.shape, piece.dest_strides))): + spans.append(later_span) + later_span += stride * (size - 1) + spans.reverse() + if remaining > later_span: + return piece.numel + + index = 0 + for axis, (size, stride) in enumerate(zip(piece.shape, piece.dest_strides)): + position = 0 + if size > 1: + # The first position from which the later axes can still reach ``remaining``. + position = max(0, -(-(remaining - spans[axis]) // stride)) + index = index * size + position + remaining -= position * stride + if remaining <= 0: + for later_size in piece.shape[axis + 1:]: + index *= later_size + return index + return index diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index a2a8eff321dc..38ad395fa3a1 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -608,6 +608,22 @@ subset of a parameter combined with a list of ranks holding this subset" — is §2.1's exact `locations` implements, and the same mechanism can later describe normalized ZeRO/offload placement without introducing another geometry IR. +`deepspeed/checkpoint/zero_affine.py` takes the first step for live training storage. It +describes, in global ranks, where ZeRO keeps each element of a parameter: the low-precision +copy (`zero_param_storage_affine_map`: replicas under ZeRO-0/1/2, `ds_tensor` partitions +under ZeRO-3, padding excluded), and the fp32 master copy with its optimizer states +(`zero_optimizer_storage_affine_maps`: ZeRO-1/2 fragments gathered from each rank's +`_hp_mapping`, which may cut a parameter mid-row, and the ZeRO-3 fp32 partition). The +matching `*_storage_shard` accessors return the flat views of ZeRO's own buffers, so a +plan's writes land in the optimizer. Offloaded states, MiCS and hpZ have no map yet. +Both parameter storage accessors reject quantized ZeRO-3 storage (`ds_quant_scale`): +its encoded values and separate quantization scales cannot be treated as ordinary +parameter elements. They reject NVMe-offloaded ZeRO-3 storage too, whose partition is a +placeholder while swapped out. `zero_param_storage_shard` also rejects a gathered ZeRO-3 +parameter: the forward reads the gathered copy, which persistent parameters keep between +steps, so a write to the partition alone would be lost. `engine.empty_partition_cache()` +releases every gathered copy first. + --- ## 9. Staging @@ -623,3 +639,60 @@ ZeRO/offload placement without introducing another geometry IR. Steps 1–3 change no behaviour: they add a representation and prove it agrees with the current one. Only step 4 changes what a checkpoint contains, and only by adding a key. + +--- + +## 10. Phase 2 — live storage maps + +Phase 2 plans `M_t ∘ M_s⁻¹` between two live topologies. A live topology is what each process +actually holds, so its map must describe storage, not only the parallel layout. + +**10.1 Two kinds of map.** +- *Layout map*: the parameter's N-D logical shape, keyed by ranks of the parallel group that + defines the layout (TP, EP). The layer builds it from its own semantics, checkpoints persist + it, and it remains the only geometry contract. +- *Storage map*: keyed by global rank; each shard is a process's flat buffer (`param.data`, + ZeRO-3 `ds_tensor`, an fp32 or optimizer-state fragment). It exists at runtime only and is + never persisted. Built from a layout map (10.3), it keeps the layout's N-D logical shape; + without a layout it describes the flattened parameter `(numel,)`. + +Both satisfy P1–P5. A plan's endpoints must be the same kind and the same logical shape; +`plan_transfer` rejects different logical shapes, so a flat storage map cannot be planned +against an N-D one of the same parameter. Matching rank namespaces is the caller's job. + +**10.2 N-D boxes in flat storage.** A ZeRO partition of an N-D layout shard is a contiguous +interval of its flattened buffer, which is generally not one box of the N-D parameter: it is +a partial row, then whole rows, then a partial row, recursively over the axes, so at most +`2 * ndim - 1` boxes. Each box is dense in the buffer (its storage strides are its own +row-major strides). `plan_transfer` accepts such a piece in a 1-D shard and treats its +storage as the interval `[dest_offset, dest_offset + numel)`; any other shard must have the +logical parameter's rank. Describing the flattened parameter instead would cut a column of a +row-major weight, such as a row-parallel shard, into one piece per row. + +**10.3 Derivation.** A storage map is derived, never hand-built: +`storage = nest_affine_maps(layout, {layout rank: inner})`, where `inner` describes how one +flattened layout shard is stored over its data-parallel group (`zero_storage_affine_map`: +replicas, or ZeRO-3 partitions `[d*n, (d+1)*n)` without padding). Nesting keeps P5: layout +pieces that describe the same block on several layout ranks are one block whose holders are +the union of every rank's holders, and the block is cut wherever any of those holders' +storage begins or ends; each resulting interval becomes the boxes of 10.2. A layout piece +that is not dense within its layout shard, such as a block narrower than the shard's rows, +cannot be merged across those rows and keeps one piece per row. Callers: +- `zero_param_storage_affine_map` and `zero_optimizer_storage_affine_maps` describe a dense + parameter (no layout, so they are already the storage map). +- `autotp_storage_affine_maps(model)` nests each parameter's published AutoTP layout map + with the storage of every TP shard. A parameter AutoTP leaves whole in a model built under + `zero.Init` stays partitioned over the whole world, so every TP rank reports the storage of + one shared copy; that storage describes the parameter on its own, without nesting. +- AutoEP expert storage nests the placement's layout map with each EP rank's EDP storage. + +**10.4 Replica choice** stays with the planner (§2.4): every holder appears in `locations`, +and the planner's default or a `replica_selector` picks one. + +**10.5 Plan lifetime.** Maps and plans depend only on layout and partitioning, not on values. +A caller builds them once and executes the plan for every sync, until either side's +placement, parallel sizes or ZeRO partitioning changes. A parameter whose buffer is +reallocated keeps its plan; only the shard views must be fetched again. + +**10.6 Not described.** Offload tiers, MiCS / hpZ, expert tensor parallelism, quantized +ZeRO-3 storage, non-elementwise transforms. diff --git a/deepspeed/checkpoint/affine_transfer.py b/deepspeed/checkpoint/affine_transfer.py index 87452bd6ec20..e38ecb819542 100644 --- a/deepspeed/checkpoint/affine_transfer.py +++ b/deepspeed/checkpoint/affine_transfer.py @@ -246,7 +246,8 @@ def _frame_of(piece: AffinePiece, rank: int, side: str, logical_shape: Tuple[int Both stride vectors must be the row-major strides of the tensor they address, which is what makes the box arithmetic below valid and what every builder in ``affine.py`` already emits. An axis the piece does not span is absent from storage, so its stride is not compared -- that is also where a - permutation would hide, and a permutation is not one box. + permutation would hide, and a permutation is not one box. A flat shard is the one exception: it may + hold a multi-axis box as a dense run, which is how ZeRO stores a tensor-parallel shard. """ where = f'{side} rank {rank} piece {piece!r}' if len(piece.shape) != len(logical_shape): @@ -271,6 +272,11 @@ def _frame_of(piece: AffinePiece, rank: int, side: str, logical_shape: Tuple[int for axis, (got, want, size) in enumerate(zip(piece.source_strides, logical_strides, piece.shape)): if size > 1 and got != want: raise UnsupportedLayoutError(f'{where}: source axis {axis} strides {got}, expected {want}') + flat_storage = len(shard_shape) == 1 and len(piece.shape) != 1 + if flat_storage: + # A flat shard, as ZeRO stores one, holds a box of a multi-axis parameter as one run of + # addresses: the box's own row-major strides. Its place in the shard is that run. + shard_strides = _row_major_strides(tuple(piece.shape)) for axis, (got, want, size) in enumerate(zip(piece.dest_strides, shard_strides, piece.shape)): if size > 1 and got != want: raise UnsupportedLayoutError(f'{where}: dest axis {axis} strides {got}, expected {want}') @@ -279,7 +285,10 @@ def _frame_of(piece: AffinePiece, rank: int, side: str, logical_shape: Tuple[int box = LogicalBox(origin, tuple(piece.shape)) if not LogicalBox(tuple([0] * len(logical_shape)), logical_shape).contains(box): raise InvalidMapError(f'{where}: covers {box!r}, outside the logical parameter of {logical_shape}') - local = LogicalBox(_origin(piece.dest_offset, shard_strides, shard_shape, where, 'dest'), tuple(piece.shape)) + if flat_storage: + local = LogicalBox((piece.dest_offset, ), (piece.numel, )) + else: + local = LogicalBox(_origin(piece.dest_offset, shard_strides, shard_shape, where, 'dest'), tuple(piece.shape)) if not LogicalBox(tuple([0] * len(shard_shape)), shard_shape).contains(local): raise InvalidMapError(f'{where}: writes {local!r}, outside the {shard_shape} shard') return _Frame(rank=rank, @@ -302,7 +311,8 @@ def _normalized(map_: ParamAffineMap, side: str, limits: PlanBudget) -> List[_Fr logical_shape = tuple(map_.logical_shape) if len(logical_shape) > limits.max_ndim: raise UnsupportedLayoutError(f'{side}: {len(logical_shape)} axes exceeds max_ndim {limits.max_ndim}') - if any(len(tuple(shape)) != len(logical_shape) for shape in map_.shard_shapes.values()): + # A flat shard is flat storage, such as ZeRO's; its pieces carry the axis correspondence. + if any(len(tuple(shape)) not in (1, len(logical_shape)) for shape in map_.shard_shapes.values()): raise UnsupportedLayoutError( f'{side}: a shard is described with a different number of axes than the logical parameter, so ' 'no axis correspondence exists to plan against') diff --git a/deepspeed/checkpoint/affine_transfer_executor.py b/deepspeed/checkpoint/affine_transfer_executor.py new file mode 100644 index 000000000000..bed36c60adc3 --- /dev/null +++ b/deepspeed/checkpoint/affine_transfer_executor.py @@ -0,0 +1,589 @@ +# 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. + +Buffers may reside in host CPU memory (e.g. for CPU-offloaded optimizer states) or on the accelerator. +Transfers crossing host and accelerator devices move through bounded pinned-memory staging arenas +bounded by the staging budget, with explicit event ordering and stream synchronization to ensure wire +payloads over NCCL remain on accelerator devices without whole-model allocations. + +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, transport_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.type != 'cpu' and buffer.device != transport_device: + raise AffineTransferExecutionError( + f'{which} location {location}: buffer is on {buffer.device}, transport requires {transport_device} or cpu' + ) + 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 target_dev, low, high in targets: + if any(target_dev == source_dev and low < source_high and source_low < high + for source_dev, source_low, source_high in sources): + raise AffineTransferExecutionError('a source and a target overlap within the transfer batch') + for index, (target_dev, low, high) in enumerate(targets): + if any(target_dev == other_dev and low < other_high and other_low < high + for other_dev, 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, transport_device): + 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 + if flat.device != transport_device: + return None + 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 + accelerator = get_accelerator() + device = torch.device(accelerator.current_device_name()) if accelerator.is_available() else torch.device('cpu') + stream_ordered = accelerator.is_available() and not 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, device) for entries, _ in rounds] + for peer, rounds in recv_rounds.items() + } + send_direct = { + peer: [_direct_payload(entries, requests, dtype, True, device) 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]) + } + + has_cpu_source = any(b.device.type == 'cpu' for req in requests for b in req.source_buffers.values()) + has_cpu_target = any(b.device.type == 'cpu' for req in requests for b in req.target_buffers.values()) + has_acc_source = any(b.device == device for req in requests for b in req.source_buffers.values()) + has_acc_target = any(b.device == device for req in requests for b in req.target_buffers.values()) + + host_send_arena = None + host_recv_arena = None + local_device_arena = None + + if stream_ordered: + accelerator = get_accelerator() + if has_cpu_source: + host_send_arena = accelerator.pin_empty(max_elements, dtype=dtype, device='cpu') + if has_cpu_target: + host_recv_arena = accelerator.pin_empty(max_elements, dtype=dtype, device='cpu') + if (has_cpu_source and has_acc_target) or (has_acc_source and has_cpu_target): + local_device_arena = torch.empty(max_elements, dtype=dtype, device=device) + except (ValueError, TypeError, KeyError, AttributeError, RuntimeError, OverflowError, MemoryError) as caught: + error = caught + rejected = 1 + recv_arenas, send_arenas = {}, {} + host_send_arena, host_recv_arena, local_device_arena = None, None, None + 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] + peer_has_cpu_entry = False + for item, operation, offset in entries: + source = _flat(operation.segment.source_rank, requests[item].source_buffers, dtype) + source_reg = _region(source, operation.source_offset, operation.chunk_shape, + operation.segment.source_strides) + if stream_ordered and source.device.type == 'cpu': + peer_has_cpu_entry = True + host_slice = host_send_arena[offset:offset + operation.elements].view(operation.chunk_shape) + host_slice.copy_(source_reg) + arena[offset:offset + operation.elements].view(operation.chunk_shape).copy_(host_slice, + non_blocking=True) + else: + arena[offset:offset + operation.elements].view(operation.chunk_shape).copy_(source_reg) + if peer_has_cpu_entry: + h2d_event = get_accelerator().Event() + h2d_event.record(get_accelerator().current_stream()) + h2d_event.synchronize() + 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) + target_reg = _region(target, operation.target_offset, operation.chunk_shape, + operation.segment.target_strides) + if stream_ordered and target.device.type == 'cpu': + host_slice = host_recv_arena[offset:offset + operation.elements].view(operation.chunk_shape) + host_slice.copy_(payload, non_blocking=True) + d2h_event = get_accelerator().Event() + d2h_event.record(get_accelerator().current_stream()) + d2h_event.synchronize() + target_reg.copy_(host_slice) + else: + target_reg.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) + source_reg = _region(source, operation.source_offset, operation.chunk_shape, operation.segment.source_strides) + target_reg = _region(target, operation.target_offset, operation.chunk_shape, operation.segment.target_strides) + + if stream_ordered and source.device.type == 'cpu' and target.device != torch.device('cpu'): + host_stage = host_send_arena[:operation.elements].view(operation.chunk_shape) + host_stage.copy_(source_reg) + dev_stage = local_device_arena[:operation.elements].view(operation.chunk_shape) + dev_stage.copy_(host_stage, non_blocking=True) + target_reg.copy_(dev_stage) + h2d_event = get_accelerator().Event() + h2d_event.record(get_accelerator().current_stream()) + h2d_event.synchronize() + elif stream_ordered and source.device != torch.device('cpu') and target.device.type == 'cpu': + dev_stage = local_device_arena[:operation.elements].view(operation.chunk_shape) + dev_stage.copy_(source_reg) + host_stage = host_recv_arena[:operation.elements].view(operation.chunk_shape) + host_stage.copy_(dev_stage, non_blocking=True) + d2h_event = get_accelerator().Event() + d2h_event.record(get_accelerator().current_stream()) + d2h_event.synchronize() + target_reg.copy_(host_stage) + else: + target_reg.copy_(source_reg) + 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[str, int, int]: + """Byte extent a tensor's elements occupy in storage, qualified by device.""" + dev_key = f'{buffer.device.type}:{buffer.device.index}' if buffer.device.index is not None else buffer.device.type + start = buffer.data_ptr() + span = buffer.numel() * buffer.element_size() + return dev_key, 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/deepspeed/checkpoint/zero_affine.py b/deepspeed/checkpoint/zero_affine.py new file mode 100644 index 000000000000..668a83f329d4 --- /dev/null +++ b/deepspeed/checkpoint/zero_affine.py @@ -0,0 +1,340 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Affine maps of the storage ZeRO keeps for a model parameter. + +A transfer planner needs every legal holder of every element. These maps describe a flattened +parameter keyed by global rank, both for the low-precision parameter the model computes with and +for the fp32 master copy and optimizer states that ZeRO partitions across the data-parallel group. +Each map comes with an accessor for the matching local shard, which is a view of ZeRO's storage, +so it can be read as a transfer source or written as a transfer target. + +Under AutoTP these maps describe one tensor-parallel shard; ``autotp_storage_affine_maps`` nests +them in each parameter's AutoTP layout to describe the whole parameter in its own shape. +""" + +import torch + +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.affine import AffinePiece, ParamAffineMap, nest_affine_maps, replicated_map + +__all__ = [ + 'zero_storage_affine_map', 'zero_param_storage_affine_map', 'zero_param_storage_shard', + 'zero_optimizer_storage_affine_maps', 'zero_optimizer_storage_shard', 'autotp_storage_affine_maps' +] + + +def _flat_piece(numel, start, rank): + return AffinePiece(shape=(numel, ), + source_offset=start, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=[rank]) + + +def _fragment_affine_map(numel, fragments): + # ``fragments`` lists (rank, start, length): each rank holds one contiguous range, possibly empty. + shard_shapes = {} + pieces_by_rank = {} + for rank, start, length in fragments: + shard_shapes[rank] = (length, ) + pieces_by_rank[rank] = [_flat_piece(length, start, rank)] if length > 0 else [] + affine_map = ParamAffineMap(logical_shape=(numel, ), shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) + affine_map.validate() + return affine_map + + +def _zero3_bounds(numel, partition_numel, index): + # The trailing padding of the last partitions holds no parameter elements, so it is not part of the shard. + start = min(index * partition_numel, numel) + stop = min(start + partition_numel, numel) + return start, stop - start + + +def zero_storage_affine_map(numel, holder_ranks, partition_numel=None): + """Describe a flattened parameter of ``numel`` elements stored by ``holder_ranks``. + + Without ``partition_numel`` every holder stores the whole parameter, as data-parallel + replicas do under ZeRO-0/1/2, so the one piece lists every holder. With it, holder ``i`` + stores elements ``[i * n, (i + 1) * n)`` as ZeRO-3 does; the trailing padding holds no + parameter elements and is not part of the shard, so a holder past the end has an empty shard. + """ + numel = int(numel) + holder_ranks = [int(rank) for rank in holder_ranks] + if numel < 0: + raise ValueError(f'Parameter numel cannot be negative: {numel}.') + if not holder_ranks or len(set(holder_ranks)) != len(holder_ranks): + raise ValueError(f'Holder ranks must be a non-empty list of distinct ranks: {holder_ranks}.') + + if partition_numel is None: + piece = AffinePiece(shape=(numel, ), + source_offset=0, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=holder_ranks) + affine_map = ParamAffineMap(logical_shape=(numel, ), + shard_shapes={rank: (numel, ) + for rank in holder_ranks}, + pieces_by_rank={rank: [piece] + for rank in holder_ranks}) + affine_map.validate() + return affine_map + + partition_numel = int(partition_numel) + if partition_numel < 0 or partition_numel * len(holder_ranks) < numel: + raise ValueError(f'{len(holder_ranks)} partitions of {partition_numel} elements cannot hold {numel} elements.') + fragments = [] + for index, rank in enumerate(holder_ranks): + start, length = _zero3_bounds(numel, partition_numel, index) + fragments.append((rank, start, length)) + return _fragment_affine_map(numel, fragments) + + +def _zero3_local_numel(param): + index = dist.get_rank(group=param.ds_process_group) if dist.is_initialized() else 0 + _, length = _zero3_bounds(param.ds_numel, param.ds_tensor.ds_numel, index) + return length + + +def _global_ranks(group): + # dist.get_all_ranks_from_group never ends on the world group: torch maps any rank of the + # world group to itself instead of raising past the end. + if not dist.is_initialized(): + return [0] + return [dist.get_global_rank(group, rank) for rank in range(dist.get_world_size(group))] + + +def _check_zero3_param_storage(param): + from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum + + if hasattr(param.ds_tensor, 'ds_quant_scale'): + raise NotImplementedError('Quantized ZeRO-3 parameter storage has no affine storage map.') + # NVMe offload swaps a partition out of memory and leaves a placeholder in its place. + if getattr(param.ds_tensor, 'final_location', None) == OffloadDeviceEnum.nvme: + raise NotImplementedError('NVMe-offloaded ZeRO-3 parameter storage has no affine storage map.') + + +def zero_param_storage_affine_map(param, replica_group=None): + """Describe the storage of a live low-precision ``param`` across global ranks. + + A ZeRO-3 parameter is described by its own partition group and ``ds_tensor`` size. + Otherwise every rank of ``replica_group`` holds a whole copy. It defaults to the group ZeRO + partitions the parameter's optimizer state over, which is the expert data-parallel group for + MoE experts, and to the data-parallel group for a parameter outside a ZeRO optimizer. Under + tensor parallelism the map describes this rank's tensor-parallel shard, which its + data-parallel peers hold too. + + Quantized and NVMe-offloaded ZeRO-3 parameter storage is not supported. + """ + if hasattr(param, 'ds_id'): + _check_zero3_param_storage(param) + holders = _global_ranks(param.ds_process_group) + return zero_storage_affine_map(param.ds_numel, holders, param.ds_tensor.ds_numel) + if replica_group is None: + replica_group = getattr(param, '_dp_group', None) + if replica_group is None: + from deepspeed.utils import groups + replica_group = groups._get_data_parallel_group() + return zero_storage_affine_map(param.numel(), _global_ranks(replica_group)) + + +def zero_param_storage_shard(param): + """This rank's shard of ``param`` as laid out by ``zero_param_storage_affine_map``: a flat view. + + A ZeRO-3 parameter must be partitioned. A gathered copy, which ZeRO-3 keeps for persistent + parameters between forwards, would go on serving the forward after the shard is written, so + release gathered copies with ``engine.empty_partition_cache()`` first. + """ + if hasattr(param, 'ds_id'): + from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + + _check_zero3_param_storage(param) + if param.ds_status != ZeroParamStatus.NOT_AVAILABLE: + raise RuntimeError(f'ZeRO-3 parameter {param.ds_id} is gathered, so its shard is not its only copy; ' + 'call engine.empty_partition_cache() first.') + return param.ds_tensor.data.view(-1).narrow(0, 0, _zero3_local_numel(param)) + return param.data.view(-1) + + +def _zero3_optimizer_offloaded(param): + return param._z3_optimizer.offload_optimizer + + +def zero_optimizer_storage_affine_maps(params): + """Describe where ZeRO keeps the fp32 master copy and optimizer states of each of ``params``. + + Returns one map per parameter, in order, or ``None`` for a parameter without an fp32 copy, + such as a frozen one. The master copy and every optimizer state of a parameter share its map. + Under ZeRO-1/2 (and the bf16 optimizer) each rank of the parameter's data-parallel group holds + one contiguous range of it, possibly empty, so the ranks exchange their ranges in one all-gather + per group: every rank must pass the same parameters in the same order. Under ZeRO-3 the + partitions follow the parameter's own, and no communication is needed. + + CPU-offloaded optimizer states are supported, and their shards are host CPU tensors. + """ + params = list(params) + maps = [None] * len(params) + if not dist.is_initialized(): + for position, param in enumerate(params): + if hasattr(param, 'ds_id'): + if not hasattr(param, '_z3_optimizer') or not param.requires_grad: + continue + if getattr(param._z3_optimizer, 'swap_optimizer', False): + raise NotImplementedError('NVMe-swapped ZeRO-3 optimizer states have no affine storage map.') + maps[position] = zero_param_storage_affine_map(param) + elif getattr(param, '_dp_group', None) is not None or hasattr(param, '_hp_mapping'): + if getattr(param, 'requires_grad', True): + maps[position] = zero_storage_affine_map(param.numel(), [0]) + return maps + + # ZeRO-1/2 parameters grouped by the group their ranges are exchanged over, in first-seen order + # so that every rank issues the all-gathers in the same order. + by_group = {} + for position, param in enumerate(params): + if hasattr(param, 'ds_id'): + # ZeRO-3 keeps no fp32 copy of a frozen parameter. + if not hasattr(param, '_z3_optimizer') or not param.requires_grad: + continue + if getattr(param._z3_optimizer, 'swap_optimizer', False): + raise NotImplementedError('NVMe-swapped ZeRO-3 optimizer states have no affine storage map.') + maps[position] = zero_param_storage_affine_map(param) + elif getattr(param, '_dp_group', None) is not None: + by_group.setdefault(param._dp_group, []).append(position) + + device = torch.device(get_accelerator().current_device_name()) + for group, positions in by_group.items(): + # Each rank reports (start, length, offloaded) for every parameter; a rank without a range of a + # parameter reports an empty one. Offload is reported rather than checked locally so that every + # rank raises together instead of leaving its peers in the all-gather. + local = torch.zeros(len(positions), 3, dtype=torch.int64, device=device) + for row, position in enumerate(positions): + mapping = params[position]._hp_mapping + if mapping is not None: + address = mapping.lp_fragment_address + local[row] = torch.tensor([address.start, address.numel, int(mapping.use_offload)]) + world_size = dist.get_world_size(group=group) + # Gloo wants the output stacked along the input's first axis, not along a new one. + gathered = torch.empty(world_size * len(positions), 3, dtype=torch.int64, device=device) + dist.all_gather_into_tensor(gathered, local, group=group) + gathered = gathered.view(world_size, len(positions), 3).cpu() + ranks = _global_ranks(group) + for row, position in enumerate(positions): + fragments = [(rank, int(gathered[index, row, 0]), int(gathered[index, row, 1])) + for index, rank in enumerate(ranks)] + maps[position] = _fragment_affine_map(params[position].numel(), fragments) + return maps + + +def zero_optimizer_storage_shard(param, optim_state_key=None): + """This rank's shard of the fp32 master copy of ``param``, or of its ``optim_state_key`` state. + + The shard is a flat view of ZeRO's storage, laid out as ``zero_optimizer_storage_affine_maps`` + describes it, so writing it updates the optimizer. Returns ``None`` where that has no map. + """ + if hasattr(param, 'ds_id'): + if not hasattr(param, '_z3_optimizer') or not param.requires_grad: + return None + _check_zero3_param_storage(param) + z3_opt = param._z3_optimizer + if getattr(z3_opt, 'swap_optimizer', False): + raise NotImplementedError('NVMe-swapped ZeRO-3 optimizer states have no affine storage map.') + group_idx, dest_offset, num_elements = z3_opt.grad_position[z3_opt.get_param_id(param)] + fp32_param = z3_opt.fp32_partitioned_groups_flat[group_idx] + if optim_state_key is not None: + if fp32_param not in z3_opt.optimizer.state or optim_state_key not in z3_opt.optimizer.state[fp32_param]: + raise RuntimeError('ZeRO links optimizer states to parameters at the first optimizer step; ' + f'call this after it to read {optim_state_key!r}.') + fp32_opt_state, _ = z3_opt._get_fp32_opt_state_partition(param, + release_swap_buffers=False, + optim_state_key=optim_state_key) + return fp32_opt_state.view(-1).narrow(0, 0, _zero3_local_numel(param)) + + if getattr(param, '_dp_group', None) is None: + return None + mapping = param._hp_mapping + if mapping is None: + # The rank holds no range of this parameter. + device = 'cpu' if getattr(param, '_is_offloaded', False) or (hasattr(param, '_zero_optimizer') and getattr( + param._zero_optimizer, 'cpu_offload', False)) else get_accelerator().current_device_name() + return torch.empty(0, dtype=torch.float32, device=device) + if optim_state_key is not None and mapping.optim_fragment is None: + raise RuntimeError('ZeRO links optimizer states to parameters at the first optimizer step; ' + f'call this after it to read {optim_state_key!r}.') + return mapping.get_hp_fragment(optim_state_key) + + +def _autotp_layout_map(name, param, tp_world_size): + from deepspeed.module_inject.layers import _get_param_uc_conversion_meta + + meta = _get_param_uc_conversion_meta(param) + if not meta: + # AutoTP leaves such a parameter whole on every tensor-parallel rank. + return replicated_map(tuple(getattr(param, 'ds_shape', param.shape)), tp_world_size) + if meta.get('unsupported_reason'): + raise NotImplementedError(f"AutoTP parameter '{name}' has no affine layout: {meta['unsupported_reason']}") + if meta.get('affine_map') is None: + raise NotImplementedError(f"AutoTP parameter '{name}' publishes no affine layout map.") + return ParamAffineMap.from_dict(meta['affine_map']) + + +def autotp_storage_affine_maps(model, optimizer_states=False, tp_group=None): + """Describe what every process stores of each parameter of an AutoTP ``model`` under ZeRO. + + Returns ``{name: map}`` over ``model.named_parameters()``. Each map nests the parameter's + AutoTP layout map with the ZeRO storage of every tensor-parallel shard, so it describes the + parameter in its own shape keyed by global rank, and each rank's shard is the one + ``zero_param_storage_shard`` returns. With ``optimizer_states`` the maps describe the fp32 + master copy and optimizer states instead, matching ``zero_optimizer_storage_shard``, and a + parameter without them maps to ``None``. A parameter AutoTP did not partition is replicated + over the tensor-parallel group, unless it kept the partition over the whole world that + ``zero.Init`` gave it, which describes it on its own. + + Collective: the tensor-parallel ranks exchange their shards' storage maps in one all-gather + over ``tp_group`` (the AutoTP group by default), and every rank must call this with the same + model. A rank's storage of a layout piece that is dense in its shard is at most + ``2 * ndim - 1`` boxes; any other layout piece keeps one piece per row. + """ + if tp_group is None: + from deepspeed.utils import groups + tp_group = groups.get_tensor_model_parallel_group() + tp_world_size = dist.get_world_size(group=tp_group) + + # Some AutoTP layers publish their metadata only when asked. + for module in model.modules(): + marker = getattr(module, '_mark_uc_metadata', None) + if marker is not None: + marker() + + named_params = list(model.named_parameters()) + params = [param for _, param in named_params] + if optimizer_states: + storage_maps = zero_optimizer_storage_affine_maps(params) + else: + storage_maps = [zero_param_storage_affine_map(param) for param in params] + local = [None if storage_map is None else storage_map.to_dict() for storage_map in storage_maps] + gathered = [None] * tp_world_size + dist.all_gather_object(gathered, local, group=tp_group) + + maps = {} + for position, (name, param) in enumerate(named_params): + storage_by_tp_rank = [rank_maps[position] for rank_maps in gathered] + if all(storage is None for storage in storage_by_tp_rank): + maps[name] = None + continue + if any(storage is None for storage in storage_by_tp_rank): + raise ValueError(f"Parameter '{name}' has storage on some tensor-parallel ranks but not on others.") + layout = _autotp_layout_map(name, param, tp_world_size) + inner = {tp_rank: ParamAffineMap.from_dict(storage) for tp_rank, storage in enumerate(storage_by_tp_rank)} + replicated = layout.to_dict() == replicated_map(layout.logical_shape, tp_world_size).to_dict() + shared = all(storage == storage_by_tp_rank[0] for storage in storage_by_tp_rank) + if tp_world_size > 1 and replicated and shared: + # A parameter AutoTP leaves whole in a model built under zero.Init keeps its partition over + # the whole world, so every tensor-parallel rank reports the storage of the same one copy. + layout = replicated_map(layout.logical_shape, 1) + inner = {0: inner[0]} + maps[name] = nest_affine_maps(layout, inner) + return maps diff --git a/deepspeed/runtime/base_optimizer.py b/deepspeed/runtime/base_optimizer.py index f04c5d35f5bd..41d00dcd0f25 100644 --- a/deepspeed/runtime/base_optimizer.py +++ b/deepspeed/runtime/base_optimizer.py @@ -379,6 +379,57 @@ def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, chec continue param_group[key] = value + self._normalize_step_for_all_param_groups() + + def _normalize_step_tensor(self, step=None, param=None, fused=False, capturable=False): + if not isinstance(self, ZeROOptimizer): + capturable = fused + fused = param + param = step + step = self + + if isinstance(fused, dict): + capturable = fused.get('capturable', False) + fused = fused.get('fused', False) + + fused = bool(fused) + capturable = bool(capturable) + + if step is None: + return None + + # Preserve scalar integer steps for built-in/C++ optimizers (like FusedAdam, CPUAdam) + # when not using PyTorch tensor-step policy. + if not (fused or capturable): + if not torch.is_tensor(step): + return step + # If it's already a tensor but was stored as scalar/cpu int, keep as int/scalar if appropriate, + # or respect standard non-fused CPU step policy. + if step.numel() == 1 and not step.is_floating_point(): + return step.item() if not torch.is_tensor(step) else step + + if param is None: + return step + + # For PyTorch fused/capturable optimizers requiring tensor steps on device: + target_device = param.device if (fused or capturable) else torch.device("cpu") + target_dtype = torch.float32 if (fused or capturable) else torch.float32 + + if not torch.is_tensor(step): + step = torch.tensor(step, dtype=target_dtype, device=target_device) + else: + step = step.to(device=target_device, dtype=target_dtype) + return step + + def _normalize_step_for_all_param_groups(self) -> None: + for param_group in self.optimizer.param_groups: + fused = bool(param_group.get('fused', False)) if param_group is not None else False + capturable = bool(param_group.get('capturable', False)) if param_group is not None else False + for param in param_group['params']: + if param in self.optimizer.state and 'step' in self.optimizer.state[param]: + self.optimizer.state[param]['step'] = self._normalize_step_tensor( + self.optimizer.state[param]['step'], param, fused, capturable) + def report_ipg_memory_usage(self, tag, param_elems, dtype=None): dtypes = self.ipg_buckets.keys() if dtype is None else [dtype] diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index cb633e47d6bb..6d9ff66d0266 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -3690,6 +3690,17 @@ def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir): else: param_group[key] = value + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + param_group_id = self.sub_group_to_group_id.get(sub_group_id, 0) + param_group = self.optimizer.param_groups[param_group_id] + fused = bool(param_group.get('fused', False)) if param_group is not None else False + capturable = bool(param_group.get('capturable', False)) if param_group is not None else False + if fp32_param in self.optimizer.state and 'step' in self.optimizer.state[fp32_param]: + self.optimizer.state[fp32_param]['step'] = self._normalize_step_tensor( + self.optimizer.state[fp32_param]['step'], fp32_param, fused, capturable) + + self._normalize_step_for_all_param_groups() + if self.swap_optimizer: # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint self.optimizer_swapper.purge_state() diff --git a/docs/code-docs/source/affine-transfer.rst b/docs/code-docs/source/affine-transfer.rst new file mode 100644 index 000000000000..a699869b7c91 --- /dev/null +++ b/docs/code-docs/source/affine-transfer.rst @@ -0,0 +1,54 @@ +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. + +Buffers may reside in host CPU memory (e.g. CPU-offloaded optimizer states) or +on the accelerator. Transfers crossing host and accelerator devices move +through bounded pinned-memory staging arenas bounded by ``StagingBudget``, +with explicit stream synchronization and event ordering ensuring that NCCL +payloads remain strictly on accelerator memory without whole-model allocations. 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..5f26f68bc041 --- /dev/null +++ b/tests/unit/checkpoint/test_affine_transfer_executor.py @@ -0,0 +1,666 @@ +# 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])) + + +class TestB2ResidencyAndOffload: + """Test CPU <-> Accelerator residency combinations, uneven partitions, padding, and offloaded optimizer states.""" + + @pytest.mark.parametrize("src_device_type, dst_device_type", [ + ("cpu", "cuda"), + ("cuda", "cpu"), + ("cpu", "cpu"), + ]) + def test_residency_combinations_with_uneven_partitions_and_guards(self, src_device_type, dst_device_type): + acc = get_accelerator() + if (src_device_type == "cuda" or dst_device_type == "cuda") and not acc.is_available(): + pytest.skip("Accelerator not available") + + src_device = torch.device(acc.current_device_name() if src_device_type == "cuda" else "cpu") + dst_device = torch.device(acc.current_device_name() if dst_device_type == "cuda" else "cpu") + + # Uneven partitions: 23 rows x 7 cols = 161 elements + # Source map: 3 uneven row partitions: [8, 8, 7] + # Target map: 3 uneven row partitions: [10, 8, 5] + logical_shape = (23, 7) + total_numel = 23 * 7 + source_map = contiguous_split_map(logical_shape, [8, 8, 7], partition_dim=0) + target_map = contiguous_split_map(logical_shape, [10, 8, 5], partition_dim=0) + plan = plan_transfer(target_map, source_map) + + full = torch.arange(total_numel, dtype=torch.float32).reshape(logical_shape) + + # Allocate storage with guard offsets to verify bounds and storage offsets + source_buffers = {} + source_guards = [] + for rank, shape in source_map.shard_shapes.items(): + numel = math.prod(shape) + storage = torch.full((numel + 18, ), -42.0, device=src_device) + # Offset 7, length numel, followed by 11 guard elements + shard = storage[7:7 + numel].view(shape) + shard.copy_(source_map.extract(full, rank).to(src_device)) + source_buffers[rank] = shard + source_guards.append((storage, numel, 7)) + + target_buffers = {} + target_guards = [] + for rank, shape in target_map.shard_shapes.items(): + numel = math.prod(shape) + storage = torch.full((numel + 14, ), -99.0, device=dst_device) + # Offset 5, length numel, followed by 9 guard elements + shard = storage[5:5 + numel].view(shape) + target_buffers[rank] = shard + target_guards.append((storage, numel, 5)) + + endpoints = TransferEndpoints({0: 0, 1: 0, 2: 0}, {0: 0, 1: 0, 2: 0}) + + # Use small budget of 32 bytes (8 fp32 elements) to force bounded chunking across pinned staging + budget = StagingBudget(bytes_per_peer=32) + stats = execute_transfers([TransferRequest("param_uneven", plan, source_buffers, target_buffers)], endpoints, + torch.float32, budget) + + assert stats.chunks >= 20, f"Expected chunking under 32-byte budget, got {stats.chunks}" + + # Verify data landed bit-exactly + for rank, target in target_buffers.items(): + expected = target_map.extract(full, rank).to(dst_device) + assert torch.equal(target, expected), f"Rank {rank} target elements mismatch" + + # Verify guard boundaries are completely untouched + for storage, numel, offset in target_guards: + assert torch.all(storage[:offset] == -99.0).item(), "Target prefix guard corrupted" + assert torch.all(storage[offset + numel:] == -99.0).item(), "Target suffix guard corrupted" + + for storage, numel, offset in source_guards: + assert torch.all(storage[:offset] == -42.0).item(), "Source prefix guard corrupted" + assert torch.all(storage[offset + numel:] == -42.0).item(), "Source suffix guard corrupted" + + def test_mixed_residency_batch(self): + acc = get_accelerator() + if not acc.is_available(): + pytest.skip("Accelerator not available") + gpu_device = torch.device(acc.current_device_name()) + cpu_device = torch.device("cpu") + + # Three requests in one batch: + # req 0: CPU -> GPU (e.g. CPU-offloaded optimizer state -> GPU model param) + # req 1: GPU -> CPU (e.g. GPU model param -> CPU-offloaded optimizer state) + # req 2: CPU -> CPU (e.g. CPU-offloaded moment -> CPU-offloaded moment resharding) + shape = (12, 4) + src_map = contiguous_split_map(shape, [6, 6], partition_dim=0) + dst_map = contiguous_split_map(shape, [4, 8], partition_dim=0) + plan = plan_transfer(dst_map, src_map) + + endpoints = TransferEndpoints({0: 0, 1: 0}, {0: 0, 1: 0}) + budget = StagingBudget(bytes_per_peer=32) # 8 elements per chunk + + configs = [ + ("cpu_to_gpu", cpu_device, gpu_device), + ("gpu_to_cpu", gpu_device, cpu_device), + ("cpu_to_cpu", cpu_device, cpu_device), + ] + + requests = [] + oracles = [] + for key, s_dev, d_dev in configs: + full = torch.randn(shape, dtype=torch.float32) + s_bufs = {r: src_map.extract(full, r).clone().to(s_dev) for r in src_map.shard_shapes} + d_bufs = { + r: torch.full(dst_map.shard_shapes[r], -1.0, dtype=torch.float32, device=d_dev) + for r in dst_map.shard_shapes + } + requests.append(TransferRequest(key, plan, s_bufs, d_bufs)) + oracles.append((d_bufs, full, d_dev)) + + execute_transfers(requests, endpoints, torch.float32, budget) + + for d_bufs, full, d_dev in oracles: + for r, buf in d_bufs.items(): + expected = dst_map.extract(full, r).to(d_dev) + assert torch.equal(buf, expected) + + @pytest.mark.parametrize("src_device_type, dst_device_type", [ + ("cpu", "cuda"), + ("cuda", "cpu"), + ("cpu", "cpu"), + ]) + def test_zero3_padding_drop_with_uneven_shards(self, src_device_type, dst_device_type): + from deepspeed.checkpoint.zero_affine import zero_storage_affine_map + acc = get_accelerator() + if (src_device_type == "cuda" or dst_device_type == "cuda") and not acc.is_available(): + pytest.skip("Accelerator not available") + + src_device = torch.device(acc.current_device_name() if src_device_type == "cuda" else "cpu") + dst_device = torch.device(acc.current_device_name() if dst_device_type == "cuda" else "cpu") + + # 19 elements under ZeRO-3 with partition size 6 across 4 ranks: + # Rank 0: [0:6] (6 elements) + # Rank 1: [6:12] (6 elements) + # Rank 2: [12:18] (6 elements) + # Rank 3: [18:19] (1 element, trailing 5 elements are padding in storage!) + source_map = zero_storage_affine_map(19, [0, 1, 2, 3], partition_numel=6) + target_map = contiguous_split_map((19, ), [7, 7, 5], partition_dim=0) + plan = plan_transfer(target_map, source_map) + + full = torch.arange(19.0, dtype=torch.float32) + + source_buffers = {} + for r, shape in source_map.shard_shapes.items(): + storage = torch.full((6, ), -777.0, device=src_device) + valid = storage[:shape[0]] + valid.copy_(source_map.extract(full, r).to(src_device)) + source_buffers[r] = valid + + target_buffers = {} + target_storage = {} + for r, shape in target_map.shard_shapes.items(): + storage = torch.full((shape[0] + 4, ), -888.0, device=dst_device) + valid = storage[2:2 + shape[0]] + target_buffers[r] = valid + target_storage[r] = storage + + endpoints = TransferEndpoints({0: 0, 1: 0, 2: 0, 3: 0}, {0: 0, 1: 0, 2: 0}) + execute_transfers([TransferRequest("z3_padding", plan, source_buffers, target_buffers)], endpoints, + torch.float32, StagingBudget(bytes_per_peer=16)) + + for r, target in target_buffers.items(): + expected = target_map.extract(full, r).to(dst_device) + assert torch.equal(target, expected) + + # Padding in target storage remained intact + for r, storage in target_storage.items(): + assert torch.all(storage[:2] == -888.0).item() + assert torch.all(storage[2 + target_map.shard_shapes[r][0]:] == -888.0).item() + + def test_live_offloaded_optimizer_states_in_place_update(self): + acc = get_accelerator() + gpu_device = torch.device(acc.current_device_name() if acc.is_available() else "cpu") + cpu_device = torch.device("cpu") + + # Simulate ZeRO-2 CPU-offloaded optimizer states: + # Master weight (fp32) and Adam moments (exp_avg, exp_avg_sq) on CPU + numel = 24 + full_exp_avg = torch.randn(numel, dtype=torch.float32) + + source_map = contiguous_split_map((numel, ), [12, 12], partition_dim=0) + target_map = contiguous_split_map((numel, ), [8, 8, 8], partition_dim=0) + + plan = plan_transfer(target_map, source_map) + endpoints = TransferEndpoints({0: 0, 1: 0}, {0: 0, 1: 0, 2: 0}) + + # Transfer CPU offloaded moments to GPU target shards + cpu_source_exp_avg = { + 0: source_map.extract(full_exp_avg, 0).clone().to(cpu_device), + 1: source_map.extract(full_exp_avg, 1).clone().to(cpu_device) + } + gpu_target_exp_avg = { + r: torch.zeros(target_map.shard_shapes[r], device=gpu_device) + for r in target_map.shard_shapes + } + + execute_transfer(plan, cpu_source_exp_avg, gpu_target_exp_avg, endpoints, torch.float32, + StagingBudget(bytes_per_peer=16)) + + for r, target in gpu_target_exp_avg.items(): + assert torch.equal(target, target_map.extract(full_exp_avg, r).to(gpu_device)) + + # Reverse transfer: GPU target back into CPU-offloaded moments, modifying source in-place + reverse_plan = plan_transfer(source_map, target_map) + reverse_endpoints = TransferEndpoints({0: 0, 1: 0, 2: 0}, {0: 0, 1: 0}) + + # Scale target on GPU + for target in gpu_target_exp_avg.values(): + target.mul_(2.0) + + execute_transfer(reverse_plan, gpu_target_exp_avg, cpu_source_exp_avg, reverse_endpoints, torch.float32, + StagingBudget(bytes_per_peer=16)) + + # Assert in-place update landed in CPU source tensors + for r, src in cpu_source_exp_avg.items(): + expected = (source_map.extract(full_exp_avg, r) * 2.0).to(cpu_device) + assert torch.equal(src, expected) diff --git a/tests/unit/checkpoint/test_affine_transfer_plan.py b/tests/unit/checkpoint/test_affine_transfer_plan.py index caff1cc8c8a1..83638bb79926 100644 --- a/tests/unit/checkpoint/test_affine_transfer_plan.py +++ b/tests/unit/checkpoint/test_affine_transfer_plan.py @@ -397,13 +397,40 @@ def test_two_pieces_claiming_one_shard_range_are_refused(): plan_transfer(replicated_map((4, 4), 1), source) -def test_a_shard_described_with_the_wrong_rank_is_refused(): - """A flat shard against a matrix parameter has no axis correspondence to plan along.""" +def test_a_flat_shard_holds_a_dense_box_of_a_matrix(): + """Flat storage, as ZeRO keeps it, may hold boxes of a matrix as dense runs of addresses.""" + full = torch.arange(32.).reshape(4, 8) + # Columns [0, 3) then columns [3, 8), each stored densely one after the other. + source = ParamAffineMap(logical_shape=(4, 8), + shard_shapes={0: (32, )}, + pieces_by_rank={ + 0: [ + AffinePiece((4, 3), 0, (8, 1), 0, (3, 1), [0]), + AffinePiece((4, 5), 3, (8, 1), 12, (5, 1), [0]), + ] + }) + storage = torch.cat([full[:, :3].reshape(-1), full[:, 3:].reshape(-1)]) + target = contiguous_split_map((4, 8), [2, 3, 3], 1) + + plan = plan_transfer(target, source) + + written = {rank: torch.full(shape, float('nan')) for rank, shape in target.shard_shapes.items()} + for segment in plan.segments: + torch.as_strided(written[segment.target_rank], segment.shape, segment.target_strides, + segment.target_offset).copy_( + torch.as_strided(storage, segment.shape, segment.source_strides, segment.source_offset)) + for rank, shard in written.items(): + assert torch.equal(shard, target.extract(full, rank)) + + +@pytest.mark.parametrize('shard_shape, dest_strides', [((32, ), (8, 1)), ((2, 2, 4), (4, 1))], + ids=['flat_but_not_dense', 'other_rank']) +def test_a_shard_without_axis_correspondence_is_refused(shard_shape, dest_strides): source = ParamAffineMap(logical_shape=(4, 4), - shard_shapes={0: (16, )}, - pieces_by_rank={0: [AffinePiece((4, 4), 0, (4, 1), 0, (4, 1), [0])]}) + shard_shapes={0: shard_shape}, + pieces_by_rank={0: [AffinePiece((4, 4), 0, (4, 1), 0, dest_strides, [0])]}) - with pytest.raises(UnsupportedLayoutError, match='number of axes'): + with pytest.raises(UnsupportedLayoutError): plan_transfer(replicated_map((4, 4), 1), source) diff --git a/tests/unit/checkpoint/test_convert_checkpoint.py b/tests/unit/checkpoint/test_convert_checkpoint.py index 3b59234832b8..a6592139c729 100644 --- a/tests/unit/checkpoint/test_convert_checkpoint.py +++ b/tests/unit/checkpoint/test_convert_checkpoint.py @@ -188,8 +188,8 @@ def test_convert_zero_checkpoint_to_fp32_state_dict(self, tmp_path): assert tensor.dtype == torch.bfloat16 torch.testing.assert_close(tensor.float(), state_dict[name], rtol=5e-3, atol=5e-3) - fp32_size = (fp32_save_dir / 'pytorch_model.bin').stat().size - bf16_size = (bf16_save_dir / 'pytorch_model.bin').stat().size + fp32_size = (fp32_save_dir / 'pytorch_model.bin').stat().st_size + bf16_size = (bf16_save_dir / 'pytorch_model.bin').stat().st_size assert bf16_size < fp32_size * 0.6 diff --git a/tests/unit/checkpoint/test_universal_checkpoint.py b/tests/unit/checkpoint/test_universal_checkpoint.py index 5a9af80d746c..94467c9f3d5f 100644 --- a/tests/unit/checkpoint/test_universal_checkpoint.py +++ b/tests/unit/checkpoint/test_universal_checkpoint.py @@ -32,7 +32,7 @@ def get_expected_mismatch_keys(): def maybe_step(t): - return not torch.is_tensor(t) or (t.device.type == 'cpu' and t.numel() == 1) + return not torch.is_tensor(t) or t.numel() == 1 def gather_opt_state(optimizer_state): @@ -275,3 +275,195 @@ def test_dp_world_size_4to2(self, baseline_ws4, tmpdir, dtype, ds_config, load_o @pytest.mark.world_size(4) def test_dp_world_size_2to4(self, baseline_ws2, tmpdir, dtype, ds_config, load_optim, use_torch_adam): self._run_test(tmpdir, dtype, ds_config, load_optim, use_torch_adam, 4) + + +@pytest.mark.parametrize("fused,capturable", [(True, False), (False, True), (False, False)]) +def test_zero_optimizer_normalize_step_tensor(fused, capturable): + from deepspeed.runtime.base_optimizer import ZeROOptimizer + from deepspeed.accelerator import get_accelerator + + acc = get_accelerator() + target_device = torch.device(acc.device_name(0) if acc.is_available() else 'cpu') + param = torch.nn.Parameter(torch.randn(2, 2, device=target_device)) + group = {"fused": fused, "capturable": capturable} + + # Step on CPU as a float tensor + cpu_step = torch.tensor(3.0, device="cpu") + normalized = ZeROOptimizer._normalize_step_tensor(cpu_step, param, group) + if fused or capturable: + assert normalized.device == param.device + assert normalized.dtype == torch.float32 + assert normalized.item() == 3.0 + else: + assert normalized.device == torch.device("cpu") + assert normalized.item() == 3.0 + + # Step on accelerator as a float tensor + if acc.is_available(): + acc_step = torch.tensor(4.0, device=acc.device_name(0)) + normalized = ZeROOptimizer._normalize_step_tensor(acc_step, param, group) + if fused or capturable: + assert normalized.device == param.device + assert normalized.dtype == torch.float32 + assert normalized.item() == 4.0 + else: + assert normalized.device == torch.device("cpu") + assert normalized.item() == 4.0 + + # Integer scalar and integer tensor steps + for raw_step in (5, torch.tensor(5, dtype=torch.int64)): + normalized = ZeROOptimizer._normalize_step_tensor(raw_step, param, group) + if fused or capturable: + assert normalized.device == param.device + assert normalized.dtype == torch.float32 + assert normalized.item() == 5.0 + else: + if isinstance(raw_step, int): + assert normalized == 5 and not torch.is_tensor(normalized) + else: + assert normalized.item() == 5 and normalized.device == torch.device("cpu") + + +@pytest.mark.parametrize("fused,capturable", [(True, False), (False, True), (False, False)]) +def test_universal_checkpoint_load_hp_step_device_policy(tmp_path, fused, capturable): + from deepspeed.runtime.base_optimizer import ZeROOptimizer + from deepspeed.accelerator import get_accelerator + + acc = get_accelerator() + if (fused or capturable) and not acc.is_available(): + pytest.skip("Accelerator required for fused/capturable optimizer test") + + # Guard against capturable=True on CPU device + if capturable and (not torch.cuda.is_available() #ignore-cuda + or acc.is_available() and acc.device_name() == "cpu"): + pytest.skip("PyTorch capturable optimizers require a CUDA/GPU accelerator device") + + device = torch.device(acc.device_name(0) if acc.is_available() else "cpu") + + zero_dir = tmp_path / "zero" + zero_dir.mkdir(parents=True) + + optim_sd = { + "param_groups": [{ + "fused": fused, + "capturable": capturable, + "lr": 0.001 + }], + "ds_version": "0.14.0", + } + torch.save(optim_sd, zero_dir / "optimizer_state.pt") + + param = torch.nn.Parameter(torch.randn(4, 4, device=device)) + client_optimizer = torch.optim.AdamW([param], lr=0.001, fused=fused, capturable=capturable) + + class DummyZeROOptimizer(ZeROOptimizer): + + def __init__(self, optimizer): + super().__init__() + self.optimizer = optimizer + self.mpu = None + self.param_names = {param: "param_0"} + self.dummy_groups = [[param]] + + def _load_global_state(self, sd): + pass + + param._hp_mapping = SimpleNamespace(get_optim_state_keys=lambda: []) + # Universal checkpoint restores steps from disk on CPU + param.load_hp_checkpoint_state = lambda *args, **kwargs: torch.tensor(7.0, device="cpu") + + zero_opt = DummyZeROOptimizer(client_optimizer) + zero_opt.load_hp_checkpoint_state_from_checkpoint_dir("dummy_groups", str(tmp_path)) + + restored_step = client_optimizer.state[param]["step"] + if fused or capturable: + assert restored_step.device == param.device + assert restored_step.dtype == torch.float32 + else: + assert restored_step.device == torch.device("cpu") + + # Verify that the restored optimizer state allows step() to execute without device errors + client_optimizer.state[param]["exp_avg"] = torch.zeros_like(param) + client_optimizer.state[param]["exp_avg_sq"] = torch.zeros_like(param) + param.grad = torch.randn_like(param) + client_optimizer.step() + assert client_optimizer.state[param]["step"].item() == 8.0 + + +@pytest.mark.parametrize("fused,capturable", [(True, False), (False, True), (False, False)]) +def test_universal_checkpoint_stage3_load_hp_step_device_policy(tmp_path, fused, capturable): + from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 + from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT + from deepspeed.accelerator import get_accelerator + + acc = get_accelerator() + if (fused or capturable) and not acc.is_available(): + pytest.skip("Accelerator required for fused/capturable optimizer test") + + # Guard against capturable=True on CPU device + if capturable and (not torch.cuda.is_available() #ignore-cuda + or acc.is_available() and acc.device_name() == "cpu"): + pytest.skip("PyTorch capturable optimizers require a CUDA/GPU accelerator device") + + device = torch.device(acc.device_name(0) if acc.is_available() else "cpu") + + zero_dir = tmp_path / "zero" + zero_dir.mkdir(parents=True) + + fp32_param = torch.nn.Parameter(torch.randn(16, device=device)) + client_optimizer = torch.optim.AdamW([fp32_param], lr=0.001, fused=fused, capturable=capturable) + + optim_sd = { + OPTIMIZER_STATE_DICT: { + "state": [{ + "step": torch.tensor(12.0, device="cpu") + }], + "param_groups": [{ + "fused": fused, + "capturable": capturable, + "lr": 0.001 + }], + }, + "ds_version": "0.14.0", + } + torch.save(optim_sd, zero_dir / "optimizer_state.pt") + + class DummyStage3Optimizer(DeepSpeedZeroOptimizer_Stage3): + + def __init__(self, optimizer): + self._backward_hook_state = SimpleNamespace() + self._is_gradient_accumulation_boundary = True + self.optimizer = optimizer + self.fp32_partitioned_groups_flat = [fp32_param] + self.fp16_partitioned_groups_flat = [fp32_param] + self.fp16_groups = [[fp32_param]] + self.fp16_partitioned_groups = [[fp32_param]] + self.param_names = {fp32_param: "param_0"} + self.sub_group_to_group_id = {0: 0} + self.swap_optimizer = False + + def _load_global_state_stage3(self, sd): + pass + + def load_hp_checkpoint_state(self, path, key, param=None): + return torch.zeros_like(param) + + def unflatten(self, flat, params): + return [flat] + + stage3_opt = DummyStage3Optimizer(client_optimizer) + stage3_opt.load_hp_checkpoint_state_from_checkpoint_dir_stage3(str(tmp_path)) + + restored_step = client_optimizer.state[fp32_param]["step"] + if fused or capturable: + assert restored_step.device == fp32_param.device + assert restored_step.dtype == torch.float32 + else: + assert restored_step.device == torch.device("cpu") + + client_optimizer.state[fp32_param]["exp_avg"] = torch.zeros_like(fp32_param) + client_optimizer.state[fp32_param]["exp_avg_sq"] = torch.zeros_like(fp32_param) + client_optimizer.param_groups[0]["params"] = [fp32_param] + fp32_param.grad = torch.randn_like(fp32_param) + client_optimizer.step() + assert client_optimizer.state[fp32_param]["step"].item() == 13.0 diff --git a/tests/unit/checkpoint/test_zero_affine.py b/tests/unit/checkpoint/test_zero_affine.py new file mode 100644 index 000000000000..b2434c4639be --- /dev/null +++ b/tests/unit/checkpoint/test_zero_affine.py @@ -0,0 +1,598 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import pytest +import torch + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.affine import (contiguous_split_map, nest_affine_maps, replicated_map, segmented_map, + sub_param_map) +from deepspeed.checkpoint.affine_transfer import PlanBudget, plan_transfer +from deepspeed.checkpoint.zero_affine import (autotp_storage_affine_maps, zero_optimizer_storage_affine_maps, + zero_optimizer_storage_shard, zero_param_storage_affine_map, + zero_param_storage_shard, zero_storage_affine_map) +from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum +from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_optimizer_state +from unit.common import DistributedTest +from unit.simple_model import random_dataloader + + +def test_replicas_list_every_holder(): + full = torch.arange(6.) + affine_map = zero_storage_affine_map(6, [4, 1, 6]) + + assert set(affine_map.shard_shapes) == {1, 4, 6} + for rank in (1, 4, 6): + assert torch.equal(affine_map.extract(full, rank), full) + + asked = [] + + def selector(request): + asked.append(sorted(candidate.source_rank for candidate in request.candidates)) + return request.candidates[-1].source_rank + + plan_transfer(contiguous_split_map((6, ), [3, 3], 0), affine_map, PlanBudget(), replica_selector=selector) + # Each target block may be read from any of the three replicas. + assert asked == [[1, 4, 6], [1, 4, 6]] + + +def test_zero3_partitions_drop_padding(): + # 10 elements in partitions of 4: holders keep 4, 4 and 2 elements; the fourth is all padding. + full = torch.arange(10.) + affine_map = zero_storage_affine_map(10, [2, 0, 3, 1], partition_numel=4) + + assert affine_map.shard_shapes == {2: (4, ), 0: (4, ), 3: (2, ), 1: (0, )} + assert torch.equal(affine_map.extract(full, 2), full[0:4]) + assert torch.equal(affine_map.extract(full, 0), full[4:8]) + assert torch.equal(affine_map.extract(full, 3), full[8:10]) + assert affine_map.extract(full, 1).numel() == 0 + + +@pytest.mark.parametrize('numel, holders, partition', [(9, [0, 1], 4), (4, [0, 0], None), (4, [], None)]) +def test_rejects_storage_that_cannot_hold_the_parameter(numel, holders, partition): + with pytest.raises(ValueError): + zero_storage_affine_map(numel, holders, partition) + + +@pytest.mark.parametrize('accessor', [zero_param_storage_affine_map, zero_param_storage_shard]) +def test_rejects_quantized_zero3_parameter_storage(accessor): + param = torch.nn.Parameter(torch.zeros(8), requires_grad=False) + param.ds_id = 0 + param.ds_numel = 8 + param.ds_tensor = torch.zeros(4, dtype=torch.int8) + param.ds_tensor.ds_numel = 4 + param.ds_tensor.ds_quant_scale = torch.ones(1) + + with pytest.raises(NotImplementedError, match='Quantized ZeRO-3 parameter storage'): + accessor(param) + + +@pytest.mark.parametrize('accessor', [zero_param_storage_affine_map, zero_param_storage_shard]) +def test_rejects_nvme_offloaded_zero3_parameter_storage(accessor): + # A partition swapped out to NVMe leaves a one-element placeholder in memory. + param = torch.nn.Parameter(torch.zeros(8), requires_grad=False) + param.ds_id = 0 + param.ds_numel = 8 + param.ds_tensor = torch.zeros(1) + param.ds_tensor.ds_numel = 4 + param.ds_tensor.final_location = OffloadDeviceEnum.nvme + + with pytest.raises(NotImplementedError, match='NVMe-offloaded ZeRO-3 parameter storage'): + accessor(param) + + +def _stored_by_rank(layout_shards, holders_by_layout_rank, partition_by_layout_rank=None): + """What each process stores, by slicing each flattened layout shard with ZeRO's rule.""" + stored = {} + for layout_rank, shard in layout_shards.items(): + flat = shard.reshape(-1) + for index, rank in enumerate(holders_by_layout_rank[layout_rank]): + if partition_by_layout_rank is None: + stored[rank] = flat.clone() + continue + partition = partition_by_layout_rank[layout_rank] + stored[rank] = flat[index * partition:(index + 1) * partition].clone() + return stored + + +def _nested(layout, holders_by_layout_rank, partition_by_layout_rank=None): + inner = {} + for layout_rank, holders in holders_by_layout_rank.items(): + numel = torch.Size(layout.shard_shapes[layout_rank]).numel() + partition = None if partition_by_layout_rank is None else partition_by_layout_rank[layout_rank] + inner[layout_rank] = zero_storage_affine_map(numel, holders, partition) + return nest_affine_maps(layout, inner) + + +def _check_storage(affine_map, full, stored): + # validate_coverage also proves every piece names exactly the processes that store its elements. + affine_map.validate_coverage() + assert affine_map.logical_shape == tuple(full.shape) + assert affine_map.shard_shapes == {rank: tuple(shard.shape) for rank, shard in stored.items()} + for rank, shard in stored.items(): + torch.testing.assert_close(affine_map.extract(full, rank), shard, rtol=0, atol=0) + torch.testing.assert_close(affine_map.rebuild(stored), full, rtol=0, atol=0) + + +# Uneven column split of a 4 x 8 weight over two tensor-parallel ranks, the layout of a row-parallel +# linear layer. Its pieces are strided in the flattened weight. +_FULL = torch.arange(32, dtype=torch.float32).reshape(4, 8) +_COLUMNS = {0: _FULL[:, :3], 1: _FULL[:, 3:]} +_HOLDERS = {0: [0, 2], 1: [1, 3]} + + +@pytest.mark.parametrize('partition', [None, {0: 6, 1: 10}, {0: 7, 1: 11}], ids=['replicas', 'even', 'mid_row']) +def test_nest_describes_zero_storage_of_a_column_split(partition): + layout = contiguous_split_map(_FULL.shape, [3, 5], partition_dim=1) + + affine_map = _nested(layout, _HOLDERS, partition) + + _check_storage(affine_map, _FULL, _stored_by_rank(_COLUMNS, _HOLDERS, partition)) + + +def test_nest_describes_a_contiguous_storage_run_as_a_few_boxes(): + # A process stores one contiguous run of its layout shard: whole rows between two partial rows. + # One piece per row would make planning a large row-parallel weight slow, so the run must stay + # at most three boxes however many rows it spans. + full = torch.arange(64 * 48, dtype=torch.float32).reshape(64, 48) + layout = contiguous_split_map(full.shape, [20, 28], partition_dim=1) + holders = {0: [0, 1, 2], 1: [3, 4]} + partition = {0: 450, 1: 900} + + affine_map = _nested(layout, holders, partition) + + _check_storage(affine_map, full, _stored_by_rank({0: full[:, :20], 1: full[:, 20:]}, holders, partition)) + for rank in range(5): + assert len(affine_map.pieces_by_rank[rank]) <= 3 + + +def test_nest_merges_holders_of_replicated_layout_shards_split_at_different_places(): + # Two layout ranks hold the whole tensor but partition it over groups of different sizes, so an + # element's holders are the union of one process from each group, cut at both groups' boundaries. + layout = replicated_map(_FULL.shape, 2) + holders = {0: [0, 1], 1: [2, 3, 4]} + partition = {0: 16, 1: 11} + + affine_map = _nested(layout, holders, partition) + + _check_storage(affine_map, _FULL, _stored_by_rank({0: _FULL, 1: _FULL}, holders, partition)) + for offset, ranks in affine_map.holders().items(): + assert ranks == {offset // 16, 2 + offset // 11} + + +def test_nest_keeps_fused_sub_parameters_apart(): + full = torch.arange(6 * 4, dtype=torch.float32).reshape(6, 4) + # Fused QKV: rows [0, 2), [2, 4) and [4, 6) are each split between two ranks. + layout = sub_param_map(full.shape, [2, 2, 2], [[1, 1], [1, 1], [1, 1]], partition_dim=0) + shards = {rank: full[[rank, 2 + rank, 4 + rank]] for rank in range(2)} + holders = {0: [0, 1], 1: [2, 3]} + partition = {0: 7, 1: 6} + + _check_storage(_nested(layout, holders, partition), full, _stored_by_rank(shards, holders, partition)) + + +def test_nest_keeps_boxes_dense_in_every_replica(): + # The last column is a strided column of rank 0's whole-weight shard but all of rank 1's shard, + # so a box spanning both rows would be one run of addresses on rank 1 only. + full = torch.arange(6, dtype=torch.float32).reshape(2, 3) + layout = segmented_map(full.shape, [(2, False), (1, True)], 1, 2, [[2, 0]]) + holders = {0: [0], 1: [1]} + + _check_storage(_nested(layout, holders), full, _stored_by_rank({0: full, 1: full[:, 2:]}, holders)) + + +def test_nest_and_plan_a_scalar_parameter(): + full = torch.tensor(5.0) + holders = {0: [0, 1], 1: [2]} + source = _nested(replicated_map((), 2), holders) + stored = _stored_by_rank({0: full, 1: full}, holders) + _check_storage(source, full, stored) + + plan = plan_transfer(replicated_map((), 1), source) + + target = torch.full((), float('nan')) + for segment in plan.segments: + source_view = torch.as_strided(stored[segment.source_rank], segment.shape, segment.source_strides, + segment.source_offset) + torch.as_strided(target, segment.shape, segment.target_strides, segment.target_offset).copy_(source_view) + assert target.item() == 5.0 + + +def test_nested_storage_plans_into_another_layout(): + source = _nested(contiguous_split_map(_FULL.shape, [3, 5], partition_dim=1), _HOLDERS, {0: 7, 1: 11}) + target_holders = {0: [4], 1: [5, 6], 2: [7]} + target = _nested(contiguous_split_map(_FULL.shape, [1, 2, 1], partition_dim=0), target_holders, {0: 8, 1: 9, 2: 8}) + source_storage = _stored_by_rank(_COLUMNS, _HOLDERS, {0: 7, 1: 11}) + expected = _stored_by_rank({0: _FULL[:1], 1: _FULL[1:3], 2: _FULL[3:]}, target_holders, {0: 8, 1: 9, 2: 8}) + + plan = plan_transfer(target, source) + + target_storage = {rank: torch.full(shape, float('nan')) for rank, shape in target.shard_shapes.items()} + writes = {rank: torch.zeros(shape, dtype=torch.int64) for rank, shape in target.shard_shapes.items()} + for segment in plan.segments: + source_view = torch.as_strided(source_storage[segment.source_rank], segment.shape, segment.source_strides, + segment.source_offset) + torch.as_strided(target_storage[segment.target_rank], segment.shape, segment.target_strides, + segment.target_offset).copy_(source_view) + torch.as_strided(writes[segment.target_rank], segment.shape, segment.target_strides, + segment.target_offset).add_(1) + for rank, shard in expected.items(): + torch.testing.assert_close(target_storage[rank], shard, rtol=0, atol=0) + assert torch.all(writes[rank] == 1) + + +def test_nest_rejects_storage_that_does_not_match_the_layout(): + layout = contiguous_split_map(_FULL.shape, [3, 5], partition_dim=1) + with pytest.raises(ValueError, match='layout ranks'): + nest_affine_maps(layout, {0: zero_storage_affine_map(12, [0])}) + with pytest.raises(ValueError, match='holds 12 elements'): + nest_affine_maps(layout, {0: zero_storage_affine_map(32, [0]), 1: zero_storage_affine_map(20, [1])}) + with pytest.raises(ValueError, match='more than one layout shard'): + nest_affine_maps(layout, {0: zero_storage_affine_map(12, [0]), 1: zero_storage_affine_map(20, [0])}) + + +class StraddlingModel(torch.nn.Module): + # 351 trainable elements split over two ranks put the boundary inside ``first.bias``, so one + # parameter is held in part by each rank. ``frozen`` has no optimizer state. + def __init__(self, hidden_dim): + super().__init__() + self.frozen = torch.nn.Linear(hidden_dim, hidden_dim) + self.frozen.requires_grad_(False) + self.first = torch.nn.Linear(hidden_dim, hidden_dim) + self.second = torch.nn.Linear(hidden_dim, hidden_dim, bias=False) + + def forward(self, x, y): + return torch.nn.functional.cross_entropy(self.second(self.first(self.frozen(x))), y) + + +def _full_param(param): + if hasattr(param, 'ds_id'): + with deepspeed.zero.GatheredParameters([param]): + return param.data.detach().clone().view(-1) + return param.data.detach().clone().view(-1) + + +def _full_optimizer_value(param, key): + if key is None: + return safe_get_full_fp32_param(param).view(-1) + return safe_get_full_optimizer_state(param, key).view(-1) + + +class TestZeroOptimizerStorageMaps(DistributedTest): + world_size = 2 + + @pytest.mark.parametrize('zero_stage, bf16_optimizer', [(1, False), (2, False), (3, False), (1, True)]) + def test_maps_describe_live_storage(self, zero_stage, bf16_optimizer): + if not get_accelerator().is_bf16_supported(): + pytest.skip('requires bf16') + hidden_dim = 13 + config = { + 'train_micro_batch_size_per_gpu': 1, + 'optimizer': { + 'type': 'Adam', + 'params': { + 'lr': 1e-2, + 'torch_adam': True + } + }, + 'bf16': { + 'enabled': True + }, + 'zero_optimization': { + 'stage': zero_stage + }, + } + if bf16_optimizer: + # fp32 gradient accumulation under ZeRO-1 selects the bf16 optimizer. + config['data_types'] = {'grad_accum_dtype': 'fp32'} + if zero_stage == 3: + with deepspeed.zero.Init(config_dict_or_path=config): + model = StraddlingModel(hidden_dim) + else: + model = StraddlingModel(hidden_dim) + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + params = list(engine.module.parameters()) + rank = dist.get_rank() + + if zero_stage != 3: + # first.bias straddles the partition boundary, so every rank holds part of it. + with pytest.raises(RuntimeError, match='first optimizer step'): + zero_optimizer_storage_shard(engine.module.first.bias, 'exp_avg') + + for x, y in random_dataloader(engine, 3, hidden_dim, engine.device, torch.bfloat16): + engine.backward(engine(x, y)) + engine.step() + + if zero_stage == 3: + # ZeRO-3 keeps small parameters gathered between steps; a write to the shard alone would be lost. + with pytest.raises(RuntimeError, match='empty_partition_cache'): + zero_param_storage_shard(engine.module.second.weight) + engine.empty_partition_cache() + + maps = zero_optimizer_storage_affine_maps(params) + assert maps[0] is None and maps[1] is None + assert zero_optimizer_storage_shard(engine.module.frozen.weight) is None + for param, affine_map in zip(params[2:], maps[2:]): + affine_map.validate_coverage() + + # The low-precision parameter is the fp32 master copy rounded, which checks the master + # copy's map against storage that does not come from the optimizer. + full_param = _full_param(param) + assert torch.equal(_full_optimizer_value(param, None).to(full_param.dtype), full_param) + param_map = zero_param_storage_affine_map(param) + assert torch.equal(zero_param_storage_shard(param), param_map.extract(full_param, rank)) + + for key in (None, 'exp_avg', 'exp_avg_sq'): + full = _full_optimizer_value(param, key) + assert torch.equal(zero_optimizer_storage_shard(param, key), affine_map.extract(full, rank)) + + # The shard is ZeRO's storage: writing each rank's id into it must show up in the gathered + # state exactly where the map places that rank's range. + zero_optimizer_storage_shard(param, 'exp_avg').fill_(rank + 1) + written = _full_optimizer_value(param, 'exp_avg') + expected = torch.zeros_like(written) + for holder, pieces in affine_map.pieces_by_rank.items(): + for piece in pieces: + expected[piece.source_offset:piece.source_offset + piece.shape[0]] = holder + 1 + assert torch.equal(written, expected) + + # A write to the parameter's shard must reach the next forward. + zero_param_storage_shard(engine.module.second.weight).zero_() + outputs = [] + hook = engine.module.second.register_forward_hook(lambda module, args, output: outputs.append(output)) + for x, y in random_dataloader(engine, 1, hidden_dim, engine.device, torch.bfloat16): + engine(x, y) + hook.remove() + assert torch.all(outputs[0] == 0) + + engine.destroy() + + +class ThreeLinears(torch.nn.Module): + + def __init__(self, hidden_dim): + super().__init__() + self.linears = torch.nn.ModuleList([torch.nn.Linear(hidden_dim, hidden_dim) for _ in range(3)]) + + def forward(self, x, y): + for linear in self.linears: + x = linear(x) + return torch.nn.functional.cross_entropy(x, y) + + +def _all_shards(shard): + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, shard.detach().cpu()) + return dict(enumerate(gathered)) + + +def _one_holder_each(layout): + # Every layout rank is a process of its own, as a rollout engine without ZeRO would be. + inner = { + rank: zero_storage_affine_map(torch.Size(shape).numel(), [rank]) + for rank, shape in layout.shard_shapes.items() + } + return nest_affine_maps(layout, inner) + + +class TestAutoTPStorageMaps(DistributedTest): + world_size = 4 + reuse_dist_env = False + + @pytest.mark.parametrize('zero_stage, zero_init', [(1, False), (2, False), (3, False), (3, True)]) + def test_maps_describe_live_storage_of_tp2_dp2(self, zero_stage, zero_init): + if not get_accelerator().is_bf16_supported(): + pytest.skip('requires bf16') + hidden_dim = 12 + config = { + 'train_micro_batch_size_per_gpu': 1, + 'optimizer': { + 'type': 'Adam', + 'params': { + 'lr': 1e-2, + 'torch_adam': True + } + }, + 'bf16': { + 'enabled': True + }, + 'zero_optimization': { + 'stage': zero_stage + }, + 'tensor_parallel': { + 'autotp_size': 2, + 'partition_config': { + 'use_default_specs': + False, + 'layer_specs': [ + { + 'patterns': [r'.*linears\.0\.weight$'], + 'partition_type': 'column' + }, + { + 'patterns': [r'.*linears\.1\.weight$'], + 'partition_type': 'row' + }, + { + 'patterns': [r'.*linears\.2\.weight$'], + 'partition_type': 'skip' + }, + ], + }, + }, + } + torch.manual_seed(0) + if zero_init: + # zero.Init partitions every parameter over the whole world before AutoTP sees the model. + with deepspeed.zero.Init(config_dict_or_path=config): + model = ThreeLinears(hidden_dim) + # The weights before AutoTP and ZeRO touch them: the oracle for what each process stores. + reference = { + name: _full_param(param).cpu().view(param.ds_shape) + for name, param in model.named_parameters() + } + else: + model = ThreeLinears(hidden_dim) + reference = {name: param.detach().to(torch.bfloat16) for name, param in model.named_parameters()} + shapes = {name: tuple(full.shape) for name, full in reference.items()} + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + params = dict(engine.module.named_parameters()) + rank = dist.get_rank() + + maps = autotp_storage_affine_maps(engine.module) + + assert set(maps) == set(reference) + for name, affine_map in maps.items(): + affine_map.validate_coverage() + assert torch.equal(zero_param_storage_shard(params[name]).cpu(), affine_map.extract(reference[name], rank)) + + # Plan from the live storage into four tensor-parallel ranks split along dim 0. + target = _one_holder_each(contiguous_split_map(shapes[name], [hidden_dim // 4] * 4, 0)) + plan = plan_transfer(target, affine_map) + source_storage = _all_shards(zero_param_storage_shard(params[name])) + written = {r: torch.full(s, float('nan'), dtype=torch.bfloat16) for r, s in target.shard_shapes.items()} + for segment in plan.segments: + torch.as_strided(written[segment.target_rank], segment.shape, segment.target_strides, + segment.target_offset).copy_( + torch.as_strided(source_storage[segment.source_rank], segment.shape, + segment.source_strides, segment.source_offset)) + for target_rank, shard in written.items(): + assert torch.equal(shard, target.extract(reference[name], target_rank)) + + for x, y in random_dataloader(engine, 2, hidden_dim, engine.device, torch.bfloat16): + engine.backward(engine(x, y)) + engine.step() + engine.empty_partition_cache() + + optimizer_maps = autotp_storage_affine_maps(engine.module, optimizer_states=True) + for name, affine_map in optimizer_maps.items(): + affine_map.validate_coverage() + # The trained low-precision parameter, rebuilt through its map, is the fp32 master copy rounded. + full_param = maps[name].rebuild(_all_shards(zero_param_storage_shard(params[name]))) + master = zero_optimizer_storage_shard(params[name]).to(torch.bfloat16).cpu() + assert torch.equal(master, affine_map.extract(full_param, rank)) + + engine.destroy() + + +def test_zero2_cpu_offloaded_optimizer_storage(): + from deepspeed.checkpoint.zero_affine import zero_optimizer_storage_shard, zero_optimizer_storage_affine_maps + from deepspeed.utils.tensor_fragment import tensor_fragment, fragment_address + + param = torch.nn.Parameter(torch.randn(12)) + param._dp_group = object() + + # 1. Readiness check: before optimizer step, optim_fragment is None + mapping = tensor_fragment( + lp_fragment=torch.randn(6), + lp_fragment_address=fragment_address(numel=6, start=0), + hp_fragment=torch.randn(6, device='cpu'), + hp_fragment_address=fragment_address(numel=6, start=0), + gradient_dict={}, + offload_gradient_dict={}, + use_offload=True, + param_group_index=0, + optim_fragment=None, + ) + param._hp_mapping = mapping + + # Master weight shard is accessible before step + master_shard = zero_optimizer_storage_shard(param, None) + assert master_shard.device.type == 'cpu' + assert torch.equal(master_shard, mapping.hp_fragment) + + # Moments raise RuntimeError before step + with pytest.raises(RuntimeError, match='first optimizer step'): + zero_optimizer_storage_shard(param, 'exp_avg') + with pytest.raises(RuntimeError, match='first optimizer step'): + zero_optimizer_storage_shard(param, 'exp_avg_sq') + + # 2. After step: optim_fragment populated + exp_avg = torch.randn(6, device='cpu') + exp_avg_sq = torch.randn(6, device='cpu') + mapping.optim_fragment = {'exp_avg': exp_avg, 'exp_avg_sq': exp_avg_sq} + + s_exp_avg = zero_optimizer_storage_shard(param, 'exp_avg') + s_exp_avg_sq = zero_optimizer_storage_shard(param, 'exp_avg_sq') + assert s_exp_avg.device.type == 'cpu' + assert s_exp_avg_sq.device.type == 'cpu' + assert torch.equal(s_exp_avg, exp_avg) + assert torch.equal(s_exp_avg_sq, exp_avg_sq) + + # 3. In-place write test + s_exp_avg.fill_(42.0) + assert torch.all(exp_avg == 42.0) + + # 4. Storage affine map generation for offloaded parameter + maps = zero_optimizer_storage_affine_maps([param]) + assert maps[0] is not None + assert maps[0].logical_shape == (12, ) + + +def test_zero3_cpu_offloaded_optimizer_storage(): + from types import SimpleNamespace + from deepspeed.checkpoint.zero_affine import zero_optimizer_storage_shard, zero_optimizer_storage_affine_maps + + param = torch.nn.Parameter(torch.randn(4), requires_grad=True) + param.ds_id = 0 + param.ds_numel = 4 + param.ds_tensor = torch.zeros(4) + param.ds_tensor.ds_numel = 4 + param.ds_process_group = None + + fp32_flat = torch.arange(4.0, device='cpu') + exp_avg_flat = torch.zeros(4, device='cpu') + exp_avg_sq_flat = torch.zeros(4, device='cpu') + + z3_opt = SimpleNamespace( + offload_optimizer=True, + swap_optimizer=False, + grad_position={0: (0, 0, 4)}, + get_param_id=lambda p: p.ds_id, + fp32_partitioned_groups_flat=[fp32_flat], + optimizer=SimpleNamespace(state={}), + ) + + def _get_fp32_opt_state_partition(p, release_swap_buffers, optim_state_key=None): + g_idx, dest_offset, numel = z3_opt.grad_position[z3_opt.get_param_id(p)] + fp = z3_opt.fp32_partitioned_groups_flat[g_idx] + if optim_state_key is None: + return fp.narrow(0, dest_offset, numel), g_idx + return z3_opt.optimizer.state[fp][optim_state_key].narrow(0, dest_offset, numel), g_idx + + z3_opt._get_fp32_opt_state_partition = _get_fp32_opt_state_partition + param._z3_optimizer = z3_opt + + # 1. Readiness check: before step, state is empty + master_shard = zero_optimizer_storage_shard(param, None) + assert master_shard.device.type == 'cpu' + assert torch.equal(master_shard, fp32_flat) + + with pytest.raises(RuntimeError, match='first optimizer step'): + zero_optimizer_storage_shard(param, 'exp_avg') + with pytest.raises(RuntimeError, match='first optimizer step'): + zero_optimizer_storage_shard(param, 'exp_avg_sq') + + # 2. After step + z3_opt.optimizer.state[fp32_flat] = { + 'exp_avg': exp_avg_flat, + 'exp_avg_sq': exp_avg_sq_flat, + } + + s_exp_avg = zero_optimizer_storage_shard(param, 'exp_avg') + s_exp_avg_sq = zero_optimizer_storage_shard(param, 'exp_avg_sq') + assert s_exp_avg.device.type == 'cpu' + assert s_exp_avg_sq.device.type == 'cpu' + + # 3. In-place write test + s_exp_avg.fill_(99.0) + assert torch.all(exp_avg_flat == 99.0) + + # 4. Storage affine map generation for offloaded ZeRO-3 parameter + maps = zero_optimizer_storage_affine_maps([param]) + assert maps[0] is not None + assert maps[0].logical_shape == (4, )