Fix vmem_chunked asymmetric allocation bug

iris/host/memory/allocators/vmem_chunked_allocator.py

@@ -12,11 +12,13 @@
 - Power-of-two free lists for O(1) alloc/free reuse
 - GC via weakref finalizers on tensor.untyped_storage()
 - Free/reuse is pure bookkeeping (no driver calls, no physical remap)
+- A block becomes reusable only once every rank has freed it (see reuse()),
+  so allocation offsets stay identical across ranks whenever each rank's GC runs
 
 Cost model:
 - Init: ~170us per chunk (create + map + set_access for 1 device)
 - Per-allocation: 0us (bump or free-list pop, no driver calls)
-- Per-free: 0us (push to free list, no driver calls)
+- Per-free: 0us (push to a local released list, no driver calls)
 - Chunk growth: ~170us (rare, every chunk_size bytes)
 """
 
@@ -125,6 +127,8 @@ class VMemChunkedAllocator(BaseAllocator):
     DEFAULT_CHUNK_SIZE = 256 * 1024 * 1024
     # Default VA size: 0 means auto-size (128 GiB)
     DEFAULT_VA_SIZE = 0
+    # Reconcile frees (SymmetricHeap.collect()) after handing out 1/COLLECT_FRACTION of the mapped heap
+    COLLECT_FRACTION = 8
 
     def __init__(
         self,
@@ -149,6 +153,10 @@ def __init__(
         self.free_lists = defaultdict(list)
         self.alloc_sizes = {}
         self._pending_free = deque()
+        # Freed on this rank but not yet reusable: offset -> size class
+        self._released = {}
+        # Bytes handed out since the last reuse(): identical on every rank, like the allocation sequence
+        self._bytes_since_collect = 0
         self.chunks: List[LocalAllocation] = []
         self._shared_regions: List[ExportableMemory] = []
         self._peer_mappings: List[PeerMapping] = []
@@ -256,10 +264,38 @@ def _grow_chunk(self):
         self.mapped_extent += self.chunk_size
 
     def _process_pending_frees(self):
-        """Process pending frees from GC finalizers. Call with lock held."""
+        """Move frees from GC finalizers to the released list. Call with lock held."""
         while self._pending_free:
             offset, size_class = self._pending_free.popleft()
-            self.free_lists[size_class].append(offset)
+            self.alloc_sizes.pop(offset, None)
+            self._released[offset] = size_class
+
+    def released_offsets(self) -> List[int]:
+        """Offsets freed on this rank that are not yet reusable, sorted."""
+        with self.lock:
+            self._process_pending_frees()
+            return sorted(self._released)
+
+    def collect_due(self) -> bool:
+        """True once 1/COLLECT_FRACTION of the mapped heap has been handed out since the last reuse()."""
+        with self.lock:
+            return self._bytes_since_collect >= self.mapped_extent // self.COLLECT_FRACTION
+
+    def reuse(self, offsets) -> None:
+        """
+        Make released blocks reusable. Call with the same offsets on every rank: those every rank has
+        released (see SymmetricHeap.collect()), so the free lists, and every later offset, stay
+        identical across ranks.
+        """
+        with self.lock:
+            self._bytes_since_collect = 0
+            offsets = sorted(offsets)
+            # Blocks at the top of the bump region go back to the bump pointer, not a size class
+            while offsets and offsets[-1] + self._released[offsets[-1]] == self.bump:
+                self.bump = offsets[-1]
+                del self._released[offsets.pop()]
+            for offset in offsets:
+                self.free_lists[self._released.pop(offset)].append(offset)
 
     def _free_callback(self, offset, size_class):
         """Called by weakref finalizer when a tensor's storage is GC'd.
@@ -294,7 +330,13 @@ def transport_tier(self) -> InterconnectLevel:
         """
         return self._interconnect
 
-    def allocate(self, num_elements: int, dtype: torch.dtype, alignment: int = 1024) -> torch.Tensor:
+    def allocate(
+        self, num_elements: int, dtype: torch.dtype, alignment: int = 1024, grow: bool = True
+    ) -> Optional[torch.Tensor]:
+        """
+        Allocate from a free list or the bump pointer. With ``grow=False``, return None instead of
+        mapping new chunks.
+        """
         with self.lock:
             if num_elements == 0:
                 return torch.empty(0, dtype=dtype, device=self.device)
@@ -320,6 +362,8 @@ def allocate(self, num_elements: int, dtype: torch.dtype, alignment: int = 1024)
                 aligned_bump = (self.bump + self.min_alignment - 1) & ~(self.min_alignment - 1)
 
                 needed = aligned_bump + aligned_size - self.mapped_extent
+                if needed > 0 and not grow:
+                    return None
                 if needed > 0:
                     n_new_chunks = math.ceil(needed / self.chunk_size)
                     if self.mapped_extent + n_new_chunks * self.chunk_size > self.va_size:
@@ -336,6 +380,7 @@ def allocate(self, num_elements: int, dtype: torch.dtype, alignment: int = 1024)
 
             # Track for free
             self.alloc_sizes[offset] = aligned_size
+            self._bytes_since_collect += aligned_size
 
             # Create tensor via CUDAArrayInterface
             va = self.base_va + offset
@@ -579,19 +624,22 @@ def release_peer_chunk(self, remote_va: int) -> None:
 
     def get_stats(self):
         """Return allocator statistics."""
-        total_free = sum(len(v) for v in self.free_lists.values())
-        free_bytes = sum(size_class * len(offsets) for size_class, offsets in self.free_lists.items())
-        return {
-            "num_chunks": len(self.chunks),
-            "mapped_bytes": self.mapped_extent,
-            "bump": self.bump,
-            "num_active_allocs": len(self.alloc_sizes),
-            "num_free_blocks": total_free,
-            "free_bytes": free_bytes,
-            "va_size": self.va_size,
-            "chunk_size": self.chunk_size,
-            "granularity": self.granularity,
-        }
+        with self.lock:
+            self._process_pending_frees()
+            total_free = sum(len(v) for v in self.free_lists.values())
+            free_bytes = sum(size_class * len(offsets) for size_class, offsets in self.free_lists.items())
+            return {
+                "num_chunks": len(self.chunks),
+                "mapped_bytes": self.mapped_extent,
+                "bump": self.bump,
+                "num_active_allocs": len(self.alloc_sizes),
+                "num_free_blocks": total_free,
+                "num_released_blocks": len(self._released),
+                "free_bytes": free_bytes,
+                "va_size": self.va_size,
+                "chunk_size": self.chunk_size,
+                "granularity": self.granularity,
+            }
 
     def close(self):
         """Release all VMem resources."""
@@ -605,6 +653,7 @@ def close(self):
         with self.lock:
             # Clear allocator bookkeeping before releasing mappings/chunks.
             self._pending_free.clear()
+            self._released.clear()
             self.free_lists.clear()
             self.alloc_sizes.clear()
 

iris/host/memory/symmetric_heap.py

@@ -8,6 +8,7 @@
 hiding the details of allocators and inter-process memory sharing.
 """
 
+import gc
 import itertools
 import logging
 import os
@@ -79,6 +80,7 @@ def _validate_local_handle(handle_bytes: bytes) -> None:
 
 
 _hostname_exchanges = itertools.count(1)
+_released_exchanges = itertools.count(1)
 
 
 def _hostnames_by_rank(num_ranks: int) -> List[str]:
@@ -91,6 +93,28 @@ def _hostnames_by_rank(num_ranks: int) -> List[str]:
     return [store.get(f"{prefix}{r}").decode("utf-8") for r in range(num_ranks)]
 
 
+def _released_by_every_rank(offsets: List[int], num_ranks: int) -> List[int]:
+    """The offsets every rank released, exchanged through the rendezvous store (a barrier, no NCCL collective)."""
+    import torch.distributed as dist
+
+    store = dist.distributed_c10d._get_default_store()
+    prefix = f"iris_released_v{next(_released_exchanges)}/"
+    store.set(f"{prefix}{dist.get_rank()}", np.asarray(offsets, dtype=np.int64).tobytes())
+    keys = [f"{prefix}{r}" for r in range(num_ranks)]
+    try:
+        store.wait(keys)
+    except Exception as exc:
+        missing = [r for r, key in enumerate(keys) if not store.check([key])]
+        raise RuntimeError(
+            f"SymmetricHeap.collect(): ranks {missing} did not reach the same allocation before the store "
+            "timed out; symmetric allocations must happen on every rank"
+        ) from exc
+    common = set(offsets)
+    for key in keys:
+        common &= set(np.frombuffer(store.get(key), dtype=np.int64).tolist())
+    return sorted(common)
+
+
 def _require_single_host(allocator_type: str, hostnames: List[str]) -> None:
     hosts = sorted(set(hostnames))
     if len(hosts) > 1:
@@ -317,7 +341,16 @@ def allocate(self, num_elements: int, dtype: torch.dtype, alignment: int = 1024)
         is_chunked = isinstance(self.allocator, VMemChunkedAllocator)
         chunks_before = self.allocator.get_num_chunks() if is_chunked else 0
 
-        tensor = self.allocator.allocate(actual_elements, dtype, alignment)
+        if is_chunked:
+            # Every rank reaches the same branches: they depend only on state identical across ranks
+            if self.allocator.collect_due():
+                self.collect()
+            tensor = self.allocator.allocate(actual_elements, dtype, alignment, grow=False)
+            if tensor is None:
+                self.collect()
+                tensor = self.allocator.allocate(actual_elements, dtype, alignment)
+        else:
+            tensor = self.allocator.allocate(actual_elements, dtype, alignment)
         tensor = tensor[:num_elements]
 
         if is_chunked:
@@ -427,6 +460,25 @@ def translate(self, ptr: int, from_rank: int, to_rank: int) -> int:
         offset = ptr - from_base
         return to_base + offset
 
+    def collect(self):
+        """
+        Collective: make the heap memory that every rank has freed reusable.
+
+        Freed symmetric tensors are only reused once every rank has freed them, so that allocations
+        land at the same offset on every rank whenever each rank's garbage collector runs.
+        ``allocate()`` calls this before growing the heap, and once 1/``COLLECT_FRACTION`` of the
+        mapped heap has been handed out since the last call; call it directly to reuse memory sooner.
+        Waits for this rank's device first, so no rank is still accessing a block it hands out.
+        """
+        if not isinstance(self.allocator, VMemChunkedAllocator):
+            return
+        gc.collect()
+        torch.cuda.synchronize(self.device_id)
+        released = self.allocator.released_offsets()
+        if self.num_ranks > 1:
+            released = _released_by_every_rank(released, self.num_ranks)
+        self.allocator.reuse(released)
+
     def refresh_peer_access(self):
         """
         Refresh peer imports for the active allocator backend.

tests/unittests/test_vmem_chunked_allocator.py

@@ -7,7 +7,7 @@
 Tests cover:
 - Basic allocation and data integrity
 - Multiple allocations and non-overlap
-- Power-of-two free list reuse via GC
+- Power-of-two free list reuse via GC, once every rank has freed a block
 - Chunk growth on overflow
 - as_symmetric (import external tensor)
 - owns_tensor detection
@@ -93,10 +93,9 @@ def test_chunked_gc_free_reuse():
     t = ctx.zeros(1024, dtype=torch.float32)
     ptr1 = t.data_ptr()
 
-    # Drop the tensor and trigger GC
+    # Drop the tensor and make freed memory reusable
     del t
-    gc.collect()
-    torch.cuda.synchronize()
+    ctx.heap.collect()
 
     # Allocate again -- should reuse the freed block
     t2 = ctx.zeros(1024, dtype=torch.float32)
@@ -144,15 +143,16 @@ def test_chunked_free_list_size_classes():
     small = ctx.zeros(size_small, dtype=torch.float32)
     medium = ctx.zeros(size_medium, dtype=torch.float32)
     large = ctx.zeros(size_large, dtype=torch.float32)
+    # Keeps the freed blocks below the top of the heap, so they go to the free lists
+    sentinel = ctx.zeros(size_small, dtype=torch.float32)  # noqa: F841
 
     small_ptr = small.data_ptr()
     medium_ptr = medium.data_ptr()
     large_ptr = large.data_ptr()
 
     # Free all
     del small, medium, large
-    gc.collect()
-    torch.cuda.synchronize()
+    ctx.heap.collect()
 
     # Re-allocate -- each should reuse from its size class
     small2 = ctx.zeros(size_small, dtype=torch.float32)
@@ -164,6 +164,109 @@ def test_chunked_free_list_size_classes():
     assert large2.data_ptr() == large_ptr
 
 
+@pytest.mark.parametrize("keep_every", [0, 16])
+def test_chunked_churn_then_large_allocation_does_not_grow(keep_every):
+    """
+    A churned scratch buffer must not fragment the heap so that a later large allocation grows it:
+    freed blocks at the top return to the bump pointer, and periodic collect() calls bound the rest
+    (also when buffers kept every ``keep_every`` iterations stop the rollback).
+    """
+    ctx = iris.iris(64 << 20, allocator_type=ALLOC_TYPE)
+    alloc = ctx.heap.allocator
+    scratch_bytes, keep_bytes = 2 * alloc.granularity, alloc.granularity  # different size classes
+    chunks = alloc.get_num_chunks()
+
+    kept = []
+    for i in range(3 * alloc.mapped_extent // scratch_bytes):
+        t = ctx.zeros(scratch_bytes // 4, dtype=torch.float32)
+        t.fill_(1.0)
+        del t
+        if keep_every and i % keep_every == keep_every - 1:
+            kept.append(ctx.zeros(keep_bytes // 4, dtype=torch.float32))
+
+    big = ctx.zeros(alloc.chunk_size // 2 // 4, dtype=torch.float32)  # noqa: F841
+    assert alloc.get_num_chunks() == chunks
+
+
+def test_chunked_freed_block_waits_for_collect():
+    """A freed block is only reused after collect(), so every rank's offsets change at the same point."""
+    ctx = iris.iris(64 << 20, allocator_type=ALLOC_TYPE)
+
+    t = ctx.zeros(1024, dtype=torch.float32)
+    ptr = t.data_ptr()
+    del t
+    gc.collect()
+
+    t2 = ctx.zeros(1024, dtype=torch.float32)
+    assert t2.data_ptr() != ptr
+    assert ctx.heap.allocator.get_stats()["num_released_blocks"] >= 1
+
+    ctx.heap.collect()
+    t3 = ctx.zeros(1024, dtype=torch.float32)
+    assert t3.data_ptr() == ptr
+
+
+def test_chunked_reuses_before_growing():
+    """An allocation that would grow the heap first reuses blocks freed on every rank."""
+    ctx = iris.iris(1 << 20, allocator_type=ALLOC_TYPE)
+    alloc = ctx.heap.allocator
+    size = alloc.granularity // 4  # float32 elements in one granule
+
+    # Fill the mapped chunks, so the next allocation of this size would grow the heap
+    held = [ctx.zeros(size, dtype=torch.float32)]
+    while alloc.bump + alloc.granularity <= alloc.mapped_extent:
+        held.append(ctx.zeros(size, dtype=torch.float32))
+    chunks = alloc.get_num_chunks()
+    ptr = held[0].data_ptr()
+    del held[0]
+
+    t = ctx.zeros(size, dtype=torch.float32)
+    assert t.data_ptr() == ptr
+    assert alloc.get_num_chunks() == chunks
+
+
+def _offsets_per_rank(ctx, t):
+    offset = torch.tensor([t.data_ptr() - int(ctx.heap.heap_bases_cpu[ctx.cur_rank])], device=t.device)
+    every = [torch.empty_like(offset) for _ in range(ctx.num_ranks)]
+    torch.distributed.all_gather(every, offset)
+    return [int(o.item()) for o in every]
+
+
+def test_chunked_offsets_symmetric_across_gc_timing():
+    """GC running on one rank before another must not change where the next allocation lands."""
+    ctx = iris.iris(64 << 20, allocator_type=ALLOC_TYPE)
+    if ctx.num_ranks < 2:
+        pytest.skip("Requires at least 2 ranks")
+
+    gc.disable()
+    try:
+        # A tensor reachable only through a reference cycle: freed when the cyclic GC runs
+        a = ctx.zeros(4096, dtype=torch.float32)
+        holder = {"t": a}
+        holder["self"] = holder
+        del a, holder
+        if ctx.cur_rank == 0:
+            gc.collect()
+        b = ctx.zeros(4096, dtype=torch.float32)
+        offsets = _offsets_per_rank(ctx, b)
+        assert len(set(offsets)) == 1, f"offsets differ across ranks: {offsets}"
+
+        # A block one rank still holds is not handed out, even after collect()
+        c = ctx.zeros(4096, dtype=torch.float32)
+        c_ptr = c.data_ptr()
+        keep = c if ctx.cur_rank == 0 else None
+        del c
+        ctx.heap.collect()
+        d = ctx.zeros(4096, dtype=torch.float32)
+        offsets = _offsets_per_rank(ctx, d)
+        assert len(set(offsets)) == 1, f"offsets differ across ranks: {offsets}"
+        assert d.data_ptr() != c_ptr
+        del keep
+    finally:
+        gc.enable()
+    ctx.barrier()
+
+
 def test_chunked_chunk_growth():
     """Test that allocator grows chunks when needed."""
     # Small chunk size to force growth