[AffineIR] Affine storage maps for live ZeRO parameters and optimizer states

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

deepspeed/checkpoint/affine_ir_spec.md

@@ -614,6 +614,22 @@ them with a target descriptor through `plan_transfer` into shard-to-shard copies
 ranks are EP-local ranks of the respective placements; the plan is not executed, and ZeRO
 and EDP storage must still be normalized before its offsets apply.
 
+`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
@@ -629,3 +645,60 @@ and EDP storage must still be normalized before its offsets apply.
 
 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.

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')