Fix vmem_chunked asymmetric allocation bug

[ariscript]

Problem

vmem_chunked frees a block when Python garbage-collects the tensor, independently on each rank, and
the next allocation of that size reuses it. A symmetric tensor that is only reachable through a
reference cycle is freed when CPython's cyclic GC next runs, and that moment differs between ranks.
Until it has run everywhere, the rank that collected reuses the freed slot while another rank
bump-allocates. The same ctx.zeros(...) call then returns different offsets on different ranks,
the bump pointers stay out of step from then on, and remote loads, stores, atomics and copies go to
the wrong peer address with no error.

Fix

Garbage collection still decides when a rank is done with a tensor; only reuse becomes collective.

• Frees are local and deferred. The finalizer queues the offset as before; the allocator keeps it
  in a per-rank released set instead of putting it straight on a free list.
• SymmetricHeap.collect() (collective) runs gc.collect() and a device sync on every rank,
  exchanges the released offsets through the rendezvous store (no NCCL collective, like
  _hostnames_by_rank), and moves the offsets every rank released onto the free lists in sorted
  order. Offsets only some ranks released stay pending.
• Allocation reads only state that is identical across ranks (free lists and bump pointer), so
  every rank takes the same branch: free list, then bump, and only when the heap would otherwise grow,
  collect() first and retry. The common path stays collective-free, as before.

The device sync before the exchange also means no rank is still accessing a block when it is handed
out again.

Behavior change

A freed block is reused only after collect(): automatically when the heap would otherwise grow,
or explicitly via ctx.heap.collect() to reuse memory sooner. Previously it was reused by the next
allocation of its size class on that rank. test_chunked_gc_free_reuse and
test_chunked_free_list_size_classes now call collect() before reallocating.

If a rank never reaches the allocation that triggers collect() (an asymmetric allocation sequence),
the exchange fails when the store times out, with an error naming the missing ranks.

Cost

• Allocation and free: unchanged (no collective) unless the heap would grow.
• At a would-be growth: one gc.collect(), one device sync and one small store exchange per rank.
  Growth already runs refresh_peer_access().

Tests

tests/unittests/test_vmem_chunked_allocator.py, new:

• test_chunked_freed_block_waits_for_collect: a freed block is not reused before collect(), and is after.
• test_chunked_reuses_before_growing: an allocation that would grow the heap reuses a freed block instead.
• test_chunked_offsets_symmetric_across_gc_timing (2+ ranks): a reference cycle collected on rank 0
  only no longer changes where the next allocation lands; a block one rank still holds is not handed
  out after collect().

Ran on 8× MI300X (gfx942, ROCm 7.1.0): the chunked allocator tests pass at 1, 2, 4 and 8 ranks, and
test_multihost_guard.py at 2 ranks.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

2 open findings

• np is used here but isn’t imported within the function. If numpy as np isn’t imported at module… · New
• alloc_sizes entries for freed blocks appear to remain until reuse() runs (where… · New

🧠 Review effort: Lite

[nirvedhmeshram]

I am still reviewing this but wanted to share one thing flagged by the agent in the time being, AI comment below

Before this PR, a churn pattern reused the same block every time and bump stayed low. Now every freed block stays out of the free lists until the heap would grow. A churn loop therefore consumes the whole mapped extent as small blocks first. The free lists are exact power-of-two size classes with no coalescing, so all that memory can then only serve that size class. The first larger allocation finds nothing to reuse after collect() and grows the heap.

Model of the PR's exact bookkeeping (256 MB heap and chunk, 1000 iterations of a freed 1 MB scratch buffer, then one 128 MB buffer):

before: bump=129 MB mapped=256 MB grows=0
after : bump=384 MB mapped=512 MB grows=1

So a workload that fit in its heap now maps another chunk, pays a refresh_peer_access(), and doubles its physical footprint. When the VA budget is tight, it can hit VA space exhausted where it used to succeed. This is a realistic pattern: a per-iteration scratch buffer plus one late large allocation.

Suggested fix (deterministic, small): in reuse(), before pushing to the free lists, walk the common offsets from the highest down. While a block ends exactly at bump, lower bump to that block's start instead of pushing it onto a free list. The common set and bump are identical on every rank, so this keeps symmetry. In the same model:

after+rollback: bump=128 MB mapped=256 MB grows=0

def reuse(self, offsets) -> None:
    with self.lock:
        offsets = sorted(offsets)
        for offset in offsets:
            self.alloc_sizes.pop(offset, None)
        # 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)

(This assumes bump is always the end of the last bump allocation, which holds today. An import_external_tensor() moves bump to mapped_extent, past the import, so blocks below an import never match and the rollback stops there, which is correct.)

This doesn't fix fragmentation in the middle of the heap, but that existed before the PR. It does stop the PR from making it systematically worse. Please also add a test: churn, then a large allocation, then assert get_num_chunks() is unchanged.

[ariscript]

Thanks for the comment, the bug has been fixed as written. Response to the raised bug by an agent:

One remaining gap: if the churn loop also makes an occasional long-lived allocation, that block stops
the rollback, and the scratch blocks below it stay in their size class until a larger allocation
needs memory. Before this PR that pattern reused the scratch block immediately. One way to bound it
deterministically: also run collect() once a fixed amount has been bump-allocated since the last
reconcile (say 1/8 of the mapped heap). Every rank sees the same bump pointer, so the decision stays
symmetric, and it costs one gc.collect() plus a small store exchange per that much allocation.

This can also be added but I wanted to check if this is worth doing.

[nirvedhmeshram]

Thanks for adding the rollback. On whether the periodic collect() is worth doing: yes. I modeled it (agent-assisted, bookkeeping only, single rank), and the remaining gap is bigger than the note suggests. The trigger needs one change, though.

Model: 256 MB heap. Each of 20,000 steps allocates and frees a 1 MB scratch buffer, every 100th step also keeps a 512 KB buffer, and a 64 MB allocation comes at the end. collect() includes the rollback from 1965d08 in every row.

	mapped at end	heap growths	collect() calls
before this PR	256 MB	0	0
current head (rollback only)	1792 MB	6	29
+ trigger on bytes bump-allocated since last collect()	1024 MB	3	78
+ trigger on bytes handed out since last collect()	256 MB	0	625

Why rollback alone grows: once the heap is cut into 1 MB blocks, a 512 KB request has nothing of its size to reuse, so it bump-allocates past the end and grows the heap. Each kept buffer also stops the next rollback, so the cycle repeats.

Why counting bump bytes isn't enough: after a collect(), the loop drains the free list without moving bump. Those blocks go back to _released, and only then does bump advance a full threshold before the next collect(). So the scratch working set grows by one threshold per cycle. Counting every byte handed out (free-list pops as well as bumps) caps it at the threshold. That count is identical on every rank, because the allocation sequence is.

Sketch:

# VMemChunkedAllocator.allocate(), once the offset is chosen (either branch)
self._bytes_since_collect += aligned_size

# VMemChunkedAllocator.reuse()
self._bytes_since_collect = 0

# SymmetricHeap.allocate(), before the grow=False attempt
if self.allocator._bytes_since_collect >= self.allocator.mapped_extent // self.allocator.COLLECT_FRACTION:
    self.collect()

On cost: besides gc.collect() and the store exchange, collect() also does torch.cuda.synchronize(), so each trigger waits for all queued GPU work on the device. In the model that's one sync every ~32 steps. The rate scales with allocation volume relative to heap size. I'd use a class constant next to DEFAULT_CHUNK_SIZE (e.g. COLLECT_FRACTION = 8) rather than a new knob. The sync has to stay, since it's what makes reuse safe.

For a test: extend the churn test so it keeps a small buffer every N iterations, then assert that the large allocation still doesn't add a chunk. With rollback only, that fails today.

[ariscript]

Also made this change per your sketch.