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

deepspeed/checkpoint/affine.py

@@ -141,6 +141,20 @@ def __hash__(self):
                      self.locations, self.scale))
 
 
+def _dest_interval(piece):
+    """The shard addresses a piece writes, when they are one contiguous run.
+
+    A piece whose destination strides match its own shape occupies
+    ``[dest_offset, dest_offset + numel)`` and nothing else, so its footprint is an interval and
+    overlap between two of them is decidable without touching elements. A piece cut along a
+    non-final dimension -- a column split, or Yuan's o_proj -- strides through the shard instead,
+    and intervals say nothing about it, so the caller gets None and checks what it can.
+    """
+    if piece.numel == 0 or 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 +193,48 @@ def holders(self):
         return holders
 
     def validate(self):
-        """Cheap structural check: every rank's pieces account for exactly its shard."""
+        """Cheap structural check: every rank's pieces account for exactly its shard.
+
+        Counting elements is not the same as placing them. Two pieces can both write the head of
+        a shard and leave its tail unwritten while summing to the shard size exactly, so the
+        totals agree and nothing here reports a problem. This also walks the destination addresses
+        of every piece that packs densely, which bounds the check by the piece count rather than
+        the shard size.
+        """
         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):
+        """Refuse two pieces of one rank that write the same shard address.
+
+        Dense pieces are the only ones an interval can describe, and a rank where every piece is
+        dense needs no separate hole check: intervals that lie inside the shard, do not overlap and
+        whose lengths sum to the shard size cannot leave a gap. Where some piece is strided the
+        hole question is not answerable from intervals, and this stays quiet rather than guess.
+        """
+        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.