From 58ee4f8f57600faaa96097d0e586bd9af0c27667 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Mon, 12 Oct 2026 18:35:34 -0500 Subject: [PATCH 1/3] Keep vmem_chunked heap offsets symmetric when ranks GC at different times The chunked allocator freed a block on the rank whose garbage collector dropped the tensor, and the next allocation of that size reused it. A symmetric tensor reachable only through a reference cycle is freed when the cyclic GC runs, which differs between ranks, so the rank that had collected reused the slot while another bump-allocated. The same allocation then landed at different offsets, the bump pointers diverged, and every later remote access went to the wrong peer address, silently. Frees still come from GC finalizers, but a freed block now only becomes reusable once every rank has freed it. SymmetricHeap.collect() runs gc.collect() and a device sync on every rank, exchanges the released offsets through the rendezvous store, and moves the common ones to the free lists in sorted order. Allocation reads only state that is identical across ranks, so every rank takes the same path: it tries the free lists and the bump pointer first, and calls collect() only when the heap would otherwise grow. The common path stays free of collectives. Reuse now waits for that point (or an explicit collect()) instead of the next allocation; the two reuse tests call collect() accordingly. New tests cover reuse before growth, and, across ranks, offsets staying symmetric when GC runs on one rank first and a block one rank still holds not being handed out. Also drop freed offsets from alloc_sizes. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../allocators/vmem_chunked_allocator.py | 42 +++++++-- iris/host/memory/symmetric_heap.py | 51 ++++++++++- .../unittests/test_vmem_chunked_allocator.py | 89 +++++++++++++++++-- 3 files changed, 169 insertions(+), 13 deletions(-) diff --git a/iris/host/memory/allocators/vmem_chunked_allocator.py b/iris/host/memory/allocators/vmem_chunked_allocator.py index e66db02df..5ecddd557 100644 --- a/iris/host/memory/allocators/vmem_chunked_allocator.py +++ b/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) """ @@ -149,6 +151,8 @@ 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 = {} self.chunks: List[LocalAllocation] = [] self._shared_regions: List[ExportableMemory] = [] self._peer_mappings: List[PeerMapping] = [] @@ -256,10 +260,28 @@ 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._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 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: + for offset in sorted(offsets): + size_class = self._released.pop(offset) + self.alloc_sizes.pop(offset, None) + self.free_lists[size_class].append(offset) def _free_callback(self, offset, size_class): """Called by weakref finalizer when a tensor's storage is GC'd. @@ -294,13 +316,17 @@ 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) - self._process_pending_frees() - elem_size = _element_size(dtype) size_bytes = num_elements * elem_size # Minimum allocation is one granule @@ -320,6 +346,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: @@ -587,6 +615,7 @@ def get_stats(self): "bump": self.bump, "num_active_allocs": len(self.alloc_sizes), "num_free_blocks": total_free, + "num_released_blocks": len(self._released) + len(self._pending_free), "free_bytes": free_bytes, "va_size": self.va_size, "chunk_size": self.chunk_size, @@ -605,6 +634,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() diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 5fab77f51..0b9177854 100644 --- a/iris/host/memory/symmetric_heap.py +++ b/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,14 @@ 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 branch: it depends only on state identical across ranks + 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 +458,24 @@ 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; 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. diff --git a/tests/unittests/test_vmem_chunked_allocator.py b/tests/unittests/test_vmem_chunked_allocator.py index 3d4fca42d..eae674a0d 100644 --- a/tests/unittests/test_vmem_chunked_allocator.py +++ b/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) @@ -151,8 +150,7 @@ def test_chunked_free_list_size_classes(): # 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 +162,85 @@ def test_chunked_free_list_size_classes(): assert large2.data_ptr() == large_ptr +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 From 1965d08106e93d195c7651d845181bbc4c7aaeb9 Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Tue, 13 Oct 2026 12:45:18 -0500 Subject: [PATCH 2/3] Return freed blocks at the top of the heap to the bump pointer Deferring reuse until the heap would grow let a churning scratch buffer fill the mapped heap with blocks of one size class first. The free lists are exact power-of-two classes without coalescing, so a later, larger allocation found nothing reusable and grew the heap, where before it fit. reuse() now walks the common offsets down from the top: while a block ends exactly at the bump pointer, the pointer moves down to its start instead of the block going onto a free list. The offsets and the bump pointer are identical on every rank, so this stays symmetric. Blocks below a live allocation or an import still go to their free list. Also drop a freed block from alloc_sizes as soon as its finalizer's free is drained (in allocate() and get_stats() as well as collect()), so num_active_allocs counts only live allocations. test_chunked_churn_then_large_allocation_does_not_grow covers the churn case (it grows the heap without this change); the size-class test keeps a live allocation above its blocks so it still exercises the free lists. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../allocators/vmem_chunked_allocator.py | 44 +++++++++++-------- .../unittests/test_vmem_chunked_allocator.py | 19 ++++++++ 2 files changed, 45 insertions(+), 18 deletions(-) diff --git a/iris/host/memory/allocators/vmem_chunked_allocator.py b/iris/host/memory/allocators/vmem_chunked_allocator.py index 5ecddd557..67c4b68e0 100644 --- a/iris/host/memory/allocators/vmem_chunked_allocator.py +++ b/iris/host/memory/allocators/vmem_chunked_allocator.py @@ -263,6 +263,7 @@ def _process_pending_frees(self): """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.alloc_sizes.pop(offset, None) self._released[offset] = size_class def released_offsets(self) -> List[int]: @@ -278,10 +279,13 @@ def reuse(self, offsets) -> None: identical across ranks. """ with self.lock: - for offset in sorted(offsets): - size_class = self._released.pop(offset) - self.alloc_sizes.pop(offset, None) - self.free_lists[size_class].append(offset) + 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. @@ -327,6 +331,8 @@ def allocate( if num_elements == 0: return torch.empty(0, dtype=dtype, device=self.device) + self._process_pending_frees() + elem_size = _element_size(dtype) size_bytes = num_elements * elem_size # Minimum allocation is one granule @@ -607,20 +613,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, - "num_released_blocks": len(self._released) + len(self._pending_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.""" diff --git a/tests/unittests/test_vmem_chunked_allocator.py b/tests/unittests/test_vmem_chunked_allocator.py index eae674a0d..00931864b 100644 --- a/tests/unittests/test_vmem_chunked_allocator.py +++ b/tests/unittests/test_vmem_chunked_allocator.py @@ -143,6 +143,8 @@ 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() @@ -162,6 +164,23 @@ def test_chunked_free_list_size_classes(): assert large2.data_ptr() == large_ptr +def test_chunked_churn_then_large_allocation_does_not_grow(): + """Freed scratch blocks at the top of the heap return to the bump pointer, so a later large allocation fits.""" + ctx = iris.iris(64 << 20, allocator_type=ALLOC_TYPE) + alloc = ctx.heap.allocator + scratch = (1 << 20) // 4 # 1 MiB of float32 + chunks = alloc.get_num_chunks() + + # Churn well past the mapped heap: freed blocks only become reusable at collect() + for _ in range(3 * alloc.mapped_extent // (1 << 20)): + t = ctx.zeros(scratch, dtype=torch.float32) + t.fill_(1.0) + del t + + 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) From cba75a556d7b1dd5bf192cbaf09c9f9575c8bf0c Mon Sep 17 00:00:00 2001 From: Ari Prakash Date: Tue, 13 Oct 2026 13:56:16 -0500 Subject: [PATCH 3/3] Reconcile frees after handing out 1/8 of the mapped heap The bump rollback only reclaims freed blocks at the top of the heap. A churn loop that also keeps an occasional buffer stops the rollback, so the scratch blocks below it stay in their size class, smaller requests of another class bump past the end, and the heap grows: in the new test, from 1 chunk to 3, where before deferred reuse it stayed at 1. Count the bytes handed out since the last reuse() (free-list pops as well as bump allocations) and call collect() from SymmetricHeap.allocate once they reach 1/COLLECT_FRACTION (8) of the mapped heap. Counting only bump bytes would let the scratch working set grow by one threshold per cycle; counting everything handed out caps it. The count follows the allocation sequence, so it is identical on every rank and the decision stays symmetric. The churn test is parametrized to also keep a buffer every 16 iterations, with scratch and kept sizes in different size classes. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../allocators/vmem_chunked_allocator.py | 11 +++++++++++ iris/host/memory/symmetric_heap.py | 7 +++++-- .../unittests/test_vmem_chunked_allocator.py | 19 +++++++++++++------ 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/iris/host/memory/allocators/vmem_chunked_allocator.py b/iris/host/memory/allocators/vmem_chunked_allocator.py index 67c4b68e0..b9b727d36 100644 --- a/iris/host/memory/allocators/vmem_chunked_allocator.py +++ b/iris/host/memory/allocators/vmem_chunked_allocator.py @@ -127,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, @@ -153,6 +155,8 @@ def __init__( 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] = [] @@ -272,6 +276,11 @@ def released_offsets(self) -> List[int]: 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 @@ -279,6 +288,7 @@ def reuse(self, offsets) -> None: 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: @@ -370,6 +380,7 @@ def allocate( # Track for free self.alloc_sizes[offset] = aligned_size + self._bytes_since_collect += aligned_size # Create tensor via CUDAArrayInterface va = self.base_va + offset diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 0b9177854..88ce1225c 100644 --- a/iris/host/memory/symmetric_heap.py +++ b/iris/host/memory/symmetric_heap.py @@ -342,7 +342,9 @@ def allocate(self, num_elements: int, dtype: torch.dtype, alignment: int = 1024) chunks_before = self.allocator.get_num_chunks() if is_chunked else 0 if is_chunked: - # Every rank reaches the same branch: it depends only on state identical across ranks + # 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() @@ -464,7 +466,8 @@ def collect(self): 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; call it directly to reuse memory sooner. + ``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): diff --git a/tests/unittests/test_vmem_chunked_allocator.py b/tests/unittests/test_vmem_chunked_allocator.py index 00931864b..98a103ead 100644 --- a/tests/unittests/test_vmem_chunked_allocator.py +++ b/tests/unittests/test_vmem_chunked_allocator.py @@ -164,18 +164,25 @@ def test_chunked_free_list_size_classes(): assert large2.data_ptr() == large_ptr -def test_chunked_churn_then_large_allocation_does_not_grow(): - """Freed scratch blocks at the top of the heap return to the bump pointer, so a later large allocation fits.""" +@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 = (1 << 20) // 4 # 1 MiB of float32 + scratch_bytes, keep_bytes = 2 * alloc.granularity, alloc.granularity # different size classes chunks = alloc.get_num_chunks() - # Churn well past the mapped heap: freed blocks only become reusable at collect() - for _ in range(3 * alloc.mapped_extent // (1 << 20)): - t = ctx.zeros(scratch, dtype=torch.float32) + 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