From 5f5af9f49c46c0859daa5b53120317c57cb54e63 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Mon, 26 Oct 2026 22:22:58 -0400 Subject: [PATCH 1/2] Check where column-split pieces write in a shard validate() checked destination bounds and overlap only for pieces whose destination is one dense interval, so pieces of a column split could write the same shard columns and pass every check. Check those pieces as boxes in the shard's grid from their metadata, and walk every shard address in validate_coverage() so the full check covers any stride. Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 95 ++++++++++++++++++- .../unit/checkpoint/test_affine_shard_map.py | 54 +++++++++++ 2 files changed, 144 insertions(+), 5 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index 312076fb8e93..d90c99598165 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -103,6 +103,10 @@ def source_offsets(self): """ return self._offsets(self.source_offset, self.source_strides) + def dest_offsets(self): + """Yield the offset into the shard of every element this piece covers.""" + return self._offsets(self.dest_offset, self.dest_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 @@ -157,6 +161,39 @@ def _dest_interval(piece): return piece.dest_offset, piece.dest_offset + piece.numel +def _dest_box(piece, shard_shape): + """Return a piece's destination as per-axis ``(starts, stops)`` in the shard, or None. + + A column split writes each piece into a few columns of every shard row, so its + destination is not one interval. It is still a box in the shard's own grid, because its + strides are the shard's row-major strides, and a box can be checked for bounds and + overlap from its corners without walking any elements. + """ + if piece.numel == 0 or len(piece.shape) != len(shard_shape) or len(piece.dest_strides) != len(shard_shape): + return None + shard_strides = _row_major_strides(shard_shape) + if any(size > 1 and got != want for size, got, want in zip(piece.shape, piece.dest_strides, shard_strides)): + return None + + starts = [] + remainder = piece.dest_offset + for stride in shard_strides: + start, remainder = divmod(remainder, stride) + starts.append(start) + stops = [start + size for start, size in zip(starts, piece.shape)] + return starts, stops + + +def _first_shared_address(box, other_box, shard_shape): + """Return the first shard address two boxes both write, or None when they are disjoint.""" + corner = [] + for start, stop, other_start, other_stop in zip(box[0], box[1], other_box[0], other_box[1]): + if max(start, other_start) >= min(stop, other_stop): + return None + corner.append(max(start, other_start)) + return sum(position * stride for position, stride in zip(corner, _row_major_strides(shard_shape))) + + class ParamAffineMap: """How one parameter is spread over a tensor-parallel group. @@ -195,21 +232,25 @@ def holders(self): return holders def validate(self): - """Check shard element counts and dense destination bounds and overlap.""" + """Check shard element counts, and destination bounds and overlap from piece metadata.""" for rank, pieces in self.pieces_by_rank.items(): held = sum(piece.numel for piece in pieces) expected = _product(self.shard_shapes[rank]) if held != expected: raise ValueError(f'Rank {rank} holds a shard of {expected} elements but its pieces ' f'account for {held}.') - self._validate_destinations(rank, pieces, expected) + self._validate_destinations(rank, pieces, self.shard_shapes[rank]) - def _validate_destinations(self, rank, pieces, shard_numel): - """Check that dense destination intervals stay within the shard and do not overlap.""" + def _validate_destinations(self, rank, pieces, shard_shape): + """Check that dense intervals and strided boxes stay within the shard and do not overlap.""" + shard_numel = _product(shard_shape) intervals = [] + has_strided_piece = False for piece in pieces: span = _dest_interval(piece) if span is None: + if piece.numel: + has_strided_piece = True continue start, stop = span if start < 0 or stop > shard_numel: @@ -225,14 +266,58 @@ def _validate_destinations(self, rank, pieces, shard_numel): 'silently overwrite another.') cursor = stop + # A column split's pieces have no interval, so the check above skips them. Check them as + # boxes against every other box-shaped piece of the rank instead. Ranks whose pieces are + # all dense are fully covered above and skip this pairwise pass. + if not has_strided_piece: + return + boxes = [] + for piece in pieces: + box = _dest_box(piece, shard_shape) + if box is None: + continue + starts, stops = box + fits = all(start >= 0 and stop <= dim for start, stop, dim in zip(starts, stops, shard_shape)) + if not fits: + if _dest_interval(piece) is None: + raise ValueError(f'Rank {rank} piece {piece} writes outside the shard of shape ' + f'{tuple(shard_shape)}.') + # A dense run may wrap past the end of a shard row. Its interval was bounds-checked + # above, and it is not a box of the shard's grid, so leave it out of the box pass. + continue + boxes.append(box) + for index, box in enumerate(boxes): + for other_box in boxes[index + 1:]: + shared = _first_shared_address(box, other_box, shard_shape) + if shared is not None: + raise ValueError(f'Rank {rank} writes shard address {shared} twice, so one piece would ' + 'silently overwrite another.') + def validate_coverage(self): - """Full check: the shards cover the parameter, and every piece is homogeneous. + """Full check: each shard address is written once, 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() + # validate() reads destinations from piece metadata, which cannot place every strided + # layout. Walking each shard address makes this check complete: every address of every + # shard is written by exactly one element. + for rank, pieces in self.pieces_by_rank.items(): + shard_numel = _product(self.shard_shapes[rank]) + written = set() + for piece in pieces: + for offset in piece.dest_offsets(): + if offset < 0 or offset >= shard_numel: + raise ValueError(f'Rank {rank} piece {piece} writes shard address {offset}, outside ' + f'the shard of {shard_numel} elements.') + if offset in written: + raise ValueError(f'Rank {rank} writes shard address {offset} twice, so one piece would ' + 'silently overwrite another.') + written.add(offset) + 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]}, ' diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index aa4bda513e42..6599e89c74af 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -505,6 +505,60 @@ def refuse_walk(self): affine_map.validate() +def _column_split_qkv_map(): + """Fused QKV split along columns, so every piece writes a few columns of each shard row.""" + return sub_param_map(shape=(4, 12), + sub_dim_sizes=[4, 4, 4], + shard_widths=[[2, 2], [2, 2], [2, 2]], + partition_dim=1) + + +def _move_destination(affine_map, rank, index, delta): + piece = affine_map.pieces_by_rank[rank][index] + affine_map.pieces_by_rank[rank][index] = AffinePiece(piece.shape, piece.source_offset, piece.source_strides, + piece.dest_offset + delta, piece.dest_strides, + piece.locations) + + +def test_strided_destination_overlap_is_refused(): + # Rank 0's value columns are moved onto its key columns. The element count still matches, + # so without a destination check the key block is overwritten and two shard columns are + # never written at all. + affine_map = _column_split_qkv_map() + _move_destination(affine_map, rank=0, index=2, delta=-2) + with pytest.raises(ValueError, match='twice'): + affine_map.validate() + + +def test_strided_destination_outside_the_shard_is_refused(): + affine_map = _column_split_qkv_map() + _move_destination(affine_map, rank=0, index=2, delta=2) + with pytest.raises(ValueError, match='outside the shard'): + affine_map.validate() + + +def test_full_check_refuses_a_destination_collision_of_any_stride(): + # The first piece writes shard addresses 0 and 2, the second writes 1 and 2, and address 3 + # is never written. Neither piece is a box of the shard's grid, so only walking the + # addresses can see it. + pieces = [AffinePiece((2, ), 0, (1, ), 0, (2, ), [0]), AffinePiece((2, ), 2, (1, ), 1, (1, ), [0])] + affine_map = ParamAffineMap((4, ), {0: (4, )}, {0: pieces}) + with pytest.raises(ValueError, match='twice'): + affine_map.validate_coverage() + + +def test_strided_destination_validation_does_not_enumerate_elements(monkeypatch): + size = 10**9 + pieces = [AffinePiece((size, 4), column, (8, 1), column, (8, 1), [0]) for column in (0, 4)] + affine_map = ParamAffineMap((size, 8), {0: (size, 8)}, {0: pieces}) + + def refuse_walk(self, base, strides): + raise AssertionError('validation must not enumerate shard elements') + + monkeypatch.setattr(AffinePiece, '_offsets', refuse_walk) + affine_map.validate() + + # 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. From efd161e963dfe1412cb978ef3ef3b22bb9109b18 Mon Sep 17 00:00:00 2001 From: Achyuthan Sivasankar Date: Mon, 26 Oct 2026 22:46:59 -0400 Subject: [PATCH 2/2] Accept strided views that wrap a shard row Bound a strided piece by its lowest and highest shard address, and use box coordinates only when the view fits the shard grid. A view that wraps a row stays valid and is checked by the address walk in validate_coverage(). Signed-off-by: Achyuthan Sivasankar --- deepspeed/checkpoint/affine.py | 25 +++++++++++++++---- .../unit/checkpoint/test_affine_shard_map.py | 16 ++++++++++++ 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/deepspeed/checkpoint/affine.py b/deepspeed/checkpoint/affine.py index d90c99598165..7aadba508961 100644 --- a/deepspeed/checkpoint/affine.py +++ b/deepspeed/checkpoint/affine.py @@ -184,6 +184,19 @@ def _dest_box(piece, shard_shape): return starts, stops +def _dest_address_range(piece): + """Return the lowest and highest shard address a piece writes.""" + low = piece.dest_offset + high = piece.dest_offset + for size, stride in zip(piece.shape, piece.dest_strides): + reach = (size - 1) * stride + if reach < 0: + low += reach + else: + high += reach + return low, high + + def _first_shared_address(box, other_box, shard_shape): """Return the first shard address two boxes both write, or None when they are disjoint.""" corner = [] @@ -273,17 +286,19 @@ def _validate_destinations(self, rank, pieces, shard_shape): return boxes = [] for piece in pieces: + if piece.numel and _dest_interval(piece) is None: + low, high = _dest_address_range(piece) + if low < 0 or high >= shard_numel: + raise ValueError(f'Rank {rank} piece {piece} writes shard addresses [{low}, {high}], ' + f'outside the shard of {shard_numel} elements.') box = _dest_box(piece, shard_shape) if box is None: continue starts, stops = box fits = all(start >= 0 and stop <= dim for start, stop, dim in zip(starts, stops, shard_shape)) if not fits: - if _dest_interval(piece) is None: - raise ValueError(f'Rank {rank} piece {piece} writes outside the shard of shape ' - f'{tuple(shard_shape)}.') - # A dense run may wrap past the end of a shard row. Its interval was bounds-checked - # above, and it is not a box of the shard's grid, so leave it out of the box pass. + # A view may wrap past the end of a shard row while staying inside the shard. It is + # not a box of the shard's grid, so the walk in validate_coverage() checks it instead. continue boxes.append(box) for index, box in enumerate(boxes): diff --git a/tests/unit/checkpoint/test_affine_shard_map.py b/tests/unit/checkpoint/test_affine_shard_map.py index 6599e89c74af..2a38d1ded648 100644 --- a/tests/unit/checkpoint/test_affine_shard_map.py +++ b/tests/unit/checkpoint/test_affine_shard_map.py @@ -537,6 +537,22 @@ def test_strided_destination_outside_the_shard_is_refused(): affine_map.validate() +def test_strided_destination_that_wraps_a_shard_row_round_trips(): + # In a (3, 3) shard, a (2, 2) view starting at address 2 with the shard's strides writes + # addresses 2, 3, 5 and 6. It wraps past the end of row 0 but stays inside the shard, so it + # is a valid destination even though it is not a box of the shard's grid. + pieces = [ + AffinePiece((2, 2), 0, (2, 1), 2, (3, 1), [0]), + AffinePiece((2, ), 4, (1, ), 0, (1, ), [0]), + AffinePiece((1, ), 6, (1, ), 4, (1, ), [0]), + AffinePiece((2, ), 7, (1, ), 7, (1, ), [0]), + ] + affine_map = ParamAffineMap((9, ), {0: (3, 3)}, {0: pieces}) + affine_map.validate_coverage() + full = torch.arange(9.) + torch.testing.assert_close(affine_map.rebuild({0: affine_map.extract(full, 0)}), full) + + def test_full_check_refuses_a_destination_collision_of_any_stride(): # The first piece writes shard addresses 0 and 2, the second writes 1 and 2, and address 3 # is never written. Neither piece is a box of the shard's grid, so only walking the