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. |
|
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 |
…3_5_moe The loader releases expert rows [K, E) for every NVFP4 family, but the disk index was constructed in one place -- qwen3_5_moe's setup_offload_expert_banks. Every other family (glm5_next, glm4_moe, gemma4, minimax_m2, minimax_m3, deepseek_v4, ...) reached that release with disk_index=None, so the engine's attach was skipped, the "disk tier: K/E pinned" line never printed, and copy_missing fed the grouped GEMM zeroed expert rows. Wrong logits, no error, no diagnostic. Each family already owns the Nvfp4ExpertSourceSpec its loader reads with, so expose it (nvfp4_expert_source_spec, re-exported from each package like the loaders are) and build the index in the shared provider, which is also where the release happens. glm5_next keeps choosing between its compressed-tensors and modelopt namings per checkpoint. Guards, so the silent path cannot come back: - _nvfp4_banks raises when a family exposes no spec (it would release and never refetch). - the engine raises when the tier is on and the provider returned no index, which also covers the non-NVFP4 providers (bf16, ds_fp4, q4_0). - qwen3_5_moe checks "nvfp4 experts only" before its fp8_block branch instead of inside the nvfp4 one, where the default fp8_block path silently dropped --moe-disk-tier. Verified on an RTX 4090: the index now builds for GLM-5.3-Flash-NVFP4 (glm5_next, 288 experts, 42 MoE layers, 118 shards) and its rows are byte-identical to the checkpoint tensors read through safetensors -- packed weights, per-row scales and the fp32 global scales, 9 of 9 spot checks. Ornith-1.5 (qwen3_5_moe) is unchanged end to end. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
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.
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.