From 7c24c08d41965b43278451df561fa32498675306 Mon Sep 17 00:00:00 2001
From: mkolodner <mkolodner@snapchat.com>
Date: Wed, 9 Sep 2026 21:23:19 +0000
Subject: [PATCH] Chunk the degree-tensor all-reduce to bound peak device
 memory

The single-shot all-reduce moved the whole padded int64 degree tensor to the
process-group device. For large graphs that tensor is several GiB, so together
with the collective buffer it can OOM small-memory GPUs during dataloader setup,
even though the reduced result is only ever used on CPU.

Reduce one fixed-size slice at a time: keep the padded tensor on CPU and, per
slice, copy it to the device, all-reduce, and copy the result back. A chunked
element-wise SUM is bit-identical to a whole-tensor SUM (every rank pads to the
same size and uses identical slice bounds); the over-counting correction is
unchanged. This runs once at loader setup, so there is no per-batch cost.
---
 gigl/distributed/utils/degree.py | 35 ++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/gigl/distributed/utils/degree.py b/gigl/distributed/utils/degree.py
index f88d3c0c6..2eb9b8d84 100644
--- a/gigl/distributed/utils/degree.py
+++ b/gigl/distributed/utils/degree.py
@@ -188,24 +188,43 @@ def _get_degree_reduce_context() -> tuple[int, torch.device]:
     return local_world_size, device
 
 
+# All-reduce the degree tensor one slice at a time. The full padded int64
+# tensor for a large graph is several GiB; materializing it (plus the collective
+# buffer) on the process-group device OOMs small GPUs. Since the result is
+# CPU-bound anyway, only one slice is moved to the device per collective. A
+# chunked element-wise SUM is bit-identical to a whole-tensor SUM (every rank
+# pads to the same size and uses identical slice bounds). This runs once at
+# loader setup, so it adds no per-batch throughput cost.
+_DEGREE_ALLREDUCE_CHUNK = 64_000_000
+
+
 def _all_reduce_single_degree_tensor(
     tensor: torch.Tensor,
     local_world_size: int,
     device: torch.device,
 ) -> torch.Tensor:
-    """All-reduce a single tensor with size sync and over-counting correction."""
+    """All-reduce a single tensor with size sync and over-counting correction.
+
+    The reduction is chunked so only one ``_DEGREE_ALLREDUCE_CHUNK``-sized slice
+    is resident on ``device`` at a time, keeping the (potentially multi-GiB)
+    degree tensor off the GPU except for the slice currently being reduced.
+    """
     # Synchronize max size across all ranks.
     local_size = torch.tensor([tensor.size(0)], dtype=torch.long, device=device)
     torch.distributed.all_reduce(local_size, op=torch.distributed.ReduceOp.MAX)
     max_size = int(local_size.item())
 
-    # Pad, convert to int64 for all_reduce, and move to the process-group device.
-    padded = _pad_to_size(tensor, max_size).to(torch.int64).to(device)
-    torch.distributed.all_reduce(padded, op=torch.distributed.ReduceOp.SUM)
-
-    # Correct for over-counting and move back to CPU. Clamp before casting so
-    # high-degree nodes saturate instead of wrapping.
-    return _clamp_to_int32(padded // local_world_size).cpu()
+    # Pad and widen on CPU; keep the full tensor off the process-group device.
+    padded = _pad_to_size(tensor, max_size).to(torch.int64)
+    for start in range(0, max_size, _DEGREE_ALLREDUCE_CHUNK):
+        end = min(start + _DEGREE_ALLREDUCE_CHUNK, max_size)
+        chunk = padded[start:end].to(device)
+        torch.distributed.all_reduce(chunk, op=torch.distributed.ReduceOp.SUM)
+        padded[start:end] = chunk.cpu()
+
+    # Correct for over-counting. Clamp before casting so high-degree nodes
+    # saturate instead of wrapping. Result stays on CPU.
+    return _clamp_to_int32(padded // local_world_size)
 
 
 @overload