[Distributed] Optimize ND shard overlap detection (pytorch#167073)

* Fixes high quadratic cost in `validate_non_overlapping_shards_metadata` when shard count is large by replacing the O(n²) nested-loop scan in `_find_nd_overlapping_shards` with a sweep-line pass giving O(n log n) behavior for ND overlap detection.
* Add test cases in `test_check_overlapping` covering 2D grid patterns, adjacent shards, and 3D multi-shard overlap scenarios to validate the optimized path.

Fixes pytorch#166941

Pull Request resolved: pytorch#167073
Approved by: https://github.com/Skylion007, https://github.com/wconstab

Author: mansiag05

test/distributed/_shard/sharding_spec/test_sharding_spec.py

@@ -490,6 +490,69 @@ def test_check_overlapping(self):
         with self.assertRaisesRegex(ValueError, "overlap"):
             validate_non_overlapping_shards_metadata(shards)
 
+        shards = [
+            ShardMetadata(
+                shard_offsets=[0, 0],
+                shard_sizes=[5, 5],
+                placement="cuda:0",
+            ),
+            ShardMetadata(
+                shard_offsets=[0, 5],
+                shard_sizes=[5, 5],
+                placement="cuda:1",
+            ),
+            ShardMetadata(
+                shard_offsets=[5, 0],
+                shard_sizes=[5, 5],
+                placement="cuda:2",
+            ),
+            ShardMetadata(
+                shard_offsets=[5, 5],
+                shard_sizes=[5, 5],
+                placement="cuda:3",
+            ),
+        ]
+        validate_non_overlapping_shards_metadata(shards)
+
+        shards = [
+            ShardMetadata(
+                shard_offsets=[0, 0],
+                shard_sizes=[5, 5],
+                placement="cuda:0",
+            ),
+            ShardMetadata(
+                shard_offsets=[5, 5],
+                shard_sizes=[5, 5],
+                placement="cuda:1",
+            ),
+        ]
+        validate_non_overlapping_shards_metadata(shards)
+
+        shards = [
+            ShardMetadata(
+                shard_offsets=[0, 0, 0],
+                shard_sizes=[5, 5, 5],
+                placement="cuda:0",
+            ),
+            ShardMetadata(
+                shard_offsets=[5, 0, 0],
+                shard_sizes=[5, 5, 5],
+                placement="cuda:1",
+            ),
+            ShardMetadata(
+                shard_offsets=[10, 0, 0],
+                shard_sizes=[5, 5, 5],
+                placement="cuda:2",
+            ),
+            ShardMetadata(
+                shard_offsets=[10, 3, 0],
+                shard_sizes=[5, 5, 5],
+                placement="cuda:3",
+            ),
+        ]
+        with self.assertRaisesRegex(ValueError, "overlap"):
+            validate_non_overlapping_shards_metadata(shards)
+
 
 # Custom ShardingSpec, an simple example to do grid sharding
 @dataclass

torch/distributed/_shard/sharding_spec/_internals.py

@@ -1,5 +1,7 @@
 # mypy: allow-untyped-defs
 import math
+import sys
+from bisect import bisect_right, insort
 from typing import Optional
 
 from torch.distributed._shard.metadata import ShardMetadata
@@ -27,31 +29,48 @@ def _check_shard_metadata_pair_overlap(shard1: ShardMetadata, shard2: ShardMetad
 def _find_nd_overlapping_shards(
     shards: list[ShardMetadata], sharded_dims: list[int]
 ) -> Optional[tuple[int, int]]:
-    # Each rank has len(sharded_dims) tuples. Each tuple represent the
-    # [begin, end] (inclusive) pair of that dimension.
-    shard_intervals = [
-        [
-            (s.shard_offsets[dim], s.shard_offsets[dim] + s.shard_sizes[dim] - 1)
-            for dim in sharded_dims
-        ]
-        for s in shards
-    ]
-
-    for i in range(len(shards)):
-        shard_i = shard_intervals[i]
-        for j in range(i + 1, len(shards)):
-            shard_j = shard_intervals[j]
-            # For each dim of each shard, check if one shard resides on the other
-            # end of second shard with respect to that dim. As an example for a 2D
-            # shard, we would check if one shard is above or on the left of the
-            # other shard.
-            overlap = True
-            for interval_i, interval_j in zip(shard_i, shard_j):
-                if interval_i[0] > interval_j[1] or interval_j[0] > interval_i[1]:
-                    overlap = False
-                    break
-            if overlap:
-                return (i, j)
+    """Find overlapping shards using sweep-line algorithm."""
+    if len(shards) <= 1:
+        return None
+
+    dims = len(sharded_dims)
+    if dims == 0:
+        return None
+
+    sweep_dim_idx = 0
+    if dims > 1:
+        max_size = 0
+        for i, dim in enumerate(sharded_dims):
+            dim_size = shards[0].shard_offsets[dim] + shards[0].shard_sizes[dim]
+            if dim_size > max_size:
+                max_size = dim_size
+                sweep_dim_idx = i
+    sweep_dim = sharded_dims[sweep_dim_idx]
+
+    sorted_indices = sorted(
+        range(len(shards)),
+        key=lambda idx: (
+            shards[idx].shard_offsets[sweep_dim],
+            *(shards[idx].shard_offsets[d] for d in sharded_dims if d != sweep_dim),
+        ),
+    )
+    active: list[tuple[int, int]] = []
+
+    for idx in sorted_indices:
+        current = shards[idx]
+        start = current.shard_offsets[sweep_dim]
+        end = start + current.shard_sizes[sweep_dim]
+
+        cutoff = bisect_right(active, (start, sys.maxsize))
+        if cutoff:
+            del active[:cutoff]
+
+        for _, other_idx in active:
+            other = shards[other_idx]
+
+            if _check_shard_metadata_pair_overlap(current, other):
+                return (other_idx, idx)
+        insort(active, (end, idx))
     return None
 
 
@@ -112,10 +131,8 @@ def validate_non_overlapping_shards_metadata(shards: list[ShardMetadata]):
         # using a O(nlogn) overlapping interval algorithm.
         pair = _find_1d_overlapping_shards(shards, sharded_dims[0])
     else:
-        # Shards are partitioned over more than one dimension. Fall back to
-        # pair-wise check. Even though O(nlogn) algorithms (line sweep) exist
-        # for 2D overlap, the implementation is not trivial and may not justify
-        # the time saving in most cases.
+        # Shards are partitioned over more than one dimension.
+        # Use sweep-line algorithm for O(n log n) complexity.
         pair = _find_nd_overlapping_shards(shards, sharded_dims)
 
     if pair: