diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py new file mode 100644 index 000000000000..8ca076ae6667 --- /dev/null +++ b/deepspeed/checkpoint/affine.py @@ -0,0 +1,482 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Affine description of how a parameter is sharded. + +Universal checkpoint currently decides how to merge a parameter by matching its *name* +against regex categories (vocabulary, row-parallel, fused sub-parameters, ...), so a +layout that no category describes cannot be converted at all. This module describes a +shard geometrically instead: each rank holds a list of affine views of the full tensor, +and the views alone say how to rebuild it. + +A piece may carry an invertible elementwise map, but never a reduction. Scaling a block by +a constant is reversible; averaging several elements into one is not, and would leave the +full tensor unrecoverable from the shards. That line is what keeps conversion invertible in +both directions. + +See ``deepspeed/checkpoint/affine_ir_spec.md`` for the full specification. +""" + +import torch + +__all__ = [ + 'AffinePiece', 'ParamAffineMap', 'AFFINE_MAP_FORMAT_VERSION', 'replicated_map', 'contiguous_split_map', + 'sub_param_map' +] + +# Encoding version of the stored map, independent of the universal checkpoint version so +# the two can move separately. A reader refuses a version it predates rather than +# misreading fields it does not know about. +AFFINE_MAP_FORMAT_VERSION = 1 + + +class AffinePiece: + """A block of elements, described where it lives in the full tensor and in the shard. + + A piece holds the same elements in the same arrangement on both sides, so ``shape`` is + shared and only the offset and strides differ. Describing the shard side explicitly is + what lets a piece land somewhere other than the end of the shard: a shard split along + a column interleaves its pieces row by row, so "append each piece in turn" is not + enough to say where the elements go. + + The offsets and strides are the arguments of ``torch.as_strided``, so a piece can be + applied to a tensor without any interpretation step. + + ``locations`` is a set of ranks rather than a single owner because a piece may be + replicated. Naming one owner here would be a scheduling decision, and the cheapest + source depends on the topology, which this description does not know about. + + ``scale`` is the factor the shard holds the block by: ``shard == full * scale``. Row + parallel layers pre-divide a replicated bias by the world size so that summing the + all-reduced outputs adds the bias exactly once, and the divisor changes with the world + size. Recording it as a number keeps that describable without a rule naming which + parameters are biases. + + A piece must be *homogeneous*: every element it covers is held by the same set of ranks + and carries the same scale. That is what makes ``locations`` exact rather than a hint, + and it is why merging two adjacent blocks is only allowed when both agree. + """ + + __slots__ = ('shape', 'source_offset', 'source_strides', 'dest_offset', 'dest_strides', 'locations', 'scale') + + def __init__(self, shape, source_offset, source_strides, dest_offset, dest_strides, locations, scale=1.0): + self.shape = tuple(int(dim) for dim in shape) + self.source_offset = int(source_offset) + self.source_strides = tuple(int(stride) for stride in source_strides) + self.dest_offset = int(dest_offset) + self.dest_strides = tuple(int(stride) for stride in dest_strides) + self.locations = frozenset(int(rank) for rank in locations) + self.scale = float(scale) + if self.scale == 0.0: + raise ValueError('A zero scale is not invertible, so the full tensor could not be rebuilt.') + + @property + def numel(self): + count = 1 + for dim in self.shape: + count *= dim + return count + + def source_view(self, full_param): + """This piece's elements as a view of the full parameter, without copying.""" + return torch.as_strided(full_param, + size=self.shape, + stride=self.source_strides, + storage_offset=self.source_offset) + + def dest_view(self, shard): + """This piece's elements as a view of the shard that holds it.""" + return torch.as_strided(shard, size=self.shape, stride=self.dest_strides, storage_offset=self.dest_offset) + + def source_offsets(self): + """Yield the offset into the full tensor of every element this piece covers. + + Used to check coverage. Pieces are small relative to the parameter, so walking + them elementwise is affordable and avoids assuming anything about their shape. + """ + return self._offsets(self.source_offset, self.source_strides) + + def _offsets(self, base, strides): + if any(dim == 0 for dim in self.shape): + # An empty piece covers no elements. Walking it anyway would report offsets it + # does not hold, and a coverage check would accept a map that is missing data. + return + if not self.shape: + yield base + return + + index = [0] * len(self.shape) + while True: + offset = base + for axis, position in enumerate(index): + offset += position * strides[axis] + yield offset + + axis = len(self.shape) - 1 + while axis >= 0: + index[axis] += 1 + if index[axis] < self.shape[axis]: + break + index[axis] = 0 + axis -= 1 + if axis < 0: + return + + def __repr__(self): + return (f'AffinePiece(shape={self.shape}, source=({self.source_offset}, {self.source_strides}), ' + f'dest=({self.dest_offset}, {self.dest_strides}), locations={sorted(self.locations)}, ' + f'scale={self.scale})') + + def __eq__(self, other): + if not isinstance(other, AffinePiece): + return NotImplemented + return (self.shape == other.shape and self.source_offset == other.source_offset + and self.source_strides == other.source_strides and self.dest_offset == other.dest_offset + and self.dest_strides == other.dest_strides and self.locations == other.locations + and self.scale == other.scale) + + def __hash__(self): + return hash((self.shape, self.source_offset, self.source_strides, self.dest_offset, self.dest_strides, + self.locations, self.scale)) + + +class ParamAffineMap: + """How one parameter is spread over a tensor-parallel group. + + ``shard_shapes`` gives each rank the shape of the tensor it holds, and + ``pieces_by_rank`` says which blocks of the full parameter make it up. + """ + + def __init__(self, logical_shape, shard_shapes, pieces_by_rank): + self.logical_shape = tuple(int(dim) for dim in logical_shape) + self.shard_shapes = {int(rank): tuple(int(dim) for dim in shape) for rank, shape in shard_shapes.items()} + self.pieces_by_rank = {int(rank): list(pieces) for rank, pieces in pieces_by_rank.items()} + + @property + def numel(self): + return _product(self.logical_shape) + + def uncovered_offsets(self): + """Return the offsets of the full tensor that no rank holds. + + A non-empty result means the parameter cannot be rebuilt from its shards, so this + is the check that has to pass before a map is used for conversion. + """ + covered = set() + for pieces in self.pieces_by_rank.values(): + for piece in pieces: + covered.update(piece.source_offsets()) + return sorted(set(range(self.numel)) - covered) + + def holders(self): + """Map each element offset of the full tensor to the set of ranks holding it.""" + holders = {} + for rank, pieces in self.pieces_by_rank.items(): + for piece in pieces: + for offset in piece.source_offsets(): + holders.setdefault(offset, set()).add(rank) + return holders + + def validate(self): + """Cheap structural check: every rank's pieces account for exactly its shard.""" + for rank, pieces in self.pieces_by_rank.items(): + held = sum(piece.numel for piece in pieces) + expected = _product(self.shard_shapes[rank]) + assert held == expected, (f'Rank {rank} holds a shard of {expected} elements but its pieces ' + f'account for {held}.') + + def validate_coverage(self): + """Full check: the shards cover the parameter, and every piece is homogeneous. + + This walks the map element by element, so it costs O(numel) and is meant for tests + and for validating a newly built map, not for every conversion of a large tensor. + """ + self.validate() + + missing = self.uncovered_offsets() + assert not missing, (f'Affine map for a parameter of shape {self.logical_shape} leaves ' + f'{len(missing)} element(s) uncovered, starting at offset {missing[0]}, ' + 'so the parameter cannot be rebuilt from its shards.') + + # A piece whose elements are not all held by the same ranks would make its own + # `locations` a lie, and a reader trusting that field would miss a replica. + holders = self.holders() + for rank, pieces in self.pieces_by_rank.items(): + for piece in pieces: + for offset in piece.source_offsets(): + assert holders[offset] == set( + piece.locations), (f'Piece {piece} on rank {rank} covers offset {offset}, which is held by ' + f'{sorted(holders[offset])}. A piece must not span elements with different ' + 'holders, or its locations cannot be trusted.') + + def rebuild(self, shards, scale_power=1): + """Rebuild the full parameter from per-rank shards, using the pieces as the plan. + + ``shards`` maps a rank to its shard. + + ``scale_power`` says how a piece's ``scale`` applies to the tensor being moved: + the shard holds ``full * scale ** scale_power``. A parameter uses 1. Optimizer + moments do not, because scaling a parameter by ``s`` scales its gradient by + ``1 / s``: Adam's first moment needs -1 and its second moment -2. The map records + the geometry; the caller knows which tensor it is moving. + """ + self.validate() + self._reject_scaled_optimizer_state(scale_power) + for rank, shard in shards.items(): + expected = _product(self.shard_shapes[rank]) + if shard.numel() != expected: + raise ValueError(f'Rank {rank} supplied a shard of {shard.numel()} elements, but the map ' + f'describes {expected}. Applying the pieces would silently read a prefix ' + 'of it and drop the rest.') + any_shard = next(iter(shards.values())) + full_param = torch.empty(self.numel, dtype=any_shard.dtype, device=any_shard.device) + + written = {} + for rank, pieces in self.pieces_by_rank.items(): + flat_shard = _flat_buffer(shards[rank]) + for piece in pieces: + target = piece.source_view(full_param) + source = piece.dest_view(flat_shard) + if len(piece.locations) > 1 and piece.source_offset in written: + # Every rank holding a replicated piece must agree. Letting the last + # writer win would hide rank drift or a corrupt shard, which the + # category-based converter guards against explicitly. + if not torch.equal(target, _scaled(source, piece.scale, scale_power)): + raise ValueError(f'Ranks {sorted(piece.locations)} hold different data for the same ' + f'block of a replicated parameter, starting at offset ' + f'{piece.source_offset}.') + continue + target.copy_(_scaled(source, piece.scale, scale_power)) + written[piece.source_offset] = True + + return full_param.view(self.logical_shape) + + def _reject_scaled_optimizer_state(self, scale_power): + """Refuse to move an optimizer state through a scaled piece. + + The moments themselves convert correctly under their powers, but Adam's update also + depends on ``lr`` and ``eps``, and those live in the coordinate the optimizer was + training in. Restoring a rescaled parameter without rescaling them + (``lr / scale``, ``eps * scale``) resumes on a different trajectory, with an error + that grows step by step and no failure to show for it. Describing that transform is + outside this module, which knows about parameter geometry and not about optimizers, + so refuse until the checkpoint contract covers it. + """ + if scale_power == 1: + return + scaled = [piece for pieces in self.pieces_by_rank.values() for piece in pieces if piece.scale != 1.0] + if scaled: + raise NotImplementedError( + 'Converting an optimizer state through a scaled piece is not supported: the moments transform ' + 'correctly, but the optimizer group hyperparameters would still be in the source coordinate, ' + 'so a resumed run would diverge.') + + def extract(self, full_param, rank, scale_power=1): + """Produce one rank's shard from the full parameter. The inverse of ``rebuild``.""" + self._reject_scaled_optimizer_state(scale_power) + flat_param = _flat_buffer(full_param) + shard_shape = self.shard_shapes[rank] + shard = torch.empty(_product(shard_shape), dtype=full_param.dtype, device=full_param.device) + + for piece in self.pieces_by_rank[rank]: + target = piece.dest_view(shard) + target.copy_(piece.source_view(flat_param)) + factor = piece.scale**scale_power + if factor != 1.0: + target.mul_(factor) + + return shard.view(shard_shape) + + def to_dict(self): + """Serialise to plain scalars, so the map can be read without importing torch.""" + return { + 'logical_shape': list(self.logical_shape), + 'ranks': { + rank: { + 'shard_shape': list(self.shard_shapes[rank]), + 'pieces': [_piece_to_dict(piece) for piece in pieces], + } + for rank, pieces in sorted(self.pieces_by_rank.items()) + }, + } + + @classmethod + def from_dict(cls, entry): + ranks = entry['ranks'] + return cls(logical_shape=entry['logical_shape'], + shard_shapes={ + int(rank): value['shard_shape'] + for rank, value in ranks.items() + }, + pieces_by_rank={ + int(rank): [_piece_from_dict(piece) for piece in value['pieces']] + for rank, value in ranks.items() + }) + + def __repr__(self): + counts = {rank: len(pieces) for rank, pieces in sorted(self.pieces_by_rank.items())} + return f'ParamAffineMap(logical_shape={self.logical_shape}, pieces_per_rank={counts})' + + +def _scaled(tensor, scale, scale_power): + """Undo the shard's scaling, giving the block as it appears in the full tensor.""" + factor = scale**scale_power + if factor == 1.0: + return tensor + return tensor / factor + + +def _flat_buffer(tensor): + """Flatten ``tensor`` into a buffer whose storage starts at its first element. + + Piece offsets are storage offsets, because that is what ``torch.as_strided`` takes. + A tensor that is a view into a larger buffer starts partway into that storage, so + applying a piece to it directly would read from the wrong place. Shards routinely + arrive this way, since splitting a tensor produces views that share one buffer. + """ + flat = tensor.reshape(-1) + if flat.storage_offset() != 0: + flat = flat.clone() + return flat + + +def _product(shape): + count = 1 + for dim in shape: + count *= dim + return count + + +def _row_major_strides(shape): + """Element stride per axis for a densely packed tensor of this shape.""" + strides = [1] * len(shape) + for axis in range(len(shape) - 2, -1, -1): + strides[axis] = strides[axis + 1] * shape[axis + 1] + return tuple(strides) + + +def replicated_map(shape, tp_degree, scale=1.0): + """Every rank holds the whole parameter. + + One piece, named by every rank, which is what lets a converter read it from whichever + rank is cheapest rather than from a designated owner. + + ``scale`` belongs here rather than on a split, because the layouts that pre-divide a + value hold it whole on every rank: a row-parallel layer replicates its bias divided by + the world size so that summing the all-reduced outputs adds it exactly once. The weight + beside it is split and unscaled. + """ + shape = tuple(shape) + strides = _row_major_strides(shape) + ranks = list(range(tp_degree)) + pieces = [ + AffinePiece(shape=shape, + source_offset=0, + source_strides=strides, + dest_offset=0, + dest_strides=strides, + locations=ranks, + scale=scale) + ] + return ParamAffineMap(logical_shape=shape, + shard_shapes={rank: shape + for rank in ranks}, + pieces_by_rank={rank: list(pieces) + for rank in ranks}) + + +def contiguous_split_map(shape, per_rank_sizes, partition_dim): + """Each rank holds one contiguous block along ``partition_dim``. + + Covers row-parallel and column-parallel layers alike: they differ only in which axis + the block is cut on, and therefore only in strides. ``per_rank_sizes`` may be uneven. + """ + shape = tuple(shape) + source_strides = _row_major_strides(shape) + pieces_by_rank = {} + shard_shapes = {} + start = 0 + for rank, size in enumerate(per_rank_sizes): + shard_shape = list(shape) + shard_shape[partition_dim] = size + shard_shape = tuple(shard_shape) + shard_shapes[rank] = shard_shape + pieces_by_rank[rank] = [ + AffinePiece(shape=shard_shape, + source_offset=start * source_strides[partition_dim], + source_strides=source_strides, + dest_offset=0, + dest_strides=_row_major_strides(shard_shape), + locations=[rank]) + ] + start += size + return ParamAffineMap(logical_shape=shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) + + +def sub_param_map(shape, sub_dim_sizes, shard_widths, partition_dim): + """A parameter that is several sub-parameters concatenated, each split across ranks. + + Fused QKV is the motivating case: the full parameter is Q then K then V along + ``partition_dim``, and a rank's shard holds its slice of each in turn. ``shard_widths[i]`` + gives the per-rank widths of sub-parameter ``i``, so the sub-parameters may be split + unevenly and need not be the same size as each other. + """ + shape = tuple(shape) + source_strides = _row_major_strides(shape) + tp_degree = len(shard_widths[0]) + + shard_shapes = {} + for rank in range(tp_degree): + shard_shape = list(shape) + shard_shape[partition_dim] = sum(widths[rank] for widths in shard_widths) + shard_shapes[rank] = tuple(shard_shape) + + pieces_by_rank = {rank: [] for rank in range(tp_degree)} + dest_starts = {rank: 0 for rank in range(tp_degree)} + sub_param_start = 0 + for sub_index, sub_size in enumerate(sub_dim_sizes): + widths = shard_widths[sub_index] + source_start = sub_param_start + for rank in range(tp_degree): + piece_shape = list(shape) + piece_shape[partition_dim] = widths[rank] + dest_strides = _row_major_strides(shard_shapes[rank]) + pieces_by_rank[rank].append( + AffinePiece(shape=tuple(piece_shape), + source_offset=source_start * source_strides[partition_dim], + source_strides=source_strides, + dest_offset=dest_starts[rank] * dest_strides[partition_dim], + dest_strides=dest_strides, + locations=[rank])) + source_start += widths[rank] + dest_starts[rank] += widths[rank] + sub_param_start += sub_size + + return ParamAffineMap(logical_shape=shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) + + +def _piece_to_dict(piece): + entry = { + 'shape': list(piece.shape), + 'source': [piece.source_offset, list(piece.source_strides)], + 'dest': [piece.dest_offset, list(piece.dest_strides)], + 'locations': sorted(piece.locations), + } + # Most pieces are unscaled, so leaving the default out keeps the stored map smaller + # and lets a reader that predates scaling still make sense of one that does not use it. + if piece.scale != 1.0: + entry['scale'] = piece.scale + return entry + + +def _piece_from_dict(entry): + source_offset, source_strides = entry['source'] + dest_offset, dest_strides = entry['dest'] + return AffinePiece(shape=entry['shape'], + source_offset=source_offset, + source_strides=source_strides, + dest_offset=dest_offset, + dest_strides=dest_strides, + locations=entry['locations'], + scale=entry.get('scale', 1.0)) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md new file mode 100644 index 000000000000..11a507db4d3d --- /dev/null +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -0,0 +1,569 @@ +# Affine IR for Universal Checkpoint + +Specification for the geometric shard description in `deepspeed/checkpoint/affine.py`. + +Design discussion: [#8252](https://github.com/deepspeedai/DeepSpeed/issues/8252), +sub-issue of [#8230](https://github.com/deepspeedai/DeepSpeed/issues/8230). Measurements +here were taken against `master` at `92843ad70`. + +Code is referred to by symbol rather than line number, so that this document does not go +stale every time a file above it shifts. + +--- + +## 1. Scope and boundary + +Universal Checkpoint stores every parameter as one full logical tensor `F`. Getting into +and out of that form currently requires knowing *what a parameter means* — is it a +vocabulary embedding, is it row-parallel, does it hold fused sub-parameters — and that +knowledge is encoded as regex lists over parameter names. This spec replaces the meaning +with geometry. + +**The rule this document follows throughout:** + +> The IR states what is true. The planner decides what to do. + +Everything that is a fact about how bytes are laid out belongs in the IR. Everything that +is a choice — which replica to read from, what order to move things in, where to place a +tensor next — belongs to the caller. Section 4 applies this test to every key in the +current schema, and two of them (`output_shape`, `target_partition_shape`) fail it: they +describe a destination, not a fact about the source, and so they are not part of the IR. + +**Non-goals.** This spec does not change how any layer partitions its weights at runtime. +It changes only how that partitioning is *described*. Every layout in tree today keeps the +exact bytes it has on every rank. + +--- + +## 2. The representation + +### 2.1 Piece + +A **piece** is a block of elements, described by where it sits in the full tensor `F` and +where it sits in the shard: + +```python +Piece = (shape, source_offset, source_strides, dest_offset, dest_strides, locations, scale) +``` + +- `shape` — the extent of the block, **shared by both sides**: a piece holds the same + elements in the same arrangement wherever it lives +- `source_offset`, `source_strides` — where the block sits in `F` +- `dest_offset`, `dest_strides` — where it sits in the rank's shard +- `locations` — the set of ranks holding this piece +- `scale` — the factor the shard holds the block by: `shard == full * scale` + +`locations` is not the same thing as which rank's shard a piece belongs to. A map lists +pieces per rank because that is what reading or writing one shard needs; `locations` says +which *other* ranks hold identical data for that same block. For an ordinary sharded +piece the two coincide. For `bigcodetype`'s replicated kv block they do not: the block +appears in every rank's list, each entry naming all four ranks, and a converter reads it +once from whichever rank is cheapest. + +Each side is exactly `torch.as_strided`'s argument list, so a piece is directly executable +with no interpretation step, in either direction. + +`scale` exists because a row-parallel layer whose output is sum-all-reduced pre-divides its +replicated bias by the world size, so that summing the outputs adds the bias exactly once. +The divisor changes with the world size, so a checkpoint that does not record it cannot be +restored at a different TP degree without a rule naming which parameters are biases — which +is the kind of semantic category this IR exists to remove. §4.3 draws the line between this +and the average op. + +**`scale` applies to the parameter, not to everything stored beside it.** A checkpoint also +carries optimizer state, and Adam's moments live in the parameter's *scaled* coordinate: +scaling a parameter by `s` scales its gradient by `1/s`, so the first moment carries `s⁻¹` +and the second `s⁻²` where the parameter carries `s`. A converter that applies the +parameter's own factor to all three corrupts the optimizer state and changes the trajectory +after a resume. The map records the geometry and the factor; the caller says which power of +it applies to the tensor being moved. + +**Worked example.** A row-parallel bias `b` at world size 4. Each rank stores `p = b/4`, so +the piece records `scale = 1/4` and every rank appears in `locations`. Writing `s` for that +factor, a shard holds `full * s**power`, and conversion recovers `full = shard / s**power`: + +| state | power | `s**power` | a rank holds | `F` recovers to | +|---|---|---|---|---| +| `fp32` | 1 | 1/4 | 2.0 | **8.0** — the logical bias | +| `exp_avg` | -1 | 4 | 8.0 | **2.0** | +| `exp_avg_sq` | -2 | 16 | 32.0 | **2.0** | + +The moments move the *opposite* way to the parameter, and the second moment twice as far. +The reason is the chain rule: the optimizer trains `p`, and `∂L/∂b = (∂L/∂p)·(1/4)`, so the +first moment of `b` is the stored moment divided by 4 and the second is divided by 16. Using +the parameter's own factor for all three would multiply `exp_avg` by 4 where it should be +divided — off by 16× — and silently change the trajectory after a resume. + +**Optimizer states are refused for now.** The moment powers above are correct — verified +end to end on real training in #8385 — but they are not sufficient. Adam's update also +depends on `lr` and `eps`, and those live in the coordinate the optimizer was training in: +resuming a rescaled parameter needs `lr / s` and `eps * s`. Keeping the source values +resumes on a different trajectory, with an error that grows every step and nothing to +signal it. Since that transform is a property of the optimizer rather than of the +parameter's geometry, it is outside this IR, and `rebuild` refuses a scaled optimizer +state until the checkpoint contract covers it. A scaled *parameter* converts normally. + +**Homogeneity.** A piece must cover elements that are all held by the same set of ranks and +all carry the same scale. This is what makes `locations` exact rather than advisory, and it +constrains merging: see §3 (P5) and §6.3. + +The shard side has to be described explicitly. It is tempting to say a shard is simply its +pieces concatenated in order, but that is false for a **column** split, where the shard +interleaves its pieces row by row rather than appending them. An implementation built on +the concatenation assumption reproduces row-split layouts correctly and silently +transposes column-split ones. + +### 2.2 Parameter map + +A parameter's map pairs each rank's shard shape with the pieces that fill it: + +```python +ParamMap = { logical_shape, { rank: (shard_shape, [Piece, ...]) } } +``` + +Because both ends of every piece are affine views, the two directions are the same loop +with the copy reversed — `extract` writes `source_view -> dest_view`, `rebuild` writes +`dest_view -> source_view`. Neither direction needs to know what the parameter *means*. + +**Implementation note.** Piece offsets are storage offsets, because that is what +`as_strided` takes. A shard handed in by a caller is frequently a *view* into a larger +buffer — splitting a tensor produces exactly that — and its elements then begin partway +into that storage. Applying a piece to such a tensor reads from the wrong place, with no +error raised. An implementation must normalise a shard to a buffer starting at its first +element before applying pieces to it. + +### 2.3 The two maps + +- `M_s` — the **source** map: the layout the checkpoint was written from. Conversion to + universal form is `M_s⁻¹`. +- `M_t` — the **target** map: the layout being restored into. + +Restore is `M_t`. Phase 2 — moving weights between two live topologies without ever +materialising `F` — is `M_t ∘ M_s⁻¹`. + +### 2.4 Why it is a relation, not a function + +Under replication, one position in `F` lives on several ranks, so `M_s⁻¹` maps one element +to a *set* of locations. This is deliberate. A function would have to name a single owner, +and naming an owner is a decision — the cheapest source depends on topology, link +bandwidth, and what else is in flight, none of which the checkpoint format knows. The IR +records that all of these ranks hold the byte; the planner picks one. + +This is the same argument that removes the average/combine op (§4.3): a reduce is not +invertible, so keeping it would make `M_s⁻¹` undefined for part of the language. + +--- + +## 3. Properties + +For a parameter map to be usable it must satisfy: + +**P1 — Coverage.** The union of all pieces over all ranks covers every element of `F` +exactly once or more. If some element is covered by no piece, `F` cannot be rebuilt. + +**P2 — Consistency under replication.** Where pieces from different ranks overlap, the +underlying data is identical. This is what makes picking any one of them safe, and it is +checkable — the current converter already asserts it for replicated parameters +(the replicated branch of `merge_tp_slices`). + +**P3 — Invertibility.** A piece may carry an invertible elementwise map — a non-zero +`scale` — but never a reduction over several source elements. Given P1–P3, `M_s⁻¹` is total +and well defined as a relation. + +**P4 — Composability.** `M_t ∘ M_s⁻¹` is again a set of affine pieces, so a transfer plan +can be computed without materialising `F`. This is not assumed: §8.2 measures it against +the real partition functions for every in-tree layout. It holds, but only for a composer +that folds pieces in N dimensions and groups them by source rank — §8.2 states both as +normative, because a composer missing either produces a correct map whose size grows with +model size. + +**P5 — Homogeneity.** Every element a piece covers is held by the same set of ranks, and +that set is the piece's `locations`. This makes `locations` authoritative: if pieces `P` +(locations `L`) and `Q` (locations `M`) share an element `e`, then P5 gives +`holders(e) = L` and `holders(e) = M`, so `L = M`. **Two pieces with different location +sets cannot overlap.** A planner can therefore trust one rank's map about what is shared, +instead of scanning every other rank to discover replication it was not told about. + +P5 does not forbid overlap. Two pieces with the *same* location set may overlap partially; +both already declare the same sharing, so nothing is hidden. + +P3 is the property that forces the average op out of the language, and the property that +vocabulary truncation currently violates (§8). + +--- + +## 4. Lowering: every current key + +This is the compatibility contract. Every key that exists today must have a defined image +in the IR, or be explicitly declared out of scope. + +### 4.1 Convert path — `merge_tp_slices` in `ds_to_universal.py` + +The current implementation is a branch chain over regex-matched name categories. Each +branch lowers to a piece list: + +| Branch (`ds_to_universal.py`) | Current behaviour | Lowers to | +|---|---|---| +| `TP_REPLICATED_PARAMETER_PATTERNS` | take `slices[0]`, assert all equal | **1 piece** covering all of `F`, `locations` = every rank | +| `PARAMETER_TO_AVERAGE_PATTERNS` | `sum(slices) / len(slices)` | **not a view** — see §4.3 | +| `PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0` | chunk each slice in 2, cat groupwise | **2 pieces/rank** | +| `PARAMETER_WITH_SUB_PARAMS` + `SUB_PARAM_SHAPE` | N sub-params, per-rank widths | **N pieces/rank** | +| default | `cat(slices, dim=cat_dim)` where `cat_dim = 1` if `PARAMETER_WITH_ROW_PARALLELISM_PATTERNS` else `0` | **1 piece/rank**, split along `cat_dim` | +| `VOCABULARY_PARAMETER_PATTERNS` | `param[:original_vocab_size, :]` | **not a view** — see §8 | + +Two keys are written *into* the produced checkpoint rather than read from metadata, and +both become redundant: + +| Key | Why it disappears | +|---|---| +| `CAT_DIM` | absorbed into `strides` — a row split and a column split differ only in stride | +| `PARAM_N_SUB_PARAMS` | absorbed into the length of the piece list | + +`SUB_PARAM_SHARD_WIDTHS` (UCP 0.4, #8185) is the closest thing in tree to the IR already: +it records the physical extent each rank holds of each sub-parameter. It lowers directly +to per-piece `shape` entries, and it is the key that makes uneven splits expressible. + +### 4.2 Restore path — `_resolve_autotp_partition` in `universal_checkpoint.py` + +| Meta key | Current use | Lowers to | +|---|---|---| +| `replicated` | return all of `full_hp_param` | 1 piece, all of `F` | +| `partition_dim` | axis to narrow on | the axis whose piece offsets vary | +| `logical_shape` | `full_hp_param.view(...)` | `shape` of `F` | +| `sub_param_sizes` / `sub_param_shape` (:82, :81) | `_narrow_sub_params` | N pieces | +| `sub_param_shard_widths` | per-rank widths inside `_narrow_sub_params` | per-piece `shape` | +| `partition_sizes` | `narrow(dim, sum(sizes[:rank]), sizes[rank])` | 1 piece, offset = prefix sum | +| *fallback* (:141-148) | `full_view.chunk(world)[rank]`, asserts divisibility | 1 piece, even split — **the assert disappears**, because uneven extents are expressible | +| `unsupported_reason` (:91-93) | raises `RuntimeError` | **eliminated** — see §5 | + +`output_shape` and `target_partition_shape` are carried in the restore meta +(by `_build_param_uc_restore_meta`) but describe the *destination*. By the boundary rule they are planner +inputs, not IR. They stay where they are; the IR does not absorb them. + +### 4.3 Reductions, and why `scale` is not one + +`PARAMETER_TO_AVERAGE_PATTERNS` is the one branch that cannot be expressed: an average is a +reduce over several source elements, and it is the one branch that would break P3 for the +whole language. + +The distinction that matters is invertibility, not arithmetic. A reduce is `N -> 1` and +destroys information, so `M_s⁻¹` is undefined on it. A scale is `1 -> 1` and is a bijection +for any non-zero constant, so `M_s⁻¹` is the reciprocal. The IR therefore admits an +invertible elementwise map and refuses a reduction. + +Dropping it is safe in tree. The constant has exactly three references repo-wide — its +declaration (in `constants.py`), its import, and its `.get` in the converter — and **no +producer anywhere in the repository** writes the key. Nothing in tree can reach that +branch. This matches delock's reading in #8230 that it papers over a historical rank-drift +issue rather than describing a real layout. + +--- + +## 5. Coverage: the layouts #8185 marks unsupported + +#8185 introduced `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` — layouts whose partitioning the +current schema cannot describe, which conversion therefore refuses. If the IR is to +replace the semantic categories, it has to cover these. This was tested rather than +assumed. + +**Method.** Per layout, per rank: run the real DeepSpeed partition function on a marker +tensor to recover which element of `F` each shard element came from; compress that index +map into maximal runs to obtain pieces; then re-run the same partition function on +**independent random data** and check that those pieces, materialised with +`torch.as_strided`, reproduce the shard bit-exactly. The last step is the actual test — if +pieces derived from markers reproduce a random tensor's shard, the mapping is a pure view +(P3), not something data-dependent. + +**Result — all four are affine-expressible.** + +| Layout | Reason string in `master` | `F` | tp | Pieces/rank | Covers `F` | +|---|---|---|---|---|---| +| `bigcodetype` (`fused_LinearLayer`, `fused_LinearLayer`) | "interleaves or replicates blocks … cannot be reassembled from the shards" | 48×8 | 4 | 2 | yes | +| `codegentype` (same class) | same | 24×8 | 2 | 11 | yes | +| Yuan value, dim 0 (`Yuan_LinearLayer`, `Yuan_LinearLayer`) | "selects noncontiguous head groups … cannot currently describe" | 32×8 | 2 | 2 | yes | +| Yuan o_proj, dim 1 (`Yuan_LinearAllreduce`, `Yuan_LinearAllreduce`) | same | 8×32 | 2 | 2, strides `(32, 1)` | yes | + +Piece count is bounded by block structure, not model size — CodeGen holds at 11 +pieces/rank while rows-per-rank grows 32×: + +``` + hidden F rows tp kv_heads rows/rank pieces/rank + 8 24 2 8 12 11 + 16 48 2 8 24 11 + 32 96 2 8 48 11 + 64 192 2 8 96 11 + 128 384 2 8 192 11 + 256 768 2 8 384 11 +``` + +Three consequences for this spec: + +1. **No information is lost in any of these cases.** All four cover `F` exactly (P1), and + Yuan's pieces are disjoint — it was always a clean partition. "Cannot be reassembled" is + a limit of the *description format*, not of the data. That is what makes it fixable by + an IR change alone. + +2. **`bigcodetype` is the concrete case for the location set.** Its kv block is + byte-identical on every rank — partial replication *within a single parameter*. A schema + with one `partition_dim` per parameter cannot say "these rows sharded, those rows + replicated"; a per-piece `locations` set says it with no extra mechanism. GPTBigCode / + StarCoder is an in-tree model this unblocks. + +3. **Strides are load-bearing.** Yuan's dim-1 case is `strides=(32, 1)` — a column + selection, not a contiguous span. The representation cannot be narrowed to "a list of + contiguous ranges." + +4. **Merging must stop at a replication boundary.** On bigcode's last rank the q slice ends + exactly where the replicated kv block begins, so merging by adjacency alone produces one + piece that is rank-private in its first half and replicated in its second — and a single + `locations` cannot describe it. Splitting there (P5) costs one extra piece across the + whole parameter, 7 to 8 at TP=4; every other layout was already homogeneous. + +The harness that produced this table is now a test — +`tests/unit/checkpoint/test_affine_shard_map.py`, 14 cases covering coverage, round-trip +against unseen data, and rebuild-inverts-extract for all four layouts. It needs no process +group or accelerator, because the partition functions take an explicit rank. + +--- + +## 6. On-disk format + +### 6.1 What is stored, and what is not + +The file stores **`M_s` only** — how the job that wrote the checkpoint had the parameter +laid out. `M_t` is not stored and must not be: the job doing the restore derives its own +map from its own layers, and a target map baked into the file would be a map for somebody +else's topology. + +Within `M_s`, §1's boundary rule splits the content in two: + +| | stable across jobs? | stored | why | +|---|---|---|---| +| geometry (`shape`, offsets, strides) | yes — a property of the parameter | **yes** | this is what makes the checkpoint portable | +| `locations` (which ranks held a piece) | no — a property of one grid | **yes, as provenance** | the converter must know which shard file to read | + +Storing locations is not a violation of the boundary. The file records *where the bytes +were*, which is a fact; it does not record where they should go next, which is a decision. +A reader that wants a different placement ignores the field entirely. + +### 6.2 The document + +Written into `UNIVERSAL_CHECKPOINT_INFO` under a new `affine_map` key, keyed by the same +exact-match patterns `collect_autotp_universal_checkpoint_info` already emits — one +`rf"^{re.escape(full_name)}$"` per parameter: + +```python +{ + "affine_map": { + "version": 1, + "params": { + "^transformer.h.0.attn.c_attn.weight$": { + "logical_shape": [48, 8], + "ranks": { + 0: {"shard_shape": [16, 8], + "pieces": [ + {"shape": [8, 8], "source": [0, [8, 1]], "dest": [0, [8, 1]], "locations": [0]}, + {"shape": [8, 8], "source": [256, [8, 1]], "dest": [64, [8, 1]], "locations": [0,1,2,3]} + ]}, + ... + } + } + } + } +} +``` + +Each piece is `shape` plus `[offset, strides]` on each side — the argument list of +`torch.as_strided` twice over, so a reader applies it with no interpretation step. `scale` +is omitted when it is 1, which is almost every piece. + +Per-parameter entries live under `params` rather than beside `version`, so that a parameter +whose name matched a metadata key could never be confused for one. + +### 6.3 Rules + +- **Additive.** A checkpoint carrying `affine_map` also carries today's keys. Old readers + ignore the new key and take the existing branch chain; new readers prefer `affine_map` + when present and fall back otherwise. Nothing written today becomes unreadable, and + nothing written by this scheme becomes unreadable by an old converter. +- **Version.** `UNIVERSAL_CHECKPOINT_VERSION_VALUE` goes 0.4 → 0.5. The inner + `"version": 1` covers the map encoding itself, so the two can move independently. +- **Pieces are canonical.** A writer must emit maximally merged pieces, folded in N + dimensions (§8.2), **subject to homogeneity** (P5): a merge may not cross a change in + `locations` or in `scale`. Both halves are normative. Skipping the merge gives a map whose + size grows with model size instead of block structure; skipping the homogeneity constraint + gives a map whose `locations` is wrong, which is worse because it still round-trips + correctly on a single topology and only misleads a phase-2 planner. +- **Plain scalars only.** No tensors, no pickled classes, no `SubparamShape` objects. The + map should be readable without importing DeepSpeed, which matters for external tooling + and for debugging a checkpoint that will not load. +- **`locations` is per piece, not per parameter.** This is the field that expresses + `bigcodetype`, whose kv block is held identically by every rank (§5). + +### 6.4 Size + +Measured on a fused-QKV parameter at TP=8: 96 pieces, 6836 bytes of compact JSON — +about **71 bytes per piece**. Extrapolated by tensor count: + +| model | tensors | `affine_map` | +|---|---|---| +| Llama-3-8B | ~291 | ~1.9 MB | +| Llama-3-70B | ~723 | ~4.7 MB | + +Megabytes, not kilobytes, and it scales with **tensor count × TP degree** rather than with +parameter size — a 70B and a 7B with the same layer count cost the same. Against +checkpoints measured in hundreds of gigabytes this is negligible, but it is large enough +that the per-rank redundancy (most ranks differ only in offset) is worth revisiting if the +map is ever loaded somewhere latency-sensitive. That is an encoding question and is +deliberately outside the semantics. + +--- + +## 7. Phase 1 — the converter + +### 7.1 The algorithm + +With the map present, `merge_tp_slices` loses its branch chain entirely: + +``` +for each parameter: + F = empty(logical_shape) + for rank, pieces in M_s: + shard = flat_buffer(load_shard(rank)) + for piece in pieces: + F[piece.source_view] = shard[piece.dest_view] + assert M_s.uncovered_offsets() == [] + if the parameter is a padded vocabulary: + F = F[:V_orig] # outside the IR, see §8.1 + write F +``` + +Restore is the same loop with the copy reversed, against the target's own `M_t`: + +``` +for each rank r of this job: + shard = empty(M_t.shard_shape[r]) + for piece in M_t.pieces[r]: + shard[piece.dest_view] = F[piece.source_view] +``` + +Both directions are driven by the same piece list. That is the structural reason the +current convert/restore asymmetry — the two builders `_build_param_uc_conversion_meta` and `_build_param_uc_restore_meta` emitting two +different key sets — stops being expressible. + +### 7.2 Two things an implementation must get right + +**Normalise the shard buffer.** Piece offsets are storage offsets, because that is what +`as_strided` takes. A loaded shard is often a *view* into a larger buffer whose elements +begin partway into that storage — splitting a tensor produces exactly that — and applying +a piece to it then reads from the wrong place, silently and with no error. `flat_buffer` +above is not decorative. + +**Take any one holder of a replicated piece.** Where `locations` has more than one rank, +every listed rank holds identical data by construction (P2), so the converter reads +whichever is cheapest and skips the rest. Phase 1 is the degenerate case of the location +set; phase 2 is where the choice becomes interesting. + +### 7.3 What changes + +Removed: the regex categories, the `cat_dim` decision, the sub-parameter narrowing +arithmetic, the `chunk` fallback with its divisibility assert +(the `chunk` fallback in `_resolve_autotp_partition`), and the `unsupported_reason` refusal path +(its `unsupported_reason` guard) — the layouts it refuses are expressible (§5). + +Added: a coverage assertion that is a real correctness property. Today's checks are shape +heuristics — they confirm the numbers line up, not that every element of the parameter has +a home. `uncovered_offsets()` answers the second question, and it is the check that fails +loudly if a future layout outgrows the language. + +Unchanged: parallelism. Pieces write to disjoint regions of `F` (or identical values where +they overlap), and parameters remain independent, so the existing `ProcessPoolExecutor` +fan-out over `merge_tp_slices` carries over untouched. + +--- + +## 8. Open cases + +**8.1 Vocabulary truncation.** The vocabulary branch of `merge_tp_slices` does +`param = param[:original_vocab_size, :]`, dropping padding rows. This is not a view, so it +violates P3 and makes `M_s⁻¹` non-total over the shards. Two options: + +- **(a)** `F` carries the padding; truncation becomes a post-step outside the IR. Every + piece stays a view and `M_s⁻¹` stays well defined for the whole language. Cost: the + padded size must be recorded. +- **(b)** The IR admits pieces that map to nothing. Cost: invertibility, i.e. P3, for the + whole language rather than for one case. + +**Resolved: (a)**, agreed with delock on #8252. One key is cheaper than one weakened +property, and the case is narrow — only legacy TP needs an even split and therefore +padding at all; AutoTP's uneven sharding does not pad `F`. + +**(a) also survives phase 2**, which was the real risk. When a source and target topology +pad to different heights (TP=4 → 3 padding rows, TP=2 → 1), the truncate/pad step is the +identity restricted to `[0, V_orig)`. Restricting an affine map to a range leaves it +affine, so `M_t ∘ truncate ∘ M_s⁻¹` composes to plain overlap arithmetic and neither +`Fsrc` nor `Ftgt` is ever materialised. Verified over all 64 source/target pairs for +TP ∈ {1,2,4,8} and vocab ∈ {13, 32, 100, 50257}, in both the truncating and padding +directions, with source padding poisoned so that reading it would fail the check. + +The composed map is **partial over the target** — target padding rows have no source. +That is much weaker than `M_s` being non-total: padding is not data, and the target has to +initialise it regardless. Piece count is bounded by the two TP degrees rather than by the vocabulary: across +vocab 13 → 262144 and TP up to 16 it never exceeds `src_tp + tgt_tp`, though the exact +count varies with how the two padded heights align. + +Consequence: legacy TP does not need excluding from phase 2, and padding need not enter +the UC file even in legacy mode — the composition needs only `V_orig`. + +Relevant: `PADDED_VOCAB_SIZE` is declared at `constants.py` and referenced **nowhere +else in the repository** (verified against `92843ad70`) — it looks like exactly this slot, +never wired up. Worth confirming before adding a new key. + +**8.2 Do we ever need explicit indices? — RESOLVED: no.** Every layout in tree is +expressible with affine pieces at bounded count (§5), and composition `M_t ∘ M_s⁻¹` between +two TP degrees stays affine for all of them, verified against the real partition functions +with `F` never materialised. No escape hatch to explicit index lists is needed. + +Two rules are **normative for any composer**, because getting either wrong yields a piece +count that grows with model size — which looks like evidence the language is insufficient: + +- **Compress in N dimensions.** A column split emits one run per row with identical strides + and constant start deltas; that is an outer dimension, and the stack is one piece. In 1-D + only, Yuan o_proj needs one piece per row (32 → 1024 as rows go 4 → 128); in N-D it is 4. +- **Group runs by source rank before folding.** A target shard interleaves blocks from + several source ranks, so runs that stack are usually not adjacent in shard order. For + Yuan o_proj 4→2 they alternate rank 0/rank 1, and a neighbour-only merge never fires: + 256 pieces instead of 4. + +Grouping is only legal because a piece carries its own destination offset (§2.1), so pieces +may be reordered freely. + +Measured piece counts are bounded by topology and block structure, not model size: across +an 8× hidden sweep, codegen 4→8 holds at 40, bigcode 2→4 at 7, Yuan o_proj at 4. + +**8.3 Scale and floating point.** `(x / N) * N` is exact for power-of-two `N` in fp32, bf16 +and fp16 alike, because dividing by `2^k` only shifts the exponent — so a bias round-trips +bit-exactly at every TP degree in normal use. It is lossy for non-power-of-two `N` (3, 6, +12), where a converted-and-restored bias may differ in the last bits from the original. + +**8.4 ZeRO and offload placement.** delock's extension in #8230 — "a subset of a parameter +combined with a list of ranks holding this subset" — is what §2.1's `locations` implements. +ZeRO-1/3 partitions and offload replicas should fall out as pieces whose `locations` +describe the DP group rather than the TP group, but this spec does not yet work through +AutoEP's expert placement, where locations are per-expert rather than per-parameter. + +--- + +## 9. Staging + +1. **Done.** IR data structure (`deepspeed/checkpoint/affine.py`: `AffinePiece`, + `ParamAffineMap`), behind no flag — pure addition, nothing reads it yet. +2. **Done.** The §5 harness as a test, asserting all four layouts lower and round-trip. +3. Converter reads `affine_map` when present, existing chain otherwise. Parity test: + both paths produce byte-identical `F` for every layout in tree. +4. Emit `affine_map` from `collect_autotp_universal_checkpoint_info`, and drop + `unsupported_reason` for the four layouts §5 covers. +5. Phase 2 — `M_t ∘ M_s⁻¹` as a transfer plan — builds on this and is out of scope here. + +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. diff --git a/deepspeed/checkpoint/constants.py b/deepspeed/checkpoint/constants.py index cdf846c1f322..536db59cfed9 100644 --- a/deepspeed/checkpoint/constants.py +++ b/deepspeed/checkpoint/constants.py @@ -62,6 +62,12 @@ # Attribute name used to store AutoTP universal-checkpoint metadata on torch Parameters. DS_AUTOTP_UC_META = "ds_autotp_universal_checkpoint_meta" AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS = "autotp_unsupported_parameter_patterns" +# Geometric description of how each parameter is sharded, keyed by parameter pattern. +# Written alongside the pattern lists below rather than replacing them, so a converter +# that predates it simply does not see the key. See checkpoint/affine_ir_spec.md. +AFFINE_MAP = 'affine_map' +AFFINE_MAP_VERSION = 'version' +AFFINE_MAP_PARAMS = 'params' # Vocabulary padding VOCAB_TENSOR = 'vocab_tensor' diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index b62a88ce7252..ec9c728db4a6 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -41,6 +41,9 @@ PIPELINE_REPLICATED_PARAMETER_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_TO_AVERAGE_PATTERNS, + AFFINE_MAP, + AFFINE_MAP_PARAMS, + AFFINE_MAP_VERSION, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, PARAMETER_WITH_SUB_PARAMS, @@ -57,6 +60,7 @@ is_autoep_zero3_partitioned_entry, validate_autoep_zero3_partitioned_metadata, ) +from deepspeed.checkpoint.affine import ParamAffineMap, AFFINE_MAP_FORMAT_VERSION def parse_arguments(): @@ -300,6 +304,13 @@ def merge_tp_slices(uc_info, dir, slice_dir, tp_degree, name_and_shapes): parameters_with_2_sub_params_cat_dim_0 = universal_checkpoint_info.get(PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, []) parameter_with_sub_params = universal_checkpoint_info.get(PARAMETER_WITH_SUB_PARAMS, []) sub_param_shard_widths = universal_checkpoint_info.get(SUB_PARAM_SHARD_WIDTHS, {}) + affine_map_info = universal_checkpoint_info.get(AFFINE_MAP, {}) + affine_map_version = affine_map_info.get(AFFINE_MAP_VERSION, AFFINE_MAP_FORMAT_VERSION) + if affine_map_version > AFFINE_MAP_FORMAT_VERSION: + raise RuntimeError( + f"Checkpoint records affine map format version {affine_map_version}, but this DeepSpeed understands " + f"up to {AFFINE_MAP_FORMAT_VERSION}. Reading it could misinterpret fields added since.") + affine_params = affine_map_info.get(AFFINE_MAP_PARAMS, {}) uc_version = universal_checkpoint_info.get(UNIVERSAL_CHECKPOINT_VERSION_KEY, 0.0) unmatched_patterns = set(replicated_parameters + parameters_to_average + parameters_with_row_parallelism + @@ -324,12 +335,30 @@ def get_matched_sub_params_pattern(name_): return subparam_shape, pattern_ return None, None + def get_matched_affine_map(name_): + """An affine map describes the layout geometrically, so it needs no category. + + Its patterns are a separate namespace from the category lists, and are not part of + `unmatched_patterns`: a map covers only the parameters it was written for, and the + rest still take the branches below. + """ + matched_ = [pattern_ for pattern_ in affine_params if re.match(pattern_, name_)] + assert len(matched_) <= 1, f'Got more than one matching affine map patterns={matched_} for {name_}' + return ParamAffineMap.from_dict(affine_params[matched_[0]]) if matched_ else None + + matched_affine_map = get_matched_affine_map(name) matched_sub_params_shape, matched_sub_params_pattern = get_matched_sub_params_pattern(name) step_merged = _merge_zero_shards(slice_base_path, "step", tp_degree, per_tp_shapes) if step_merged: _save_checkpoint(os.path.join(param_base_path, "step.pt"), step_merged[0]) + # How a piece's scale applies to each state. Scaling a parameter by `s` scales its + # gradient by `1 / s`, so Adam's first moment carries the inverse and its second moment + # the inverse square. Using the parameter's factor for all three would corrupt the + # optimizer state and change the trajectory after a resume. + scale_powers = {"fp32": 1, "exp_avg": -1, "exp_avg_sq": -2} + for state in ("fp32", "exp_avg", "exp_avg_sq"): slices = _merge_zero_shards(slice_base_path, state, tp_degree, per_tp_shapes) final_path = os.path.join(param_base_path, f"{state}.pt") @@ -337,7 +366,14 @@ def get_matched_sub_params_pattern(name_): #print(f"Expected shape: {shape}") #print(f"Fragment sizes:", list(frag.shape for frag in slices)) ckpt_dict = {} - if get_matched_pattern(replicated_parameters, name): + if matched_affine_map is not None: + # The pieces say where every element of the parameter lives, so none of the + # category branches below are consulted. This branch only decides `param`; it + # writes none of the per-category keys those branches add to `ckpt_dict`, + # because the geometry is what a restoring job needs and it is not tied to a + # category. + param = matched_affine_map.rebuild(dict(enumerate(slices)), scale_powers[state]) + elif get_matched_pattern(replicated_parameters, name): if len(slices) > 1: assert all([slices[0].equal(other_slice) for other_slice in slices[1:]]) param = slices[0] diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 2932335b09bb..ce47a9de1946 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -4557,8 +4557,15 @@ def _load_checkpoint(self, is_pipe_parallel = isinstance(self.module, PipelineModule) - load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, - self.checkpoint_mp_rank, + checkpoint_mp_world_size = self.mp_world_size + checkpoint_mp_rank = self.checkpoint_mp_rank + if self.load_universal_checkpoint(): + # UC restores weights from zero/. Read metadata from the corresponding + # source rank without invoking Megatron's model-specific weight merger. + checkpoint_mp_rank = checkpoint_mp_rank * len(ckpt_list) // checkpoint_mp_world_size + checkpoint_mp_world_size = len(ckpt_list) + load_path, checkpoint, _ = sd_loader.load(checkpoint_mp_world_size, + checkpoint_mp_rank, is_pipe_parallel=is_pipe_parallel) if checkpoint is None: diff --git a/deepspeed/utils/tensor_fragment.py b/deepspeed/utils/tensor_fragment.py index 1947ec3d8853..e21de5c3a235 100644 --- a/deepspeed/utils/tensor_fragment.py +++ b/deepspeed/utils/tensor_fragment.py @@ -79,7 +79,7 @@ def map_to_flat_opt_states(flat_hp_tensor, lp_tensors, optim_state, opt_keys): hp_fragment_address = lp._hp_mapping.get_hp_fragment_address() hp_fragment = buffer.narrow(0, hp_fragment_address.start, hp_fragment_address.numel) hp_fragment.data.copy_(lp._hp_mapping.get_hp_fragment(optim_state_key=key).data) - lp._hp_mapping.hp_fragment = hp_fragment + lp._hp_mapping.optim_fragment[key] = hp_fragment optim_state[hp_param][key] = buffer diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py new file mode 100644 index 000000000000..9275ca14a690 --- /dev/null +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -0,0 +1,587 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Check that AutoTP's fused and shared-QK layouts are describable as affine views. + +These are the layouts `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` currently refuses to +convert, on the grounds that the parameter "cannot be reassembled from the shards". The +tests below show the shards do cover the full tensor, and that affine pieces reproduce +them exactly, so what is missing is a way to describe the layout rather than the data. + +Pieces are recovered by running the real partition function on a marker tensor and then +checked against *independent random data*. That second step is what makes this a test of +geometry: if pieces derived from markers reproduce a random tensor's shard bit-exactly, +the layout is a pure view and not something that depends on the values. + +The partition functions used here take an explicit rank, so nothing in this file needs a +process group or an accelerator. +""" + +import pytest +import torch + +from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, AFFINE_MAP_FORMAT_VERSION, replicated_map, + contiguous_split_map, sub_param_map, _scaled) +from deepspeed.module_inject.fusedqkv_utils import prepare_tp_fused_qkvw, shard_value_with_share_qk +from deepspeed.module_inject.tp_shard import AutoTPMeta + + +class _NamedModule(torch.nn.Module): + """`get_fused_qkv_type` picks a layout by matching the module's class name as text.""" + + def __init__(self, name): + super().__init__() + self._name = name + + def __str__(self): + return self._name + + +def _row_markers(rows, cols): + return torch.arange(rows, dtype=torch.float64).reshape(rows, 1).repeat(1, cols) + + +def _col_markers(rows, cols): + return torch.arange(cols, dtype=torch.float64).reshape(1, cols).repeat(rows, 1) + + +def _index_runs(indices): + """Group consecutive indices into (start_index, position_in_shard, length) runs. + + Compressing runs is what bounds the number of pieces: a contiguous block of the full + tensor needs one piece however long it is. + """ + runs = [] + start = 0 + for j in range(1, len(indices) + 1): + if j == len(indices) or indices[j] != indices[j - 1] + 1: + runs.append((indices[start], start, j - start)) + start = j + return runs + + +def _source_indices(markers, shard, axis): + """Recover which row (or column) of the full tensor each row of the shard came from.""" + if axis == 'row': + lookup = {tuple(markers[i].tolist()): i for i in range(markers.shape[0])} + keys = [tuple(shard[j].tolist()) for j in range(shard.shape[0])] + else: + lookup = {tuple(markers[:, i].tolist()): i for i in range(markers.shape[1])} + keys = [tuple(shard[:, j].tolist()) for j in range(shard.shape[1])] + + indices = [] + for key in keys: + assert key in lookup, 'a shard slice is not a slice of the full tensor, so the layout is not a view' + indices.append(lookup[key]) + return indices + + +def _split_by_holders(indices, holders): + """Break runs wherever the set of ranks holding the data changes. + + Merging across such a boundary would produce a piece whose own `locations` is wrong for + part of it. GPTBigCode's last rank is the case that forces this: its q slice ends exactly + where the replicated kv block begins, so unconstrained merging fuses a rank-private block + onto a fully replicated one. + """ + split = [] + for source_index, dest_index, length in _index_runs(indices): + start = 0 + for step in range(1, length + 1): + at_end = step == length + if at_end or holders[source_index + step] != holders[source_index + start]: + split.append((source_index + start, dest_index + start, step - start)) + start = step + return split + + +def _pieces_for_rank(markers, shard, rank, axis, cols, holders): + """Turn one rank's shard into affine pieces of the full tensor.""" + pieces = [] + indices = _source_indices(markers, shard, axis) + for source_index, dest_index, length in _split_by_holders(indices, holders): + if axis == 'row': + shape = (length, cols) + source_offset = source_index * cols + dest_offset = dest_index * cols + dest_strides = (cols, 1) + else: + shard_cols = shard.shape[1] + shape = (shard.shape[0], length) + source_offset = source_index + dest_offset = dest_index + dest_strides = (shard_cols, 1) + pieces.append( + AffinePiece(shape=shape, + source_offset=source_offset, + source_strides=(cols, 1), + dest_offset=dest_offset, + dest_strides=dest_strides, + locations=sorted(holders[source_index]))) + return pieces + + +def _build_map(shard_fn, rows, cols, mp_size, axis): + """Derive the affine map of a layout by probing the real partition function.""" + markers = _row_markers(rows, cols) if axis == 'row' else _col_markers(rows, cols) + shards = {rank: shard_fn(markers.clone(), rank) for rank in range(mp_size)} + + # Which ranks hold each slice of the full tensor. Needed before pieces can be cut, + # because a piece may not span slices with different holders. + holders = {} + for rank, shard in shards.items(): + for index in _source_indices(markers, shard, axis): + holders.setdefault(index, set()).add(rank) + + pieces_by_rank = {} + shard_shapes = {} + for rank, shard in shards.items(): + shard_shapes[rank] = tuple(shard.shape) + pieces_by_rank[rank] = _pieces_for_rank(markers, shard, rank, axis, cols, holders) + return ParamAffineMap(logical_shape=(rows, cols), shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) + + +def _fused_qkv_shard_fn(layout_name, mp_size, meta): + module = _NamedModule(layout_name) + + def shard_fn(full_param, rank): + return prepare_tp_fused_qkvw(module, full_param, mp_size, rank, meta) + + return shard_fn + + +def _shared_qk_shard_fn(shard_value, meta): + + def shard_fn(full_param, rank): + return shard_value_with_share_qk(full_param, None, rank, 2, shard_value, meta)[0].data + + return shard_fn + + +def _meta(num_kv_heads, n_embd=None, num_attention_heads=None): + """The per-model shard metadata AutoTP derives from a model config.""" + return AutoTPMeta(num_kv_heads=num_kv_heads, n_embd=n_embd, num_attention_heads=num_attention_heads) + + +# (id, rows, cols, mp_size, axis, make_meta, build_shard_fn) +LAYOUTS = [ + ('bigcode', 48, 8, 4, 'row', lambda: _meta(4, n_embd=32, num_attention_heads=4), + lambda mp, meta: _fused_qkv_shard_fn('GPTBigCodeBlock', mp, meta)), + ('codegen', 24, 8, 2, 'row', lambda: _meta(8, n_embd=8, num_attention_heads=8), + lambda mp, meta: _fused_qkv_shard_fn('CodeGenBlock', mp, meta)), + ('yuan_value', 32, 8, 2, 'row', lambda: _meta(8), lambda mp, meta: _shared_qk_shard_fn(True, meta)), + ('yuan_oproj', 8, 32, 2, 'col', lambda: _meta(8), lambda mp, meta: _shared_qk_shard_fn(False, meta)), +] + + +@pytest.mark.parametrize('name, rows, cols, mp_size, axis, make_meta, build_shard_fn', + LAYOUTS, + ids=[layout[0] for layout in LAYOUTS]) +class TestAffineShardMap: + + def test_shards_cover_the_full_parameter(self, name, rows, cols, mp_size, axis, make_meta, build_shard_fn): + """Every element of the full tensor is held by some rank, so it can be rebuilt.""" + affine_map = _build_map(build_shard_fn(mp_size, make_meta()), rows, cols, mp_size, axis) + assert affine_map.uncovered_offsets() == [] + + def test_every_piece_is_homogeneous(self, name, rows, cols, mp_size, axis, make_meta, build_shard_fn): + """No piece spans elements held by different sets of ranks, so locations is exact.""" + affine_map = _build_map(build_shard_fn(mp_size, make_meta()), rows, cols, mp_size, axis) + affine_map.validate_coverage() + + def test_pieces_reproduce_shards_of_unseen_data(self, name, rows, cols, mp_size, axis, make_meta, build_shard_fn): + """Pieces derived from markers must reproduce shards of data they were not derived from.""" + shard_fn = build_shard_fn(mp_size, make_meta()) + affine_map = _build_map(shard_fn, rows, cols, mp_size, axis) + + torch.manual_seed(0) + full_param = torch.randn(rows, cols, dtype=torch.float64) + for rank in range(mp_size): + expected = shard_fn(full_param.clone(), rank) + assert torch.equal(affine_map.extract(full_param, rank), expected) + + def test_rebuild_inverts_extract(self, name, rows, cols, mp_size, axis, make_meta, build_shard_fn): + """Rebuilding from the shards returns the parameter the shards were cut from.""" + shard_fn = build_shard_fn(mp_size, make_meta()) + affine_map = _build_map(shard_fn, rows, cols, mp_size, axis) + + torch.manual_seed(1) + full_param = torch.randn(rows, cols, dtype=torch.float64) + shards = {rank: shard_fn(full_param.clone(), rank) for rank in range(mp_size)} + assert torch.equal(affine_map.rebuild(shards), full_param) + + +def test_piece_count_does_not_grow_with_model_size(): + """Piece count follows the block structure of the layout, not the size of the tensor. + + This is what keeps the description small: if piece count tracked the tensor it would + be no cheaper than storing an index per element. + """ + piece_counts = set() + for hidden in (8, 16, 32, 64, 128): + meta = _meta(8, n_embd=hidden, num_attention_heads=8) + rows = 3 * hidden + affine_map = _build_map(_fused_qkv_shard_fn('CodeGenBlock', 2, meta), rows, 8, 2, 'row') + piece_counts.add(len(affine_map.pieces_by_rank[0])) + + assert len(piece_counts) == 1, f'piece count varied with model size: {sorted(piece_counts)}' + + +def test_replicated_piece_is_shared_by_every_rank(): + """GPTBigCode gives every rank the same kv block, which one owner per parameter cannot say. + + The kv rows appear in every rank's shard, so describing this layout needs replication + to be expressible for part of a parameter rather than all of it. + """ + meta = _meta(4, n_embd=32, num_attention_heads=4) + mp_size = 4 + affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size, meta), 48, 8, mp_size, 'row') + + kv_offsets = set(range(32 * 8, 48 * 8)) + for rank in range(mp_size): + held = set() + for piece in affine_map.pieces_by_rank[rank]: + held.update(piece.source_offsets()) + assert kv_offsets <= held, f'rank {rank} does not hold the whole kv block' + + +def test_pieces_never_span_a_replication_boundary(): + """GPTBigCode's last rank is where an unconstrained merge would cross one. + + Its q slice ends exactly where the replicated kv block begins, so merging by adjacency + alone yields one piece that is rank-private in its first half and replicated in its + second. Splitting at the boundary costs one extra piece and keeps locations honest. + """ + meta = _meta(4, n_embd=32, num_attention_heads=4) + mp_size = 4 + affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size, meta), 48, 8, mp_size, 'row') + + last_rank_pieces = affine_map.pieces_by_rank[mp_size - 1] + assert len(last_rank_pieces) == 2 + assert {len(piece.locations) for piece in last_rank_pieces} == {1, mp_size} + affine_map.validate_coverage() + + +def test_scale_round_trips_through_both_directions(): + """A row-parallel bias is replicated pre-divided by the world size, so a piece scales it. + + `shard == full * scale`, so extracting multiplies and rebuilding divides. Recording the + factor keeps the layout describable without a rule that names which parameters are biases. + """ + world_size = 4 + full_bias = torch.randn(16, dtype=torch.float64) + pieces_by_rank = { + rank: [ + AffinePiece(shape=(16, ), + source_offset=0, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=range(world_size), + scale=1.0 / world_size) + ] + for rank in range(world_size) + } + affine_map = ParamAffineMap(logical_shape=(16, ), + shard_shapes={rank: (16, ) + for rank in range(world_size)}, + pieces_by_rank=pieces_by_rank) + affine_map.validate_coverage() + + shards = {rank: affine_map.extract(full_bias, rank) for rank in range(world_size)} + assert torch.equal(shards[0], full_bias / world_size) + assert torch.equal(affine_map.rebuild(shards), full_bias) + + +def test_scale_round_trip_is_exact_for_power_of_two_world_sizes(): + """Dividing by a power of two only shifts the exponent, so no bias precision is lost.""" + full_bias = torch.randn(1024, dtype=torch.float32) + for world_size in (2, 4, 8, 16): + assert torch.equal(full_bias / world_size * world_size, full_bias) + + +# The universal-checkpoint converter merges tp slices with a chain of category-specific +# `torch.cat` arithmetic. The tests below rebuild the same parameter through the affine map +# and require the two to agree exactly, so the map can replace a branch without changing +# what any checkpoint converts to. + + +def _random_slices(shapes): + torch.manual_seed(0) + return [torch.randn(shape, dtype=torch.float64) for shape in shapes] + + +def test_parity_replicated(): + """`merge_tp_slices` takes slices[0] after asserting every rank matches.""" + tp_degree = 4 + full = torch.randn(6, 8, dtype=torch.float64) + slices = [full.clone() for _ in range(tp_degree)] + + expected = slices[0] + affine_map = replicated_map((6, 8), tp_degree) + assert torch.equal(affine_map.rebuild(dict(enumerate(slices))), expected) + + +@pytest.mark.parametrize('cat_dim', [0, 1]) +def test_parity_contiguous_split(cat_dim): + """The default branch: `cat(slices, dim=1)` for row parallelism, `dim=0` otherwise.""" + per_rank = [3, 3, 2, 2] # deliberately uneven; `chunk` would disagree here + shapes = [(size, 8) if cat_dim == 0 else (8, size) for size in per_rank] + slices = _random_slices(shapes) + + expected = torch.cat(slices, dim=cat_dim) + affine_map = contiguous_split_map(tuple(expected.shape), per_rank, cat_dim) + assert torch.equal(affine_map.rebuild(dict(enumerate(slices))), expected) + + +def test_parity_two_sub_params_cat_dim_0(): + """The 2-sub-param branch chunks each slice, merges each half, then concatenates.""" + tp_degree, half = 4, 3 + slices = _random_slices([(2 * half, 8)] * tp_degree) + + chunked = [torch.chunk(tp_slice, 2, dim=0) for tp_slice in slices] + expected = torch.cat([ + torch.cat([chunk[0] for chunk in chunked], dim=0), + torch.cat([chunk[1] for chunk in chunked], dim=0), + ], + dim=0) + + total = half * tp_degree + affine_map = sub_param_map(shape=(2 * total, 8), + sub_dim_sizes=(total, total), + shard_widths=[[half] * tp_degree, [half] * tp_degree], + partition_dim=0) + assert torch.equal(affine_map.rebuild(dict(enumerate(slices))), expected) + + +def test_parity_sub_params_with_uneven_widths(): + """The sub-parameter branch, with the per-rank widths #8185 added for uneven splits. + + Q, K and V are different sizes and none divides evenly by the tp degree, which is the + case the pre-0.4 metadata could not describe at all. + """ + tp_degree = 3 + shard_widths = [[3, 2, 2], [2, 2, 1], [1, 1, 1]] + sub_dim_sizes = [sum(widths) for widths in shard_widths] + rows_per_rank = [sum(widths[rank] for widths in shard_widths) for rank in range(tp_degree)] + slices = _random_slices([(rows, 8) for rows in rows_per_rank]) + + # Exactly the arithmetic in `merge_tp_slices`: for each sub-parameter, take every + # rank's block of it in turn, then concatenate the sub-parameters. + offsets = [0] * tp_degree + merged_chunks = [] + for widths in shard_widths: + blocks = [] + for rank, tp_slice in enumerate(slices): + blocks.append(tp_slice.narrow(0, offsets[rank], widths[rank])) + offsets[rank] += widths[rank] + merged_chunks.append(torch.cat(blocks, dim=0)) + expected = torch.cat(merged_chunks, dim=0) + + affine_map = sub_param_map(shape=(sum(sub_dim_sizes), 8), + sub_dim_sizes=sub_dim_sizes, + shard_widths=shard_widths, + partition_dim=0) + affine_map.validate_coverage() + assert torch.equal(affine_map.rebuild(dict(enumerate(slices))), expected) + + +def test_parity_round_trips_back_to_the_original_slices(): + """Extracting from the merged parameter returns the slices it was built from.""" + per_rank = [3, 3, 2, 2] + slices = _random_slices([(size, 8) for size in per_rank]) + affine_map = contiguous_split_map((10, 8), per_rank, 0) + + full = affine_map.rebuild(dict(enumerate(slices))) + for rank, tp_slice in enumerate(slices): + assert torch.equal(affine_map.extract(full, rank), tp_slice) + + +def test_serialisation_round_trips(): + """The stored form must rebuild the same parameter as the map it came from.""" + affine_map = sub_param_map(shape=(18, 8), + sub_dim_sizes=[9, 9], + shard_widths=[[4, 3, 2], [3, 3, 3]], + partition_dim=0) + restored = ParamAffineMap.from_dict(affine_map.to_dict()) + + torch.manual_seed(0) + slices = {rank: torch.randn(shape, dtype=torch.float64) for rank, shape in affine_map.shard_shapes.items()} + assert torch.equal(restored.rebuild(slices), affine_map.rebuild(slices)) + + +def test_serialisation_holds_only_plain_scalars(): + """Nothing torch-specific may reach the file, so the map stays readable on its own.""" + stored = contiguous_split_map((10, 8), [3, 3, 2, 2], 0).to_dict() + + def check(value): + if isinstance(value, dict): + for key, item in value.items(): + assert isinstance(key, (str, int)), f'unexpected key type {type(key)}' + check(item) + elif isinstance(value, list): + for item in value: + check(item) + else: + assert isinstance(value, (int, float, str)), f'unexpected value type {type(value)}' + + check(stored) + + +def test_serialisation_preserves_scale(): + """A scaled piece must survive the file, or a restored bias would be off by the divisor.""" + piece = AffinePiece(shape=(4, ), + source_offset=0, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=[0, 1], + scale=0.25) + affine_map = ParamAffineMap(logical_shape=(4, ), + shard_shapes={ + 0: (4, ), + 1: (4, ) + }, + pieces_by_rank={ + 0: [piece], + 1: [piece] + }) + restored = ParamAffineMap.from_dict(affine_map.to_dict()) + assert restored.pieces_by_rank[0][0] == piece + + +def test_stored_version_is_the_format_version(): + """A map written now must declare the version this code writes, not the checkpoint's.""" + assert AFFINE_MAP_FORMAT_VERSION >= 1 + + +# Guards added in review. Each one covers a way a map could produce a plausible but wrong +# parameter rather than failing, which is the failure mode that matters for a checkpoint. + + +def test_zero_extent_piece_covers_nothing(): + """A rank may hold none of a sub-parameter, and an empty piece must not claim coverage.""" + piece = AffinePiece(shape=(0, 8), + source_offset=0, + source_strides=(8, 1), + dest_offset=0, + dest_strides=(8, 1), + locations=[0]) + assert piece.numel == 0 + assert list(piece.source_offsets()) == [] + + +def test_disagreeing_replicas_are_refused(): + """Replicas must match. Letting the last writer win would hide a corrupt shard.""" + affine_map = _replicated_bias_map(world_size=2, scale=1.0) + shards = {0: torch.zeros(4, dtype=torch.float64), 1: torch.ones(4, dtype=torch.float64)} + with pytest.raises(ValueError, match='different data'): + affine_map.rebuild(shards) + + +def test_shard_that_disagrees_with_the_map_is_refused(): + """A shard longer than the map describes would be silently truncated to its prefix.""" + affine_map = contiguous_split_map((10, 8), [3, 3, 2, 2], 0) + shards = {rank: torch.randn(shape, dtype=torch.float64) for rank, shape in affine_map.shard_shapes.items()} + shards[0] = torch.randn(5, 8, dtype=torch.float64) + with pytest.raises(ValueError, match='describes'): + affine_map.rebuild(shards) + + +def test_zero_scale_is_refused(): + """Zero cannot be inverted, so the full tensor could never be recovered.""" + with pytest.raises(ValueError, match='not invertible'): + AffinePiece(shape=(4, ), + source_offset=0, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=[0], + scale=0.0) + + +def _replicated_bias_map(world_size, scale): + piece = AffinePiece(shape=(4, ), + source_offset=0, + source_strides=(1, ), + dest_offset=0, + dest_strides=(1, ), + locations=range(world_size), + scale=scale) + return ParamAffineMap(logical_shape=(4, ), + shard_shapes={rank: (4, ) + for rank in range(world_size)}, + pieces_by_rank={rank: [piece] + for rank in range(world_size)}) + + +def test_scale_power_gives_each_state_its_own_factor(): + """The moments' factors are the inverse and inverse square of the parameter's. + + Verified at the arithmetic level because `rebuild` refuses a scaled optimizer state + until the checkpoint contract covers the optimizer's own hyperparameters. The factors + themselves are correct, and are what that contract will build on. + """ + scale = 0.25 + shard = torch.full((4, ), 8.0, dtype=torch.float64) + assert torch.equal(_scaled(shard, scale, 1), shard * 4) # parameter: divided by 1/4 + assert torch.equal(_scaled(shard, scale, -1), shard / 4) # first moment: the inverse + assert torch.equal(_scaled(shard, scale, -2), shard / 16) # second moment: inverse square + + +def test_scaled_parameter_still_round_trips(): + """The parameter itself converts through a scaled piece, which is the supported case.""" + world_size = 4 + affine_map = _replicated_bias_map(world_size, scale=1.0 / world_size) + full = torch.randn(4, dtype=torch.float64) + + shards = {rank: affine_map.extract(full, rank) for rank in range(world_size)} + assert torch.equal(shards[0], full / world_size) + assert torch.equal(affine_map.rebuild(shards), full) + + +# Guards added in review. Each one covers a way a map could produce a plausible but wrong +# parameter rather than failing, which is the failure mode that matters for a checkpoint. + + +def test_scaled_optimizer_state_is_refused(): + """Correct moments are not enough: Adam's lr and eps live in the source coordinate. + + Converting them without rescaling `lr / scale` and `eps * scale` resumes a run on a + different trajectory, with an error that grows each step and nothing to signal it. Until + the checkpoint contract covers those hyperparameters, refuse rather than convert. + """ + affine_map = _replicated_bias_map(world_size=4, scale=0.25) + shards = {rank: torch.ones(4, dtype=torch.float64) for rank in range(4)} + + affine_map.rebuild(shards, scale_power=1) # the parameter itself is fine + for scale_power in (-1, -2): + with pytest.raises(NotImplementedError, match='source coordinate'): + affine_map.rebuild(shards, scale_power) + + +def test_unscaled_optimizer_state_still_converts(): + """The refusal is about scaling, not about optimizer states.""" + affine_map = _replicated_bias_map(world_size=2, scale=1.0) + shards = {rank: torch.ones(4, dtype=torch.float64) for rank in range(2)} + for scale_power in (1, -1, -2): + assert torch.equal(affine_map.rebuild(shards, scale_power), torch.ones(4, dtype=torch.float64)) + + +def test_replicated_map_carries_the_scale(): + """The layouts that pre-divide a value replicate it whole, so the scale belongs here. + + A row-parallel layer divides its bias by the world size and gives every rank the whole + thing, so summing the all-reduced outputs adds the bias once. The weight beside it is + split and unscaled, which is why the split constructors take no scale. + """ + world_size = 4 + full_bias = torch.randn(5, dtype=torch.float64) + affine_map = replicated_map((5, ), world_size, scale=1.0 / world_size) + affine_map.validate_coverage() + + shards = {rank: affine_map.extract(full_bias, rank) for rank in range(world_size)} + assert torch.equal(shards[0], full_bias / world_size) + assert torch.equal(affine_map.rebuild(shards), full_bias) + assert affine_map.pieces_by_rank[0][0].locations == frozenset(range(world_size)) diff --git a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py index af615fa2cc27..58bbe065c6e8 100644 --- a/tests/unit/checkpoint/test_autotp_uc_checkpoint.py +++ b/tests/unit/checkpoint/test_autotp_uc_checkpoint.py @@ -32,7 +32,7 @@ from deepspeed.utils import RepeatingLoader, groups from deepspeed.module_inject.tp_shard import get_shard_size_list -from unit.common import DistributedTest +from unit.common import DistributedTest, DistributedFixture class _DummyAddress: @@ -1567,3 +1567,181 @@ def test_strict_correctness(self, tmpdir): # 2c: optimizer state usable. _train_steps(restored_engine, hidden_dim, steps=1) + + +class AffineResumeModel(torch.nn.Module): + """A column/row pair with a replicated, post-reduction row bias.""" + + def __init__(self): + super().__init__() + self.fc1 = torch.nn.Linear(16, 16) + self.fc2 = torch.nn.Linear(16, 16) + + def forward(self, x): + return self.fc2(torch.tanh(self.fc1(x))) + + +def _affine_resume_engine(tp_size, load_universal=False): + torch.manual_seed(1234) + model = AffineResumeModel() + config = { + "train_micro_batch_size_per_gpu": 2, + "zero_optimization": { + "stage": 1 + }, + "zero_allow_untested_optimizer": True, + "checkpoint": { + "load_universal": load_universal + }, + } + if tp_size > 1: + config["tensor_parallel"] = { + "autotp_size": tp_size, + "partition_config": { + "use_default_specs": + False, + "layer_specs": [ + { + "patterns": [r".*fc1\.weight$"], + "partition_type": "column" + }, + { + "patterns": [r".*fc2\.weight$"], + "partition_type": "row" + }, + ], + }, + } + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.8, 0.95), eps=1e-6) + engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config) + return engine + + +def _affine_resume_batch(engine, step): + generator = torch.Generator().manual_seed(9000 + step) + x = torch.randn(2, 16, generator=generator).to(engine.device) + target = torch.randn(2, 16, generator=generator).to(engine.device) + return x, target + + +def _affine_resume_full_tensor(name, value, tp_size): + if tp_size == 1: + return value.detach().cpu().clone() + shards = [torch.empty_like(value) for _ in range(tp_size)] + dist.all_gather(shards, value.contiguous(), group=groups.get_tensor_model_parallel_group()) + if name == "fc2.bias": + for shard in shards[1:]: + torch.testing.assert_close(shard, shards[0]) + full = shards[0] + else: + full = torch.cat(shards, dim=1 if name == "fc2.weight" else 0) + return full.detach().cpu().clone() + + +def _affine_resume_state(engine, tp_size, gradients=False): + from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state + + result = {} + for name, param in engine.module.named_parameters(): + if gradients: + values = {"grad": safe_get_full_grad(param)} + else: + values = { + "fp32": safe_get_full_fp32_param(param), + "exp_avg": safe_get_full_optimizer_state(param, "exp_avg"), + "exp_avg_sq": safe_get_full_optimizer_state(param, "exp_avg_sq"), + } + for key, value in values.items(): + assert value is not None, (name, key) + result[f"{name}/{key}"] = _affine_resume_full_tensor(name, value, tp_size) + result["global_steps"] = engine.global_steps + if not gradients: + steps = [state["step"].item() for state in engine.optimizer.optimizer.state.values()] + assert steps and all(step == engine.global_steps for step in steps) + result["optimizer_step"] = steps[0] + return result + + +def _affine_resume_step(engine, tp_size, step): + x, target = _affine_resume_batch(engine, step) + logits = engine(x) + loss = torch.nn.functional.mse_loss(logits, target) + engine.backward(loss) + result = _affine_resume_state(engine, tp_size, gradients=True) + result["logits"] = logits.detach().cpu().clone() + result["loss"] = loss.detach().cpu().clone() + engine.step() + result.update(_affine_resume_state(engine, tp_size)) + return result + + +class affine_resume_checkpoint(DistributedFixture): + world_size = 2 + + def run(self, tmpdir, affine_layout): + from deepspeed.checkpoint.affine import (AFFINE_MAP_FORMAT_VERSION, contiguous_split_map, replicated_map) + from deepspeed.checkpoint.constants import AFFINE_MAP, AFFINE_MAP_PARAMS, AFFINE_MAP_VERSION + + engine = _affine_resume_engine(self.world_size) + for step in range(4): + _affine_resume_step(engine, self.world_size, step) + + # Phase 1 does not yet emit AutoTP affine metadata. Supply the known fixture + # layout through the checkpoint API; all weights and moments come from training. + # LinearAllreduce adds this bias AFTER reduction, so its real scale is one. + maps = { + r"^fc1\.weight$": contiguous_split_map((16, 16), [8, 8], 0), + r"^fc1\.bias$": contiguous_split_map((16, ), [8, 8], 0), + r"^fc2\.weight$": contiguous_split_map((16, 16), [8, 8], 1), + r"^fc2\.bias$": replicated_map((16, ), self.world_size), + } + uc_info = { + UNIVERSAL_CHECKPOINT_VERSION_KEY: UNIVERSAL_CHECKPOINT_VERSION_VALUE, + AFFINE_MAP: { + AFFINE_MAP_VERSION: AFFINE_MAP_FORMAT_VERSION, + AFFINE_MAP_PARAMS: { + name: layout.to_dict() + for name, layout in maps.items() + }, + }, + } + engine.save_checkpoint(tmpdir, + tag="affine_resume", + client_state={UNIVERSAL_CHECKPOINT_INFO: uc_info} if affine_layout else {}) + dist.barrier() + if dist.get_rank() == 0: + _convert_to_universal(os.path.join(tmpdir, "affine_resume"), os.path.join(tmpdir, "affine_universal")) + dist.barrier() + reference = [_affine_resume_state(engine, self.world_size)] + for step in range(4, 8): + reference.append(_affine_resume_step(engine, self.world_size, step)) + if dist.get_rank() == 0: + torch.save(reference, os.path.join(tmpdir, "affine_reference.pt")) + dist.barrier() + engine.destroy() + + +@pytest.mark.parametrize("affine_layout", [False, True], ids=["legacy", "affine"]) +@pytest.mark.parametrize("world_size", [1, 2], ids=["tp1", "tp2"]) +class TestAffineUniversalCheckpointResume(DistributedTest): + + def test_resume_matches_uninterrupted_training(self, affine_resume_checkpoint, tmpdir, affine_layout, world_size): + tp_size = dist.get_world_size() + reference = torch.load(os.path.join(tmpdir, "affine_reference.pt"), weights_only=False) + engine = _affine_resume_engine(tp_size, load_universal=True) + load_path, _ = engine.load_checkpoint(tmpdir, tag="affine_universal", load_optimizer_states=True) + assert load_path is not None + actual = [_affine_resume_state(engine, tp_size)] + for step in range(4, 8): + actual.append(_affine_resume_step(engine, tp_size, step)) + # The uninterrupted job is independent of converter/loader geometry. Resetting + # moments, losing the row bias, or slicing along the wrong axis changes this trace. + for index, (restored, expected) in enumerate(zip(actual, reference)): + assert restored.keys() == expected.keys() + for name in expected: + torch.testing.assert_close(restored[name], + expected[name], + atol=2e-6, + rtol=2e-5, + msg=lambda message: f"snapshot {index}, {name}: {message}") + engine.destroy()