diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index e855bcd32f..bc902ad7ab 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -681,6 +681,9 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: ) except PinFailed as exc: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc + resident_bytes = sum(t.nbytes for per_layer in banks.sources.values() for t in per_layer if t.is_cuda) + self._weights_bytes += resident_bytes + self._post_weights_free -= resident_bytes if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) object.__setattr__(config, "moe_cache_size", size) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d0..6e6e757a6e 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -93,6 +93,19 @@ def build_expert_banks( specs = {role: ((E, *spec.shape), spec.dtype) for role, spec in layout.items() if not spec.resident} hb = alloc_layer_banks(specs, num_layers) banks = {role: [b.tensor for b in hb[role]] for role in specs} + gpu_layers = set() + fit_host_ram = layer_sink is None and not dummy and device.type == "cuda" and method.kind is QuantKind.FP8_BLOCK + layer_bytes = sum(per_layer[0].nbytes for per_layer in hb.values()) + if fit_host_ram: + available = _host_available_bytes() + if available is not None and layer_bytes: + # Keep one layer of host staging space while packing; CMA cannot back registered banks. + gpu_layers = set(range(min(num_layers, max(0, (layer_bytes * (num_layers + 1) - available + layer_bytes - 1) // layer_bytes)))) + if gpu_layers: + for role, per_layer in banks.items(): + for layer_id in gpu_layers: + per_layer[layer_id] = torch.zeros_like(per_layer[layer_id], device=device) + logger.info(f"expert banks: keeping {len(gpu_layers)} FP8 layers ({len(gpu_layers) * layer_bytes / 2**30:.2f} GiB) on GPU to fit host RAM") alphas = { role: torch.empty(num_layers * E, dtype=spec.dtype, device=device) for role, spec in layout.items() if spec.resident @@ -122,6 +135,14 @@ def _fill(sink) -> None: # refuse before writing: a duplicate row would also complete the layer early and hand the sink a half-filled bank if written[layer_id, e0:e1].any(): raise ValueError(f"expert rows written more than once: layer {layer_id}, experts {e0}:{e1}") + if fit_host_ram and layer_id not in gpu_layers and not written[layer_id].any(): + # Reclaim and host registration change usable RAM during a long load. + available = _host_available_bytes() + if available is not None and available < 2 * layer_bytes: + for per_layer in banks.values(): + per_layer[layer_id] = torch.zeros_like(per_layer[layer_id], device=device) + gpu_layers.add(layer_id) + logger.info(f"expert banks: keeping FP8 layer {layer_id} on GPU after host RAM changed") written[layer_id, e0:e1] = 1 out = {role: banks[role][layer_id][e0:e1] for role in specs} got = method.pack(piece, out) @@ -138,7 +159,7 @@ def _fill(sink) -> None: _fill(layer_sink) elif torch.cuda.is_available(): with PinPipeline() as pins: - _fill(pins) + _fill(lambda layer_id, layer_banks: pins(layer_id, layer_banks) if layer_id not in gpu_layers else None) else: _fill(None) @@ -152,6 +173,15 @@ def _fill(sink) -> None: _PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader +def _host_available_bytes() -> int | None: + try: + with open("/proc/meminfo") as f: + mem = {key: int(value.split()[0]) * 1024 for key, value in (line.split(":", 1) for line in f)} + return max(0, mem["MemAvailable"] - mem.get("CmaFree", 0)) + except (OSError, KeyError, ValueError): + return None + + def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: if parallel: raise NotImplementedError( @@ -209,15 +239,7 @@ def _host_ram_fits_parallel(model_path: str) -> bool: extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True, i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure.""" - avail = None - try: - with open("/proc/meminfo") as f: - for line in f: - if line.startswith("MemAvailable:"): - avail = int(line.split()[1]) * 1024 - break - except OSError: - pass + avail = _host_available_bytes() if avail is None: return True try: # resolve a hub id to its local cache dir (no-op for a local path) so glob sees the shards diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc773..405df2cd36 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -76,6 +76,63 @@ def _bound(quant): assert torch.all(banks.sources["gate_up_global"][0].float() > 0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("available_layers, gpu_layers", [(4, ()), (3, (0,)), (1, (0, 1, 2)), ((4, 4, 1, 1), (1, 2))]) +def test_fp8_banks_fit_host_ram_without_changing_weight_bytes(monkeypatch, available_layers, gpu_layers): + from freetoken.layers.quantization.moe.base import MoEConfig + from freetoken.layers.quantization.moe.fp8_block import Fp8BlockMoEMethod + from freetoken.moe import expert_banks + from freetoken.moe.offload_cache import OffloadMoeCache + + method = Fp8BlockMoEMethod(MoEConfig(4, 256, 128, 2, strategy="offload")) + specs = method.layout() + layer_bytes = sum(4 * torch.empty(s.shape, dtype=s.dtype).nbytes for s in specs.values()) + if isinstance(available_layers, tuple): + availability = iter(available_layers) + monkeypatch.setattr(expert_banks, "_host_available_bytes", lambda: next(availability) * layer_bytes) + else: + monkeypatch.setattr(expert_banks, "_host_available_bytes", lambda: available_layers * layer_bytes) + pieces = [] + for layer in range(3): + piece = { + "gate_up": torch.full((4, 256, 256), layer + 1, dtype=torch.float32).to(torch.float8_e4m3fn), + "down": torch.full((4, 256, 128), layer + 2, dtype=torch.float32).to(torch.float8_e4m3fn), + "gate_up_scale": torch.full((4, 2, 2), layer + 1, dtype=torch.bfloat16), + "down_scale": torch.full((4, 2, 1), layer + 2, dtype=torch.bfloat16), + } + pieces.append((layer, 0, 4, piece)) + banks = expert_banks.build_expert_banks(method, 3, iter(pieces), device=torch.device("cuda")) + cache = OffloadMoeCache(3, 4, 8, device=torch.device("cuda"), quant_format="fp8_block", prefill_overlap=True) + cache.set_bank_sources(banks.sources) + cache.begin_prefill() + for layer, _, _, piece in pieces: + got = cache.wait_prefill_layer(layer) + for (name, per_layer), copied in zip(cache.bank_sources.items(), got): + source = per_layer[layer] + assert source.is_cuda == (layer in gpu_layers) + reference = torch.zeros(source.shape, dtype=source.dtype) + value = piece[name] + reference[..., :value.shape[-1]].copy_(value) + assert torch.equal(copied.cpu().view(torch.uint8), reference.view(torch.uint8)) + cache.release_prefill_layer(layer) + cache.reset() + for layer, _, _, _ in pieces: + ids = torch.tensor([3, 1], dtype=torch.int32, device="cuda") + cache.ensure_experts(layer, ids) + cache.copy_missing() + for name, per_layer in cache.bank_sources.items(): + got = cache.bank_caches[name][ids.long()].cpu().view(torch.uint8) + assert torch.equal(got, per_layer[layer].cpu()[[3, 1]].view(torch.uint8)) + + +def test_host_available_bytes_excludes_cma(monkeypatch): + from io import StringIO + from freetoken.moe.expert_banks import _host_available_bytes + + monkeypatch.setattr("builtins.open", lambda *a, **kw: StringIO("MemAvailable: 2048 kB\nCmaFree: 1792 kB\n")) + assert _host_available_bytes() == 256 * 1024 + + def test_offload_moe_layer_prefill_forward_uses_single_layer_cache_view(monkeypatch): layer, cache = _make_layer_and_cache() topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32)