From b6b0d505464d8e2834e6b3dc34fc959df56c62b4 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Fri, 28 Aug 2026 08:24:45 -0400 Subject: [PATCH 01/12] Add affine shard map for universal checkpoint Describe a parameter's tensor-parallel layout as affine views of the full tensor instead of matching its name against regex categories. Each piece records where a block sits in the full parameter and in the shard, so both conversion directions are the same copy with the ends swapped. Covers the fused QKV and Yuan shared-QK layouts that AutoTP currently marks unsupported, with tests showing their shards do cover the full parameter. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 205 ++++++++++++++++ .../unit/checkpoint/test_affine_shard_map.py | 218 ++++++++++++++++++ 2 files changed, 423 insertions(+) create mode 100644 deepspeed/checkpoint/affine.py create mode 100644 tests/unit/checkpoint/test_affine_shard_map.py diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py new file mode 100644 index 000000000000..2010d67d43c5 --- /dev/null +++ b/deepspeed/checkpoint/affine.py @@ -0,0 +1,205 @@ +# Copyright (c) Microsoft Corporation. +# 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. + +The representation is deliberately view-only. A piece never combines several elements of +the full tensor, which is what keeps the mapping invertible: given the full tensor a rank's +shard can be produced, and given the shards the full tensor can be rebuilt. +""" + +import torch + +__all__ = ['AffinePiece', 'ParamAffineMap'] + + +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. + """ + + __slots__ = ('shape', 'source_offset', 'source_strides', 'dest_offset', 'dest_strides', 'locations') + + def __init__(self, shape, source_offset, source_strides, dest_offset, dest_strides, locations): + 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) + + @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 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)})') + + 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) + + def __hash__(self): + return hash( + (self.shape, self.source_offset, self.source_strides, self.dest_offset, self.dest_strides, self.locations)) + + +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 validate(self): + 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.') + + 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 rebuild(self, shards): + """Rebuild the full parameter from per-rank shards, using the pieces as the plan. + + ``shards`` maps a rank to its shard. Where pieces overlap the data is identical by + construction, so writing them in any order gives the same result. + """ + self.validate() + any_shard = next(iter(shards.values())) + full_param = torch.empty(self.numel, dtype=any_shard.dtype, device=any_shard.device) + + for rank, pieces in self.pieces_by_rank.items(): + flat_shard = _flat_buffer(shards[rank]) + for piece in pieces: + piece.source_view(full_param).copy_(piece.dest_view(flat_shard)) + + return full_param.view(self.logical_shape) + + def extract(self, full_param, rank): + """Produce one rank's shard from the full parameter. The inverse of ``rebuild``.""" + 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]: + piece.dest_view(shard).copy_(piece.source_view(flat_param)) + + return shard.view(shard_shape) + + 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 _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 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..265c5a283345 --- /dev/null +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -0,0 +1,218 @@ +# Copyright (c) Microsoft Corporation. +# 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 +from deepspeed.module_inject.fusedqkv_utils import prepare_tp_fused_qkvw, shard_value_with_share_qk +from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads + + +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 _pieces_for_rank(markers, shard, rank, axis, cols): + """Turn one rank's shard into affine pieces of the full tensor.""" + pieces = [] + for source_index, dest_index, length in _index_runs(_source_indices(markers, shard, axis)): + 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=[rank])) + 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) + pieces_by_rank = {} + shard_shapes = {} + for rank in range(mp_size): + shard = shard_fn(markers.clone(), rank) + shard_shapes[rank] = tuple(shard.shape) + pieces_by_rank[rank] = _pieces_for_rank(markers, shard, rank, axis, cols) + 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): + module = _NamedModule(layout_name) + + def shard_fn(full_param, rank): + return prepare_tp_fused_qkvw(module, full_param, mp_size, rank) + + return shard_fn + + +def _shared_qk_shard_fn(shard_value): + + def shard_fn(full_param, rank): + return shard_value_with_share_qk(full_param, None, rank, 2, shard_value)[0].data + + return shard_fn + + +def _configure_heads(num_kv_heads, n_embd=None, num_attention_heads=None): + set_num_kv_heads(num_kv_heads) + if n_embd is not None: + set_n_embd(n_embd) + if num_attention_heads is not None: + set_num_attention_heads(num_attention_heads) + + +# (id, rows, cols, mp_size, axis, configure, build_shard_fn) +LAYOUTS = [ + ('bigcode', 48, 8, 4, 'row', lambda: _configure_heads(4, n_embd=32, num_attention_heads=4), + lambda mp: _fused_qkv_shard_fn('GPTBigCodeBlock', mp)), + ('codegen', 24, 8, 2, 'row', lambda: _configure_heads(8, n_embd=8, num_attention_heads=8), + lambda mp: _fused_qkv_shard_fn('CodeGenBlock', mp)), + ('yuan_value', 32, 8, 2, 'row', lambda: _configure_heads(8), lambda mp: _shared_qk_shard_fn(True)), + ('yuan_oproj', 8, 32, 2, 'col', lambda: _configure_heads(8), lambda mp: _shared_qk_shard_fn(False)), +] + + +@pytest.mark.parametrize('name, rows, cols, mp_size, axis, configure, 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, configure, build_shard_fn): + """Every element of the full tensor is held by some rank, so it can be rebuilt.""" + configure() + affine_map = _build_map(build_shard_fn(mp_size), rows, cols, mp_size, axis) + assert affine_map.uncovered_offsets() == [] + + def test_pieces_reproduce_shards_of_unseen_data(self, name, rows, cols, mp_size, axis, configure, build_shard_fn): + """Pieces derived from markers must reproduce shards of data they were not derived from.""" + configure() + shard_fn = build_shard_fn(mp_size) + 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, configure, build_shard_fn): + """Rebuilding from the shards returns the parameter the shards were cut from.""" + configure() + shard_fn = build_shard_fn(mp_size) + 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): + _configure_heads(8, n_embd=hidden, num_attention_heads=8) + rows = 3 * hidden + affine_map = _build_map(_fused_qkv_shard_fn('CodeGenBlock', 2), 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. + """ + _configure_heads(4, n_embd=32, num_attention_heads=4) + mp_size = 4 + affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size), 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' From 13e2c649136e320d27743ff855afdaab9eb3a0d2 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:22:28 -0400 Subject: [PATCH 02/12] Add scale and homogeneity to the affine shard map A row-parallel layer pre-divides its replicated bias by the world size, so a piece carries the factor its shard holds the block by. Scaling is invertible where a reduction is not, so this keeps conversion reversible in both directions. A piece must also cover elements held by the same set of ranks. Merging by adjacency alone fuses a rank-private block onto a replicated one where they happen to be neighbours, leaving a piece whose own locations is wrong for half of it. Adds the specification the module implements. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 82 ++- deepspeed/checkpoint/affine_ir_spec.md | 530 ++++++++++++++++++ .../unit/checkpoint/test_affine_shard_map.py | 101 +++- 3 files changed, 692 insertions(+), 21 deletions(-) create mode 100644 deepspeed/checkpoint/affine_ir_spec.md diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 2010d67d43c5..9718e5ad8f4e 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -10,9 +10,12 @@ shard geometrically instead: each rank holds a list of affine views of the full tensor, and the views alone say how to rebuild it. -The representation is deliberately view-only. A piece never combines several elements of -the full tensor, which is what keeps the mapping invertible: given the full tensor a rank's -shard can be produced, and given the shards the full tensor can be rebuilt. +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 @@ -35,17 +38,29 @@ class AffinePiece: ``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') + __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): + 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) + assert self.scale != 0.0, 'A zero scale is not invertible, so the full tensor could not be rebuilt.' @property def numel(self): @@ -97,18 +112,20 @@ def _offsets(self, base, strides): 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'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.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)) + return hash((self.shape, self.source_offset, self.source_strides, self.dest_offset, self.dest_strides, + self.locations, self.scale)) class ParamAffineMap: @@ -139,17 +156,46 @@ def uncovered_offsets(self): 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(): - 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}.') + 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): """Rebuild the full parameter from per-rank shards, using the pieces as the plan. @@ -164,7 +210,10 @@ def rebuild(self, shards): for rank, pieces in self.pieces_by_rank.items(): flat_shard = _flat_buffer(shards[rank]) for piece in pieces: - piece.source_view(full_param).copy_(piece.dest_view(flat_shard)) + target = piece.source_view(full_param) + target.copy_(piece.dest_view(flat_shard)) + if piece.scale != 1.0: + target.div_(piece.scale) return full_param.view(self.logical_shape) @@ -175,7 +224,10 @@ def extract(self, full_param, rank): shard = torch.empty(_product(shard_shape), dtype=full_param.dtype, device=full_param.device) for piece in self.pieces_by_rank[rank]: - piece.dest_view(shard).copy_(piece.source_view(flat_param)) + target = piece.dest_view(shard) + target.copy_(piece.source_view(flat_param)) + if piece.scale != 1.0: + target.mul_(piece.scale) return shard.view(shard_shape) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md new file mode 100644 index 000000000000..840f9e5a86bb --- /dev/null +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -0,0 +1,530 @@ +# 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) +``` + +- `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. + +**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 +(`collect_autotp_universal_checkpoint_info` produces `rf"^{re.escape(full_name)}$"`, one per parameter): + +```python +{ + "affine_map": { + "version": 1, + "^transformer.h.0.attn.c_attn.weight$": { + "logical_shape": [48, 8], + "ranks": { + 0: {"shard_shape": [24, 8], + "pieces": [ + {"shape": [8, 8], "source": [0, [8, 1]], "dest": [0, [8, 1]], "locations": [0]}, + {"shape": [16, 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. + +### 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/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 265c5a283345..cab1f4d1ee47 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -76,10 +76,30 @@ def _source_indices(markers, shard, axis): return indices -def _pieces_for_rank(markers, shard, rank, axis, cols): +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 = [] - for source_index, dest_index, length in _index_runs(_source_indices(markers, shard, axis)): + 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 @@ -97,19 +117,27 @@ def _pieces_for_rank(markers, shard, rank, axis, cols): source_strides=(cols, 1), dest_offset=dest_offset, dest_strides=dest_strides, - locations=[rank])) + 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 in range(mp_size): - shard = shard_fn(markers.clone(), rank) + for rank, shard in shards.items(): shard_shapes[rank] = tuple(shard.shape) - pieces_by_rank[rank] = _pieces_for_rank(markers, shard, rank, axis, cols) + 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) @@ -160,6 +188,12 @@ def test_shards_cover_the_full_parameter(self, name, rows, cols, mp_size, axis, affine_map = _build_map(build_shard_fn(mp_size), rows, cols, mp_size, axis) assert affine_map.uncovered_offsets() == [] + def test_every_piece_is_homogeneous(self, name, rows, cols, mp_size, axis, configure, build_shard_fn): + """No piece spans elements held by different sets of ranks, so locations is exact.""" + configure() + affine_map = _build_map(build_shard_fn(mp_size), rows, cols, mp_size, axis) + affine_map.validate_coverage() + def test_pieces_reproduce_shards_of_unseen_data(self, name, rows, cols, mp_size, axis, configure, build_shard_fn): """Pieces derived from markers must reproduce shards of data they were not derived from.""" configure() @@ -216,3 +250,58 @@ def test_replicated_piece_is_shared_by_every_rank(): 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. + """ + _configure_heads(4, n_embd=32, num_attention_heads=4) + mp_size = 4 + affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size), 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) From 1d1d6039fec2e923863761da22676974399abbc7 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:25:24 -0400 Subject: [PATCH 03/12] Lower the current checkpoint metadata to affine maps Add constructors for the layouts the universal checkpoint converter already handles: replicated parameters, contiguous splits along either axis, and parameters holding several sub-parameters split unevenly across ranks. Row and column parallelism differ only in stride, so one constructor covers both and the recorded concat dimension becomes redundant. Tests require each constructor to reproduce merge_tp_slices' own arithmetic exactly, so a map can replace a branch without changing what a checkpoint converts to. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 107 +++++++++++++++++- .../unit/checkpoint/test_affine_shard_map.py | 100 +++++++++++++++- 2 files changed, 205 insertions(+), 2 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 9718e5ad8f4e..d91d0b7aa3b1 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -20,7 +20,9 @@ import torch -__all__ = ['AffinePiece', 'ParamAffineMap'] +__all__ = [ + 'AffinePiece', 'ParamAffineMap', 'row_major_strides', 'replicated_map', 'contiguous_split_map', 'sub_param_map' +] class AffinePiece: @@ -255,3 +257,106 @@ def _product(shape): 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): + """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. + """ + 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) + ] + 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, scale=1.0): + """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], + scale=scale) + ] + 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) diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index cab1f4d1ee47..65ea7a79cd28 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -21,7 +21,8 @@ import pytest import torch -from deepspeed.checkpoint.affine import AffinePiece, ParamAffineMap +from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, replicated_map, contiguous_split_map, + sub_param_map) from deepspeed.module_inject.fusedqkv_utils import prepare_tp_fused_qkvw, shard_value_with_share_qk from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads @@ -305,3 +306,100 @@ def test_scale_round_trip_is_exact_for_power_of_two_world_sizes(): 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) From 3cf12341cf2630a2f9d3e13ead593ded0334f22e Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:34:50 -0400 Subject: [PATCH 04/12] Convert tp slices through an affine map when one is present Read the geometric description from universal checkpoint info and rebuild the parameter from it, falling back to the existing category branches when a parameter has no map. Nothing writes a map yet, so this changes no conversion. Add the on-disk form, which holds plain scalars so the map can be read without importing DeepSpeed, and omits the scale factor where it is 1. Tests require each constructor to reproduce merge_tp_slices' own arithmetic, including the uneven sub-parameter widths that earlier metadata could not describe. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 52 ++++++++++++++++++ deepspeed/checkpoint/affine_ir_spec.md | 30 ++++++----- deepspeed/checkpoint/constants.py | 6 +++ deepspeed/checkpoint/ds_to_universal.py | 25 ++++++++- .../unit/checkpoint/test_affine_shard_map.py | 53 +++++++++++++++++++ 5 files changed, 153 insertions(+), 13 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index d91d0b7aa3b1..39a199a52a42 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -233,6 +233,32 @@ def extract(self, full_param, rank): 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})' @@ -360,3 +386,29 @@ def sub_param_map(shape, sub_dim_sizes, shard_widths, partition_dim): 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 index 840f9e5a86bb..5d38c959ae60 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -316,22 +316,24 @@ 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 -(`collect_autotp_universal_checkpoint_info` produces `rf"^{re.escape(full_name)}$"`, one per parameter): +exact-match patterns `collect_autotp_universal_checkpoint_info` already emits — one +`rf"^{re.escape(full_name)}$"` per parameter: ```python { "affine_map": { "version": 1, - "^transformer.h.0.attn.c_attn.weight$": { - "logical_shape": [48, 8], - "ranks": { - 0: {"shard_shape": [24, 8], - "pieces": [ - {"shape": [8, 8], "source": [0, [8, 1]], "dest": [0, [8, 1]], "locations": [0]}, - {"shape": [16, 8], "source": [256, [8, 1]], "dest": [64, [8, 1]], "locations": [0,1,2,3]} - ]}, - ... + "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]} + ]}, + ... + } } } } @@ -339,7 +341,11 @@ exact-match patterns `collect_autotp_universal_checkpoint_info` already emits ``` 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. +`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 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..ec76737911ca 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -41,6 +41,8 @@ PIPELINE_REPLICATED_PARAMETER_PATTERNS, TP_REPLICATED_PARAMETER_PATTERNS, PARAMETER_TO_AVERAGE_PATTERNS, + AFFINE_MAP, + AFFINE_MAP_PARAMS, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS, PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0, PARAMETER_WITH_SUB_PARAMS, @@ -57,6 +59,7 @@ is_autoep_zero3_partitioned_entry, validate_autoep_zero3_partitioned_metadata, ) +from deepspeed.checkpoint.affine import ParamAffineMap def parse_arguments(): @@ -300,6 +303,7 @@ 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_params = universal_checkpoint_info.get(AFFINE_MAP, {}).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,6 +328,19 @@ 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. + """ + for pattern_, entry_ in affine_params.items(): + if re.match(pattern_, name_): + return ParamAffineMap.from_dict(entry_) + return 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) @@ -337,7 +354,13 @@ 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. A checkpoint converted this way carries + # the map rather than the per-category keys those branches write, because the + # geometry is what the restoring side needs and it is not tied to a category. + param = matched_affine_map.rebuild(dict(enumerate(slices))) + 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/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 65ea7a79cd28..80fb80cf97dc 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -403,3 +403,56 @@ def test_parity_round_trips_back_to_the_original_slices(): 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 From 166a3d2f9daebd9b89a26dc94bf138c28e16f364 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:37:39 -0400 Subject: [PATCH 05/12] Build shard metadata per model in the affine map tests AutoTP now derives kv-head and grain values into an AutoTPMeta per model instead of process-wide globals, so the tests construct one and pass it to the partition functions they exercise. Signed-off-by: Achyuthan Sivasankar --- .../unit/checkpoint/test_affine_shard_map.py | 67 +++++++++---------- 1 file changed, 30 insertions(+), 37 deletions(-) diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 80fb80cf97dc..b5b1f21c5280 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -24,7 +24,7 @@ from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, replicated_map, contiguous_split_map, sub_param_map) from deepspeed.module_inject.fusedqkv_utils import prepare_tp_fused_qkvw, shard_value_with_share_qk -from deepspeed.module_inject.tp_shard import set_num_kv_heads, set_n_embd, set_num_attention_heads +from deepspeed.module_inject.tp_shard import AutoTPMeta class _NamedModule(torch.nn.Module): @@ -142,63 +142,57 @@ def _build_map(shard_fn, rows, cols, mp_size, axis): 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): +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) + return prepare_tp_fused_qkvw(module, full_param, mp_size, rank, meta) return shard_fn -def _shared_qk_shard_fn(shard_value): +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)[0].data + return shard_value_with_share_qk(full_param, None, rank, 2, shard_value, meta)[0].data return shard_fn -def _configure_heads(num_kv_heads, n_embd=None, num_attention_heads=None): - set_num_kv_heads(num_kv_heads) - if n_embd is not None: - set_n_embd(n_embd) - if num_attention_heads is not None: - set_num_attention_heads(num_attention_heads) +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, configure, build_shard_fn) +# (id, rows, cols, mp_size, axis, make_meta, build_shard_fn) LAYOUTS = [ - ('bigcode', 48, 8, 4, 'row', lambda: _configure_heads(4, n_embd=32, num_attention_heads=4), - lambda mp: _fused_qkv_shard_fn('GPTBigCodeBlock', mp)), - ('codegen', 24, 8, 2, 'row', lambda: _configure_heads(8, n_embd=8, num_attention_heads=8), - lambda mp: _fused_qkv_shard_fn('CodeGenBlock', mp)), - ('yuan_value', 32, 8, 2, 'row', lambda: _configure_heads(8), lambda mp: _shared_qk_shard_fn(True)), - ('yuan_oproj', 8, 32, 2, 'col', lambda: _configure_heads(8), lambda mp: _shared_qk_shard_fn(False)), + ('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, configure, build_shard_fn', +@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, configure, build_shard_fn): + 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.""" - configure() - affine_map = _build_map(build_shard_fn(mp_size), rows, cols, mp_size, axis) + 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, configure, build_shard_fn): + 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.""" - configure() - affine_map = _build_map(build_shard_fn(mp_size), rows, cols, mp_size, axis) + 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, configure, build_shard_fn): + 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.""" - configure() - shard_fn = build_shard_fn(mp_size) + shard_fn = build_shard_fn(mp_size, make_meta()) affine_map = _build_map(shard_fn, rows, cols, mp_size, axis) torch.manual_seed(0) @@ -207,10 +201,9 @@ def test_pieces_reproduce_shards_of_unseen_data(self, name, rows, cols, 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, configure, build_shard_fn): + 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.""" - configure() - shard_fn = build_shard_fn(mp_size) + shard_fn = build_shard_fn(mp_size, make_meta()) affine_map = _build_map(shard_fn, rows, cols, mp_size, axis) torch.manual_seed(1) @@ -227,9 +220,9 @@ def test_piece_count_does_not_grow_with_model_size(): """ piece_counts = set() for hidden in (8, 16, 32, 64, 128): - _configure_heads(8, n_embd=hidden, num_attention_heads=8) + meta = _meta(8, n_embd=hidden, num_attention_heads=8) rows = 3 * hidden - affine_map = _build_map(_fused_qkv_shard_fn('CodeGenBlock', 2), rows, 8, 2, 'row') + 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)}' @@ -241,9 +234,9 @@ def test_replicated_piece_is_shared_by_every_rank(): 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. """ - _configure_heads(4, n_embd=32, num_attention_heads=4) + meta = _meta(4, n_embd=32, num_attention_heads=4) mp_size = 4 - affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size), 48, 8, mp_size, 'row') + 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): @@ -260,9 +253,9 @@ def test_pieces_never_span_a_replication_boundary(): 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. """ - _configure_heads(4, n_embd=32, num_attention_heads=4) + meta = _meta(4, n_embd=32, num_attention_heads=4) mp_size = 4 - affine_map = _build_map(_fused_qkv_shard_fn('GPTBigCodeBlock', mp_size), 48, 8, mp_size, 'row') + 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 From 91499088f8dae5c449be8ed869bf84d5575522ba Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:41:26 -0400 Subject: [PATCH 06/12] Refuse an affine map format newer than the reader understands The stored map declares its own encoding version, separate from the universal checkpoint version so the two can move independently. A reader that predates a version would otherwise misinterpret fields added since, so refuse instead. Keep the stride helper private, since nothing outside the module builds a piece by hand yet. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 8 +++++++- deepspeed/checkpoint/ds_to_universal.py | 10 ++++++++-- tests/unit/checkpoint/test_affine_shard_map.py | 9 +++++++-- 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 39a199a52a42..156955cfbf80 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -21,9 +21,15 @@ import torch __all__ = [ - 'AffinePiece', 'ParamAffineMap', 'row_major_strides', 'replicated_map', 'contiguous_split_map', 'sub_param_map' + 'AffinePiece', 'ParamAffineMap', 'AFFINE_MAP_FORMAT_VERSION', 'row_major_strides', '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. diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index ec76737911ca..0343894cffb9 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -43,6 +43,7 @@ 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, @@ -59,7 +60,7 @@ is_autoep_zero3_partitioned_entry, validate_autoep_zero3_partitioned_metadata, ) -from deepspeed.checkpoint.affine import ParamAffineMap +from deepspeed.checkpoint.affine import ParamAffineMap, AFFINE_MAP_FORMAT_VERSION def parse_arguments(): @@ -303,7 +304,12 @@ 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_params = universal_checkpoint_info.get(AFFINE_MAP, {}).get(AFFINE_MAP_PARAMS, {}) + affine_map_info = universal_checkpoint_info.get(AFFINE_MAP, {}) + affine_map_version = affine_map_info.get(AFFINE_MAP_VERSION, AFFINE_MAP_FORMAT_VERSION) + assert affine_map_version <= AFFINE_MAP_FORMAT_VERSION, ( + 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 + diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index b5b1f21c5280..b27156216e58 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -21,8 +21,8 @@ import pytest import torch -from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, replicated_map, contiguous_split_map, - sub_param_map) +from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, AFFINE_MAP_FORMAT_VERSION, replicated_map, + contiguous_split_map, sub_param_map) from deepspeed.module_inject.fusedqkv_utils import prepare_tp_fused_qkvw, shard_value_with_share_qk from deepspeed.module_inject.tp_shard import AutoTPMeta @@ -449,3 +449,8 @@ def test_serialisation_preserves_scale(): }) 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 From d63251bb34235939ed12d4cefd8b89a2ed80df5f Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 11:47:08 -0400 Subject: [PATCH 07/12] Keep the stride helper private Nothing outside the module builds a piece by hand yet, so exporting the helper widens the public surface for no caller. Step 4 can promote it when it needs it. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 156955cfbf80..bd6a9b7b5d75 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -21,8 +21,8 @@ import torch __all__ = [ - 'AffinePiece', 'ParamAffineMap', 'AFFINE_MAP_FORMAT_VERSION', 'row_major_strides', 'replicated_map', - 'contiguous_split_map', 'sub_param_map' + '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 @@ -291,7 +291,7 @@ def _product(shape): return count -def row_major_strides(shape): +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): @@ -306,7 +306,7 @@ def replicated_map(shape, tp_degree): rank is cheapest rather than from a designated owner. """ shape = tuple(shape) - strides = row_major_strides(shape) + strides = _row_major_strides(shape) ranks = list(range(tp_degree)) pieces = [ AffinePiece(shape=shape, @@ -330,7 +330,7 @@ def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0): 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) + source_strides = _row_major_strides(shape) pieces_by_rank = {} shard_shapes = {} start = 0 @@ -344,7 +344,7 @@ def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0): source_offset=start * source_strides[partition_dim], source_strides=source_strides, dest_offset=0, - dest_strides=row_major_strides(shard_shape), + dest_strides=_row_major_strides(shard_shape), locations=[rank], scale=scale) ] @@ -361,7 +361,7 @@ def sub_param_map(shape, sub_dim_sizes, shard_widths, partition_dim): unevenly and need not be the same size as each other. """ shape = tuple(shape) - source_strides = row_major_strides(shape) + source_strides = _row_major_strides(shape) tp_degree = len(shard_widths[0]) shard_shapes = {} @@ -379,7 +379,7 @@ def sub_param_map(shape, sub_dim_sizes, shard_widths, partition_dim): for rank in range(tp_degree): piece_shape = list(shape) piece_shape[partition_dim] = widths[rank] - dest_strides = row_major_strides(shard_shapes[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], From 2e846d91153bc41c663f99f289121735e26927ee Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Tue, 1 Sep 2026 13:35:11 -0400 Subject: [PATCH 08/12] Address review: scale optimizer states by their own power Adam's moments live in the parameter's scaled coordinate, so scaling a parameter by s scales its gradient by 1/s. Applying the parameter's factor to the moments would corrupt the optimizer state and change the trajectory after a resume, so the caller now says which power of the scale applies. Refuse rather than guess in four more places: replicas that disagree, a shard whose size contradicts the map, a map format newer than this reader, and a zero scale. Empty pieces no longer report covering elements they do not hold. Raise instead of asserting where the check guards checkpoint compatibility, since asserts are stripped under python -O. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 55 +++++++++-- deepspeed/checkpoint/affine_ir_spec.md | 8 ++ deepspeed/checkpoint/ds_to_universal.py | 29 +++--- .../unit/checkpoint/test_affine_shard_map.py | 91 +++++++++++++++++++ 4 files changed, 162 insertions(+), 21 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index bd6a9b7b5d75..85bfd7a68db2 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -68,7 +68,8 @@ def __init__(self, shape, source_offset, source_strides, dest_offset, dest_strid self.dest_strides = tuple(int(stride) for stride in dest_strides) self.locations = frozenset(int(rank) for rank in locations) self.scale = float(scale) - assert self.scale != 0.0, 'A zero scale is not invertible, so the full tensor could not be rebuilt.' + 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): @@ -97,6 +98,10 @@ def source_offsets(self): 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 @@ -205,27 +210,48 @@ def validate_coverage(self): f'{sorted(holders[offset])}. A piece must not span elements with different ' 'holders, or its locations cannot be trusted.') - def rebuild(self, shards): + 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. Where pieces overlap the data is identical by - construction, so writing them in any order gives the same result. + ``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() + 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) - target.copy_(piece.dest_view(flat_shard)) - if piece.scale != 1.0: - target.div_(piece.scale) + 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 extract(self, full_param, rank): + def extract(self, full_param, rank, scale_power=1): """Produce one rank's shard from the full parameter. The inverse of ``rebuild``.""" flat_param = _flat_buffer(full_param) shard_shape = self.shard_shapes[rank] @@ -234,8 +260,9 @@ def extract(self, full_param, rank): for piece in self.pieces_by_rank[rank]: target = piece.dest_view(shard) target.copy_(piece.source_view(flat_param)) - if piece.scale != 1.0: - target.mul_(piece.scale) + factor = piece.scale**scale_power + if factor != 1.0: + target.mul_(factor) return shard.view(shard_shape) @@ -270,6 +297,14 @@ def __repr__(self): 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. diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 5d38c959ae60..8545d1e411f6 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -70,6 +70,14 @@ restored at a different TP degree without a rule naming which parameters are bia 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. + **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. diff --git a/deepspeed/checkpoint/ds_to_universal.py b/deepspeed/checkpoint/ds_to_universal.py index 0343894cffb9..ec9c728db4a6 100755 --- a/deepspeed/checkpoint/ds_to_universal.py +++ b/deepspeed/checkpoint/ds_to_universal.py @@ -306,9 +306,10 @@ def merge_tp_slices(uc_info, dir, slice_dir, tp_degree, name_and_shapes): 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) - assert affine_map_version <= AFFINE_MAP_FORMAT_VERSION, ( - 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.") + 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) @@ -341,10 +342,9 @@ def get_matched_affine_map(name_): `unmatched_patterns`: a map covers only the parameters it was written for, and the rest still take the branches below. """ - for pattern_, entry_ in affine_params.items(): - if re.match(pattern_, name_): - return ParamAffineMap.from_dict(entry_) - return None + 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) @@ -353,6 +353,12 @@ def get_matched_affine_map(name_): 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") @@ -362,10 +368,11 @@ def get_matched_affine_map(name_): ckpt_dict = {} 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. A checkpoint converted this way carries - # the map rather than the per-category keys those branches write, because the - # geometry is what the restoring side needs and it is not tied to a category. - param = matched_affine_map.rebuild(dict(enumerate(slices))) + # 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:]]) diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index b27156216e58..71801f2590ec 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -454,3 +454,94 @@ def test_serialisation_preserves_scale(): 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 _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_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) + + +@pytest.mark.parametrize('scale_power', [1, -1, -2]) +def test_scale_power_round_trips_each_optimizer_state(scale_power): + """Adam's moments live in the parameter's scaled coordinate, so they undo differently. + + Scaling a parameter by `s` scales its gradient by `1 / s`, so the first moment carries + the inverse of the parameter's factor and the second moment the inverse square. Using + the parameter's own factor for all three would corrupt the optimizer state. + """ + 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, scale_power) for rank in range(world_size)} + expected = full * (1.0 / world_size)**scale_power + assert torch.equal(shards[0], expected) + assert torch.equal(affine_map.rebuild(shards, scale_power), full) + + +def test_optimizer_moments_do_not_use_the_parameter_factor(): + """The three states must not come out of the same shard with the same value.""" + world_size = 4 + affine_map = _replicated_bias_map(world_size, scale=1.0 / world_size) + shard = torch.full((4, ), 8.0, dtype=torch.float64) + shards = {rank: shard.clone() for rank in range(world_size)} + + rebuilt = {power: affine_map.rebuild(shards, power) for power in (1, -1, -2)} + assert torch.equal(rebuilt[1], shard * world_size) + assert torch.equal(rebuilt[-1], shard / world_size) + assert torch.equal(rebuilt[-2], shard / world_size**2) From fcec0621812e69c0e4fd377ec5a43649dfd9490c Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Fri, 4 Sep 2026 12:46:29 -0400 Subject: [PATCH 09/12] Address review: header, scale in the piece tuple, scaling example Use the DeepSpeed Team copyright line that new files carry; the license check does not ask for the Microsoft one. Record scale in the piece definition itself, not only in the notes below it, and add a worked example of how each optimizer state recovers from a scaled shard. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 2 +- deepspeed/checkpoint/affine_ir_spec.md | 18 +++++++++++++++++- tests/unit/checkpoint/test_affine_shard_map.py | 2 +- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 85bfd7a68db2..d2723e8d9728 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -1,4 +1,4 @@ -# Copyright (c) Microsoft Corporation. +# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 8545d1e411f6..72d1721b5381 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -43,7 +43,7 @@ A **piece** is a block of elements, described by where it sits in the full tenso where it sits in the shard: ```python -Piece = (shape, source_offset, source_strides, dest_offset, dest_strides, locations) +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 @@ -78,6 +78,22 @@ parameter's own factor to all three corrupts the optimizer state and changes the 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. + **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. diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 71801f2590ec..bd8b26f12a22 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -1,4 +1,4 @@ -# Copyright (c) Microsoft Corporation. +# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team From 10bde3e76428822cc3372c86c1f0b9235744caa6 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Wed, 9 Sep 2026 16:26:16 -0400 Subject: [PATCH 10/12] Refuse to convert a scaled optimizer state The moment powers are right, but Adam's update also depends on lr and eps, and those live in the coordinate the optimizer trained in. Restoring a rescaled parameter without rescaling them resumes on a different trajectory, with an error that grows each step and nothing to signal it. That transform belongs to the optimizer rather than to parameter geometry, so refuse until the checkpoint contract covers it. A scaled parameter still converts. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 22 +++++ deepspeed/checkpoint/affine_ir_spec.md | 9 ++ .../unit/checkpoint/test_affine_shard_map.py | 94 ++++++++++++------- 3 files changed, 89 insertions(+), 36 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index d2723e8d9728..7a4eb8b0fe86 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -222,6 +222,7 @@ def rebuild(self, shards, scale_power=1): 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: @@ -251,8 +252,29 @@ def rebuild(self, shards, scale_power=1): 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) diff --git a/deepspeed/checkpoint/affine_ir_spec.md b/deepspeed/checkpoint/affine_ir_spec.md index 72d1721b5381..11a507db4d3d 100644 --- a/deepspeed/checkpoint/affine_ir_spec.md +++ b/deepspeed/checkpoint/affine_ir_spec.md @@ -94,6 +94,15 @@ first moment of `b` is the stored moment divided by 4 and the second is divided 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. diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index bd8b26f12a22..96d96462af57 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -22,7 +22,7 @@ import torch from deepspeed.checkpoint.affine import (AffinePiece, ParamAffineMap, AFFINE_MAP_FORMAT_VERSION, replicated_map, - contiguous_split_map, sub_param_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 @@ -460,21 +460,6 @@ def test_stored_version_is_the_format_version(): # parameter rather than failing, which is the failure mode that matters for a checkpoint. -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_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), @@ -516,32 +501,69 @@ def test_zero_scale_is_refused(): scale=0.0) -@pytest.mark.parametrize('scale_power', [1, -1, -2]) -def test_scale_power_round_trips_each_optimizer_state(scale_power): - """Adam's moments live in the parameter's scaled coordinate, so they undo differently. +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)}) + - Scaling a parameter by `s` scales its gradient by `1 / s`, so the first moment carries - the inverse of the parameter's factor and the second moment the inverse square. Using - the parameter's own factor for all three would corrupt the optimizer state. +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, scale_power) for rank in range(world_size)} - expected = full * (1.0 / world_size)**scale_power - assert torch.equal(shards[0], expected) - assert torch.equal(affine_map.rebuild(shards, scale_power), full) + 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) -def test_optimizer_moments_do_not_use_the_parameter_factor(): - """The three states must not come out of the same shard with the same value.""" - world_size = 4 - affine_map = _replicated_bias_map(world_size, scale=1.0 / world_size) - shard = torch.full((4, ), 8.0, dtype=torch.float64) - shards = {rank: shard.clone() for rank in range(world_size)} +# 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. - rebuilt = {power: affine_map.rebuild(shards, power) for power in (1, -1, -2)} - assert torch.equal(rebuilt[1], shard * world_size) - assert torch.equal(rebuilt[-1], shard / world_size) - assert torch.equal(rebuilt[-2], shard / world_size**2) + 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)) From d02c56bf13888f984839706ad34d8be83a70d6f2 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Thu, 10 Sep 2026 09:59:08 -0400 Subject: [PATCH 11/12] Carry the scale on the replicated map, not the split 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 the all-reduced sum adds it once. The weight beside it is split and unscaled, so no in-tree layout scales a split and the split constructors no longer take the argument. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 15 ++++++++++----- tests/unit/checkpoint/test_affine_shard_map.py | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 7a4eb8b0fe86..8ca076ae6667 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -356,11 +356,16 @@ def _row_major_strides(shape): return tuple(strides) -def replicated_map(shape, tp_degree): +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) @@ -371,7 +376,8 @@ def replicated_map(shape, tp_degree): source_strides=strides, dest_offset=0, dest_strides=strides, - locations=ranks) + locations=ranks, + scale=scale) ] return ParamAffineMap(logical_shape=shape, shard_shapes={rank: shape @@ -380,7 +386,7 @@ def replicated_map(shape, tp_degree): for rank in ranks}) -def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0): +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 @@ -402,8 +408,7 @@ def contiguous_split_map(shape, per_rank_sizes, partition_dim, scale=1.0): source_strides=source_strides, dest_offset=0, dest_strides=_row_major_strides(shard_shape), - locations=[rank], - scale=scale) + locations=[rank]) ] start += size return ParamAffineMap(logical_shape=shape, shard_shapes=shard_shapes, pieces_by_rank=pieces_by_rank) diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 96d96462af57..9275ca14a690 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -567,3 +567,21 @@ def test_unscaled_optimizer_state_still_converts(): 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)) From dd5c57f1580e8380cdc26b1e3febce73960529e9 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 9 Sep 2026 16:26:07 +0800 Subject: [PATCH 12/12] Fix universal checkpoint resume across AutoTP sizes Read universal checkpoint metadata without invoking the Megatron weight merger when tensor parallelism changes. Keep FP32 parameter fragments attached to master weights while rebinding restored Adam moment fragments to their optimizer buffers. Exercise real ZeRO-1 Adam training, save, legacy/affine conversion, and four-step TP2-to-TP1/TP2 resume against uninterrupted training. Compare logits, losses, gradients, FP32 weights, moments, and optimizer steps. Signed-off-by: 0z5a (cherry picked from commit e753a02b72dd18bc220f31ecac5ce5f6c1a89c54) --- deepspeed/runtime/engine.py | 11 +- deepspeed/utils/tensor_fragment.py | 2 +- .../checkpoint/test_autotp_uc_checkpoint.py | 180 +++++++++++++++++- 3 files changed, 189 insertions(+), 4 deletions(-) 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_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()