diff --git a/docs/reference/host/tensor-creation.md b/docs/reference/host/tensor-creation.md index 7f7e3963b..70dbcd965 100644 --- a/docs/reference/host/tensor-creation.md +++ b/docs/reference/host/tensor-creation.md @@ -21,4 +21,5 @@ APIs on `Iris` that create and initialize tensors on the Iris symmetric heap. ```{eval-rst} .. automethod:: iris.host.iris.Iris.as_symmetric .. automethod:: iris.host.iris.Iris.is_symmetric +.. automethod:: iris.host.iris.Iris.reclaim_freed ``` diff --git a/iris/host/iris.py b/iris/host/iris.py index aa1674334..24dd02f01 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -673,6 +673,31 @@ def is_symmetric(self, tensor: torch.Tensor) -> bool: """ return self.heap.is_symmetric(tensor) + def reclaim_freed(self): + """ + Make symmetric heap memory that every rank has freed reusable. + + Collective: all ranks must call it together. With `allocator_type='vmem_chunked'`, a freed + symmetric tensor is only reused once every rank has freed it, so that allocations land at the + same offset on every rank. Iris reclaims freed memory itself before growing the heap and + periodically as memory is allocated; call this to reuse freed memory sooner. Other + allocators never reuse freed memory, so for them this does nothing. + + Waits for the work queued on this rank's HIP streams before reusing memory. Copy-engine + transfers are not covered: call :meth:`quiet` before dropping their tensors. + + Each call runs ``gc.collect()``, a device-wide ``torch.cuda.synchronize()`` that breaks + CPU/GPU overlap, and a round-trip with every rank. Do not call it every step. + + Example: + >>> ctx = iris.iris(allocator_type='vmem_chunked') + >>> scratch = ctx.zeros(1 << 20, dtype=torch.float32) + >>> del scratch + >>> ctx.reclaim_freed() # on every rank + >>> reused = ctx.zeros(1 << 20, dtype=torch.float32) # same memory as scratch + """ + self.heap.collect() + def full(self, size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False): """ Creates a tensor of size size filled with fill_value. The tensor's dtype is inferred from fill_value. diff --git a/iris/host/memory/allocators/vmem_chunked_allocator.py b/iris/host/memory/allocators/vmem_chunked_allocator.py index e66db02df..f1e33fb02 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) """ @@ -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 only via SymmetricHeap.collect(), which passes the same + offsets on every rank: those every rank has released. That keeps the free lists, and every + later offset, 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() diff --git a/iris/host/memory/symmetric_heap.py b/iris/host/memory/symmetric_heap.py index 5fab77f51..9edee6e9a 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,33 @@ 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() + rank = dist.get_rank() + exchange = next(_released_exchanges) + prefix = f"iris_released_v{exchange}/" + store.set(f"{prefix}{rank}", np.asarray(offsets, dtype=np.int64).tobytes()) + keys = [f"{prefix}{r}" for r in range(num_ranks)] + try: + store.wait(keys) + except dist.DistStoreError as exc: + missing = [r for r, key in enumerate(keys) if not store.check([key])] + raise TimeoutError( + f"SymmetricHeap.collect(): timed out after {store.timeout.total_seconds():.0f}s (the store's timeout) " + f"waiting for ranks {missing}; symmetric allocations and collect() must happen on every rank" + ) from exc + # Every rank has published for this exchange, so all of them finished reading the previous one + if exchange > 1: + store.delete_key(f"iris_released_v{exchange - 1}/{rank}") + 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 +346,19 @@ 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 + collected = self.allocator.collect_due() + if collected: + self.collect() + tensor = self.allocator.allocate(actual_elements, dtype, alignment, grow=False) + if tensor is None: + # A collect() just above found everything there is to reuse; nothing was freed since + if not collected: + 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 +468,31 @@ 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 the work queued on this rank's HIP streams first, so no kernel is still accessing a + block it hands out. Copy-engine (SDMA) transfers are not covered: quiet them with + ``Iris.quiet()`` before dropping their tensors. + + Each call runs ``gc.collect()``, a device-wide ``torch.cuda.synchronize()`` that breaks CPU/GPU + overlap, and a store round-trip with every rank. Do not call it every step. + """ + 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..95b9057fc 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.reclaim_freed() # Allocate again -- should reuse the freed block t2 = ctx.zeros(1024, dtype=torch.float32) @@ -144,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() @@ -151,8 +152,7 @@ def test_chunked_free_list_size_classes(): # Free all del small, medium, large - gc.collect() - torch.cuda.synchronize() + ctx.reclaim_freed() # Re-allocate -- each should reuse from its size class small2 = ctx.zeros(size_small, dtype=torch.float32) @@ -164,6 +164,114 @@ 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.reclaim_freed() + 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.reclaim_freed() + 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 + # d reuses the block freed everywhere; had c been reused, the bump pointer would sit at c + e = ctx.zeros(4096, dtype=torch.float32) + offsets = _offsets_per_rank(ctx, e) + assert len(set(offsets)) == 1, f"offsets differ across ranks: {offsets}" + assert e.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