open-interpretability 0.0.20__py3-none-any.whl

openinterp/harness/manifold.py

@@ -0,0 +1,200 @@
+"""Measure the multi-dimensional feature check (:class:`~openinterp.core.manifold.ManifoldCheck`).
+
+The axes are the audited pole and its top partner poles. Each axis is divided by its
+peak over every token of the top examples. So a partner with large activations doesn't
+fill the first component. Unit variance would be worse: it would blow a partner that is
+silent on a family, with tiny jitter, up into a full axis.
+
+* **The activating set.** Each top example is run again, and the poles are read at the
+  audited pole's peak.
+* **The families.** Each word goes into every carrier sentence of
+  :data:`openinterp.harness.specificity.CARRIERS`. A pole's strength on the word is its
+  peak over the word's tokens, as a mean over the carriers.
+
+A circle's score follows Engels et al. (2024). Their first component is often an
+intensity direction and the circle lies in components 2 and 3. So the score takes the
+better of the planes 1–2 and 2–3. The p-value repeats that choice on each shuffled
+order, so it pays for the choice.
+"""
+
+from __future__ import annotations
+
+import cmath
+import math
+from collections.abc import Sequence
+
+import torch
+
+from openinterp.core.manifold import FamilyManifold, ManifoldCheck, ManifoldPole
+from openinterp.core.schemas import ACTIVATION_THRESHOLD, ActivatingExample, CorrelatedNeuron
+from openinterp.core.stats import spearman
+from openinterp.harness.hooks import NeuronHookRunner, pole_strength
+from openinterp.harness.specificity import CARRIERS
+
+PARTNERS = 5  # partner poles beside the audited pole
+SHOWN = 3  # principal components whose explained variance is recorded
+SHUFFLES = 999  # shuffled orders for each p-value
+FAMILIES: dict[str, tuple[str, tuple[str, ...]]] = {
+    "weekdays": ("circle", ("Monday", "Tuesday", "Wednesday", "Thursday", "Friday",
+                            "Saturday", "Sunday")),
+    "months": ("circle", ("January", "February", "March", "April", "May", "June", "July",
+                          "August", "September", "October", "November", "December")),
+    "numbers": ("line", ("one", "two", "three", "four", "five", "six", "seven", "eight",
+                         "nine", "ten", "eleven", "twelve")),
+}  # fmt: skip
+
+
+def axes(neuron_idx: int, pole, partners: Sequence[CorrelatedNeuron]) -> list[ManifoldPole]:
+    """The audited pole, then its top :data:`PARTNERS` partner poles."""
+    own = ManifoldPole(neuron_idx=neuron_idx, pole=pole)
+    return [own, *(ManifoldPole(neuron_idx=c.neuron_idx, pole=c.pole) for c in partners)][
+        : 1 + PARTNERS
+    ]
+
+
+def _strengths(acts: torch.Tensor, poles: Sequence[ManifoldPole]) -> torch.Tensor:
+    """``[seq, len(poles)]``: each pole's strength at each token, from layer ``acts``."""
+    columns = [pole_strength(acts[:, p.neuron_idx], p.pole) for p in poles]
+    return torch.stack(columns, dim=1).float().cpu()
+
+
+def _layer(runner: NeuronHookRunner, text: str, layer: int) -> torch.Tensor:
+    """``[seq, width]``: the MLP hidden activations of ``layer`` on ``text``."""
+    return runner.layer_activations(text, [layer])[layer][0]
+
+
+def word_vector(runner, layer: int, poles, word: str) -> torch.Tensor | None:
+    """Each pole's peak over ``word``'s tokens, as a mean over the carriers.
+
+    ``None`` when no carrier measures a token of the word.
+    """
+    found = []
+    for carrier in CARRIERS:
+        text, start = carrier.format(word), carrier.index("{}")
+        offsets = runner.tokenizer(text, return_offsets_mapping=True)["offset_mapping"]
+        values = _strengths(_layer(runner, text, layer), poles)
+        held = [i for i, (a, b) in enumerate(offsets) if a < start + len(word) and b > start]
+        if len(offsets) == len(values) and held:
+            found.append(values[held].max(dim=0).values)
+    return torch.stack(found).mean(dim=0) if found else None
+
+
+def peak_vector(runner, layer: int, poles, example: ActivatingExample, skip_first: bool):
+    """The poles at the audited pole's peak, and each pole's peak over the example.
+
+    The example's text is run again.
+    """
+    values = _strengths(_layer(runner, example.text, layer), poles)
+    start = int(skip_first and len(values) > 1)
+    return values[start + int(values[start:, 0].argmax())], values[start:].max(dim=0).values
+
+
+def scaled(rows: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    """``rows`` with each pole divided by its scale; a pole with no scale is 0."""
+    return torch.where(scales > 0, rows / scales.clamp_min(1e-12), torch.zeros_like(rows))
+
+
+def pca(rows: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """The coordinates, the variances and the unit components of the centred ``rows``."""
+    z = rows.double() - rows.double().mean(dim=0)
+    u, s, vt = torch.linalg.svd(z, full_matrices=False)
+    return u * s, s**2, vt
+
+
+def ratios(variances: torch.Tensor) -> list[float]:
+    """The first :data:`SHOWN` shares of variance, or none when nothing varies."""
+    total = float(variances.sum())
+    return [round(float(v) / total, 4) for v in variances[:SHOWN]] if total > 0 else []
+
+
+def participation_ratio(variances: torch.Tensor) -> float:
+    """``(Σλ)² / Σλ²``: the effective number of dimensions; 0 when nothing varies."""
+    squares = float((variances**2).sum())
+    return round(float(variances.sum()) ** 2 / squares, 4) if squares > 0 else 0.0
+
+
+def circle_score(coords: torch.Tensor, places: Sequence[int], n: int) -> float:
+    """The mean resultant length of angle minus cycle place, in the better direction."""
+    angles = [math.atan2(float(y), float(x)) for x, y in coords]
+    cycle = [2 * math.pi * k / n for k in places]
+    forward = abs(sum(cmath.exp(1j * (a - c)) for a, c in zip(angles, cycle, strict=True)))
+    backward = abs(sum(cmath.exp(1j * (a + c)) for a, c in zip(angles, cycle, strict=True)))
+    return max(forward, backward) / len(angles)
+
+
+def _best(shape: str, coords: torch.Tensor, order: Sequence[int]) -> tuple[float, list[int]]:
+    """The best score over the allowed components, and those components (1-based).
+
+    On a tie the first components win. With no components the score is 0.
+    """
+    if shape == "circle":
+        planes = [[0, 1], [1, 2]][: max(coords.shape[1] - 1, 0)]
+        found = [(circle_score(coords[:, p], order, len(order)), p) for p in planes]
+    else:
+        values = torch.tensor(order, dtype=torch.float64)
+        found = [
+            (abs(spearman(coords[:, k], values)), [k]) for k in range(min(SHOWN, coords.shape[1]))
+        ]
+    score, picked = max(found, key=lambda f: f[0], default=(0.0, []))
+    return score, [k + 1 for k in picked]
+
+
+def family_manifold(
+    family: str, rows: torch.Tensor, threshold: float, scales: torch.Tensor | None = None
+) -> FamilyManifold:
+    """The PCA and shape of one family's ``rows`` (members in order, poles as columns).
+
+    ``threshold`` is in the audited pole's units. The PCA sees ``rows`` divided by
+    ``scales`` (:func:`scaled`), when given.
+    """
+    shape = FAMILIES[family][0]
+    coords, variances, vt = pca(rows if scales is None else scaled(rows, scales))
+    order = list(range(len(rows)))
+    if not float(variances.sum()) > 0:  # nothing varies: no shape and no components
+        coords = coords[:, :0]
+    score, components = _best(shape, coords, order)
+    rng = torch.Generator().manual_seed(len(rows))  # fixed: a record always replays
+    beaten = sum(
+        _best(shape, coords, torch.randperm(len(rows), generator=rng).tolist())[0] >= score - 1e-9
+        for _ in range(SHUFFLES)
+    )
+    share = float((vt[[c - 1 for c in components], 0] ** 2).sum()) if components else 0.0
+    plane = sorted((float(variances[c - 1]) for c in components), reverse=True)
+    round_ = plane[1] / plane[0] if shape == "circle" and len(plane) == 2 and plane[0] > 0 else None
+    return FamilyManifold(
+        family=family, shape=shape, members=len(rows),
+        active_members=int((rows[:, 0] >= threshold).sum()),
+        explained_variance=ratios(variances), components=components,
+        score=round(min(score, 1.0), 4), p_value=round((1 + beaten) / (1 + SHUFFLES), 4),
+        pole_share=round(min(share, 1.0), 4),
+        roundness=None if round_ is None else round(round_, 4),
+    )  # fmt: skip
+
+
+def manifold_check(
+    runner: NeuronHookRunner,
+    layer: int,
+    poles: list[ManifoldPole],
+    examples: Sequence[ActivatingExample],
+    skip_first: bool = False,
+) -> ManifoldCheck:
+    """Measure the check for ``poles[0]`` and its partners in ``layer``."""
+    max_act = max((e.max_activation for e in examples), default=0.0)
+    measured = [peak_vector(runner, layer, poles, e, skip_first) for e in examples]
+    peaks = [p for p, _ in measured]
+    scales = torch.stack([s for _, s in measured]).max(dim=0).values if measured else None
+    variances = pca(scaled(torch.stack(peaks), scales))[1] if len(peaks) > 1 else torch.zeros(0)
+    families = []
+    for family, (_, words) in FAMILIES.items():
+        vectors = [word_vector(runner, layer, poles, w) for w in words]
+        if all(v is not None for v in vectors):
+            rows = torch.stack(vectors)
+            threshold = ACTIVATION_THRESHOLD * max_act
+            families.append(family_manifold(family, rows, threshold, scales))
+    return ManifoldCheck(
+        poles=poles,
+        activating_examples=len(peaks),
+        activating_explained_variance=ratios(variances),
+        participation_ratio=participation_ratio(variances),
+        families=families,
+    )

openinterp/harness/memory.py

@@ -0,0 +1,471 @@
+"""Streaming mode: run a model larger than this machine's memory, slowly.
+
+In streaming mode the decoder layers stay in the checkpoint's safetensors files on disk.
+``accelerate`` reads a layer's weights when the forward pass reaches it, and frees them
+after. The embeddings, the final norm and the output head stay in memory. Peak memory is
+then about one layer, those modules and one batch's activations. The cost is time: every
+forward pass reads its layers from disk again.
+
+It is never the default for a model that fits. ``--memory auto`` (the default) streams
+only when the weights at the chosen dtype, plus :data:`WORKING_MARGIN_BYTES`, exceed free
+RAM plus free swap. ``--memory stream`` forces it, and ``--memory normal`` never streams.
+The maths is the same, so a streamed record reproduces in normal mode.
+
+A streamed scan reads its layers once per batch, so it uses larger batches
+(:func:`stream_batch_tokens`) to read each layer fewer times.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator, Mapping
+from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import click
+import torch
+from torch import nn
+
+MemoryMode = Literal["auto", "normal", "stream"]
+MEMORY_MODES: tuple[MemoryMode, ...] = ("auto", "normal", "stream")
+
+# What a scan or review holds beyond the weights, measured on CPU in float32. It is about
+# 2.5–3 GB for GPT-2 and Qwen2.5-0.5B and 1.5B (docs/experiments.md, Memory).
+WORKING_MARGIN_BYTES = 3 * 2**30
+
+# A streamed scan's batch: as many tokens as its activations fit in, within these bounds.
+# The floor is the normal scan's batch (harness/scan.py, BATCH_TOKENS).
+STREAM_BATCH_TOKENS_MAX = 65536
+
+# The most rows a streamed pass of single prompts runs at once (:func:`stream_rows`).
+STREAM_ROWS_MAX = 64
+
+MEMORY_ENV = "OPENINTERP_MEMORY_GB"  # the memory to plan for, in place of what is free
+
+_requested: MemoryMode = "auto"
+
+
+def request(mode: MemoryMode) -> None:
+    """Set the mode that the next model load uses (the ``--memory`` option)."""
+    global _requested
+    _requested = mode
+
+
+def requested() -> MemoryMode:
+    """The mode that ``--memory`` asked for."""
+    return _requested
+
+
+@dataclass(frozen=True)
+class MemoryPlan:
+    """Whether to stream a model, and the numbers that decided it."""
+
+    stream: bool
+    need_bytes: int | None  # weights at the chosen dtype, plus the working margin
+    available_bytes: int | None  # free RAM plus free swap
+    reason: str
+
+
+def available_bytes(device: str | torch.device | None = None) -> int | None:
+    """Free memory to plan for on ``device``, or ``None`` when it can't be read.
+
+    On a CUDA device that is the GPU's free memory. Elsewhere it is free RAM plus free
+    swap. ``device`` defaults to the one the last model loaded on. ``$OPENINTERP_MEMORY_GB``
+    overrides it, to run as if only that much were free: on a shared machine, or to
+    test streaming on a large one.
+    """
+    budget = os.environ.get(MEMORY_ENV)
+    if budget:
+        return int(float(budget) * 2**30)
+    device = torch.device(device if device is not None else _budget_device)
+    try:
+        if device.type == "cuda":
+            return int(torch.cuda.mem_get_info(device)[0])
+        import psutil
+
+        return int(psutil.virtual_memory().available + psutil.swap_memory().free)
+    except Exception:  # psutil missing, an unsupported platform, or no CUDA
+        return None
+
+
+_budget_device = "cpu"  # where the last model was loaded: what available_bytes plans for
+
+
+def checkpoint_bytes(snapshot: Path) -> int:
+    """Total size of the safetensors (else PyTorch) weight files in a model snapshot."""
+    files = list(snapshot.glob("*.safetensors")) or list(snapshot.glob("*.bin"))
+    return sum(f.stat().st_size for f in files)
+
+
+def weights_bytes(snapshot: Path, config, dtype: torch.dtype | str) -> int:
+    """The weights' size in memory at ``dtype``: the checkpoint scaled from its own dtype."""
+    stored = getattr(config, "torch_dtype", None) or getattr(config, "dtype", None)
+    size = checkpoint_bytes(snapshot)
+    if dtype == "auto" or stored is None:
+        return size
+    stored = getattr(torch, stored) if isinstance(stored, str) else stored
+    return size * _itemsize(dtype) // _itemsize(stored)
+
+
+def plan(
+    mode: MemoryMode, need: int | None, have: int | None, where: str = "RAM and swap"
+) -> MemoryPlan:
+    """Decide whether to stream: always for ``stream``, never for ``normal``.
+
+    For ``auto``, stream only when ``need`` is known to exceed ``have``, the free
+    memory ``where`` names.
+    """
+    if mode == "stream":
+        return MemoryPlan(True, need, have, "--memory stream")
+    if mode == "normal" or need is None or have is None:
+        return MemoryPlan(False, need, have, f"--memory {mode}")
+    if need <= have:
+        return MemoryPlan(False, need, have, "the model fits")
+    reason = f"it needs about {_gb(need)} and {_gb(have)} is free ({where})"
+    return MemoryPlan(True, need, have, reason)
+
+
+def warn(memory: MemoryPlan, model_name: str) -> None:
+    """Say on stderr that ``model_name`` streams from disk, and why."""
+    click.echo(
+        f"Streaming {model_name} from disk, because {memory.reason}. This is much "
+        "slower: every forward pass reads its layers from disk. Pass --memory normal "
+        "to load it whole.",
+        err=True,
+    )
+
+
+def load_causal_lm(model_name: str, device: str, dtype, revision: str | None, **extra):
+    """Load ``model_name`` whole on ``device``, or streamed from disk (see the module doc).
+
+    A streamed model runs on the CPU. Its offloaded layers read the checkpoint files in
+    place, so streaming writes no copy of them. A mixture-of-experts checkpoint that
+    stores each expert apart is the exception. transformers stacks its experts as it
+    loads, and streaming writes the stacks to :func:`offload_dir`.
+    """
+    from transformers import AutoModelForCausalLM
+
+    from openinterp.harness.moe import use_openinterp_experts
+
+    global _budget_device
+    memory = _plan_load(model_name, dtype, revision, device)
+    if not memory.stream:
+        kwargs = {"dtype": dtype, "revision": revision, **extra}
+        model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs).to(device)
+        use_openinterp_experts(model)
+        _budget_device = device
+        return model
+    warn(memory, model_name)
+    resident = stream_device(device)
+    if resident != device:
+        click.echo(f"Streaming runs on the CPU, not {device}.", err=True)
+    empty = _empty_model(model_name, revision)
+    device_map = stream_device_map(empty, _layers_path(empty, model_name), resident)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        dtype=dtype,
+        revision=revision,
+        device_map=device_map,
+        offload_folder=str(offload_dir()),
+        **extra,
+    )
+    keep_checkpoint_open(model, resident)
+    use_openinterp_experts(model)  # reads only the routed experts (harness/moe.py)
+    _budget_device = resident
+    return model
+
+
+def stream_device(device: str) -> str:
+    """Where a streamed model runs: ``device`` if it is CUDA, else the CPU.
+
+    On CUDA the embeddings, the final norm and the head stay on the GPU. Each layer is
+    read from disk and copied to the GPU when a pass reaches it. ``from_pretrained``
+    with a device map on MPS ends in a segmentation fault (transformers 5.20), so MPS
+    streams on the CPU.
+    """
+    return device if torch.device(device).type == "cuda" else "cpu"
+
+
+class OpenCheckpoint(Mapping):
+    """``accelerate``'s offloaded weights, read from safetensors files that stay open.
+
+    ``accelerate`` opens the checkpoint file again for each tensor it reads. On
+    Qwen2.5-0.5B that took about 19 ms a tensor, and a pass reads 288 tensors. So it was
+    5.6 s of a 5.9 s pass over 16 tokens. With each file opened once, the pass took
+    0.18 s, and 0.14 s whole. The tensors read are the same. An open file is
+    memory-mapped. The pages it reads are in the file cache, which the system can free.
+    They are not the process's own memory.
+    """
+
+    def __init__(self, loader, device: str = "cpu") -> None:
+        self.loader = loader  # accelerate's OffloadedWeightsLoader
+        self.files: dict[str, object] = {}
+        self.device = torch.device(device)  # where the layers run: tensors are sent there
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        info = self.loader.index.get(key)
+        if key in self.loader.state_dict or not info or not info.get("safetensors_file"):
+            return self.loader[key]
+        path = info["safetensors_file"]
+        if path not in self.files:
+            from safetensors import safe_open
+
+            self.files[path] = safe_open(path, framework="pt", device="cpu")
+        tensor = self.files[path].get_tensor(info.get("weight_name", key))
+        return self._sent(tensor, info)
+
+    def slice(self, key: str, index: int) -> torch.Tensor:
+        """Row ``index`` of the stacked tensor ``key``, read without reading the rest.
+
+        A mixture-of-experts layer stacks its experts in one tensor. A pass reads only
+        the experts its tokens are routed to (:mod:`openinterp.harness.moe`).
+        """
+        info = self.loader.index.get(key)
+        if key in self.loader.state_dict or not info or not info.get("safetensors_file"):
+            return self.loader[key][index]
+        path = info["safetensors_file"]
+        if path not in self.files:
+            from safetensors import safe_open
+
+            self.files[path] = safe_open(path, framework="pt", device="cpu")
+        row = self.files[path].get_slice(info.get("weight_name", key))[index]
+        return self._sent(row, info)
+
+    def _sent(self, tensor: torch.Tensor, info: dict) -> torch.Tensor:
+        """``tensor`` on the layers' device, then at its dtype.
+
+        On a GPU it is sent in the checkpoint's dtype and converted there. Take a float32
+        run of Qwen2.5-0.5B, whose checkpoint is bfloat16. A streamed pass then took 0.66 s
+        on an RTX 4070 Ti, against 2.8 s converted on the CPU first.
+        """
+        tensor = tensor.to(self.device)
+        return tensor.to(getattr(torch, info["dtype"])) if "dtype" in info else tensor
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.loader)
+
+    def __len__(self) -> int:
+        return len(self.loader)
+
+
+def keep_checkpoint_open(model: nn.Module, device: str = "cpu") -> None:
+    """Make every streamed module of ``model`` read through one :class:`OpenCheckpoint`.
+
+    ``device`` is where the layers run (:func:`stream_device`).
+    """
+    from accelerate.utils import PrefixedDataset
+
+    opened: dict[int, OpenCheckpoint] = {}
+    for module in model.modules():
+        weights = getattr(getattr(module, "_hf_hook", None), "weights_map", None)
+        if isinstance(weights, PrefixedDataset) and hasattr(weights.dataset, "index"):
+            key = id(weights.dataset)
+            weights.dataset = opened.setdefault(key, OpenCheckpoint(weights.dataset, device))
+
+
+def offload_dir() -> Path:
+    """Where ``accelerate`` may write weights it can't read in place (a ``.bin`` checkpoint).
+
+    ``$XDG_CACHE_HOME/openinterp/offload``, else ``~/.cache/openinterp/offload``. It stays
+    empty for a safetensors checkpoint.
+    """
+    root = os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache"
+    return Path(root) / "openinterp" / "offload"
+
+
+def _plan_load(model_name: str, dtype, revision: str | None, device: str = "cpu") -> MemoryPlan:
+    """The streaming decision for loading ``model_name`` at ``dtype`` on ``device``."""
+    mode = requested()
+    if mode == "normal":
+        return plan(mode, None, None)
+    need = None
+    # Offline, or an unknown layout: don't stream unless asked.
+    with suppress(Exception):
+        snapshot = _snapshot(model_name, revision)
+        from transformers import AutoConfig
+
+        config = AutoConfig.from_pretrained(model_name, revision=revision)
+        need = weights_bytes(snapshot, config, dtype) + WORKING_MARGIN_BYTES
+    on_gpu = torch.device(device).type == "cuda"
+    return plan(mode, need, available_bytes(device), "on the GPU" if on_gpu else "RAM and swap")
+
+
+def _snapshot(model_name: str, revision: str | None) -> Path:
+    """The local folder that holds ``model_name``'s weights, downloading them if needed."""
+    if Path(model_name).is_dir():
+        return Path(model_name)
+    from huggingface_hub import snapshot_download
+
+    patterns = ["*.json", "*.safetensors"]
+    return Path(snapshot_download(model_name, revision=revision, allow_patterns=patterns))
+
+
+def _empty_model(model_name: str, revision: str | None) -> nn.Module:
+    """``model_name``'s architecture with every weight on the meta device (no memory)."""
+    from accelerate import init_empty_weights
+    from transformers import AutoConfig, AutoModelForCausalLM
+
+    config = AutoConfig.from_pretrained(model_name, revision=revision)
+    with init_empty_weights():
+        return AutoModelForCausalLM.from_config(config)
+
+
+def _layers_path(model: nn.Module, model_name: str) -> str:
+    """The dotted path of ``model``'s decoder layers."""
+    from openinterp.harness.hooks import LAYER_PATHS, _try_path
+
+    for path in LAYER_PATHS:
+        layers = _try_path(model, path)
+        if isinstance(layers, nn.ModuleList) and len(layers) > 0:
+            return path
+    from openinterp.core.exceptions import ArchitectureNotSupportedError
+
+    raise ArchitectureNotSupportedError(f"no decoder layers found on {model_name}")
+
+
+def stream_device_map(model: nn.Module, layers_path: str, resident: str = "cpu") -> dict[str, str]:
+    """An ``accelerate`` device map with each decoder layer on disk, the rest on ``resident``.
+
+    ``model`` may be on the meta device. ``layers_path`` names its decoder layers.
+    ``accelerate`` runs the layers on the first device in the map that isn't the CPU or
+    the disk. So they run on the GPU when ``resident`` is one.
+    """
+    device_map: dict[str, str] = {}
+
+    def assign(module: nn.Module, prefix: str) -> None:
+        for name, child in module.named_children():
+            path = f"{prefix}{name}"
+            if path == layers_path:
+                device_map.update({f"{path}.{i}": "disk" for i in range(len(child))})
+            elif layers_path.startswith(path + "."):
+                assign(child, path + ".")
+            else:
+                device_map[path] = resident
+
+    assign(model, "")
+    return device_map
+
+
+def stream_batch_tokens(config, dtype: torch.dtype, have: int | None, floor: int) -> int:
+    """A streamed scan's batch: the tokens whose activations fit in half of ``have``.
+
+    One token of one layer holds about 4 residual-width and 3 MLP-width vectors at once
+    (the residual, attention and MLP inputs and outputs). The result is between
+    ``floor`` and :data:`STREAM_BATCH_TOKENS_MAX`.
+    """
+    width = getattr(config, "hidden_size", None) or getattr(config, "n_embd", None)
+    inner = getattr(config, "intermediate_size", None) or getattr(config, "n_inner", None)
+    if not width or have is None:
+        return floor
+    per_token = _itemsize(dtype) * (4 * width + 3 * (inner or 4 * width))
+    tokens = (have // 2) // per_token
+    return int(max(floor, min(STREAM_BATCH_TOKENS_MAX, tokens)))
+
+
+def stream_rows(runner, tokens_per_row: int, default: int, *, logits: bool) -> int:
+    """Rows per forward pass: ``default``, or more when ``runner``'s model streams.
+
+    Every streamed pass reads its layers from disk again. So a streamed model runs as
+    many rows at once as a quarter of the planned memory holds (:func:`available_bytes`).
+    A pass that returns ``logits`` holds them and one log-softmax copy: 1.2 MB a token at
+    Qwen2.5's 151,936 tokens. Rows are independent, so only float rounding changes. The
+    result is at least ``default`` and at most :data:`STREAM_ROWS_MAX`.
+    """
+    have = available_bytes()
+    if not is_streamed(runner.model) or have is None:
+        return default
+    have //= _sharing  # passes that run at once share the memory (layer_major)
+    config = runner.model.config
+    dtype = next(runner.model.parameters()).dtype
+    width = getattr(config, "hidden_size", None) or getattr(config, "n_embd", 1)
+    inner = getattr(config, "intermediate_size", None) or getattr(config, "n_inner", None)
+    per_token = _itemsize(dtype) * (4 * width + 3 * (inner or 4 * width))
+    if logits:
+        per_token += 2 * 4 * getattr(config, "vocab_size", 0)  # float32 logits and log-probs
+    rows = (have // 4) // max(1, per_token * max(1, tokens_per_row))
+    return int(max(default, min(STREAM_ROWS_MAX, rows)))
+
+
+_sharing = 1  # how many passes run at once (openinterp.harness.layer_major)
+
+
+@contextmanager
+def shared_by(passes: int) -> Iterator[None]:
+    """Plan each pass's memory as a share of the budget while ``passes`` run at once."""
+    global _sharing
+    previous, _sharing = _sharing, max(1, passes)
+    try:
+        yield
+    finally:
+        _sharing = previous
+
+
+def loaded(module: nn.Module) -> AbstractContextManager:
+    """A context in which ``module``'s own weights are real tensors, also when streamed.
+
+    A streamed module keeps its weights on the meta device between passes. Read them
+    inside this context, and copy what must outlive it. Child modules are not loaded.
+    """
+    if getattr(module, "_hf_hook", None) is None:
+        return nullcontext()
+    from accelerate.utils import align_module_device
+
+    return align_module_device(module)
+
+
+@contextmanager
+def recompute_below(runner, layer: int) -> Iterator[None]:
+    """In a streamed model, recompute the layers below ``layer`` in a backward pass.
+
+    Autograd keeps the weights a layer used until the backward pass. So a gradient
+    through streamed layers holds all of them at once. Inside this block each layer
+    below ``layer`` keeps only its input. The backward pass runs it again, which reads
+    its weights from disk again. The values are the same. Take Qwen2.5-1.5B at layer
+    11, with frozen weights (``cli._load``). Integrated gradients over 1,024 tokens
+    peaked at 1.8 GB with this, and at 4.5 GB without. It took 5.2 s, not 4.2 s.
+
+    ``layer`` itself isn't recomputed, since a hook there may end the forward pass. The
+    key-value cache is off, since a recomputed layer would add to it twice. A model that
+    isn't streamed is left as it is.
+    """
+    if not is_streamed(runner.model):
+        yield
+        return
+    from torch.utils.checkpoint import checkpoint
+
+    def recomputed(forward):
+        def run(*args, **kwargs):
+            if not torch.is_grad_enabled():
+                return forward(*args, **kwargs)
+            return checkpoint(forward, *args, use_reentrant=False, **kwargs)
+
+        return run
+
+    config = runner.model.config
+    use_cache, config.use_cache = getattr(config, "use_cache", None), False
+    blocks = list(runner.layers)[:layer]
+    forwards = [block.forward for block in blocks]
+    for block, forward in zip(blocks, forwards, strict=True):
+        block.forward = recomputed(forward)
+    try:
+        yield
+    finally:
+        for block, forward in zip(blocks, forwards, strict=True):
+            block.forward = forward
+        config.use_cache = use_cache
+
+
+def is_streamed(model: nn.Module) -> bool:
+    """Whether some of ``model``'s weights stay on disk between passes."""
+    return any(p.device.type == "meta" for p in model.parameters())
+
+
+def _itemsize(dtype: torch.dtype | str) -> int:
+    dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
+    return torch.empty((), dtype=dtype).element_size()
+
+
+def _gb(n: int) -> str:
+    return f"{n / 2**30:.1f} GB"