[checkpoint] Refuse a map whose pieces write the same shard address

deepspeed/checkpoint/affine.py

@@ -141,6 +141,16 @@ def __hash__(self):
                      self.locations, self.scale))
 
 
+def _dest_interval(piece):
+    """Return a dense piece's destination interval, or None for a strided or empty piece."""
+    if piece.numel == 0 or len(piece.dest_strides) != len(piece.shape):
+        return None
+    if any(size > 1 and got != want
+           for size, got, want in zip(piece.shape, piece.dest_strides, _row_major_strides(piece.shape))):
+        return None
+    return piece.dest_offset, piece.dest_offset + piece.numel
+
+
 class ParamAffineMap:
     """How one parameter is spread over a tensor-parallel group.
 
@@ -179,12 +189,35 @@ def holders(self):
         return holders
 
     def validate(self):
-        """Cheap structural check: every rank's pieces account for exactly its shard."""
+        """Check shard element counts and dense destination bounds and overlap."""
         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}.')
+            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)
+
+    def _validate_destinations(self, rank, pieces, shard_numel):
+        """Check that dense destination intervals stay within the shard and do not overlap."""
+        intervals = []
+        for piece in pieces:
+            span = _dest_interval(piece)
+            if span is None:
+                continue
+            start, stop = span
+            if start < 0 or stop > shard_numel:
+                raise ValueError(f'Rank {rank} piece {piece} writes shard addresses [{start}, {stop}), '
+                                 f'outside the shard of {shard_numel} elements.')
+            intervals.append(span)
+
+        intervals.sort()
+        cursor = 0
+        for start, stop in intervals:
+            if start < cursor:
+                raise ValueError(f'Rank {rank} writes shard address {start} twice, so one piece would '
+                                 'silently overwrite another.')
+            cursor = stop
 
     def validate_coverage(self):
         """Full check: the shards cover the parameter, and every piece is homogeneous.

tests/unit/checkpoint/test_affine_shard_map.py

@@ -456,6 +456,55 @@ def test_stored_version_is_the_format_version():
     assert AFFINE_MAP_FORMAT_VERSION >= 1
 
 
+def test_singleton_stride_overlap_is_refused():
+    pieces = [AffinePiece((1, 2), offset, (4, 1), 0, (4, 1), [0]) for offset in (0, 2)]
+    affine_map = ParamAffineMap((1, 4), {0: (1, 4)}, {0: pieces})
+    with pytest.raises(ValueError, match='twice'):
+        affine_map.validate()
+
+
+def test_singleton_stride_out_of_bounds_is_refused():
+    pieces = [AffinePiece((1, 2), source, (4, 1), dest, (4, 1), [0]) for source, dest in ((0, 0), (2, 3))]
+    affine_map = ParamAffineMap((1, 4), {0: (1, 4)}, {0: pieces})
+    with pytest.raises(ValueError, match='outside the shard'):
+        affine_map.validate()
+
+
+def test_singleton_stride_contiguous_layout_round_trips():
+    pieces = [AffinePiece((1, 2), offset, (4, 1), offset, (4, 1), [0]) for offset in (0, 2)]
+    affine_map = ParamAffineMap((1, 4), {0: (1, 4)}, {0: pieces})
+    shard = torch.arange(4).reshape(1, 4)
+    affine_map.validate()
+    torch.testing.assert_close(affine_map.rebuild({0: shard}), shard)
+
+
+def test_strided_column_pieces_remain_valid():
+    pieces = [AffinePiece((2, 2), offset, (4, 1), offset, (4, 1), [0]) for offset in (0, 2)]
+    affine_map = ParamAffineMap((2, 4), {0: (2, 4)}, {0: pieces})
+    shard = torch.arange(8).reshape(2, 4)
+    affine_map.validate()
+    torch.testing.assert_close(affine_map.rebuild({0: shard}), shard)
+
+
+def test_shard_element_count_mismatch_is_value_error():
+    piece = AffinePiece((2, ), 0, (1, ), 0, (1, ), [0])
+    affine_map = ParamAffineMap((4, ), {0: (4, )}, {0: [piece]})
+    with pytest.raises(ValueError, match='account for 2'):
+        affine_map.validate()
+
+
+def test_destination_validation_does_not_enumerate_elements(monkeypatch):
+    size = 10**9
+    piece = AffinePiece((size, 8), 0, (8, 1), 0, (8, 1), [0])
+    affine_map = ParamAffineMap((size, 8), {0: (size, 8)}, {0: [piece]})
+
+    def refuse_walk(self):
+        raise AssertionError('validation must not enumerate parameter elements')
+
+    monkeypatch.setattr(AffinePiece, 'source_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.