diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index bcbe6bcf2..b5588a81d 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -41,6 +41,13 @@ class EngineConfig: # (cudaMemcpyBatchAsync); no-op unless moe_cache_size > 2 * num_experts. moe_prefill_hit_d2d: bool = False moe_collect_stats: bool = False # capture decode miss-rate counters into the cuda graph + # Disk tier (--moe-disk-tier, see moe/disk_tier.py): "off" = classic behavior. + # When "on", experts [0, expert_ram_experts) per layer stay pinned in RAM and the + # rest are fetched from the original checkpoint on slot-cache miss. Requires the + # native NVFP4 layout, gpu decode target, no prefill overlap, no cuda graphs. + moe_disk_tier: str = "off" + expert_ram_experts: int = 0 + disk_fetch_workers: int = 8 # CPU MoE backend (--moe-backend cpu): number of CPU worker threads computing # the decode experts. 0 = auto (physical cores). Ignored by other backends. moe_cpu_threads: int = 0 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b5a6fa3b0..e103c6a3b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -565,6 +565,29 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: "(locked layers prefill via synchronous pageable copies)" ) object.__setattr__(config, "moe_prefill_overlap", False) + disk_tier = None + if config.moe_disk_tier == "on": + from freetoken.moe.disk_tier import DiskTierSpec + + E = config.model_config.num_experts + # Collect ALL unmet preconditions and raise once: each used to surface as a + # separate boot-time ValueError, costing a full boot per missing flag. + problems = [] + if not 0 < config.expert_ram_experts < E: + problems.append( + f"--expert-ram-experts must be in (0, {E}) with --moe-disk-tier on") + if decode_target != "gpu": + problems.append( + "--moe-disk-tier v0 requires the gpu decode path (--moe-backend offload)") + if config.moe_prefill_overlap: + problems.append("--moe-disk-tier v0 requires --disable-moe-prefill-overlap") + if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs >= 1: + problems.append( + "--moe-disk-tier v0 requires --cuda-graph-max-bs 0 (cuda graphs disabled)") + if problems: + raise ValueError( + "--moe-disk-tier on: unmet preconditions:\n - " + "\n - ".join(problems)) + disk_tier = DiskTierSpec(ram_experts=config.expert_ram_experts) if cache_factory is None: # Fast path: an FTW checkpoint loads its repacked banks directly. # Slow path: load_expert_banks auto-picks parallel vs serial baseline by @@ -590,6 +613,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: parallel=expert_parallel, decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"), layer_residency=requested_residency, + disk_tier=disk_tier, ) if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks) @@ -626,6 +650,21 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency) + if banks.disk_index is not None: + cache.attach_disk_tier( + banks.disk_index, banks.disk_ram_experts, + workers=config.disk_fetch_workers) + logger.info_rank0( + f"disk tier: {banks.disk_ram_experts}/{config.model_config.num_experts} " + f"experts/layer pinned in RAM; the rest fetched from " + f"{config.model_path} on slot-cache miss") + elif disk_tier is not None: + # The loader released experts [K, E) but no fetcher came back: serving would + # multiply by zeroed rows and log nothing. Fail where the flag was set. + raise NotImplementedError( + "--moe-disk-tier on: this checkpoint's expert provider returned no disk " + "index, so experts released at load would never be refetched " + f"(quant_format={banks.quant_format!r})") cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) else: cache = cache_factory(config, self.device) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded5..dc6dc37b5 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -311,6 +311,9 @@ def _decode_routed( return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) cache.ensure_experts(self.layer_id, topk_ids) cache.copy_missing() + if (cache.disk_tier_enabled and self.layer_id == 0 + and os.environ.get("FT_DISK_TIER_VERIFY")): + cache._disk_tier.verify_decode_mapping(cache, self.layer_id, topk_ids) return self._expert_gemm( cache, hidden_states, @@ -397,7 +400,12 @@ def _prefill_routed( ) cache.release_prefill_layer(self.layer_id) return out - cache.materialize_layer(self.layer_id) + if cache.disk_tier_enabled: + # Disk tier: stream only the RAM-resident prefix + fetch the routed + # disk-resident experts (identity slots, so topk_ids pass through). + cache.materialize_layer(self.layer_id, topk_ids) + else: + cache.materialize_layer(self.layer_id) cache.copy_missing() return self._expert_gemm( cache, diff --git a/python/freetoken/models/gemma4/__init__.py b/python/freetoken/models/gemma4/__init__.py index 37a8dcb39..5ee8b87a5 100644 --- a/python/freetoken/models/gemma4/__init__.py +++ b/python/freetoken/models/gemma4/__init__.py @@ -12,11 +12,12 @@ from .weight import ( iter_weights, iter_weights_parallel, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Gemma4Attention", "Gemma4ForCausalLM", "Gemma4MultimodalEmbedder", diff --git a/python/freetoken/models/gemma4/weight.py b/python/freetoken/models/gemma4/weight.py index 43c2355d6..af4b8f705 100644 --- a/python/freetoken/models/gemma4/weight.py +++ b/python/freetoken/models/gemma4/weight.py @@ -275,8 +275,17 @@ def _expert_name(raw_name: str) -> str | None: yield _expert_name(raw_name), tensor + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -286,11 +295,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -304,6 +314,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/glm4_moe/__init__.py b/python/freetoken/models/glm4_moe/__init__.py index a07134c34..14db4d12e 100644 --- a/python/freetoken/models/glm4_moe/__init__.py +++ b/python/freetoken/models/glm4_moe/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import Glm4MoeForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Glm4MoeForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py index 7efc6e6be..3f983d0e3 100644 --- a/python/freetoken/models/glm4_moe/weight.py +++ b/python/freetoken/models/glm4_moe/weight.py @@ -191,7 +191,16 @@ def _iter_resident_weights(reader, config, primary) -> Iterator[tuple[str, torch # -------------------------------------------------------------------------------------- # Routed expert host banks (NVFP4) for the offload cache. # -------------------------------------------------------------------------------------- -def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> dict[str, torch.Tensor]: + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict[str, torch.Tensor]: """Build the pinned CPU NVFP4 banks for GLM-4's routed experts. experts exist only for layers [first_k_dense_replace, num_layers) and pack by MoE layer @@ -205,11 +214,12 @@ def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> di drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -223,6 +233,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/glm5_next/__init__.py b/python/freetoken/models/glm5_next/__init__.py index c1c05140a..927b7a9e8 100644 --- a/python/freetoken/models/glm5_next/__init__.py +++ b/python/freetoken/models/glm5_next/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import Glm5NextForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Glm5NextForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/glm5_next/weight.py b/python/freetoken/models/glm5_next/weight.py index fda852256..8a349a351 100644 --- a/python/freetoken/models/glm5_next/weight.py +++ b/python/freetoken/models/glm5_next/weight.py @@ -95,7 +95,19 @@ def _select_expert_source_spec(model_path: str) -> Nvfp4ExpertSourceSpec: _KDA_IN_PROJ = ("q_proj", "k_proj", "v_proj", "b_proj", "f_a_proj", "g_a_proj") -def load_nvfp4_expert_sources(model_path: str, config, layer_sink=None): + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _select_expert_source_spec(model_path) + +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None): + # disk_tier is threaded through the same way qwen3_5_moe does it: the caller builds the + # NVMe tier and every NVFP4 family's loader has to pass it down, or the tail of the bank + # is never registered and load fails with an unexpected-kwarg TypeError. return load_nvfp4_expert_source_banks( model_path, config, @@ -103,6 +115,7 @@ def load_nvfp4_expert_sources(model_path: str, config, layer_sink=None): drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/minimax_m2/__init__.py b/python/freetoken/models/minimax_m2/__init__.py index e4e1211ce..2f39f6bcb 100644 --- a/python/freetoken/models/minimax_m2/__init__.py +++ b/python/freetoken/models/minimax_m2/__init__.py @@ -1,8 +1,9 @@ from .config import parse_config from .model import MiniMaxM2ForCausalLM from .weight import iter_weights, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel +from .weight import nvfp4_expert_source_spec -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "MiniMaxM2ForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/minimax_m2/weight.py b/python/freetoken/models/minimax_m2/weight.py index 103f8ec56..aedba6bf4 100644 --- a/python/freetoken/models/minimax_m2/weight.py +++ b/python/freetoken/models/minimax_m2/weight.py @@ -88,11 +88,20 @@ def raw() -> Iterator[tuple[str, torch.Tensor]]: yield from iter_merged_tensors(raw(), _MERGE_RULES, model_name="minimax_m2") + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( model_path: str, config, *, - layer_sink=None, + layer_sink=None, disk_tier=None, ) -> dict[str, torch.Tensor]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -102,11 +111,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -120,6 +130,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/minimax_m3/__init__.py b/python/freetoken/models/minimax_m3/__init__.py index 71c0bb51f..926aeca96 100644 --- a/python/freetoken/models/minimax_m3/__init__.py +++ b/python/freetoken/models/minimax_m3/__init__.py @@ -2,11 +2,12 @@ from .model import MiniMaxM3ForCausalLM from .weight import ( iter_weights, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "MiniMaxM3ForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/minimax_m3/weight.py b/python/freetoken/models/minimax_m3/weight.py index e10a7237e..1b96c3b68 100644 --- a/python/freetoken/models/minimax_m3/weight.py +++ b/python/freetoken/models/minimax_m3/weight.py @@ -248,8 +248,17 @@ def iter_weights( reader.close() + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, list[torch.Tensor]]: """CPU NVFP4 expert source banks for the offload cache; see load_nvfp4_expert_source_banks.""" return load_nvfp4_expert_source_banks( @@ -259,11 +268,12 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ): """parallel: same NVFP4 source banks via the common chunked multi-threaded O_DIRECT reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -277,6 +287,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/nvfp4_banks.py b/python/freetoken/models/nvfp4_banks.py index 6b933ff1d..a23adcb38 100644 --- a/python/freetoken/models/nvfp4_banks.py +++ b/python/freetoken/models/nvfp4_banks.py @@ -86,6 +86,7 @@ def load_nvfp4_expert_source_banks( drop_page_cache: DropPageCache, primary: bool, layer_sink=None, + disk_tier=None, ) -> dict[str, list[torch.Tensor]]: """Build the 6 native NVFP4 source banks by streaming checkpoint shards (serial per-shard read). @@ -150,6 +151,16 @@ def load_nvfp4_expert_source_banks( drop_page_cache(path) _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill + K = disk_tier.ram_experts if disk_tier is not None else None + # Rows this loader actually materializes per layer. The banks are still allocated + # at the full [E, ...] shape -- rows K..E-1 are the disk tier's fetch destination and + # must exist as address space -- but they are never read or written here, so they + # stay unbacked. Filling them and releasing afterwards (the previous order) made the + # load-time peak the FULL expert set, which is exactly the case the tier exists for: + # GLM-5.3-Flash (166 GiB of experts) swapped a 61 GB box to a halt at shard 73/118. + rows_per_layer = E if K is None else min(K, E) + if disk_tier is not None and layer_sink is not None: + raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] gate_up_global = [b.tensor for b in _hb["gate_up_global"]] @@ -160,7 +171,7 @@ def load_nvfp4_expert_source_banks( from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline def _load(sink) -> int: - tracker = LayerCompletionTracker(E * 6, _hb, sink) + tracker = LayerCompletionTracker(rows_per_layer * 6, _hb, sink) placed = 0 for shard in tqdm(sorted(weight_shards), desc=f"Loading {spec.desc}", disable=not primary): path = os.path.join(folder, shard) @@ -168,6 +179,8 @@ def _load(sink) -> int: for name, match, bank_layer_id in weight_shards[shard]: layer = int(match.group("layer")) expert = int(match.group("expert")) + if expert >= rows_per_layer: + continue # disk-resident row: fetched on demand, never loaded here proj = match.group("proj") role = spec.proj_to_role[proj] kind = _canon_kind(spec, match.group("kind")) @@ -202,10 +215,15 @@ def _load(sink) -> int: if layer_sink is not None: placed = _load(layer_sink) else: - with PinPipeline() as pins: + with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) + if K is not None: + from freetoken.moe.disk_tier import check_tail_unbacked, release_bank_tails - expected = num_layers * E * 6 + release_bank_tails(_hb, E, K) + check_tail_unbacked(_hb, E, K) + + expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" return { "gate_up_packed": gate_up_packed, @@ -227,6 +245,7 @@ def load_nvfp4_expert_source_banks_parallel( workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ) -> dict[str, list[torch.Tensor]]: """parallel counterpart of :func:`load_nvfp4_expert_source_banks`, byte-for-byte same placement. bulk weight/weight_scale read via chunked multi-threaded O_DIRECT reader @@ -274,6 +293,16 @@ def load_nvfp4_expert_source_banks_parallel( drop_page_cache(path) _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I) # unpinned; pinned after fill + K = disk_tier.ram_experts if disk_tier is not None else None + # Rows this loader actually materializes per layer. The banks are still allocated + # at the full [E, ...] shape -- rows K..E-1 are the disk tier's fetch destination and + # must exist as address space -- but they are never read or written here, so they + # stay unbacked. Filling them and releasing afterwards (the previous order) made the + # load-time peak the FULL expert set, which is exactly the case the tier exists for: + # GLM-5.3-Flash (166 GiB of experts) swapped a 61 GB box to a halt at shard 73/118. + rows_per_layer = E if K is None else min(K, E) + if disk_tier is not None and layer_sink is not None: + raise NotImplementedError("disk tier: the converter (layer_sink) path is not supported yet") gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]] gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]] gate_up_global = [b.tensor for b in _hb["gate_up_global"]] @@ -285,10 +314,15 @@ def load_nvfp4_expert_source_banks_parallel( # Pass 2: bulk weight/weight_scale via the common parallel reader; place by name. def _load(sink) -> int: - tracker = LayerCompletionTracker(E * 6, _hb, sink) + tracker = LayerCompletionTracker(rows_per_layer * 6, _hb, sink) placed = 0 + def _wanted(n: str) -> bool: + info = weight_info.get(n) + # Filter at the reader so disk-resident rows cost no I/O at all, not just no write. + return info is not None and int(info[0].group("expert")) < rows_per_layer + for name, tensor in iter_expert_tensors_parallel( - folder, lambda n: n in weight_info, workers=workers, chunk=chunk + folder, _wanted, workers=workers, chunk=chunk ): match, bank_layer_id = weight_info[name] layer = int(match.group("layer")) @@ -321,10 +355,15 @@ def _load(sink) -> int: if layer_sink is not None: placed = _load(layer_sink) else: - with PinPipeline() as pins: + with PinPipeline(prefix_rows=K) as pins: placed = _load(pins) + if K is not None: + from freetoken.moe.disk_tier import check_tail_unbacked, release_bank_tails + + release_bank_tails(_hb, E, K) + check_tail_unbacked(_hb, E, K) - expected = num_layers * E * 6 + expected = num_layers * rows_per_layer * 6 assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}" return { "gate_up_packed": gate_up_packed, diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f2..589fc0ea7 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -3,12 +3,13 @@ from .weight import ( iter_weights, iter_weights_parallel, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, setup_offload_expert_banks, ) -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Qwen3_5MoEForCausalLM", "parse_config", "iter_weights", diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index d07f18cd7..a27f7c552 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -834,7 +834,7 @@ def close(self) -> None: def setup_offload_expert_banks( model_path: str, model_config, *, device: torch.device, dtype: torch.dtype, dummy: bool = False, parallel: bool = False, workers: int = 8, chunk: int = 8 << 20, - decode_target: str = "gpu", layer_sink=None, + decode_target: str = "gpu", layer_sink=None, disk_tier=None, ): """Build the routed-expert offload banks. The qwen3_5_moe module always exports this hook, so it intercepts *every* qwen3_5_moe offload load -- defer non-block-fp8 checkpoints (plain @@ -849,16 +849,35 @@ def setup_offload_expert_banks( providers for non-block-fp8 checkpoints. ``decode_target`` is forwarded so the cpu backend gets CPU-readable (native, non- - GPU-tiled) bank layouts -- e.g. native ``nvfp4`` rows rather than marlin/b12x.""" + GPU-tiled) bank layouts -- e.g. native ``nvfp4`` rows rather than marlin/b12x. + + ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``): pin only the first + ``ram_experts`` experts per layer, release the rest, and attach a + :class:`~freetoken.moe.disk_tier.Nvfp4DiskIndex` so the offload cache can + fetch the disk-resident experts on miss.""" eq = getattr(model_config, "expert_quant", "none") + # Checked before the branch below, not inside it: the fp8_block path never passes + # disk_tier on, so a guard inside the nvfp4 branch silently accepted the flag there. + if disk_tier is not None and eq != "nvfp4": + raise NotImplementedError( + f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") if eq != "fp8_block": from freetoken.moe.expert_banks import _PROVIDERS # nvfp4 -> _nvfp4_banks, none -> _bf16_banks - return _PROVIDERS[eq](model_path, model_config, device, dtype, dummy, - parallel=parallel, workers=workers, chunk=chunk, - decode_target=decode_target, layer_sink=layer_sink) + banks = _PROVIDERS[eq](model_path, model_config, device, dtype, dummy, + parallel=parallel, workers=workers, chunk=chunk, + decode_target=decode_target, layer_sink=layer_sink, + disk_tier=disk_tier) + # The disk index rides back on the banks: _nvfp4_banks builds it for every family + # through nvfp4_expert_source_spec, so there is nothing family-specific left here. + return banks if get_tp_info().size > 1: raise NotImplementedError("qwen3_5_moe fp8 expert banks support TP=1 only") + if disk_tier is not None: + # the fp8_block path builds no disk index; without this the flag would be + # accepted and silently dropped (nothing released, nothing fetched) + raise NotImplementedError( + f"disk tier: only nvfp4 experts are supported (got expert_quant={eq!r})") from freetoken.moe.expert_banks import ExpertBanks mode = os.environ.get("FREETOKEN_FP8_EXPERTS", "fp8").strip().lower() @@ -1060,8 +1079,17 @@ def _load(sink) -> None: return ExpertBanks("bf16", banks, streamed=layer_sink is not None) + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + def load_nvfp4_expert_sources( - model_path: str, config, *, layer_sink=None + model_path: str, config, *, layer_sink=None, disk_tier=None ) -> dict[str, torch.Tensor]: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" @@ -1072,11 +1100,13 @@ def load_nvfp4_expert_sources( drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ): """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -1090,6 +1120,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 04c29d778..e18b007b9 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -13,6 +13,7 @@ from .model import Qwen4ExpForCausalLM from .weight import ( iter_weights, + nvfp4_expert_source_spec, load_nvfp4_expert_sources, load_nvfp4_expert_sources_parallel, load_ple_table, @@ -24,7 +25,7 @@ # load_nvfp4_expert_sources via the model spec. from freetoken.models.qwen3_5_moe.weight import setup_offload_expert_banks -__all__ = [ +__all__ = ["nvfp4_expert_source_spec", "Qwen4ExpForCausalLM", "iter_weights", "load_nvfp4_expert_sources", diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index f8d2a7494..476929f08 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -289,7 +289,16 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> dict: + +def nvfp4_expert_source_spec(model_path: str, config): + """The source spec the disk tier must index this checkpoint with. + + Same object the loader passes to ``load_nvfp4_expert_source_banks``, exposed so the + shared provider can build the disk index without knowing the family: the index has to + read the rows the loader placed, so one spec has to serve both.""" + return _NVFP4_SOURCE_SPEC + +def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None, disk_tier=None) -> dict: """Build the CPU NVFP4 expert source banks for the offload cache (gate/up fused on the output-row axis, down separate; weight_scale_2 carried as the per-row global scale).""" return load_nvfp4_expert_source_banks( model_path, @@ -298,11 +307,12 @@ def load_nvfp4_expert_sources(model_path: str, config, *, layer_sink=None) -> di drop_page_cache=drop_page_cache, primary=get_tp_info().is_primary(), layer_sink=layer_sink, + disk_tier=disk_tier, ) def load_nvfp4_expert_sources_parallel( - model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None + model_path: str, config, *, workers: int = 8, chunk: int = 8 << 20, layer_sink=None, disk_tier=None ) -> dict: """parallel: same NVFP4 source banks via the common chunked multi-threaded reader.""" from freetoken.models.nvfp4_banks import load_nvfp4_expert_source_banks_parallel @@ -316,6 +326,7 @@ def load_nvfp4_expert_sources_parallel( workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier, ) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9b..075904c9c 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -293,6 +293,20 @@ def load_moe_expert_sources( ) +def nvfp4_moe_expert_source_spec(model_path: str, model_config): + """The family's ``Nvfp4ExpertSourceSpec`` for this checkpoint, or None if it has none. + + The disk tier indexes the same rows the loader placed, so it needs the very spec the + family loader used -- which only the family knows (glm5_next picks between a + compressed-tensors and a modelopt naming per checkpoint). None means the family has no + NVFP4 expert source spec at all, which the caller must treat as "no disk tier here" + rather than silently continue. + """ + _config, spec = _spec_for_model_path(model_path) + getter = _model_override(spec, "nvfp4_expert_source_spec") + return None if getter is None else getter(model_path, model_config) + + def load_nvfp4_moe_expert_sources( model_path: str, model_config, @@ -302,12 +316,14 @@ def load_nvfp4_moe_expert_sources( workers: int = 8, chunk: int = 8 << 20, layer_sink=None, + disk_tier=None, ) -> dict: """Load (or fabricate, with ``dummy=True``) packed NVFP4 expert source banks. ``parallel=True`` uses the model's ``load_nvfp4_expert_sources_parallel`` (common chunked multi-threaded O_DIRECT reader). ``layer_sink``: see ``models.nvfp4_banks.load_nvfp4_expert_source_banks``; forwarded to the per-model - loader, which forwards it on.""" + loader, which forwards it on. ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``) + pins only the first ``ram_experts`` experts per layer and releases the rest.""" _config, spec = _spec_for_model_path(model_path) if dummy: builder = ( @@ -319,9 +335,10 @@ def load_nvfp4_moe_expert_sources( if loader is None: # no parallel reader -> let the caller fall back to serial raise NotImplementedError( f"{spec.module} provides no load_nvfp4_expert_sources_parallel") - return loader(model_path, model_config, workers=workers, chunk=chunk, layer_sink=layer_sink) + return loader(model_path, model_config, workers=workers, chunk=chunk, layer_sink=layer_sink, + disk_tier=disk_tier) loader = _load_attr(spec.module, "load_nvfp4_expert_sources") - return loader(model_path, model_config, layer_sink=layer_sink) + return loader(model_path, model_config, layer_sink=layer_sink, disk_tier=disk_tier) def load_q4_0_moe_expert_sources( diff --git a/python/freetoken/moe/disk_tier.py b/python/freetoken/moe/disk_tier.py new file mode 100644 index 000000000..959ce319a --- /dev/null +++ b/python/freetoken/moe/disk_tier.py @@ -0,0 +1,685 @@ +"""Disk tier: NVMe-backed MoE experts (VRAM <- RAM <- NVMe). + +Lets the offload backend serve experts that do NOT fit in pinned RAM: the RAM +bank holds only the first ``ram_experts`` experts per layer (pinned), the rest +stay on disk in the original checkpoint. When the GPU slot cache misses a +disk-resident expert, :class:`DiskTier` fetches its rows with O_DIRECT preadv +into a small pinned staging buffer and H2D-copies them into the slot the LRU +kernel already assigned, then shrinks the miss list so the existing PCIe +``copy_missing`` path only moves the RAM-resident misses. + +v0 scope (prototype): +* native NVFP4 layout only (the "triton" backend banks -- what sm_120 picks); +* ``decode_target == "gpu"`` (offload) only -- the CPU executor reads banks + directly and would read released pages; +* synchronous fetch (the layer waits for its disk misses); no CUDA-graph + capture (the miss-list D2H/H2D round trip is host-side and variable); +* prefill_overlap off (the double-buffer prefill path bypasses the slot cache). + +The bank rows are read from the ORIGINAL safetensors shards: every expert +tensor is a contiguous per-expert tensor, so a bank row is one (or two, for +the gate|up-fused banks) aligned super-block preads. No FTW conversion needed. +""" + +from __future__ import annotations + +import ctypes +import json +import os +import struct +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import torch + +from freetoken.moe.host_banks import HostBank + +_ALIGN = 4096 + + +@dataclass(frozen=True) +class DiskTierSpec: + """Engine -> loader: how many experts per layer stay pinned in RAM. + + Experts ``[0, ram_experts)`` are pinned as usual; ``[ram_experts, E)`` keep + their bank rows allocated but their pages are released after load and are + served from disk by :class:`DiskTier`.""" + + ram_experts: int + + +def release_bank_tails(banks_by_name: dict[str, list[HostBank]], num_experts: int, + ram_experts: int) -> None: + """MADV_DONTNEED the unpinned tail rows of every bank layer (post-load). + + The release is an optimization, not an invariant: the tail rows were never + written at load, so when a row boundary is not page-aligned (the small scale + banks) we warn and skip that bank instead of failing the boot.""" + _PAGE = 4096 + for layer_banks in banks_by_name.values(): + for bank in layer_banks: + row_bytes = bank.nbytes // num_experts + offset = ram_experts * row_bytes + size = bank.nbytes - offset + if offset % _PAGE or size % _PAGE: + print(f"[disk-tier] WARNING: bank row boundary not page-aligned " + f"(ram_experts={ram_experts}, row_bytes={row_bytes}); skipping " + f"the release for this bank -- the tail rows were never written, " + f"so nothing is lost", flush=True) + continue + bank.release_range(offset, size) + + +def tail_resident_bytes(bank: HostBank, num_experts: int, ram_experts: int) -> int: + """Bytes the kernel currently backs in the released tail rows ``[ram_experts, E)``. + + mincore(2) over the tail's byte range: one syscall, per-page residency. + Conservative -- mincore also reports private-anon pages mapped from the + shared zero page (a plain READ of a tail row), so this overcounts what + actually costs RAM (cgroup memory.stat shmem is the real number). + Returns -1 if mincore itself fails.""" + row_bytes = bank.nbytes // num_experts + off = ram_experts * row_bytes + size = bank.nbytes - off + if size <= 0: + return 0 + _PAGE = 4096 + vec = (ctypes.c_ubyte * (size // _PAGE))() + libc = ctypes.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_ubyte)] + if libc.mincore(ctypes.c_void_p(bank.addr + off), size, vec) != 0: + return -1 + return sum(vec) * _PAGE + + +def check_tail_unbacked(banks_by_name: dict[str, list[HostBank]], num_experts: int, + ram_experts: int) -> None: + """Startup sanity check for the lazy-tail invariant, right after + ``release_bank_tails``: nothing reads or writes the tail rows, so the kernel + should be backing ~none of them. The expected worst case is one 2 MiB THP + 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 and the disk-tier RAM math no longer holds. Always logs, warns + above the bound.""" + _HUGE = 2 << 20 + resident = tail = n_banks = 0 + for layer_banks in banks_by_name.values(): + for bank in layer_banks: + row_bytes = bank.nbytes // num_experts + t = bank.nbytes - ram_experts * row_bytes + if t <= 0: + continue + r = tail_resident_bytes(bank, num_experts, ram_experts) + if r < 0: + continue # mincore failed: skip rather than warn on our own probe + resident += r + tail += t + n_banks += 1 + if n_banks == 0: + return + from freetoken.distributed import try_get_tp_info + tp = try_get_tp_info() + rank = getattr(tp, "rank", "?") + size_ = getattr(tp, "size", "?") + bound = n_banks * _HUGE + print(f"[disk-tier] tail check rank={rank}/{size_}: resident {resident >> 20} MiB " + f"of {tail >> 20} MiB (warn bound {bound >> 20} MiB = 1x2MiB per bank layer)", + flush=True) + if resident > bound: + print(f"[disk-tier] WARNING: tail rows more resident than the THP bound -- " + f"something is reading the released rows; the disk-tier RAM math no longer holds", + flush=True) + +# Native NVFP4 bank order (== _BANK_SCHEMAS["nvfp4"]) and, per bank, the +# checkpoint segments that make up one expert row: (proj, kind, dst_row_start, +# dst_row_end). The gate|up-fused banks splice gate rows then up rows on the +# output-row axis; down banks are a single segment. Row ends are None = rest. +_NVP4_BANK_SEGS = ( + (("gate_proj", "weight", 0, None), ("up_proj", "weight", None, None)), + (("gate_proj", "weight_scale", 0, None), ("up_proj", "weight_scale", None, None)), + (("gate_proj", "weight_scale_2", 0, None), ("up_proj", "weight_scale_2", None, None)), + (("down_proj", "weight", 0, None),), + (("down_proj", "weight_scale", 0, None),), + (("down_proj", "weight_scale_2", 0, None),), +) + + + +def _preadv_error(tier, staging, shard_idx: int, off: int, a0: int, slen: int, + direct: bool) -> OSError: + vma = "?" + try: + for line in open("/proc/self/maps"): + lo, hi = line.split()[0].split("-") + if int(lo, 16) <= staging.addr < int(hi, 16): + vma = line.strip()[:120] + break + except OSError: + pass + return OSError( + f"disk-tier preadv failed: shard={shard_idx} off={off} a0={a0} " + f"slen={slen} direct={direct} buf={hex(staging.addr)} " + f"staging_size={tier._staging_size} thread={threading.current_thread().name} " + f"vma={vma}" + ) + + +def _read_safetensors_offsets(path: str) -> dict[str, tuple[int, int]]: + """{tensor_name: (start, end)} from a shard's safetensors header, as ABSOLUTE + file offsets (data_offsets are relative to the data section, i.e. after the + 8-byte length + header JSON).""" + with open(path, "rb") as f: + (hlen,) = struct.unpack(" per-segment (shard_idx, offset, nbytes) locations. + + Built from the original checkpoint: the HF index json (name -> shard) plus + each referenced shard's safetensors header (name -> byte range). Expert + tensors are per-expert and contiguous, so a row is exactly one byte range + per segment. + """ + + def __init__(self, model_dir: str, config, spec) -> None: + from freetoken.models.nvfp4_banks import _num_moe_layers + from freetoken.utils.hf import download_hf_weight + + model_dir = download_hf_weight(model_dir) # hub id -> local cache dir; no-op if local + index_path = os.path.join(model_dir, "model.safetensors.index.json") + with open(index_path, encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + + num_layers = _num_moe_layers(config) + # (bank_layer, expert, proj, kind) -> (tensor_name, shard) + loc: dict[tuple[int, int, str, str], tuple[str, str]] = {} + for name, shard in weight_map.items(): + m = spec.key_pattern.match(name) + if m is None: + continue + bank_layer = spec.layer_to_bank(int(m.group("layer")), config) + if bank_layer is None: + continue + loc[(bank_layer, int(m.group("expert")), m.group("proj"), m.group("kind"))] = ( + name, shard) + + shards = sorted(set(shard for _, shard in loc.values())) + self.shard_paths = [os.path.join(model_dir, s) for s in shards] + offsets = {s: _read_safetensors_offsets(os.path.join(model_dir, s)) for s in shards} + shard_idx = {s: i for i, s in enumerate(shards)} + + E = config.num_experts + seg_size = struct.calcsize(" packed segments per expert + for bank_idx in range(len(_NVP4_BANK_SEGS)): + per_layer = [] + for layer in range(num_layers): + rows = bytearray() + for e in range(E): + for proj, kind, _, _ in _NVP4_BANK_SEGS[bank_idx]: + key = (layer, e, proj, kind) + entry = loc.get(key) + if entry is None: + raise KeyError( + f"disk tier: no {proj}.{kind} tensor for layer {layer} expert {e} " + f"(bank {bank_idx}) in {index_path}" + ) + name, shard = entry + start, end = offsets[shard][name] + rows += struct.pack(" list[tuple[int, int, int]]: + """[(shard_idx, offset, nbytes)] for one expert row, in segment order.""" + base = expert * self._seg_size * len(_NVP4_BANK_SEGS[bank_idx]) + raw = self.entries[bank_idx][layer][base:base + self._seg_size * len(_NVP4_BANK_SEGS[bank_idx])] + return [ + struct.unpack_from(" staging -> GPU slot.""" + + def __init__(self, index: Nvfp4DiskIndex, cache, ram_experts: int, workers: int = 8) -> None: + self._index = index + self._ram = ram_experts + self._banks = list(cache.banks) # [(per_layer_host, gpu_cache)] in schema order + self._row_bytes = [ + b[0][0][0].numel() * b[0][0][0].element_size() for b in self._banks + ] # full expert-row bytes per bank (staging must hold the biggest one) + # Per-bank destination row slices (gate|up split at the row midpoint). + self._dst_slices: list[list[tuple[int, int]]] = [] + for bank_idx, (host_layer, _gpu) in enumerate(self._banks): + row = host_layer[0][0] + if len(_NVP4_BANK_SEGS[bank_idx]) == 2: + mid = row.shape[0] // 2 + self._dst_slices.append([(0, mid), (mid, row.shape[0])]) + else: + self._dst_slices.append([(0, row.shape[0])]) + if os.environ.get("FT_DISK_TIER_VERIFY"): + # TP=2 debug: prove per-rank whether the host bank rows are full or + # TP-sharded, and that the disk index's full-row segments match them. + from freetoken.distributed import try_get_tp_info + tp = try_get_tp_info() + host_shapes = [tuple(b[0][0][0].shape) for b in self._banks] + disk_bytes = [ + sum(nb for _, _, nb in self._index.row_segments(bi, 0, 0)) + for bi in range(len(self._banks)) + ] + print(f"[disk-tier-init] tp_rank={getattr(tp, 'rank', '?')}/{getattr(tp, 'size', '?')} " + f"ram={ram_experts} host_row_shapes={host_shapes} " + f"host_row_bytes={self._row_bytes} disk_row_bytes={disk_bytes}", flush=True) + max_row = max(self._row_bytes) + self._staging_size = ((max_row + _ALIGN - 1) // _ALIGN + 2) * _ALIGN + self._staging = threading.local() + self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="disk-tier") + self._fd_lock = threading.Lock() + self._fds: dict[int, tuple[int, bool]] = {} + self._fetches = 0 + self._fetch_bytes = 0 + self._decode_verify_steps = 0 + self._map_verify_steps = 0 + self._cache = cache + + # ------------------------------------------------------------------ fds + def _fd(self, shard_idx: int) -> tuple[int, bool]: + """(fd, o_direct) for a shard; O_DIRECT falls back to plain preadv where the + filesystem refuses it (tmpfs/overlayfs -- tests).""" + ent = self._fds.get(shard_idx) + if ent is None: + with self._fd_lock: + ent = self._fds.get(shard_idx) + if ent is None: + path = self._index.shard_paths[shard_idx] + try: + fd = os.open(path, os.O_RDONLY | os.O_DIRECT) + direct = True + except OSError: + fd = os.open(path, os.O_RDONLY) + direct = False + ent = (fd, direct) + self._fds[shard_idx] = ent + return ent + + # --------------------------------------------------------------- staging + # Staging ring depth per worker thread. A buffer must not be overwritten by the + # next preadv until the async H2D copy that read it has finished (pinned-memory + # reuse race -- the copy is DMA, still reading host bytes after copy_ returns). + # The depth only needs to cover one copy's DMA time in host-side preadv time; + # the per-slot CUDA event below makes any shallower lap correct, just slower. + _STAGING_RING = 8 + + def _staging_ring(self) -> list: + ring = getattr(self._staging, "ring", None) + if ring is None: + # Fresh worker threads default to CUDA device 0, but the rank may + # live on another device (TP>1). The H2D copies land on the + # destination tensor's device stream, while ev.record() below uses + # the thread's CURRENT stream -- without this, the ring's reuse + # guard waits on an idle stream and a preadv can overwrite the + # buffer mid-DMA (corrupted slot rows on TP=2 rank 1). + dev = self._banks[0][1].device + if dev.type == "cuda": + torch.cuda.set_device(dev) + ring = [] + for _ in range(self._STAGING_RING): + buf = HostBank((self._staging_size,), torch.uint8) + buf.pin() # pin once per worker thread + ev = torch.cuda.Event() if torch.cuda.is_available() else None + if ev is not None: + ev.record() # start "complete"; re-recorded after each copy + ring.append([buf, ev]) + self._staging.ring = ring + return ring + + # ---------------------------------------------------------------- fetch + def _fetch_expert(self, layer: int, expert: int, slot: int) -> None: + # The server runs under inference_mode; the fetch pool threads do not, + # so the H2D writes into the (inference) slot cache need their own scope. + with torch.inference_mode(): + self._fetch_expert_inner(layer, expert, slot) + + def _fetch_expert_inner(self, layer: int, expert: int, slot: int) -> None: + ring = self._staging_ring() + ri = getattr(self._staging, "ri", 0) + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + row = gpu_cache[slot] + segs = self._index.row_segments(bank_idx, layer, expert) + for (d0, d1), (shard_idx, off, nbytes) in zip(self._dst_slices[bank_idx], segs): + staging, ev = ring[ri] + if ev is not None: + # This buffer's last async H2D copy must be done before the + # preadv below overwrites it (pinned-memory reuse race). + ev.synchronize() + ri = (ri + 1) % len(ring) + fd, direct = self._fd(shard_idx) + if direct: + a0 = off & ~(_ALIGN - 1) + slen = (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) + else: + a0, slen = off, nbytes + mv = (ctypes.c_char * slen).from_address(staging.addr) + try: + os.preadv(fd, [mv], a0) + except OSError: + raise _preadv_error(self, staging, shard_idx, off, a0, slen, direct) + row_off = off - a0 + if bank_idx in (2, 5): + # Global-scale banks: the checkpoint stores a per-expert fp32 + # SCALAR (weight_scale_2); the bank row is that value as fp16 + # broadcast across the row -- convert + fill, no byte copy. + val = staging.tensor[row_off:row_off + 4].view(torch.float32)[0].to( + torch.float16) + row[d0:d1].fill_(val) + continue + src = staging.tensor[row_off:row_off + nbytes] + dst = row[d0:d1] + dst.copy_(src.view(dst.dtype).view(dst.shape), non_blocking=True) + if ev is not None: + # Arm: the next reuse of this buffer waits for this copy. + ev.record() + self._staging.ri = ri + self._fetches += 1 + self._fetch_bytes += sum(self._row_bytes) + + def _sync_fetches(self) -> None: + """Wait for the pool threads' async H2D copies to land. + + The copies are enqueued on the pool threads' default stream; f.result() only + waits for them to be ENQUEUED. The GEMM's stream is not ordered with that + stream, so sync the default stream before the GEMM reads the slots.""" + # Key this on where the BANKS live, not on whether the machine has a GPU: the + # disk-tier unit tests build CPU banks on a CUDA box, and default_stream() rejects + # a CPU device outright. There is nothing to order for a CPU->CPU copy anyway. + device = self._banks[0][1].device + if device.type != "cuda": + return # CPU banks: the copies are synchronous CPU->CPU + torch.cuda.default_stream(device).synchronize() + + def _verify_slot(self, cache, layer: int, expert: int, slot: int | None = None, + phase: str = "prefill") -> None: + """One-shot debug: read back an expert's slot rows and compare against the + checkpoint bytes (ground truth). Gated on FT_DISK_TIER_VERIFY.""" + import torch + if slot is None: + slot = expert # identity mapping (prefill) + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[slot].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + ref = self._ref_row(bank_idx, layer, expert, flat.numel(), + slot_row.element_size(), + slot_row.numel() // slot_row.shape[0] if slot_row.dim() > 1 else 1) + try: + ref = ref.to(flat.device) + match = bool(torch.equal(flat, ref)) + if match: + print(f"[verify] {phase} L{layer} bank={bank_idx} expert={expert} " + f"slot={slot} match=True", flush=True) + continue + diff = (flat != ref) + nz = torch.nonzero(diff).flatten() + print(f"[verify] {phase} L{layer} bank={bank_idx} expert={expert} " + f"slot={slot} match=False n_diff={int(diff.sum())}/{flat.numel()} " + f"first_off={int(nz[0])} last_off={int(nz[-1])} " + f"slot_norm={slot_row.float().norm().item():.4f} " + f"ref_norm={ref.view(slot_row.dtype).view(slot_row.shape).float().norm().item():.4f} " + f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", flush=True) + self._identify_overwriter(layer, bank_idx, expert, flat) + except Exception as exc: # never crash the server in debug + print(f"[verify] bank={bank_idx} expert={expert} ERROR {exc!r}", flush=True) + + def _ref_row(self, bank_idx: int, layer: int, expert: int, row_bytes: int, + row_el: int, row_leading: int) -> torch.Tensor: + """Reference row bytes for (bank, layer, expert) straight from the checkpoint.""" + ref = torch.zeros(row_bytes, dtype=torch.uint8) + segs = self._index.row_segments(bank_idx, layer, expert) + for (d0, d1), (shard_idx, off, nbytes) in zip(self._dst_slices[bank_idx], segs): + fd, direct = self._fd(shard_idx) + a0 = off if not direct else (off & ~(_ALIGN - 1)) + slen = nbytes if not direct else (off + nbytes - a0 + _ALIGN - 1) & ~(_ALIGN - 1) + buf = os.pread(fd, slen, a0) + row_off = off - a0 + seg = buf[row_off:row_off + nbytes] + if bank_idx in (2, 5): + import struct as _st + import numpy as _np + f16 = _np.float16(_st.unpack(" None: + """Debug: on a verify mismatch, find whose row the slot actually holds. + + Compares the slot's first 64 bytes against (a) every other expert of the + same layer and (b) the same expert in every other layer, straight from the + checkpoint. A hit names the overwriter (e.g. a staging-buffer reuse race + landing a neighbour's preadv); no hit means a partial mix. Only runs on + mismatch, so the ~300 extra preads are free otherwise.""" + head = flat[:64].cpu() + host_row = self._banks[bank_idx][0][0][0] # expert-0 row (all rows share its shape) + row_el = host_row.element_size() + row_leading = host_row.numel() // host_row.shape[0] if host_row.dim() > 1 else 1 + num_experts = self._cache.num_experts + num_layers = len(self._banks[bank_idx][0]) + hits = [] + for e in range(num_experts): + if e == expert: + continue + ref = self._ref_row(bank_idx, layer, e, flat.numel(), row_el, row_leading) + if bool(torch.equal(ref[:64], head)): + hits.append(f"L{layer}_e{e}") + for L in range(num_layers): + if L == layer: + continue + ref = self._ref_row(bank_idx, L, expert, flat.numel(), row_el, row_leading) + if bool(torch.equal(ref[:64], head)): + hits.append(f"L{L}_e{expert}") + print(f"[overwriter] L{layer} B{bank_idx} e{expert} head64 matches: " + f"{hits if hits else 'NONE (partial mix?)'}", flush=True) + + def verify_ram(self, cache, layer: int) -> None: + """One-shot debug: after the PCIe copy, check a RAM-resident expert's slot rows + against the checkpoint reference. Gated on FT_DISK_TIER_VERIFY.""" + expert = min(10, self._ram - 1) # a RAM-resident expert + print(f"[verify-ram] layer={layer} expert={expert} (RAM prefix)", flush=True) + self._verify_slot(cache, layer, expert) + # Check ALL RAM experts: slot vs host row (host correctness established + # separately). Count mismatches; identify the source of the first one. + n_bad = 0 + identified = False + for e in range(self._ram): + for bank_idx, (host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[e].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + hflat = host_layer[layer][e].contiguous().view(torch.uint8).reshape(-1) + if flat.numel() != hflat.numel() or not bool(torch.equal(flat.cpu(), hflat)): + n_bad += 1 + if n_bad <= 12: + print(f"[verify-ram] MISMATCH e={e} bank={bank_idx} " + f"slot_head={flat[:8].tolist()} host_head={hflat[:8].tolist()}", + flush=True) + if not identified: + identified = True + self._identify_source(layer, bank_idx, flat) + print(f"[verify-ram] layer={layer} mismatches={n_bad}/{self._ram * len(self._banks)}", + flush=True) + + def _identify_source(self, layer: int, bank_idx: int, flat: torch.Tensor) -> None: + """Debug: find where a corrupted slot's bytes came from (GPU slot, host row, + or checkpoint row).""" + flat_cpu = flat.cpu() + head = flat_cpu[:16] + found = [] + _host_layer, gpu_cache = self._banks[bank_idx] + gpu_flat = gpu_cache.view(torch.uint8).reshape(gpu_cache.shape[0], -1) + if gpu_flat.shape[1] == flat_cpu.numel(): + cand = torch.nonzero( + (gpu_flat[:, :16].cpu() == head.unsqueeze(0)).all(dim=1)).flatten().tolist() + for s in cand[:16]: + if bool(torch.equal(gpu_flat[s].cpu(), flat_cpu)): + found.append(f"gpu_slot={s}(L{layer},B{bank_idx})") + for e in range(self._ram): + hflat = _host_layer[layer][e].contiguous().view(torch.uint8).reshape(-1) + if hflat.numel() == flat_cpu.numel() and bool(torch.equal(hflat, flat_cpu)): + found.append(f"host_row_e{e}(L{layer},B{bank_idx})") + import itertools + num_layers = len(self._banks[bank_idx][0]) + targets = sorted(set(itertools.product([layer], range(256))) + | set(itertools.product(range(num_layers), [10]))) + for (L, e) in targets: + row_el, row_leading = None, None + hrow = _host_layer[layer][0] + row_el = hrow.element_size() + row_leading = hrow.numel() // hrow.shape[0] if hrow.dim() > 1 else 1 + try: + ref = self._ref_row(bank_idx, L, e, flat_cpu.numel(), row_el, row_leading) + except Exception: + continue + if bool(torch.equal(ref, flat_cpu)): + found.append(f"checkpoint_L{L}_e{e}") + print(f"[identify] L{layer} B{bank_idx} n={flat_cpu.numel()} " + f"source={found if found else 'UNKNOWN'}", flush=True) + + def verify_decode_mapping(self, cache, layer_id: int, topk_ids: torch.Tensor) -> None: + """Debug: after the LRU rewrite + fetch/copy, check that every slot the GEMM + will read actually holds the expert the bookkeeping says it holds. Gated on + FT_DISK_TIER_VERIFY; first 4 decode steps only.""" + if self._map_verify_steps >= 4: + return + self._map_verify_steps += 1 + slots = torch.unique(topk_ids.reshape(-1)) + nbad = 0 + for s in slots.tolist(): + s = int(s) + flat_id = int(cache.id_of_slot[s].item()) + if flat_id < 0: + print(f"[verify-map] step={self._map_verify_steps} slot={s} id_of_slot=-1", + flush=True) + nbad += 1 + continue + expert = flat_id % cache.num_experts + for bank_idx, (_host_layer, gpu_cache) in enumerate(self._banks): + slot_row = gpu_cache[s].contiguous() + flat = slot_row.view(torch.uint8).reshape(-1) + ref = self._ref_row(bank_idx, layer_id, expert, flat.numel(), + slot_row.element_size(), + slot_row.numel() // slot_row.shape[0] + if slot_row.dim() > 1 else 1) + if not bool(torch.equal(flat.cpu(), ref)): + nbad += 1 + if nbad <= 8: + print(f"[verify-map] step={self._map_verify_steps} slot={s} " + f"expert={expert} bank={bank_idx} MISMATCH " + f"slot_head={flat[:8].tolist()} ref_head={ref[:8].tolist()}", + flush=True) + print(f"[verify-map] step={self._map_verify_steps} slots={slots.numel()} bad={nbad}", + flush=True) + + def materialize_layer(self, cache, layer_id: int, expert_ids: torch.Tensor) -> None: + """Disk-tier prefill: materialize the RAM-resident prefix into identity slots + (the normal kernel restricted to K experts; the following ``copy_missing`` + streams it over PCIe), then fetch the routed disk-resident experts into + THEIR identity slots. The identity mapping (position == expert id) is + preserved, so the prefill GEMM is unchanged.""" + from freetoken.moe.offload_kernels import _materialize_layer_gpu + + # Prefill identity mapping owns ALL of slots [0, E) for this layer, but + # the kernel only scans slots < materialize_count, so the disk slots + # [ram, E) that still hold a previous layer's experts (previous prefill + # layer or decode LRU) would keep their slot_for_id entries -- phantom + # decode hits that read another layer's weights. Clear them first + # (device-side, no sync). + seg = cache.id_of_slot[self._ram:cache.num_experts] + valid = seg >= 0 + cache.slot_for_id.view(-1)[seg[valid].long()] = -1 + seg[valid] = -1 + cache.usage[self._ram:cache.num_experts][valid] = 0 + + _materialize_layer_gpu(cache, layer_id, materialize_count=self._ram) + routed = expert_ids.reshape(-1) + disk = torch.unique(routed[routed >= self._ram]) + if os.environ.get("FT_DISK_TIER_DEBUG") and layer_id < 3: + print(f"[disk-tier dbg] layer={layer_id} routed={routed.numel()} " + f"unique_disk={disk.numel()} disk={disk.tolist()[:12]}", flush=True) + if disk.numel() == 0: + return + # cache.step was already incremented by the kernel; assign the 0-d tensor + # device-side (same dtype/device as usage) instead of .item()-ing it, which + # would sync the stream once per layer on the prefill/decode path. + futures = [ + self._pool.submit(self._fetch_expert, layer_id, int(e), int(e)) + for e in disk.tolist() + ] + for f in futures: + f.result() + self._sync_fetches() + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id in (0, 20) and disk.numel() > 0: + limit = disk.numel() if layer_id == 0 else 6 # layer 0: ALL experts (race hunt) + for e in disk.tolist()[:limit]: + self._verify_slot(cache, layer_id, int(e), phase="prefill") + # Same bookkeeping the materialize kernel writes, per fetched expert. + flat = layer_id * cache.num_experts + disk + cache.slot_for_id[layer_id, disk] = disk + cache.id_of_slot[disk] = flat + cache.usage[disk] = cache.step + + def fetch_pending(self, cache, layer_id: int) -> None: + """Fetch this layer's disk-resident misses into their slots; shrink the miss + list to the RAM-resident remainder for the existing PCIe copy path.""" + n = int(cache.num_indices.item()) + if n == 0: + return + src = cache.src_indices[:n].cpu() + slots = cache.evict_slots[:n].cpu() + disk = [i for i in range(n) if int(src[i]) >= self._ram] + if os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0: + print(f"[fetch-pend] layer=0 n={n} ndisk={len(disk)} " + f"src_head={src[:4].tolist()}", flush=True) + if not disk: + return + futures = [ + self._pool.submit(self._fetch_expert, layer_id, int(src[i]), int(slots[i])) + for i in disk + ] + for f in futures: + f.result() + self._sync_fetches() + if (os.environ.get("FT_DISK_TIER_VERIFY") and layer_id == 0 + and self._decode_verify_steps < 3): + self._decode_verify_steps += 1 + for i in disk[:8]: + print(f"[verify-decode] step={self._decode_verify_steps} " + f"expert={int(src[i])} slot={int(slots[i])} ndisk={len(disk)}", flush=True) + self._verify_slot(cache, layer_id, int(src[i]), int(slots[i]), + phase="decode") + disk_set = set(disk) + ram = [i for i in range(n) if i not in disk_set] + if ram: + sel = torch.tensor(ram, dtype=torch.long) + cache.src_indices[:len(ram)].copy_(src[sel].to(cache.src_indices.dtype)) + cache.evict_slots[:len(ram)].copy_(slots[sel].to(cache.evict_slots.dtype)) + cache.num_indices.fill_(len(ram)) + + def refresh(self, cache) -> None: + """Rebind the slot-cache references after a runtime cache rebuild.""" + self._banks = list(cache.banks) + self._cache = cache + + def stats(self) -> dict: + return {"experts_fetched": self._fetches, "bytes_fetched": self._fetch_bytes} diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba8..13bc4f38d 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -49,6 +49,11 @@ class ExpertBanks: # streamed straight to its sink instead of staying materialized here) -- set by # convert.py's per-format streaming gate; ``sources`` may hold released tensors. streamed: bool = False + # Disk tier (None when off): a moe.disk_tier.Nvfp4DiskIndex over the original + # checkpoint plus how many experts per layer are RAM-resident (the rest are + # disk-resident and fetched on slot-cache miss). + disk_index: object | None = field(default=None) + disk_ram_experts: int = 0 _PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader @@ -160,7 +165,7 @@ def assemble(self, num_layers: int): return sources, gate_up_alpha, down_alpha -def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None, disk_tier=None) -> ExpertBanks: from freetoken.models.weight import load_nvfp4_moe_expert_sources from freetoken.moe.nvfp4_backends import ( b12x_repack_layer, @@ -183,6 +188,27 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, getattr(model_config, "nvfp4_backend", "auto"), activation=getattr(model_config, "hidden_act", "silu")) native = backend == "triton" + if disk_tier is not None and (not native or dummy): + raise NotImplementedError( + "disk tier requires the native NVFP4 layout (triton backend, decode_target=gpu, " + "not dummy)") + # Build the index HERE, not in a per-family setup: the loader below releases expert rows + # [K, E) for every family, so a family that reaches the release without an index would + # serve zeroed experts. Resolving the spec through the family hook keeps the index and + # the loader reading the same rows. + disk_kw = {} + if disk_tier is not None: + from freetoken.models.weight import nvfp4_moe_expert_source_spec + from freetoken.moe.disk_tier import Nvfp4DiskIndex + + source_spec = nvfp4_moe_expert_source_spec(model_path, model_config) + if source_spec is None: + raise NotImplementedError( + f"--moe-disk-tier on: {type(model_config).__name__} exposes no " + "nvfp4_expert_source_spec, so the tier cannot locate expert rows in the " + "checkpoint (the loader would release them and never refetch)") + disk_kw = {"disk_index": Nvfp4DiskIndex(model_path, model_config, source_spec), + "disk_ram_experts": disk_tier.ram_experts} repack_sink = None if not native and not dummy and layer_sink is not None: @@ -197,14 +223,14 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, # parallel only parallelizes the source read; the backend repack below is reused unchanged. sources = load_nvfp4_moe_expert_sources( model_path, model_config, dummy=dummy, parallel=parallel, workers=workers, chunk=chunk, - layer_sink=sink, + layer_sink=sink, disk_tier=disk_tier, ) # CPU-compute decode (cpu/hybrid) reads the native ModelOpt rows directly (its # dequant-in-GEMV kernel), so keep the native "nvfp4" layout and skip the GPU-tiled # marlin/b12x repacks (which only the GPU W4A16 kernels can read). if decode_target == "cpu": return ExpertBanks("nvfp4", {name: sources[name] for name in _BANK_SCHEMAS["nvfp4"]}, - streamed=sink is not None) + streamed=sink is not None, **disk_kw) # Pick the expert-GEMM backend by compute capability (and MoE width: auto keeps # small-I MoE on the Triton M=1 GEMV, which beats b12x's tensor cores at single-stream # decode) and repack the banks (in place; the tiled blocks are byte-identical per @@ -212,7 +238,7 @@ def _nvfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, logger.info(f"NVFP4 expert backend: {backend}") if backend == "triton": return ExpertBanks("nvfp4", {name: sources[name] for name in _BANK_SCHEMAS["nvfp4"]}, - streamed=sink is not None) + streamed=sink is not None, **disk_kw) quant_format = f"nvfp4_{backend}" if repack_sink is not None: # Streamed conversion: each layer was already repacked + written by the wrapper. @@ -304,13 +330,14 @@ def _model_setup_override(model_config): } -def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None, disk_tier=None) -> ExpertBanks: """Dispatch to the model's setup-override or the per-quant provider. ``parallel=True`` is the parallel read; a provider that hasn't implemented it raises NotImplementedError (the caller falls back to serial). ``decode_target`` lets the cpu backend force CPU-readable (native, non-GPU-tiled) bank layouts. ``layer_sink`` (converter only) is forwarded to setups/providers that declare the parameter; the rest ignore it and stay on the - materialize-and-write path (``ExpertBanks.streamed`` reports which happened).""" + materialize-and-write path (``ExpertBanks.streamed`` reports which happened). + ``disk_tier`` (a ``moe.disk_tier.DiskTierSpec``) is forwarded the same way.""" setup = _model_setup_override(model_config) if setup is not None: import inspect @@ -330,6 +357,8 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel kw["decode_target"] = decode_target if "layer_sink" in params and layer_sink is not None: kw["layer_sink"] = layer_sink + if "disk_tier" in params and disk_tier is not None: + kw["disk_tier"] = disk_tier return setup(model_path, model_config, **kw) expert_quant = model_config.expert_quant @@ -341,7 +370,7 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel return _PROVIDERS[expert_quant]( model_path, model_config, device, dtype, dummy, parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, - layer_sink=layer_sink, + layer_sink=layer_sink, disk_tier=disk_tier, ) @@ -418,6 +447,7 @@ def load_expert_banks( decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, + disk_tier=None, ) -> ExpertBanks: """Load (or fabricate, with ``dummy=True``) the expert banks. Two paths, both returning the same normalized ``ExpertBanks`` and both pinning after fill: @@ -442,6 +472,10 @@ def load_expert_banks( from freetoken.checkpoint.ftw import is_ftw_checkpoint, load_ftw_banks if model_path and is_ftw_checkpoint(model_path) and not dummy: + if disk_tier is not None: + raise NotImplementedError( + "disk tier v0 reads the original safetensors checkpoint; FTW checkpoints " + "are not supported yet (serve from the source path)") banks = load_ftw_banks( model_path, num_layers=model_config.num_moe_layers, workers=workers, chunk=chunk, layer_residency=layer_residency, @@ -483,13 +517,13 @@ def load_expert_banks( with requested_residency(layer_residency) as residency_plan: try: banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, - decode_target, layer_sink) + decode_target, layer_sink, disk_tier) except NotImplementedError as exc: if not parallel: raise logger.warning_rank0(f"parallel reader unavailable ({exc}); falling back to serial build") banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, False, workers, chunk, - decode_target, layer_sink) + decode_target, layer_sink, disk_tier) return _echo_residency(banks, layer_residency, residency_plan) diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index 436f52d2d..08217ed14 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 @@ def __init__(self, shape: tuple[int, ...], dtype: torch.dtype, 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,11 +148,52 @@ def pin(self) -> None: 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). + + The unpinned tail keeps its filled pages until :meth:`release_range` drops + them; nothing may DMA from the tail (the disk tier's miss filter guarantees + the GPU never copies those rows).""" + if self._pinned: + return + from freetoken.kernel.pinned import host_register + + row_bytes = self.nbytes // self.tensor.shape[0] + nbytes = nrows * row_bytes + try: + host_register(self.addr, nbytes) + except RuntimeError as exc: + raise RuntimeError( + 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 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). + + 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") + 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) @@ -289,7 +338,8 @@ class PinPipeline: A clean context-manager exit drains the queue and re-raises the first settle failure. """ - def __init__(self) -> None: + def __init__(self, prefix_rows: int | None = None) -> None: + self._prefix_rows = prefix_rows self._q: queue.SimpleQueue = queue.SimpleQueue() self._exc: BaseException | None = None # the current device is thread-local: a fresh thread sits on device 0 and cudaHostRegister would build its context there -- carry the creator's (bound) device into the worker @@ -308,9 +358,16 @@ def _run(self) -> None: continue # drain without settling after a failure bank, residency, plan, layer_id = item try: - _settle(bank, residency) - if plan is not None and residency == HostResidency.LOCKED.value: - plan.record(layer_id, bank.residency.value) + if self._prefix_rows is not None: + # Disk tier: pin only the RAM-resident expert prefix; the + # disk-resident tail is released by the caller. Takes + # precedence over the residency label (H2D needs the + # prefix page-locked regardless). + bank.pin_prefix(self._prefix_rows) + else: + _settle(bank, residency) + if plan is not None and residency == HostResidency.LOCKED.value: + plan.record(layer_id, bank.residency.value) except BaseException as exc: # surfaced by wait()/__exit__ self._exc = exc diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index e1f20dd2f..e12f4fd37 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -154,6 +154,9 @@ def __post_init__(self) -> None: # offload/PCIe path. Set by the engine after construction (empty = all-GPU, # all layers = the plain --moe-backend cpu case). self.cpu_layer_ids: frozenset = frozenset() + # Disk tier (None when off): a moe.disk_tier.DiskTier that fetches + # disk-resident slot-cache misses before the PCIe copy path. + self._disk_tier = None # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path. self.validate_rebuild(self.cache_size) assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, ( @@ -509,6 +512,8 @@ def rebuild(self, cache_size: int) -> None: self.prefill_overlap = False if self.prefill_overlap: self._init_prefill_overlap_buffers() + if self._disk_tier is not None: + self._disk_tier.refresh(self) # slot caches were reallocated def set_alphas( self, gate_up_alpha: torch.Tensor | None, down_alpha: torch.Tensor | None @@ -850,7 +855,19 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction ) - def materialize_layer(self, layer_id: int) -> None: + @property + def disk_tier_enabled(self) -> bool: + return self._disk_tier is not None + + def materialize_layer(self, layer_id: int, expert_ids: torch.Tensor | None = None) -> None: + if self._disk_tier is not None: + # Disk tier: stream only the RAM-resident prefix, then fetch the routed + # disk-resident experts into their identity slots (needs the routing). + assert expert_ids is not None, "disk-tier prefill needs the routed expert ids" + self._pending_src_layer = layer_id + self._pending_whole_layer = True + self._disk_tier.materialize_layer(self, layer_id, expert_ids) + return from freetoken.moe.offload_kernels import materialize_layer self._pending_src_layer = layer_id @@ -985,11 +1002,25 @@ def decode_routing_stats(self) -> dict: "norm_entropy": norm_ent, } + def attach_disk_tier(self, index, ram_experts: int, workers: int = 8) -> None: + """Enable the NVMe tier: disk-resident slot-cache misses are fetched from the + original checkpoint before the PCIe copy path (see moe/disk_tier.py).""" + from freetoken.moe.disk_tier import DiskTier + + assert self.decode_target == "gpu", "disk tier v0 supports the gpu (offload) path only" + assert self.quant_format == "nvfp4", f"disk tier v0 supports native nvfp4 banks (got {self.quant_format!r})" + assert not self.prefill_overlap, "disk tier v0 does not support prefill overlap" + self._disk_tier = DiskTier(index, self, ram_experts, workers=workers) + def copy_missing(self) -> None: assert self.banks, "set_bank_sources must register the banks first" layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" - if layer_id in self._unpinned_layers: + if self._disk_tier is not None: + # Fetch this layer's disk-resident misses into their slots, then shrink the + # miss list to the RAM-resident remainder for the PCIe copy below. + self._disk_tier.fetch_pending(self, layer_id) + elif layer_id in self._unpinned_layers: if not self._pending_whole_layer: raise RuntimeError( f"layer {layer_id} is unpinned: its only copy is the whole-layer " @@ -1001,6 +1032,13 @@ def copy_missing(self) -> None: for per_layer, cache in self.banks: cache[: self.num_experts].copy_(per_layer[layer_id]) return + 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()): + 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) if self._copy_fused_ok: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit @@ -1016,18 +1054,21 @@ def copy_missing(self) -> None: self.src_indices, self.num_indices, ) - return - - from freetoken.kernel import fast_index_copy_jit + else: + from freetoken.kernel import fast_index_copy_jit - for per_layer, cache in self.banks: - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + for per_layer, cache in self.banks: + fast_index_copy_jit( + cache, + self.evict_slots, + per_layer[layer_id], + self.src_indices, + self.num_indices, + ) + if (self._disk_tier is not None and layer_id == 0 + and self._pending_whole_layer + and os.environ.get("FT_DISK_TIER_VERIFY")): + self._disk_tier.verify_ram(self, layer_id) def iter_offload_moe_layers(model) -> Iterator: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d..3b6d67fb7 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -188,7 +188,10 @@ def _ensure_experts_hybrid_cpu( flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item()) -def _materialize_layer_gpu(cache, layer_id: int) -> None: +def _materialize_layer_gpu(cache, layer_id: int, materialize_count: int | None = None) -> None: + # materialize_count < num_experts: the disk tier's RAM-resident prefix only; the + # flat-id base still uses the full num_experts (the id space is layer * E + expert). + count = cache.num_experts if materialize_count is None else materialize_count block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size)) _materialize_layer_kernel[(1,)]( cache.slot_for_id, @@ -200,6 +203,7 @@ def _materialize_layer_gpu(cache, layer_id: int) -> None: cache.num_indices, layer_id, cache.num_experts, + count, cache.cache_size, BLOCK=block, ) @@ -257,11 +261,12 @@ def _materialize_layer_kernel( num_indices_ptr, layer_id: tl.constexpr, num_experts: tl.constexpr, + materialize_count: tl.constexpr, cache_size: tl.constexpr, BLOCK: tl.constexpr, ): off = tl.arange(0, BLOCK) - expert_mask = off < num_experts + expert_mask = off < materialize_count slot_mask = off < cache_size slot = off @@ -282,7 +287,7 @@ def _materialize_layer_kernel( tl.store(usage_ptr + slot, step, mask=expert_mask) tl.store(evict_slots_ptr + off, slot, mask=expert_mask) tl.store(src_indices_ptr + off, off, mask=expert_mask) # layer-local row - tl.store(num_indices_ptr, num_experts) + tl.store(num_indices_ptr, materialize_count) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 5b4db587d..89ea3b005 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -547,6 +547,39 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="The unified MoE cache eviction policy.", ) + parser.add_argument( + "--moe-disk-tier", + default=ServerArgs.moe_disk_tier, + choices=["off", "on"], + help=( + "NVMe tier for MoE experts (see moe/disk_tier.py): experts beyond " + "--expert-ram-experts per layer stay on disk in the original checkpoint " + "and are fetched on slot-cache miss. Requires native NVFP4 banks. " + "v0 preconditions (all enforced at once at boot): --moe-backend offload " + "(gpu decode), --disable-moe-prefill-overlap, --cuda-graph-max-bs 0, " + "and 0 < --expert-ram-experts < num_experts." + ), + ) + parser.add_argument( + "--expert-ram-experts", + type=int, + default=ServerArgs.expert_ram_experts, + help=( + "With --moe-disk-tier on: experts per layer kept pinned in RAM " + "(0 < N < num_experts; the rest are disk-resident). Keep " + "N * (smallest bank row bytes) page-aligned (a multiple of 4096) or " + "the small scale banks' tail rows stay resident instead of released " + "(warns, does not abort). The rule is per-model: e.g. Qwen3.8-Flash-Next " + "needs a multiple of 8, Ornith-1.5-35B a multiple of 2." + ), + ) + parser.add_argument( + "--disk-fetch-workers", + type=int, + default=ServerArgs.disk_fetch_workers, + help="Disk-tier O_DIRECT fetch threads (default 8).", + ) + parser.add_argument( "--moe-cpu-threads", type=int, diff --git a/tests/moe/test_disk_tier.py b/tests/moe/test_disk_tier.py new file mode 100644 index 000000000..aea439d14 --- /dev/null +++ b/tests/moe/test_disk_tier.py @@ -0,0 +1,352 @@ +"""CPU tests for the NVMe disk tier (moe/disk_tier.py). + +A synthetic NVFP4 MoE checkpoint (2 layers, 4 experts) is written as safetensors +shards with deterministic per-tensor content; the tests verify that +:class:`Nvfp4DiskIndex` resolves the right byte ranges and that +:class:`DiskTier` places the right bytes into slot-cache rows and rewrites the +miss list. No CUDA needed -- the "GPU" banks here are CPU tensors and the +staging pin is stubbed. +""" + +import json +import re +import struct +import threading +import types + +import pytest +import torch + +from freetoken.moe.disk_tier import DiskTier, Nvfp4DiskIndex +from freetoken.moe.host_banks import HostBank +from freetoken.models.nvfp4_banks import Nvfp4ExpertSourceSpec + +H, I, E, L = 16, 32, 4, 2 +SHARDS = ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors") + +SPEC = Nvfp4ExpertSourceSpec( + key_pattern=re.compile( + r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." + r"(?Pgate_proj|up_proj|down_proj)\." + r"(?Pweight|weight_scale|weight_scale_2)$" + ), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, + desc="disk-tier test", +) + +# (proj, kind, shape, dtype) -- the native NVFP4 per-expert tensor layout. +TENSOR_SPECS = ( + ("gate_proj", "weight", (I, H // 2), torch.uint8), + ("up_proj", "weight", (I, H // 2), torch.uint8), + ("down_proj", "weight", (H, I // 2), torch.uint8), + ("gate_proj", "weight_scale", (I, H // 16), torch.uint8), + ("up_proj", "weight_scale", (I, H // 16), torch.uint8), + ("down_proj", "weight_scale", (H, I // 16), torch.uint8), + # weight_scale_2 is a per-expert fp32 SCALAR in the real checkpoints; the + # bank row is its fp16 value broadcast across the row. + ("gate_proj", "weight_scale_2", (), torch.float32), + ("up_proj", "weight_scale_2", (), torch.float32), + ("down_proj", "weight_scale_2", (), torch.float32), +) + +BANK_SHAPES = ( + (2 * I, H // 2), # gate_up_packed + (2 * I, H // 16), # gate_up_scale + (2 * I,), # gate_up_global + (H, I // 2), # down_packed + (H, I // 16), # down_scale + (H,), # down_global +) +BANK_DTYPES = (torch.uint8, torch.uint8, torch.float16, torch.uint8, torch.uint8, torch.float16) + + +def _name(layer, expert, proj, kind): + return f"model.language_model.layers.{layer}.mlp.experts.{expert}.{proj}.{kind}" + + +def _tensor_for(layer, expert, proj, kind): + """Deterministic content: a base offset per (layer, expert, proj, kind) so any + misplacement is visible.""" + proj_i = ("gate_proj", "up_proj", "down_proj").index(proj) + kind_i = ("weight", "weight_scale", "weight_scale_2").index(kind) + base = layer * 100000 + expert * 1000 + proj_i * 100 + kind_i * 10 + for p, k, shape, dtype in TENSOR_SPECS: + if p == proj and k == kind: + if shape == (): # per-expert fp32 scalar (kept in fp16 range) + return torch.tensor(float(base % 50000), dtype=torch.float32) + n = int(torch.tensor(shape).prod()) + if dtype == torch.uint8: + return torch.arange(n, dtype=torch.uint8).add_(base % 251).view(shape) + return (torch.arange(n, dtype=torch.float32) + base).to(dtype).view(shape) + raise AssertionError((proj, kind)) + + +@pytest.fixture() +def checkpoint(tmp_path): + import safetensors.torch + + by_shard = {s: {} for s in SHARDS} + weight_map = {} + for layer in range(L): + for expert in range(E): + for proj, kind, _shape, _dtype in TENSOR_SPECS: + shard = SHARDS[(layer * E + expert) % 2] + name = _name(layer, expert, proj, kind) + by_shard[shard][name] = _tensor_for(layer, expert, proj, kind) + weight_map[name] = shard + for shard, tensors in by_shard.items(): + safetensors.torch.save_file(tensors, str(tmp_path / shard), metadata={"format": "pt"}) + with open(tmp_path / "model.safetensors.index.json", "w", encoding="utf-8") as f: + json.dump({"weight_map": weight_map, "metadata": None}, f) + config = types.SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, + num_layers=L, first_k_dense_replace=0) + return tmp_path, config + + +def _index(checkpoint): + path, config = checkpoint + return Nvfp4DiskIndex(str(path), config, SPEC) + + +def test_index_segments_match_file_bytes(checkpoint): + import safetensors + + path, config = checkpoint + index = _index(checkpoint) + assert len(index.shard_paths) == 2 + # Every (bank, layer, expert) row segment must point at the exact tensor bytes. + for bank_idx in range(6): + for layer in range(L): + for expert in range(E): + segs = index.row_segments(bank_idx, layer, expert) + expected = { + 0: [("gate_proj", "weight"), ("up_proj", "weight")], + 1: [("gate_proj", "weight_scale"), ("up_proj", "weight_scale")], + 2: [("gate_proj", "weight_scale_2"), ("up_proj", "weight_scale_2")], + 3: [("down_proj", "weight")], + 4: [("down_proj", "weight_scale")], + 5: [("down_proj", "weight_scale_2")], + }[bank_idx] + assert len(segs) == len(expected) + for (shard_idx, off, nbytes), (proj, kind) in zip(segs, expected): + shard_path = index.shard_paths[shard_idx] + with open(shard_path, "rb") as f: + (hlen,) = struct.unpack(" EFAULT/segfault at fetch time). + max_row = max( + b[0][0][0].numel() * b[0][0][0].element_size() for b in cache.banks) + assert tier._staging_size >= max_row + 2 * 4096, (tier._staging_size, max_row) + # Stub the pinned staging (HostBank.pin needs CUDA) with the same per-thread + # ring semantics as production (threading.local, no CUDA events on CPU). + local = threading.local() + + def _staging_ring(): + ring = getattr(local, "ring", None) + if ring is None: + ring = [[HostBank((tier._staging_size,), torch.uint8), None] + for _ in range(tier._STAGING_RING)] + local.ring = ring + return ring + + tier._staging_ring = _staging_ring + return tier + + +def _expected_rows(layer, expert): + """The 6 bank rows for one expert, as flat uint8, in schema order.""" + rows = [] + for bank_idx, (shape, dtype) in enumerate(zip(BANK_SHAPES, BANK_DTYPES)): + if bank_idx == 2: + gate = _tensor_for(layer, expert, "gate_proj", "weight_scale_2").to(torch.float16) + up = _tensor_for(layer, expert, "up_proj", "weight_scale_2").to(torch.float16) + row = torch.cat([gate.expand(I), up.expand(I)]).view(shape) + elif bank_idx == 5: + row = (_tensor_for(layer, expert, "down_proj", "weight_scale_2") + .to(torch.float16).expand(H).view(shape)) + elif bank_idx < 3: + kind = ("weight", "weight_scale")[bank_idx] + gate = _tensor_for(layer, expert, "gate_proj", kind) + up = _tensor_for(layer, expert, "up_proj", kind) + row = torch.cat([gate.reshape(-1), up.reshape(-1)]).view(shape) + else: + row = _tensor_for(layer, expert, "down_proj", ("weight", "weight_scale")[bank_idx - 3]) + rows.append(row.contiguous().view(torch.uint8).reshape(-1)) + return rows + + +def test_fetch_expert_places_all_banks(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache) + for layer in range(L): + for expert in range(E): + slot = (layer * E + expert) % 8 + tier._fetch_expert(layer, expert, slot) + expected = _expected_rows(layer, expert) + for bank_idx, (host_layer, gpu_cache) in enumerate(cache.banks): + got = gpu_cache[slot].contiguous().view(torch.uint8).reshape(-1) + assert torch.equal(got, expected[bank_idx]), (bank_idx, layer, expert) + stats = tier.stats() + assert stats["experts_fetched"] == L * E + + +def test_fetch_pending_filters_and_rewrites(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache, ram_experts=2) # experts 0,1 RAM; 2,3 disk + layer = 1 + # Miss list: expert 0 (RAM), 2 (disk), 3 (disk) -> slots 5, 6, 7. + cache.src_indices[:3] = torch.tensor([0, 2, 3], dtype=torch.int32) + cache.evict_slots[:3] = torch.tensor([5, 6, 7], dtype=torch.int32) + cache.num_indices.fill_(3) + + tier.fetch_pending(cache, layer) + + # Disk misses fetched into their slots... + expected = _expected_rows(layer, 2) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[6].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + expected = _expected_rows(layer, 3) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[7].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + # ...and the miss list shrank to the RAM-resident remainder. + assert cache.num_indices.item() == 1 + assert cache.src_indices[0].item() == 0 + assert cache.evict_slots[0].item() == 5 + + +def test_fetch_pending_all_ram_is_noop(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache, ram_experts=2) + cache.src_indices[:2] = torch.tensor([0, 1], dtype=torch.int32) + cache.evict_slots[:2] = torch.tensor([0, 1], dtype=torch.int32) + cache.num_indices.fill_(2) + tier.fetch_pending(cache, 0) + assert cache.num_indices.item() == 2 + assert tier.stats()["experts_fetched"] == 0 + + +def test_fetch_pending_all_disk_clears_list(checkpoint): + cache = _fake_cache() + tier = _tier(checkpoint, cache, ram_experts=2) + cache.src_indices[:1] = torch.tensor([3], dtype=torch.int32) + cache.evict_slots[:1] = torch.tensor([4], dtype=torch.int32) + cache.num_indices.fill_(1) + tier.fetch_pending(cache, 0) + assert cache.num_indices.item() == 0 + expected = _expected_rows(0, 3) + for bank_idx, (_host, gpu_cache) in enumerate(cache.banks): + assert torch.equal( + gpu_cache[4].contiguous().view(torch.uint8).reshape(-1), expected[bank_idx]) + + +def test_release_range_frees_pages(): + """release_range must actually drop the resident pages. The bank is a + MAP_PRIVATE anonymous mapping, so MADV_DONTNEED frees for real; mincore + verifies the pages are gone (a MAP_SHARED mapping would keep them).""" + import ctypes as ct + + size = 4 * 1024 * 1024 + bank = HostBank((size,), torch.uint8) + bank.tensor.fill_(7) # fault every page in + libc = ct.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ct.c_void_p, ct.c_size_t, ct.POINTER(ct.c_ubyte)] + libc.mincore.restype = ct.c_int + + def resident_pages(addr, nbytes): + vec = (ct.c_ubyte * ((nbytes + 4095) // 4096))() + assert libc.mincore(addr, nbytes, vec) == 0 + return sum(1 for b in vec if b & 1) + + pages = size // 4096 + assert resident_pages(bank.addr, size) == pages + bank.release_range(0, size) + assert resident_pages(bank.addr, size) == 0 + # The mapping stays valid: the refaulted pages read back as zeros. + assert bank.tensor[0] == 0 + + +def test_tail_unbacked_after_release(): + """Lazy-tail invariant (the disk-tier RAM math): after release_range, the + tail rows [K, E) back NO pages -- until something writes them. mincore over + the tail is the cheap startup check check_tail_unbacked() runs for real.""" + from freetoken.moe.disk_tier import release_bank_tails, tail_resident_bytes + + E, K = 8, 4 + bank = HostBank((E, 4096), torch.uint8) # page-sized rows + bank.tensor[:K].fill_(1) # touch only the prefix + release_bank_tails({"b": [bank]}, E, K) + assert tail_resident_bytes(bank, E, K) == 0 + # The mapping still works: one tail write backs exactly one page (the + # invariant is "nothing touches the tail", not "the tail refuses to back"). + bank.tensor[K].fill_(2) + assert tail_resident_bytes(bank, E, K) == 4096 + + +def test_release_bank_tails_unaligned_row_boundary(): + """A row boundary that is not page-aligned (the small scale banks) must not + fail the boot: release_bank_tails warns and skips that bank instead of + asserting in release_range. The tail rows were never written, so skipping + loses nothing. Aligned boundaries still release.""" + import ctypes as ct + + from freetoken.moe.disk_tier import release_bank_tails + + libc = ct.CDLL("libc.so.6", use_errno=True) + libc.mincore.argtypes = [ct.c_void_p, ct.c_size_t, ct.POINTER(ct.c_ubyte)] + libc.mincore.restype = ct.c_int + + def resident_pages(addr, nbytes): + vec = (ct.c_ubyte * ((nbytes + 4095) // 4096))() + assert libc.mincore(addr, nbytes, vec) == 0 + return sum(1 for b in vec if b & 1) + + # A real Ornith gate_up_scale row size: 2048 bytes/row, NOT page-aligned. + E, K = 256, 127 + bank = HostBank((E, 2048), torch.uint8) + bank.tensor.fill_(7) # fault every page in + assert resident_pages(bank.addr, bank.nbytes) == bank.nbytes // 4096 + # K=127: offset = 127*2048 = 259072, not % 4096 -> warn+skip, no AssertionError. + release_bank_tails({"gate_up_scale": [bank]}, E, K) + # Skipped: the (already resident) tail pages are untouched, not freed. + assert resident_pages(bank.addr, bank.nbytes) == bank.nbytes // 4096 + + # Aligned K on the same bank shape: offset = 128*2048 = 262144 (% 4096) -> releases. + bank2 = HostBank((E, 2048), torch.uint8) + bank2.tensor.fill_(7) + release_bank_tails({"gate_up_scale": [bank2]}, E, 128) + assert resident_pages(bank2.addr + 128 * 2048, bank2.nbytes - 128 * 2048) == 0 diff --git a/tests/moe/test_disk_tier_families.py b/tests/moe/test_disk_tier_families.py new file mode 100644 index 000000000..c7a0f349e --- /dev/null +++ b/tests/moe/test_disk_tier_families.py @@ -0,0 +1,64 @@ +"""Every NVFP4 family must hand the disk tier a source spec. + +The loader releases expert rows ``[K, E)`` for every family, so a family that reaches +that release without an index would serve zeroed experts silently. These tests pin the +hook that keeps the index and the loader reading the same rows. +""" +from __future__ import annotations + +import importlib + +import pytest + +# Families whose loader passes an Nvfp4ExpertSourceSpec to load_nvfp4_expert_source_banks. +NVFP4_FAMILIES = [ + "qwen3_5_moe", "qwen4_exp", "glm4_moe", "glm5_next", "gemma4", "minimax_m2", "minimax_m3", +] + + +@pytest.mark.parametrize("family", NVFP4_FAMILIES) +def test_family_exposes_a_source_spec_hook(family): + mod = importlib.import_module(f"freetoken.models.{family}.weight") + getter = getattr(mod, "nvfp4_expert_source_spec", None) + assert callable(getter), ( + f"{family} defines _NVFP4_SOURCE_SPEC but exposes no nvfp4_expert_source_spec, so " + "the shared provider cannot build a disk index for it") + + +@pytest.mark.parametrize("family", [f for f in NVFP4_FAMILIES if f != "glm5_next"]) +def test_hook_returns_the_spec_the_loader_uses(family): + # glm5_next is excluded here only because its hook reads the checkpoint config to pick + # between the compressed-tensors and modelopt namings; it is covered by the test below. + mod = importlib.import_module(f"freetoken.models.{family}.weight") + spec = mod.nvfp4_expert_source_spec("unused/for/these/families", None) + assert spec is mod._NVFP4_SOURCE_SPEC + assert spec.key_pattern.groupindex.keys() >= {"layer", "expert", "proj", "kind"} + assert set(spec.proj_to_role.values()) == {"gate", "up", "down"} + + +def test_glm5_next_hook_follows_the_checkpoint_quant_method(monkeypatch): + mod = importlib.import_module("freetoken.models.glm5_next.weight") + + class _Cfg: + def __init__(self, method): + self.quantization_config = {"quant_method": method} + + monkeypatch.setattr(mod, "cached_load_hf_config", lambda path: _Cfg("compressed-tensors")) + assert mod.nvfp4_expert_source_spec("p", None) is mod._NVFP4_CT_SOURCE_SPEC + monkeypatch.setattr(mod, "cached_load_hf_config", lambda path: _Cfg("modelopt")) + assert mod.nvfp4_expert_source_spec("p", None) is mod._NVFP4_SOURCE_SPEC + + +def test_provider_refuses_a_family_without_a_spec(monkeypatch): + """A family with no hook must fail at load, not release rows and serve zeros.""" + from freetoken.moe import expert_banks + from freetoken.moe.disk_tier import DiskTierSpec + + monkeypatch.setattr("freetoken.models.weight.nvfp4_moe_expert_source_spec", + lambda path, config: None) + # select_nvfp4_backend is reached before the spec lookup; decode_target="cpu" keeps the + # path "native" so the earlier NotImplementedError does not mask the one under test. + with pytest.raises(NotImplementedError, match="nvfp4_expert_source_spec"): + expert_banks._nvfp4_banks( + "does/not/matter", object(), None, None, False, + decode_target="cpu", disk_tier=DiskTierSpec(ram_experts=1)) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 422ca8675..3b346c1d3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -868,3 +868,22 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def test_copy_miss_verify_probe_gated_on_disk_tier(monkeypatch, capsys): + """FT_DISK_TIER_VERIFY must not fire the [copy-miss] probe when the disk tier + is off: the probe's .item()/.cpu() are device-to-host syncs, which crash any + CUDA graph capture (PR #337 issuecomment-5519070434 -- every graph-capturing + boot died with the env var left over from a tier session). Gated on the tier + like its neighbours, and skipped while a stream is capturing.""" + layer, cache = _make_layer_and_cache() + cache._pending_src_layer = 0 + cache.evict_slots = torch.tensor([0, 1], dtype=torch.int32) + cache.src_indices = torch.tensor([2, 3], dtype=torch.int32) + cache.num_indices = torch.tensor(2) + cache._copy_fused_ok = False + monkeypatch.setattr("freetoken.kernel.fast_index_copy_jit", lambda *a, **k: None) + monkeypatch.setenv("FT_DISK_TIER_VERIFY", "1") + + cache.copy_missing() + assert "[copy-miss]" not in capsys.readouterr().out