Measure the 7B's VRAM peak instead of bracketing it

pipeline/config.py

@@ -156,6 +156,7 @@ def require_mimic() -> None:
 RPT_TOOL_GATE_PIVOT = REPORTS / "tool_gate_pivot.json"
 RPT_TOOL_TARGETS = REPORTS / "tool_target_candidates.json"
 RPT_TOOL_CAUSAL_PARITY = REPORTS / "tool_causal_parity.json"
+RPT_TOOL_VRAM_PROBE = REPORTS / "tool_vram_probe.json"
 
 # --------------------------------------------------------------------------
 # Keep every temporary file on D:. The C: drive is tight (~13 GB), and a
@@ -604,6 +605,53 @@ def _flag(name: str, default: bool) -> bool:
 LLM_REVISION = os.getenv("PM_LLM_REVISION", "main")
 LLM_QUANT = os.getenv("PM_LLM_QUANT", "nf4")          # nf4 | none
 LLM_PREFLIGHT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"    # ~350 MB toolchain probe
+# WHERE THE EMBEDDING TABLE LIVES ONCE LOADED.
+#
+# This model unties its embeddings (`tie_word_embeddings: false`) over a 152,064
+# token vocabulary, so `embed_tokens` and `lm_head` are 545 M parameters EACH and
+# bitsandbytes quantises neither -- `nn.Embedding` is not `nn.Linear`, and an
+# untied `lm_head` is what `get_keys_to_not_convert` returns. 1040 MiB apiece, in
+# 16-bit, on the scarcest resource in the machine. Only the embedding is a
+# candidate to move: `lm_head` runs a 152,064-way matmul every decode step.
+#
+# ⚠️ `cpu` PAYS, BUT NOT WHERE IT LOOKS LIKE IT DOES. Moving the table frees
+# 1039 MiB in torch's books and returns only 130 MiB to the driver -- the rest is
+# a hole inside a partially-used segment, and `empty_cache()` cannot hand that
+# back without `expandable_segments`, a no-op on Windows. Measured 2026-09-07:
+#
+#   after load          driver_free=1022  alloc=5309  reserved=6050
+#   after emb.to(cpu)   driver_free=1024  alloc=4270  reserved=6050
+#   after empty_cache   driver_free=1154  alloc=4270  reserved=5920
+#
+# Read straight after the load that says "130 MiB, not worth it". It is the wrong
+# place to read. The 909 MiB stays as REUSABLE ARENA: generation allocates its KV
+# cache and activations inside it instead of asking the driver for new segments,
+# so the card ends a generation with room on it. Three loads each way:
+#
+#                 free after generation        generation
+#   embed=cuda     119 /  97 / 117 MiB     25.6 / 17.8 / 21.6 s
+#   embed=cpu      770 / 692 / 678 MiB     14.4 / 13.2 / 13.3 s
+#
+# ~600 MiB more headroom and generation 37% faster, because generation time swings
+# ~5x with what is left on the card and this is what leaves something on it. Output
+# is byte-identical across both placements -- one sha256 over six runs -- which is
+# what makes the swap free rather than a quality trade.
+#
+# ⚠️ It does NOT move the load PEAK, so it does not lower the VRAM gate: the table
+# is on the card while `from_pretrained` runs whatever happens to it afterwards.
+LLM_EMBED_DEVICE = os.getenv("PM_LLM_EMBED_DEVICE", "cpu")   # cpu | cuda
+#: Free VRAM the 7B needs, from the DRIVER, before a load may start. One
+#: definition, two readers: `s19_generate`'s capability check and the demo
+#: service's `explanation.generator()`. It lived in the demo service alone until
+#: 2026-09-07, which left the pipeline stage asserting `vram_free_gb > 5.5` --
+#: a figure that had drifted below a level already known to segfault.
+#:
+#: MEASURED, not bracketed. `tools/vram_probe.py` samples the driver at 5 Hz
+#: across the load; three consecutive runs peaked at 6059, 6171 and 6239 MiB,
+#: spread 180. 6420 = max peak + spread. The full table, the accepted risk and
+#: the reason the peak does not move with LLM_EMBED_DEVICE are on
+#: `MIN_FREE_VRAM_MIB` in `pulsemind_demo/back-end/pythonService/explanation.py`.
+LLM_MIN_FREE_VRAM_MIB = 6420
 LLM_MAX_NEW_TOKENS = 220
 # Greedy, fixed seed. Non-determinism would make a prompt change unattributable,
 # and attributing changes is the entire purpose of the grounding checker.

pipeline/core/generate.py

@@ -69,11 +69,103 @@ def _quant_config(quant: str):
                               bnb_4bit_use_double_quant=True)
 
 
+def _host_embedding(model, compute_device) -> None:
+    """Move the embedding table into host RAM and keep it there.
+
+    ⚠️ NOT VIA `device_map`, and the reason matters.
+
+    `{"model.embed_tokens": "cpu", "": 0}` resolves correctly -- accelerate walks
+    a parameter name from the longest prefix down (`utils/modeling.py:1991`) -- so
+    it looks exactly like the right answer. It is not. In a map whose main device
+    is a GPU, `cpu` means OFFLOAD, not "run there":
+
+        offloaded_devices = ["disk"] if main_device == "cpu" ... else ["cpu", "disk"]
+        offload = {name: device in offloaded_devices ...}   # big_modeling.py:405
+
+    The module is left on `meta`, its execution device is set to the GPU, and
+    `AlignDevicesHook` streams the weight ONTO the card for every forward and
+    frees it after. For a 1040 MiB table across a prefill and 220 decode steps
+    that is the wrong trade twice over: the copies dominate generation, and the
+    table is resident again at exactly the moment the card is fullest. Measured
+    here it did not finish loading a 0.5B in seven minutes.
+
+    So: load normally, then move the module and let two hooks handle the crossing.
+    Public torch API, no accelerate internals, and the table crosses the bus once
+    at startup instead of once per token.
+
+    What travels instead is the ACTIVATIONS -- `seq_len x 3584 x 2` bytes on the
+    prefill, 7 KB per decode step. Three orders of magnitude less.
+
+    ⚠️ This lowers what the model HOLDS, not what the load PEAKS at. The table is
+    on the card while `from_pretrained` runs, so the VRAM gate -- which protects a
+    load -- cannot come down on the strength of this alone.
+
+    ⚠️ AND THE DRIVER BARELY NOTICES, WHICH IS NOT THE SAME AS IT NOT WORKING.
+    Measured: torch's `allocated` falls by 1039 MiB and the driver gets 130 back,
+    because a freed block inside a partially-used segment needs
+    `expandable_segments` to be returned and that is a no-op on Windows. Read
+    there and this looks worthless. The other 909 MiB is REUSABLE ARENA:
+    generation puts its KV cache and activations inside it instead of asking the
+    driver for new segments, so the card ends a generation with room on it.
+    Three loads each way -- `embed=cuda` finished with 119/97/117 MiB free in
+    25.6/17.8/21.6 s; `embed=cpu` with 770/692/678 MiB free in 14.4/13.2/13.3 s.
+    Generation swings ~5x with what is left on the card, and this is what leaves
+    something on it.
+
+    ⚠️ AFTER THIS RUNS, `model.device` IS A LIE. `PreTrainedModel.device` reports
+    the device of the FIRST parameter, and `embed_tokens` is the first parameter,
+    so a model whose compute lives entirely on the GPU starts answering "cpu".
+    Anything that sends inputs to `model.device` then sends them to the host: the
+    ids land on the CPU, `cache_position` and `position_ids` are derived from
+    THEIR device, and the rotary embedding dies on `mat2 is on cpu`. Measured --
+    the load succeeded and generation failed on the first forward.
+
+    Hence `compute_device`, threaded explicitly and carried on `Generator`. The
+    load knows the answer before the move makes the question ambiguous.
+    """
+    import torch
+
+    emb = model.get_input_embeddings()
+    out = model.get_output_embeddings()
+
+    # ⚠️ TIED WEIGHTS SHARE ONE STORAGE, so moving the input embedding moves
+    # `lm_head` with it -- silently, because they are the same Parameter object.
+    # `lm_head` runs a 152,064-way matmul every decode step and would then run it
+    # on the CPU. This model unties (`tie_word_embeddings: false`), but the
+    # documented fallback in config is "a 3B model", and every smaller Qwen2.5
+    # ties. That is a reachable path, so it refuses rather than degrades.
+    if out is not None and out.weight is emb.weight:
+        # Named from the LOADED model, not from config: they agree in production
+        # and diverge exactly when someone is testing an override, which is the
+        # one moment a wrong model name in the message costs real time.
+        raise RuntimeError(
+            f"{getattr(model, 'name_or_path', None) or C.LLM_MODEL_ID} ties its "
+            f"input and output embeddings, so moving the table would put lm_head "
+            f"on the CPU as well. Set PM_LLM_EMBED_DEVICE=cuda for this model.")
+
+    emb.to("cpu")
+    # The ids are built on the GPU by `generate`; the gather has to happen beside
+    # the table, and the activations have to come back.
+    emb.register_forward_pre_hook(lambda _m, args: (args[0].to("cpu"),) + args[1:])
+    emb.register_forward_hook(
+        lambda _m, _args, out: out.to(compute_device, non_blocking=True))
+    # Without this the allocator keeps the 1040 MiB and the driver still reports
+    # it used -- the saving would be real in torch's books and invisible in the
+    # only figure the gate reads.
+    torch.cuda.empty_cache()
+
+
 @dataclass(slots=True)
 class Generator:
     model: object
     tokenizer: object
     provenance: dict
+    #: Where the compute lives. NOT `model.device` -- with the embedding table on
+    #: the host that property reports the first parameter's device, which is the
+    #: host, and every input would follow it there. Defaults to None so an older
+    #: caller constructing a Generator positionally still works; `__call__` then
+    #: falls back to the old behaviour, which is correct whenever nothing moved.
+    device: object = None
     max_new_tokens: int = C.LLM_MAX_NEW_TOKENS
     latencies: list = field(default_factory=list)
 
@@ -82,7 +174,8 @@ def __call__(self, payload: dict) -> str:
         msgs = E.render_prompt(payload)
         enc = self.tokenizer.apply_chat_template(
             msgs, add_generation_prompt=True, tokenize=True,
-            return_dict=True, return_tensors="pt").to(self.model.device)
+            return_dict=True, return_tensors="pt").to(
+                self.device if self.device is not None else self.model.device)
         n_in = enc["input_ids"].shape[1]
         t0 = time.time()
         with torch.inference_mode():
@@ -96,6 +189,49 @@ def __call__(self, payload: dict) -> str:
         return self.tokenizer.decode(out[0][n_in:], skip_special_tokens=True).strip()
 
 
+def vram_status() -> dict:
+    """Free and total VRAM, from the driver rather than from CUDA's view of it.
+
+    ⚠️ `torch.cuda.mem_get_info()` reports what the CUDA runtime believes a new
+    allocation could get, which on this WDDM setup is not what the device
+    actually has free. Measured 2026-08-19 with the 7B already resident:
+
+        nvidia-smi   260 MiB free of 8151
+        torch       6759 MiB free of 8151
+
+    A 6.3 GB disagreement, and torch is the optimistic one. That is precisely how
+    a preflight passes and the load then segfaults -- `s19_generate` gates on
+    `vram_free_gb > 5.5`, which torch would have cleared with a quarter of a
+    gigabyte actually available. NVML talks to the driver; CUDA talks to its own
+    bookkeeping. Prefer the driver, and say which one answered.
+
+    Totals agree exactly and always did: 8151 MiB is 7.96 GiB is 8.55 decimal GB.
+    The card is an 8 GB card; every other figure in the notes was that same
+    number with its units mangled.
+    """
+    import shutil
+    import subprocess
+
+    exe = shutil.which("nvidia-smi")
+    if exe:
+        try:
+            raw = subprocess.run(
+                [exe, "--query-gpu=memory.total,memory.free",
+                 "--format=csv,noheader,nounits"],
+                capture_output=True, text=True, timeout=10, check=True).stdout
+            total_mib, free_mib = (int(v) for v in raw.splitlines()[0].split(","))
+            return {"free_mib": free_mib, "total_mib": total_mib, "source": "nvidia-smi"}
+        except Exception:  # noqa: BLE001
+            pass  # fall through -- an unreadable driver is not a reason to stop
+
+    import torch
+    free, total = torch.cuda.mem_get_info()
+    return {"free_mib": round(free / 1024**2), "total_mib": round(total / 1024**2),
+            # Named so it shows up in the report: this number is the optimistic
+            # one, and a reader should know the driver was not available.
+            "source": "torch.cuda.mem_get_info (nvidia-smi unavailable)"}
+
+
 def capability_check() -> dict:
     """The cheap check the stage runs every time: is the toolchain present and
     is there room?
@@ -130,9 +266,15 @@ def capability_check() -> dict:
             f"HF_HUB_CACHE) explicitly before importing transformers.")
 
     p = torch.cuda.get_device_properties(0)
-    free, total = torch.cuda.mem_get_info()
+    # From the driver, not from CUDA's bookkeeping -- see `vram_status`. The
+    # `_gb` keys keep their names and decimal-GB units because `s19_generate`
+    # gates on `vram_free_gb`; what changed is that they are now true.
+    vram = vram_status()
     out = {"gpu": p.name, "compute_capability": f"sm_{p.major}{p.minor}",
-           "vram_free_gb": round(free / 1e9, 2), "vram_total_gb": round(total / 1e9, 2),
+           "vram_free_gb": round(vram["free_mib"] * 1024**2 / 1e9, 2),
+           "vram_total_gb": round(vram["total_mib"] * 1024**2 / 1e9, 2),
+           "vram_free_mib": vram["free_mib"], "vram_total_mib": vram["total_mib"],
+           "vram_source": vram["source"],
            "hub_cache": str(hub),
            "quantisation": C.LLM_QUANT, "torch": torch.__version__}
     if C.LLM_QUANT != "none":
@@ -187,14 +329,47 @@ def load_generator() -> Generator:
         C.LLM_MODEL_ID, revision=C.LLM_REVISION,
         quantization_config=_quant_config(C.LLM_QUANT), device_map={"": 0})
     model.eval()
+
+    # Read BEFORE anything moves: every parameter is still on the card here, so
+    # this is unambiguous. After `_host_embedding` it would not be.
+    compute_device = model.device
+
+    if C.LLM_EMBED_DEVICE == "cpu":
+        _host_embedding(model, compute_device)
+    elif C.LLM_EMBED_DEVICE != "cuda":
+        raise ValueError(
+            f"unknown LLM_EMBED_DEVICE {C.LLM_EMBED_DEVICE!r}; expected 'cpu' or 'cuda'")
+
+    # WHERE THE EMBEDDING ACTUALLY LANDED, asserted rather than assumed.
+    #
+    # A move that silently failed produces a model that loads, generates
+    # correctly, and quietly holds the 1040 MiB this setting exists to free. A
+    # saving that did not happen must not look like one that did.
+    embed = model.get_input_embeddings().weight
+    want = "cpu" if C.LLM_EMBED_DEVICE == "cpu" else "cuda"
+    if embed.device.type != want:
+        raise RuntimeError(
+            f"LLM_EMBED_DEVICE={C.LLM_EMBED_DEVICE} but embed_tokens is on "
+            f"{embed.device}; the move did not take.")
     gpu = torch.cuda.get_device_properties(0)
 
     prov = {
         "model": C.LLM_MODEL_ID,
         "revision": C.LLM_REVISION,
         "quantisation": C.LLM_QUANT,
         "dtype": str(model.dtype),
-        "device": str(model.device),
+        # The COMPUTE device, read before the embedding moved. `model.device`
+        # would now name the host and misreport where this model actually runs.
+        "device": str(compute_device),
+        # Part of model identity for the same reason `quantisation` is: it says
+        # which machine configuration wrote this line. It is deliberately NOT in
+        # `FP_GENERATE` -- an embedding lookup is a gather, and bf16 round-trips
+        # through the host exactly, so moving it is provably output-neutral and
+        # fingerprinting it would force a re-run that changes nothing. That
+        # argument rests on a hash: `tools/vram_probe.py` reports one per run and
+        # the two placements must agree. If they ever stop agreeing, this is no
+        # longer a placement detail and belongs in the fingerprint.
+        "embed_device": str(model.get_input_embeddings().weight.device),
         "compute_capability": f"sm_{gpu.major}{gpu.minor}",
         "max_new_tokens": C.LLM_MAX_NEW_TOKENS,
         "seed": C.LLM_SEED,
@@ -205,4 +380,5 @@ def load_generator() -> Generator:
     if C.LLM_QUANT != "none":
         import bitsandbytes
         prov["bitsandbytes"] = bitsandbytes.__version__
-    return Generator(model=model, tokenizer=tok, provenance=prov)
+    return Generator(model=model, tokenizer=tok, provenance=prov,
+                     device=compute_device)

pipeline/stages/s19_generate.py

@@ -145,10 +145,19 @@ def _run(sample: int | None = None, with_evidence: bool = True) -> None:
         # context and fragmented segments behind, and the 7B load then fails on
         # fragmentation rather than capacity -- which is exactly what happened.
         # The full probe is `--preflight`, a separate command and process.
-        assert pf["vram_free_gb"] > 5.5, (
-            f"only {pf['vram_free_gb']:.2f} GB VRAM free; a 7B in NF4 needs a "
-            f"~4.8 GB block plus activations. Close whatever is holding the GPU, "
-            f"or set PM_LLM_QUANT=none with a 3B model.")
+        #
+        # MEASURED, in the same unit the driver reports. The old form asserted
+        # `vram_free_gb > 5.5`, written when the figure came from
+        # `torch.cuda.mem_get_info()` and was optimistic by gigabytes; once
+        # `vram_status()` started telling the truth, 5.5 GB (5245 MiB) sat well
+        # BELOW a level that had already segfaulted, so the gate protected
+        # nothing. The peak is now 6239 MiB across three runs -- see
+        # `tools/vram_probe.py` and `reports/tool_vram_probe.json`.
+        assert pf["vram_free_mib"] >= C.LLM_MIN_FREE_VRAM_MIB, (
+            f"only {pf['vram_free_mib']} MiB VRAM free of "
+            f"{pf['vram_total_mib']}; the 7B peaks at ~6239 MiB and this stage "
+            f"needs {C.LLM_MIN_FREE_VRAM_MIB}. Close whatever is holding the GPU, or "
+            f"set PM_LLM_QUANT=none with a 3B model.")
 
     with stage(f"Load {C.LLM_MODEL_ID}"):
         gen = G.load_generator()

pipeline/tools/check_env.py

@@ -54,9 +54,12 @@ def main() -> None:
         if ok:
             name = torch.cuda.get_device_name(0)
             cap = torch.cuda.get_device_capability(0)
-            free, total = torch.cuda.mem_get_info(0)
+            # Driver figures, not CUDA's -- torch over-reports free VRAM on this
+            # WDDM setup by gigabytes. See pipeline/core/generate.vram_status.
+            from pipeline.core.generate import vram_status
+            v = vram_status()
             log(f"[green]PyTorch  CUDA : OK[/green] {name} sm_{cap[0]}{cap[1]} | "
-                f"{free/1e9:.1f} GB free of {total/1e9:.1f} GB")
+                f"{v['free_mib']} MiB free of {v['total_mib']} ({v['source']})")
             t = torch.randn(4096, 4096, device="cuda")
             torch.cuda.synchronize()
             log(f"  matmul check: {float((t @ t).sum()):.1f}")