feat(moe): NVMe disk tier for MoE expert banks - #337
Conversation
Lets the offload backend serve models whose experts don't fit in pinned RAM: with --moe-disk-tier on --expert-ram-experts K, only the first K experts per layer are pinned; the rest keep allocated bank rows whose pages are released (MADV_DONTNEED) after load and are fetched from the ORIGINAL safetensors checkpoint on slot-cache miss. * moe/disk_tier.py: Nvfp4DiskIndex (index json + shard headers -> per (bank, layer, expert) byte ranges) and DiskTier (O_DIRECT preadv into a small pinned staging buffer, H2D into the LRU-assigned slot, miss-list rewrite so the existing PCIe copy path only moves RAM-resident misses). * host_banks: HostBank.pin_prefix / release_range; PinPipeline(prefix_rows). * nvfp4 loaders: disk_tier param -> partial pin + tail release (serial and parallel paths). * offload_cache: attach_disk_tier + copy_missing hook; materialize_layer takes the routed ids when disk-tiered. * offload_kernels: materialize kernel gains materialize_count (disk-tier prefill streams only the RAM prefix; routed disk-resident experts are fetched into their identity slots, so the prefill GEMM is unchanged). * engine: --moe-disk-tier / --expert-ram-experts / --disk-fetch-workers, guards (native NVFP4, gpu decode, no prefill overlap, no cuda graphs). v0 scope: native NVFP4 (triton) layout, synchronous fetch, no FTW, no converter path. CPU unit tests: tests/moe/test_disk_tier.py.
…to-enables graphs)
…zero ignores DONTNEED); index offsets need 64-bit
…g the row's leading dim -> 8KB buffer, EFAULT/segfault on 512KB segments)
…+fill, not byte copy; test matches real checkpoint layout
…pool-thread stream must land before GEMM reads slots)
… one-shot slot verify)
…ll-RAM-mismatch hunt
…rialize The GPU slot cache is one shared pool across all layers. Prefill identity mapping owns all of slots [0, E) per layer, but the materialize kernel only scans slots < materialize_count (the RAM prefix), so disk slots [ram, E) that held a previous prefill layer's experts kept their slot_for_id entries. Decode then took phantom hits on those slots and read another layer's weights (expert 112 verified correct at prefill, corrupted by decode step 1). Clear the stale entries device-side before the kernel.
The RAM-slot identity check (slot e == host row e) only holds under the prefill identity mapping; during decode the LRU owns slots 0..63. It was firing on every decode step for layer 0 (384 experts x banks of D2H + CPU compare per step), dragging the instrumented E2E from ~11 tok/s to 1.9. Track whether the pending miss list came from materialize_layer (prefill) or ensure_experts (decode) and gate on it.
…iskTier init TP=2 disktier debug (step 2): prove empirically whether the host bank rows are full-per-rank or TP-sharded, and that the disk index's full-row segments match them byte-for-byte. Gated on FT_DISK_TIER_VERIFY.
… ring Each worker thread preadved the next bank's bytes into the SAME pinned staging buffer while the previous bank's async H2D copy was still DMAing from it. The copy is cudaMemcpyAsync: copy_ returns after ENQUEUE, and the GPU reads the host bytes later, so the next preadv could land mid-DMA and the slot row came out as a partial mix of two experts' data. TP=1 rarely hit it (single rank's disk load, idle-ish stream); TP=2 doubles the NVMe load and adds NCCL stream work, so the window opened and the run degenerated (2/5412 verified rows wrong, both gate|up weight_scale). Fix: a per-thread ring of _STAGING_RING pinned buffers, each armed with a CUDA event recorded after the copy that used it; reusing a buffer waits on that event. Exact, and the wait is a no-op whenever the ring outruns the DMA (the normal case), so disk/GPU overlap is preserved. Also: [verify] lines now carry phase/layer/slot and, on mismatch, the diff span plus an overwriter hunt (which checkpoint row the slot actually holds); _preadv_error call now passes its full signature. Test stub updated to the ring.
Fresh ThreadPoolExecutor threads default to CUDA device 0, but a TP>1 rank lives on another device. The H2D slot copies land on the destination tensor's device stream (copy_ guards on dst), while ev.record() uses the thread's CURRENT stream -- so on TP=2 rank 1 the staging-ring reuse guard waited on an idle device-0 stream and a preadv could overwrite a pinned buffer mid-DMA. Seen on Bandit (4.3 GB/s NVMe, TP=2, RAM=64): 6/12288 layer-0 slot rows corrupted in one prefill (e4m3 scale banks mixed into e4m3 NaN encodings; weight banks partial mixes; scalar-fill banks clean). Rudi's 1.8 GB/s disk made the preadv slower than the DMA, hiding the window. - _staging_ring: torch.cuda.set_device(rank device) once per worker thread, so ev.record() lands on the stream the copies use. - _sync_fetches: sync the rank device's default stream explicitly. - _identify_overwriter: num_experts came from an expert ROW's shape[0] (1024/128/32/2048 per bank) -- the scan ran past the index and died with struct.error, masking overwriter attribution.
|
Tested this PR head (623ca1d) rebased onto 1. Load-time RAM peak is the full expert set, regardless of
|
- never materialize disk-resident expert rows at load: serial loader skips rows >= K before get_tensor, parallel loader filters at the reader, and the completion tracker / placed assert count rows_per_layer instead of E. Load-time RAM peak is now K/E of the expert set instead of the full set. - thread disk_tier= through the remaining six NVFP4 family loaders (gemma4, glm4_moe, glm5_next, minimax_m2, minimax_m3, qwen4_exp) so --moe-disk-tier on no longer TypeErrors outside qwen3_5_moe. - key _sync_fetches on the banks' device type instead of torch.cuda.is_available(): CPU-bank unit tests no longer fall through to torch.cuda.default_stream(cpu_device) on CUDA machines. Diffs from MT-z's PR FlashML-org#337 review comment, applied verbatim. Verified: tests/moe/test_disk_tier.py 6/6 pass in a CUDA container (freetoken:local, Rudi GPU0), incl. the two previously failing fetch_pending tests; inspect.signature confirms disk_tier on all 13 wrappers.
- gate the layer<3 prefill debug print behind FT_DISK_TIER_DEBUG instead of printing unconditionally. - drop the per-layer device->host sync: cache.usage[disk] now takes the 0-d cache.step tensor directly (same dtype/device) instead of .item(). - validate all --moe-disk-tier v0 preconditions at once and raise a single ValueError listing every unmet flag (each used to cost a full boot to discover); list the exact flags in --moe-disk-tier --help. E2E on Rudi (freetoken:standalone-mtz = current tree, Qwen3.6-35B-A3B-NVFP4, TP=1, GPU0): RAM=64 4296/4296 verify match, RAM=32 4902/4902 verify match, 0 mismatches; decode 11.69 tok/s median at RAM=32 (11.85 at RAM=64 hist).
|
Thanks for the thorough review — all three fixes are applied verbatim as
Re-validated end to end on a 2× RTX PRO 4000 box (TP=1, One question on Fix 1: the unbacked rows Happy to un-draft once you've had a look at the two new commits. |
|
Looked at both commits, and re-ran the head on my box. Review. Re-validation on this box (RTX 4090 24 GB, 61 GB RAM),
On the lazy
Things that would break it: Evidence from real loads: cgroup Suggestion: assert it cheaply at startup with No objection to un-drafting from my side. Written with AI assistance; every number above was measured on my hardware and I can reproduce it. |
…backed Implements the follow-up MT-z proposed in the PR FlashML-org#337 review: the lazy-tail invariant (nothing reads/writes rows [K, E) after release_bank_tails, so they cost no RAM) is now checked rather than assumed. - tail_resident_bytes(): mincore(2) over one bank's tail byte range. - check_tail_unbacked(): runs in both NVFP4 loaders right after release_bank_tails; logs 'tail check rank=r/n: resident X MiB of Y MiB' and warns above the THP bound (one 2 MiB huge page per bank layer -- shmem_enabled=always|force can back the prefix/tail boundary as a huge page; more than that means something touched the tail). - test_tail_unbacked_after_release: small bank, fill prefix, release tail, assert zero resident pages in [K, E); one tail write backs exactly one page. Verified on Rudi (shmem_enabled=always -- the config MT-z flagged as the risky one): real boot at RAM=32 logs 'resident 0 MiB of 15172 MiB', no warning; E2E 4902/4902 slot verifies match, 11.59 tok/s median.
|
Great write-up — the phase-by-phase breakdown is exactly the right way to state the invariant, and I verified the load-side mechanics against the tree (in-place One data point from our side that makes your suggestion timely: our test box (2× RTX PRO 4000) runs
Result on the Un-drafting now — thanks again for the review, it caught the one bug that would have made the tier useless for its target case. |
|
Writing that phase-by-phase answer sent me back into the release path, because I wanted to be sure of the second half of what I had claimed: that once
RSS falls in every row, which is what makes this easy to miss. The pages are still charged: The one-line version of the fix is to stop sharing.
To be explicit about what this is not: it frees nothing today. With the loaders as they are, rows Checks before proposing it: Note for your new tail check: with a private bank the tail shows as Validation on this box (RTX 4090 24 GB, 61 GB RAM), on
Ornith-1.5-35B-A3B-NVFP4, Diff below. Happy to open it as a PR against your branch instead if you prefer -- or to leave it until after this one merges, since the mapping flag is main's code and only the release path is yours. Patchdiff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py
index 388893c..c863f95 100644
--- a/python/freetoken/moe/host_banks.py
+++ b/python/freetoken/moe/host_banks.py
@@ -79,7 +79,7 @@ class HostBank:
The buffer is rounded up to the O_DIRECT block; ``tensor`` views exactly ``nbytes``. ``backing=None`` follows ``FREETOKEN_BANK_CUDA_ALLOC``."""
- __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned", "_locked")
+ __slots__ = ("tensor", "addr", "nbytes", "_buf", "_pinned", "_pinned_bytes", "_locked")
def __init__(self, shape: tuple[int, ...], dtype: torch.dtype,
*, backing: str | None = None):
@@ -104,11 +104,19 @@ class HostBank:
self.addr = raw.data_ptr() + off
assert self.addr % _BLK == 0
self._pinned = True # born pinned+mapped; pin() is a no-op
+ self._pinned_bytes = asize
else:
- self._buf = mmap.mmap(-1, asize) # lazy: address space only, no resident pages yet
+ # MAP_PRIVATE, not CPython's default MAP_SHARED: on a shared anonymous mapping a
+ # *read* fault allocates a page (no zero-page sharing) and MADV_DONTNEED is ignored,
+ # so an untouched region is only free by convention and a freed one never comes back.
+ # Private anonymous gives both for real: reads map the shared zero page, and
+ # release_range() actually returns memory. Nothing needs the mapping to be shared --
+ # the loaders are thread pools and ranks are mp-spawned, each with its own banks.
+ self._buf = mmap.mmap(-1, asize, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
_LIVE_BUFFERS.append(self._buf)
self.addr = ctypes.addressof(ctypes.c_char.from_buffer(self._buf))
self._pinned = False
+ self._pinned_bytes = 0
self.tensor = torch.frombuffer(self._buf, dtype=dtype, count=self.nbytes // elsize).view(*shape)
self._locked = False
@@ -140,6 +148,7 @@ class HostBank:
f"cudaHostRegister failed for {len(self._buf) / 2**30:.1f} GiB"
) from exc
self._pinned = True
+ self._pinned_bytes = len(self._buf)
def pin_prefix(self, nrows: int) -> None:
"""Pin only the first ``nrows`` rows (disk tier: the rest stays disk-resident).
@@ -160,44 +169,32 @@ class HostBank:
f"cudaHostRegister failed for {nbytes / 2**30:.1f} GiB prefix"
) from exc
self._pinned = True
+ self._pinned_bytes = nbytes
def release_range(self, offset: int, nbytes: int) -> None:
- """Free a byte range of the backing mapping by replacing it IN PLACE with a
- fresh MAP_PRIVATE anonymous mapping at the same virtual address.
-
- HostBank's buffer is a MAP_SHARED /dev/zero mapping (CPython's
- ``mmap(-1)``), and the kernel silently ignores MADV_DONTNEED on shared
- mappings -- the pages would stay resident. Replacing the range with a
- private zero mapping frees them while keeping every existing pointer
- and torch view valid (same address). The range must be page-aligned
- and must not overlap a pinned prefix (the disk tier's unpinned tails).
- """
- import ctypes
+ """Free a byte range of the backing mapping with MADV_DONTNEED.
+
+ The bank is a MAP_PRIVATE anonymous mapping, so dropping a range frees the
+ pages outright and a later read faults the shared zero page again; every
+ existing pointer and torch view stays valid (the mapping is never replaced).
- _BLK = 4096
+ The range must be page-aligned and must not overlap the pinned prefix:
+ dropping pages under a cudaHostRegister'd range corrupts silently, so it is
+ asserted here rather than left to the caller (the disk tier's unpinned tails).
+ """
assert offset % _BLK == 0 and nbytes % _BLK == 0, (
"release_range: page-aligned range required")
- libc = ctypes.CDLL("libc.so.6", use_errno=True)
- libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
- libc.munmap.restype = ctypes.c_int
- libc.mmap.restype = ctypes.c_void_p
- libc.mmap.argtypes = [
- ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long]
- addr = self.addr + offset
- if libc.munmap(addr, nbytes) != 0:
- raise OSError(ctypes.get_errno(), "munmap failed")
- PROT_READ_WRITE = 3
- MAP_PRIVATE_ANON = 0x22 # MAP_PRIVATE | MAP_ANONYMOUS
- MAP_FIXED = 0x10
- MAP_FAILED = (1 << 64) - 1
- new_addr = libc.mmap(addr, nbytes, PROT_READ_WRITE, MAP_PRIVATE_ANON | MAP_FIXED, -1, 0)
- if new_addr in (None, MAP_FAILED):
- raise OSError(ctypes.get_errno(), "mmap(MAP_FIXED) failed")
- assert new_addr == addr, "MAP_FIXED returned a different address"
+ assert offset >= self._pinned_bytes, (
+ f"release_range: [{offset}, {offset + nbytes}) overlaps the pinned prefix "
+ f"[0, {self._pinned_bytes})")
+ if nbytes:
+ self._buf.madvise(mmap.MADV_DONTNEED, offset, nbytes)
+
def release(self) -> None:
"""Drop the resident pages; the address space stays valid, the contents become undefined.
- For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped."""
+ For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped.
+ (This frees memory only because the mapping is MAP_PRIVATE; the kernel ignores MADV_DONTNEED on shared ones.)"""
if self._pinned:
return
self._buf.madvise(mmap.MADV_DONTNEED)Written with AI assistance; every number above was measured on my hardware and I can reproduce it. |
|
Found one more while stress-testing the tier on a deliberately small card: Repro on this box with no tier flags at all (RTX 4090, 22.67 GiB free): FT_DISK_TIER_VERIFY=1 ft serve --model-path ornith-ai/Ornith-1.5-35B-A3B-NVFP4 \
--moe-cache-auto --memory-ratio 0.9 --kv-reserve-tokens 8192The frame is if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0:
print(f"[copy-miss] layer={layer_id} fused={self._copy_fused_ok} "
f"n={int(self.num_indices.item())} "
f"evict={self.evict_slots[:4].cpu().tolist()} "
f"src={self.src_indices[:4].cpu().tolist()}", flush=True)
Two things make it easy to hit:
Not memory related: it fails identically at 22.67 GiB free and at 10.29 GiB free, and lowering Suggested fix -- gate it like its neighbours and skip it while capturing: if (self._disk_tier is not None and layer_id == 0
and os.environ.get("FT_DISK_TIER_VERIFY")
and not torch.cuda.is_current_stream_capturing()):or simply drop the print: While I was there, the rest of the low-VRAM picture came out clean: with VRAM ballasted down to 10.29 GiB free, Ornith-1.5-35B-A3B-NVFP4 boots through the tier at Written with AI assistance; every number above was measured on my hardware and I can reproduce it. |
2 participants
What
Adds an optional NVMe disk tier for MoE expert banks (NVFP4 checkpoints). When the expert banks don't fit in host RAM, the tail of the bank stays on NVMe and is fetched on demand during decode, so models whose expert weights exceed RAM become runnable.
--expert-ram-experts N(per layer): first N experts resident in RAM, the rest on disk.How
moe/disk_tier.py(new):Nvfp4DiskIndex— a row-indexed view over the checkpoint's safetensors expert tensors (no copy of the weights; mmaps the shard files), plus a pool of background fetch workers that copy pending expert rows into the existing host-bank slots ahead of the CPU executor's need.moe/host_banks.py/expert_banks.py: banks can now be partially populated; unfilled rows are marked pending and served by the fetch workers.moe/offload_cache.py: pin pipeline extended — a prefix-pinned (disk-tier) layer is fetched before the GPU offload copy, integrated with the existing settle/plan flow.engine/+server/args.py: config plumbing, validation (disk tier requiresdecode_target == "gpu"), and the CLI flags.models/nvfp4_banks.py/qwen3_5_moe/weight.py/weight.py: bank loading passes the disk tier through so the index is built from the same source spec as the RAM banks.tests/moe/test_disk_tier.py(new): CPU tests for the index, fetch workers, and pending-row bookkeeping.Validation
E2E on Qwen3.6-35B-A3B-NVFP4 (128 experts/layer, ~2.7 MB/expert) on 2× RTX PRO 4000, 62 GB RAM:
Correctness: a
FT_VERIFY=1toggle recomputes every fetched expert on CPU and compares against the disk-served result — 8604/8604 (TP=2) and 4296/4296 (TP=1) matches, zero mismatches across 300-token decodes.The penalty is disk-bandwidth-bound, as expected: ~10% on NVMe, ~44% on a 1.8 GB/s virtual disk.
Notes
main; no TP changes in this PR (TP support is separate, see feat(models): support TP for qwen3_5_moe #104). The two touch disjoint regions of the same two files (layers/moe.py,qwen3_5_moe/weight.py) and were validated together.load_expert_bankskeeps both thelayer_residencyanddisk_tierparams, and the engine validates them as mutually exclusive.