feat(moe): NVMe disk tier for MoE expert banks - #337
Conversation
|
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. |
|
Went back over the provider -> engine path on the merged head, and two of the four things I found need fixing before this lands. Neither showed up in my runs, because this box only exercises the two families that happen to work. 1. Only
|
| families | |
|---|---|
| resolves a setup (fetcher attaches) | qwen3_5_moe, qwen4_exp (re-exported at qwen4_exp/__init__.py:25), gpt_oss |
falls through to _nvfp4_banks (no index) |
glm5_next, glm4_moe, gemma4, minimax_m2, minimax_m3, deepseek_v4, glm_moe_dsa, muse_glimmer |
That table also explains why my earlier runs looked clean: Ornith-1.5 is qwen3_5_moe and Qwen3.8-Flash-Next is qwen4_exp through the re-export, so both log the attach line and both verify (16302 slot checks on Qwen3.8 at K=224, 414 on Ornith at K=128, zero mismatches). The GLM-5.3-Flash figures in my first report were load-phase only -- shard count, load time, cgroup shmem -- which is exactly the part that looks right whether or not a fetcher exists.
The cheap fix is to fail where the code already fails loudly (_nvfp4_banks's own NotImplementedError, the eq != "nvfp4" guard):
if banks.disk_index is not None:
cache.attach_disk_tier(...)
logger.info_rank0(...)
elif disk_tier is not None:
raise NotImplementedError(
"--moe-disk-tier on: this model family builds no disk index; experts [K, E) "
"are released at load and would never be refetched")Building the index inside _nvfp4_banks instead would fix the families rather than reject them, since what qwen3_5_moe adds is a per-family source spec -- but the guard is what stops silent corruption today.
2. release_bank_tails asserts on any --expert-ram-experts whose row offset is not page-aligned
release_bank_tails (disk_tier.py:57-58) passes ram_experts * row_bytes and bank.nbytes - offset to release_range, which hard-asserts 4 KiB alignment on both (host_banks.py:178-179). row_bytes = bank.nbytes // num_experts is not page-aligned in general -- the scale banks are small. Real values from the [disk-tier-init] line here: Ornith host_row_bytes=[1048576, 131072, 2048, 524288, 65536, 4096], Qwen3.8-Flash-Next [1638400, 204800, 2560, 819200, 102400, 5120]. So whether a boot survives depends on K:
| E, row bytes, K | release_bank_tails |
|---|---|
| 256, 2048, 128 | ok |
| 256, 2048, 127 | AssertionError: release_range: page-aligned range required |
| 512, 2560, 224 | ok |
| 512, 2560, 225 | AssertionError |
Repro, no GPU and no checkpoint needed:
import torch
from freetoken.moe.host_banks import HostBank
from freetoken.moe.disk_tier import release_bank_tails
banks = {"gate_up_scale": [HostBank((256, 2048), torch.uint8)]} # a real Ornith bank row size
release_bank_tails(banks, 256, 127) # AssertionError; 128 is fineIn other words any odd --expert-ram-experts on Ornith, or anything not a multiple of 8 on Qwen3.8, fails at load -- after the wait for the weights, and with a message that points at page alignment rather than at the flag. Since rows [K, E) were never written, the release is an optimization and not an invariant: warning and skipping (or rounding the offset up to the next page and releasing the remainder) turns this back into a no-op instead of a boot failure.
3. Minor: qwen3_5_moe's default fp8_block path drops the flag silently
setup_offload_expert_banks only enters the disk-tier branch inside if eq != "fp8_block":, and the eq != "nvfp4" guard (qwen3_5_moe/weight.py:891-893) lives inside that branch, so it never fires for the default path. _build_fp8_expert_banks(...) is called without disk_tier. Net effect on a default fp8_block checkpoint: the flag is accepted, nothing is released, nothing is fetched, no diagnostic. Raising the same NotImplementedError there would be consistent.
4. Trivial: the FT_DISK_TIER_VERIFY init probe indexes expert 100
disk_tier.py:264 sums self._index.row_segments(bi, 0, 100) to print disk_row_bytes. The intent is one expert's row, so it wants row_segments(bi, 0, 0); as written it slices past the per-layer buffer and struct.unpack_from raises for any model with E < 100. Gated on the verify flag, and my three checkpoints have 256/288/512 experts, so it never fired here.
Happy to send 1 and 2 as a patch if you want them in this PR rather than in your own words.
Written with AI assistance; every number above was measured on my hardware and I can reproduce it.
- host_banks: bank mmap is now MAP_PRIVATE|MAP_ANONYMOUS (CPython's default MAP_SHARED is backed by an internal shmem object, so release_range's munmap+MAP_FIXED gave back the address range but not the pages -- a touched tail stays charged for the life of the process). release_range is one madvise(MADV_DONTNEED) + an assert that the range does not overlap the cudaHostRegister'd prefix (_pinned_bytes tracking). Patch verbatim from MT-z (issuecomment-5518521155). - offload_cache: gate the [copy-miss] FT_DISK_TIER_VERIFY print on self._disk_tier is not None and not-capturing -- .item()/.cpu() inside a captured CUDA graph crashed every graph-capturing boot with the env var set, tier off (issuecomment-5519070434). - engine: --moe-disk-tier on a family that builds no disk index now raises NotImplementedError instead of silently serving zeroed experts (issuecomment-5520770513 item 1). - disk_tier: release_bank_tails warns and skips banks whose row boundary is not page-aligned (odd --expert-ram-experts used to AssertionError at load); init probe reads expert 0, not expert 100 (items 2 and 4). - qwen3_5_moe: the default fp8_block path raises instead of silently dropping the disk-tier flag (item 3). - tests: regression test for the unaligned row boundary; updated the release-range test docstring (MAP_PRIVATE, not the private remap).
|
All three findings verified against the tree and applied — new head 5518521155 (MAP_SHARED bank, released tails never free): applied your patch to 5519070434 ( 5520770513:
Test results (Rudi): |
|
The cgroup line is the one number the old logs could not show, and it now says the same thing on two kernels with opposite THP settings. One thing worth putting in
And an offer on the family guard. The What it does NOT have is an end-to-end serving run on any newly supported family: GLM's 18.0 GiB of non-expert weights plus one layer's 3.89 GiB expert floor do not fit this 24 GB card, so I can verify where the index points but not that the model then answers correctly. If that gap is why you would rather keep the guard as the contract, that is a good reason and I would not argue with it. Say the word either way. Written with AI assistance; every number above was measured on my hardware and I can reproduce it. |
Per MT-z (PR FlashML-org#337 issuecomment-5527183880): which K values release the tail cleanly is a per-model arithmetic rule decided by the smallest bank row (the fp32 global scales) -- K * min(row_bytes) must be a multiple of 4096, else the small scale bands stay resident (warns, does not abort). Qwen3.8-Flash-Next needs a multiple of 8; Ornith-1.5-35B a multiple of 2.
|
Thanks for the second-hardware confirmation — the cgroup line matching across opposite THP settings ( The On the family guard: we'd like to lift it — send the commit. I don't see a public fork/branch for it; push the branch anywhere I can fetch (or paste the diff) and I'll take it in. I checked the premise against the tree: the NVFP4 families that currently hit the guard are exactly gemma4, glm4_moe, glm5_next, minimax_m2, minimax_m3, and each already owns its The E2E gap is closable here: Rudi has 62 GB RAM + 2× RTX PRO 4000 (24 GB each), so GLM-5.3-Flash-NVFP4 (18 GiB non-expert + the per-layer expert floor) should fit with the tier on — that's the plan for validating the newly enabled families once the commit lands. |
|
Pushed — it sits directly on your current head, so it should be a clean cherry-pick: git fetch https://github.com/MT-z/FreeToken.git feat/moe-disk-tier-all-nvfp4-families
git cherry-pick FETCH_HEAD
It lands where you said it shouldYour reading was right, and it is the reading the commit already had: the index is built in
The seam is one new hook, The guard stays, with a corrected messageThis is the one conflict against your head, and it is worth a look. The current wording is the one
That parenthetical becomes false with this commit, since every NVFP4 family then attaches one.
The guard itself is unchanged in spirit — it is now a last-resort net for a checkpoint whose family Tested hereAnd against real checkpoints, resolving the spec end to end (no GPU needed for this part):
The first two are on your list of newly enabled families, so that is two of the five with a real Shout if you would rather have it shaped differently — it is your PR and I am happy to redo it. Written with AI assistance; the test results and checkpoint resolutions above are from my own |
|
Cherry-picked cleanly onto CPU suite (his tree): Byte-for-byte, re-run here against the real checkpoint (GLM-5.3-Flash-NVFP4, glm5_next, on Rudi): the index resolves through The E2E gap is closed — on gemma4.
Regression on the refactored path (qwen4_exp): Qwen3.8-Flash-Next at K=224 re-run on the new head — 1284/1284 slot verifies, GLM-5.3-Flash E2E: confirmed your fit analysis on Rudi — TP=1 leaves only 646 MB of cache budget after the ~18 GiB non-expert weights (the 288-slot floor needs 4.18 GB), and TP=2 is rejected by The guard message change is the right call — the old "only qwen3_5_moe and re-exporters" parenthetical was already stale. |
The probe's .item()/.cpu() syncs crash CUDA graph capture when FT_DISK_TIER_VERIFY is left set with the tier off (PR FlashML-org#337 issuecomment-5519070434). Verified on Rudi (test_offload.py 25/25); this test was developed against a6bd5c0 but never committed.
…ike text 037f102 narrowed the rule from "the whole prompt in one chunk" to "the image span in one chunk", which is what a 196-token sprite in a 166k-token turn needs. The span is [first image token, last+1) because ``mm_embeds`` is one concatenated tensor scattered in one forward -- so it grows with the TEXT between two screenshots, not just with the pictures. An agent conversation reaches the limit by talking: 400 prompt with images needs 10392 contiguous tokens in one prefill chunk (the image tokens span [160334, 170726) and cannot be split) Nothing configurable moves that. Cheaper images (~490 tokens each after the clamp) only buy more turns before the gap between the first and last one exceeds a chunk, and raising --max-prefill-length OOMs long before it helps: a 32k chunk's activations do not fit beside a 5 GiB KV pool on a 24 GiB card (measured -- it took the worker down twice today). So the concatenated tensor stops being scattered whole. ``_merge_multimodal`` takes the rows belonging to the placeholders inside ITS OWN forward -- the ones an earlier chunk or a prefix-cache hit already consumed sit in front of the window -- and the adder chunks an image prompt exactly like a text one. ``Req.mm_scatter`` and the whole pull-back / reject path go away with it, ~90 lines. Both families that carry a tower here are converted; the approach is gdevenyi's, from FlashML-org#386 (28fd56d). The span cap 09ea814 put in ``match_req`` goes too. It existed because a hit landing inside a placeholder run left half the run cached and half to forward, which the all-in-one-forward scatter could not represent; the window skips the cached half instead. Without the cap a prompt that ends with its image keeps its prefix -- 20,800 of 20,840 tokens on the repeat here, 6.0 s -> 1.2 s, and a different image at the same position still misses (answered "Green" where the cached one answers "Blue"). Measured on Ornith-1.5-35B-A3B-NVFP4, one 4090, --max-prefill-length left at its 8192 default: 2 images with 9k of text between span ~19k 10,186 tokens, 3.2 s (was a 400) 6 images with 9k between each span ~50k 55,360 tokens, 18.1 s (was a 400) A(blue) 9k B(green), and reversed "Blue, Green" / "green blue" -- read across the boundary, in order tests/tokenizer 58, tests/scheduler 90, tests/kvcache/radix 142: all passed. Twelve tests pinning the removed rule are gone and three cover the window (a span wider than a chunk now admits; a chunk scatters only its own rows; a chunk holding no placeholder scatters nothing). The ``_NoSwa`` stub gained the ``page_size`` the reservation math has been reading, which is what had six of these failing on this branch already. A cold system-test run is character-identical to the same branch without this commit, all seven cases. Assisted-by: Claude Opus 5 Re-verified on this branch (no FlashML-org#337/FlashML-org#354/FlashML-org#287 under it): tests/tokenizer 58, tests/scheduler 88, tests/kvcache/radix 142 all passed; a cold system-test run is character-identical to the same change on the daily branch, all seven cases; the two shapes that used to 400 (spans of ~19k and ~50k tokens) answer at the 8192 default.
Re-land of PR FlashML-org#337 onto the post-FlashML-org#418 quantization refactor: the disk tier serves experts that do not fit in pinned RAM -- the RAM bank holds the first --expert-ram-experts per layer (pinned), the rest stay in the original safetensors checkpoint and are fetched O_DIRECT -> pinned staging -> H2D into the slot the LRU kernel already assigned, shrinking the miss list so the existing PCIe copy_missing path only moves the RAM-resident misses. Adapted to the new architecture: * Bank layout is owned by the expert kernel (BankSpec per role); the tier speaks the native NVFP4 (triton) layout, which is unchanged by the refactor. The index and fetch path are written against it. * The per-family nvfp4_expert_source_spec hook from round 4 is subsumed by upstream's own nvfp4_expert_spec hook (added in the refactor); the index resolves the spec through the same hook the reader uses, so index and loader read the same rows. Built in _method_expert_banks next to the release, so a family that releases rows without an index fails loudly instead of serving zeroed experts. * Load-time RAM peak is the K/E prefix, not the full expert set: the NVFP4 reader never reads rows [K, E) (serial: no get_tensor; parallel: filtered at the reader) and yields an empty piece per skipped expert so build_expert_banks completes the layer, pins only the prefix (PinPipeline(prefix_rows=K)) and releases the tail (release_bank_tails). * Host banks are MAP_PRIVATE|MAP_ANONYMOUS (MT-z round 3): release_range is one madvise(MADV_DONTNEED) that really frees; the mincore startup check (check_tail_unbacked) warns if anything backed the released tail. * FT_DISK_TIER_VERIFY [copy-miss] probe gated on the tier + not capturing (round 3 finding); --moe-disk-tier preconditions collected into one error. CPU suite: test_disk_tier.py 8/8, test_disk_tier_families.py 15/15, test_offload.py green (incl. the new probe-gate regression test).
…ot the tail The serial/parallel readers read experts [0, K), so the completeness check must expect L*K*9 tensors, not L*(E-K)*9. Symmetric K (K == E-K) masks it; Qwen3.8-Flash-Next (E=512, K=224) hit it at boot.
c3abff6 to
badf5f4
Compare
|
@MT-z — I rebased the branch onto the latest
Re-validated on 2× RTX PRO 4000 Blackwell, 62 GB RAM:
The rebase also caught a real bug in the reader's completeness check (it expected |
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, so models whose expert weights exceed RAM become runnable.
--expert-ram-experts N(per layer): first N experts resident in RAM (pinned), the rest on disk in the original checkpoint.main(2026-09-10) on top of the quantization refactor (refactor(quant): config, scheme and method layers for quantization #418/fix(qwen4_exp): support modelopt MIXED_PRECISION checkpoints #426/refactor(quant): hand the checkpoint QuantConfig to the weight readers #427): the index now resolves each family'snvfp4_expert_spechook — the same spec the reader uses — so index and loader read the same rows.How
moe/disk_tier.py(new):Nvfp4DiskIndex— a row-indexed view over the checkpoint's safetensors expert tensors (no copy of the weights), plusDiskTier, a pool of O_DIRECT fetch workers that copy disk-resident slot-cache misses into a small pinned staging ring and H2D them into the slot the LRU kernel already assigned, shrinking the miss list so the existing PCIecopy_missingpath only moves the RAM-resident misses.moe/host_banks.py: banks areMAP_PRIVATE|MAP_ANONYMOUS(a shared anonymous mmap would never return released pages);pin_prefix(K)pins only the RAM-resident rows;release_rangeis onemadvise(MADV_DONTNEED)that really frees; amincorestartup check (check_tail_unbacked) warns if anything backed the released tail.moe/expert_banks.py/expert_pieces.py/models/nvfp4_banks.py: load-time RAM peak is the K/E prefix, not the full expert set — the NVFP4 reader never reads rows[K, E)(serial: noget_tensor; parallel: filtered at the reader) and yields an empty piece per skipped expert sobuild_expert_bankscompletes the layer, pins only the prefix, and releases the tail next to where the index is built.moe/offload_cache.py/layers/moe.py: prefill materializes the RAM prefix into identity slots and fetches the routed disk experts; decode fetches disk misses before the PCIe copy.FT_DISK_TIER_VERIFY=1gives slot-by-slot byte verification against the checkpoint (all verify probes are gated on the tier and skip CUDA-graph capture).engine/+server/args.py: config plumbing, all v0 preconditions collected into one boot error (--moe-strategy offload,--disable-moe-prefill-overlap,--cuda-graph-max-bs 0,0 < N < num_experts).tests/moe/test_disk_tier.py,test_disk_tier_families.py(new) +test_offload.py: CPU tests for the index, fetch path, release/tail invariants, the per-family spec hook, and the verify-probe gate.Validation (2026-09-10, rebased head)
Rudi: 2× RTX PRO 4000 Blackwell, 62 GB RAM,
shmem_enabled=always,FT_DISK_TIER_VERIFY=1:0 MiB of 36551 MiBCPU suite:
test_disk_tier.py8/8,test_disk_tier_families.py15/15,test_offload.pygreen.Notes
decode_target == "gpu", no prefill overlap, no CUDA graphs, original (non-FTW) checkpoints.N × (smallest bank row bytes)a multiple of 4096 (noted in--help) or the small scale banks' tails stay resident (warns, does not abort).