logogram 0.1.0__py3-none-any.whl

logogram/steering.py

@@ -0,0 +1,258 @@
+"""Steering: add a direction to the residual stream at several strengths, and measure the effect.
+
+At each steered site (a residual stream state at one token), the direction is the mean of
+(other prompt - receiver prompt) over a training split of the pairs. Adding it at strength 1 to a
+held-out receiver prompt moves that prompt's activation by the average difference between the two
+prompts; the effect is normalized like patching, so 1 means the output moved as far as switching
+to the other prompt. Directions are never applied to the prompts that computed them. A random
+direction of the same length, at the same strengths, is the control: a direction that matters
+should beat it.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import torch
+
+from logogram.backends.base import Cancelled, ModelBackend, ModelInfo, Patch
+from logogram.engine import (
+    EngineError,
+    EngineResult,
+    LayerFn,
+    ProgressFn,
+    _answer_tensors,
+    _chunks,
+    _LayerSources,
+    _metric,
+    check_gap,
+    compute_baselines,
+)
+from logogram.prompts import PreparedPrompt, group_by_length
+from logogram.sites import (
+    ResolvedSite,
+    ScopeError,
+    _check_site,
+    resolve_position,
+    site_label,
+)
+from logogram.spec import (
+    AllPositions,
+    LayerComponentsScope,
+    Site,
+    SitesScope,
+    Spec,
+    Steering,
+    strength_text,
+)
+
+RESIDUAL = ("resid_pre", "resid_mid", "resid_post")
+
+
+def split_pairs(n: int, train_fraction: float, seed: int) -> tuple[list[int], list[int]]:
+    """Positions of the pairs that train the direction and of the held-out pairs it is measured
+    on: a seeded shuffle from the raw PCG64 stream, so the split is the same on every machine."""
+    n_train = int(round(n * train_fraction))
+    if n_train < 1 or n - n_train < 2:
+        raise EngineError(
+            f"Steering needs at least one prompt pair to compute the direction and two to measure "
+            f"it on, but {n} pairs with a training share of {train_fraction:.0%} leaves "
+            f"{n_train} and {n - n_train}. Use more prompts or another training share."
+        )
+    order = list(range(n))
+    raw = np.random.PCG64(seed).random_raw(n)
+    for t in range(n - 1, 0, -1):
+        j = int(raw[t] % np.uint64(t + 1))
+        order[t], order[j] = order[j], order[t]
+    return sorted(order[:n_train]), sorted(order[n_train:])
+
+
+def steered_sites(spec: Spec, info: ModelInfo, prompts: list[PreparedPrompt]) -> list[Site]:
+    """The residual-stream sites a steering spec adds its direction at, one token per prompt."""
+    scope = spec.scope
+    if isinstance(scope, LayerComponentsScope):
+        if len(scope.components) != 1 or scope.components[0] not in RESIDUAL:
+            raise ScopeError(
+                "Steering adds a direction to the residual stream: choose exactly one residual "
+                "site (before, between or after the layer's blocks)."
+            )
+        sites = [
+            Site(kind=scope.components[0], layer=layer, position=scope.position)
+            for layer in range(info.n_layers)
+        ]
+    elif isinstance(scope, SitesScope):
+        sites = list(scope.sites)
+        if any(s.kind not in RESIDUAL for s in sites):
+            raise ScopeError(
+                "Steering adds a direction to the residual stream; heads, attention outputs and "
+                "MLP outputs can't be steered here. Choose residual sites."
+            )
+    else:
+        raise ScopeError(
+            "Steering runs at one residual site in every layer, or at chosen residual sites. "
+            "Choose “Attention and MLP per layer” with a single residual component, or pick sites."
+        )
+    for site in sites:
+        if isinstance(site.position, AllPositions):
+            raise ScopeError(
+                "Steering adds the direction at one token of each prompt. Choose the last token, a "
+                "named position or a token index."
+            )
+        _check_site(site, info)
+        for prompt in prompts:
+            resolve_position(site.position, prompt)
+    return sites
+
+
+def run_steering(
+    spec: Spec,
+    backend: ModelBackend,
+    prompts: list[PreparedPrompt],
+    *,
+    on_progress: ProgressFn | None = None,
+    on_layer: LayerFn | None = None,
+    cancel: threading.Event | None = None,
+    on_start: Callable[[list[ResolvedSite], dict[str, Any]], None] | None = None,
+) -> EngineResult:
+    exp = spec.experiment
+    assert isinstance(exp, Steering)
+    info = backend.info
+    batch_size = spec.execution.batch_size
+    base = steered_sites(spec, info, prompts)
+    train_at, test_at = split_pairs(len(prompts), exp.train_fraction, exp.seed)
+    train = [prompts[i] for i in train_at]
+    test = [prompts[i] for i in test_at]
+    receiver = exp.apply_to
+    reference = "corrupt" if receiver == "clean" else "clean"
+
+    # The variants of each site: every strength along the direction, then along the control.
+    variants = [(c, False) for c in exp.coefficients]
+    if exp.control:
+        variants += [(c, True) for c in exp.coefficients]
+    keys = [
+        f"random {strength_text(c)}" if control else strength_text(c) for c, control in variants
+    ]
+    uniform = len({(s.kind, s.position.model_dump_json()) for s in base}) == 1
+    sites: list[ResolvedSite] = []
+    for b, site in enumerate(base):
+        for v, ((coefficient, control), key) in enumerate(zip(variants, keys, strict=True)):
+            sites.append(
+                ResolvedSite(
+                    index=len(sites),
+                    site=site,
+                    row=b,
+                    col=v,
+                    label=f"{site_label(site)} {key}",
+                    variant={"coefficient": coefficient, "control": control},
+                    variant_key=key,
+                )
+            )
+    layout = {
+        "kind": "steering",
+        "site": base[0].kind if uniform else None,
+        "row_title": "Layer" if uniform else "Site",
+        "col_title": "Strength",
+        "rows": [
+            {"key": str(b), "label": str(s.layer) if uniform else site_label(s)}
+            for b, s in enumerate(base)
+        ],
+        "cols": [
+            {"key": key, "label": key, "coefficient": c, "control": control}
+            for (c, control), key in zip(variants, keys, strict=True)
+        ],
+    }
+    if on_start is not None:
+        on_start(sites, layout)
+
+    test_groups = group_by_length(test)
+    train_groups = group_by_length(train)
+    baselines = compute_baselines(backend, test, test_groups, batch_size, cancel)
+    warnings = check_gap(spec, baselines, test, receiver, reference)
+    n = len(test)
+    patched_ld = np.zeros((len(sites), n))
+    patched_prob = np.zeros((len(sites), n))
+    result = EngineResult(
+        sites=sites,
+        layout=layout,
+        prompts=test,
+        baselines=baselines,
+        receiver=receiver,
+        reference=reference,
+        patched_ld=patched_ld,
+        patched_prob=patched_prob,
+        warnings=warnings,
+        extra={
+            "steering": {
+                "train": [p.index for p in train],
+                "test": [p.index for p in test],
+                "norms": [],
+            }
+        },
+    )
+    model_dtype = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
+    dtype = model_dtype[info.dtype]
+    local_of = {i: li for g in test_groups for li, i in enumerate(g.members)}
+    total = len(sites) * n
+    done = 0
+    for b, site in enumerate(base):
+        probe = ResolvedSite(0, site, 0, 0, "")
+        kinds = [site.kind]
+        # The direction: mean (reference - receiver) over the training pairs, at the site's token.
+        toward = _LayerSources(
+            backend, train, train_groups, site.layer, kinds, reference, batch_size, cancel
+        )
+        away = _LayerSources(
+            backend, train, train_groups, site.layer, kinds, receiver, batch_size, cancel
+        )
+        direction = (
+            toward.table(site.kind, probe).double() - away.table(site.kind, probe).double()
+        ).mean(0)
+        norm = float(direction.norm())
+        result.extra["steering"]["norms"].append(norm)
+        control_direction = None
+        if exp.control:
+            generator = torch.Generator().manual_seed(exp.seed * 1_000_003 + b)
+            noise = torch.randn(direction.shape, generator=generator, dtype=torch.float64)
+            control_direction = noise / noise.norm() * norm
+        # The held-out receivers' own activations, to which the direction is added.
+        own = _LayerSources(
+            backend, test, test_groups, site.layer, kinds, receiver, batch_size, cancel
+        ).table(site.kind, probe)
+        positions = [resolve_position(site.position, p) for p in test]
+        first = b * len(variants)
+        for group in test_groups:
+            tokens = group.clean if receiver == "clean" else group.corrupt
+            rows = [(v, p) for v in range(len(variants)) for p in group.members]
+            for sl in _chunks(len(rows), batch_size):
+                if cancel is not None and cancel.is_set():
+                    raise Cancelled()
+                chunk = rows[sl]
+                values = torch.stack(
+                    [
+                        own[p].double()
+                        + variants[v][0]
+                        * (control_direction if variants[v][1] else direction).to(own.device)
+                        for v, p in chunk
+                    ]
+                ).to(dtype)
+                patch = Patch(
+                    kind=site.kind,
+                    layer=site.layer,
+                    values=values,
+                    positions=torch.tensor([positions[p] for _, p in chunk], dtype=torch.long),
+                )
+                local = torch.tensor([local_of[p] for _, p in chunk], dtype=torch.long)
+                logits = backend.final_logits(tokens[local], patch)
+                ld, prob = _metric(logits, *_answer_tensors(test, [p for _, p in chunk]))
+                for (v, p), value_ld, value_prob in zip(chunk, ld, prob, strict=True):
+                    patched_ld[first + v, p] = value_ld
+                    patched_prob[first + v, p] = value_prob
+                done += len(chunk)
+                if on_progress is not None:
+                    on_progress(done, total, site.layer)
+        if on_layer is not None:
+            on_layer(site.layer, list(range(first, first + len(variants))), result)
+    return result

logogram/system.py

@@ -0,0 +1,379 @@
+"""Hardware and environment report, with concrete fixes for common problems.
+
+Used by the first-run system check, by ``logogram doctor`` and by the memory estimate. Nothing
+here is written to disk or sent anywhere.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+import os
+import platform
+import shutil
+import subprocess
+import sys
+from dataclasses import asdict, dataclass, field
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+from typing import Any
+
+import psutil
+
+GB = 1024**3
+
+
+@dataclass
+class Issue:
+    severity: str  # "error" | "warning" | "info"
+    title: str
+    detail: str
+    fix: str | None = None  # a command or a short instruction
+
+
+@dataclass
+class GPU:
+    name: str
+    memory_total: int
+    memory_free: int | None
+    bf16: bool
+
+
+@dataclass
+class SystemReport:
+    os: str
+    machine: str
+    python: str
+    cpu: str
+    cpu_cores: int | None
+    cpu_threads: int | None
+    memory_total: int
+    memory_available: int
+    gpus: list[GPU]
+    backend: str  # cuda | mps | cpu
+    torch: str
+    torch_cuda: str | None
+    transformer_lens: str | None
+    transformers: str | None
+    recommended_dtype: str
+    precision_note: str
+    hf_cache_free: int | None
+    issues: list[Issue] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def _pkg(name: str) -> str | None:
+    try:
+        return version(name)
+    except PackageNotFoundError:
+        return None
+
+
+@functools.cache
+def cpu_name() -> str:
+    system = platform.system()
+    try:
+        if system == "Linux":
+            for line in Path("/proc/cpuinfo").read_text().splitlines():
+                if line.lower().startswith("model name"):
+                    return line.split(":", 1)[1].strip()
+        elif system == "Darwin":
+            out = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.brand_string"],
+                capture_output=True,
+                text=True,
+                timeout=2,
+            )
+            if out.stdout.strip():
+                return out.stdout.strip()
+        elif system == "Windows":
+            name = platform.processor()
+            if name:
+                return name
+    except Exception:  # noqa: BLE001
+        pass
+    return platform.processor() or platform.machine() or "Unknown CPU"
+
+
+def _os_label() -> str:
+    system = platform.system()
+    if system == "Darwin":
+        return f"macOS {platform.mac_ver()[0]}"
+    if system == "Windows":
+        return f"Windows {platform.release()}"
+    if system == "Linux":
+        try:
+            info = platform.freedesktop_os_release()
+            return info.get("PRETTY_NAME") or info.get("NAME") or "Linux"
+        except OSError:
+            return "Linux"
+    return system or "Unknown"
+
+
+def nvidia_driver_present() -> bool:
+    if Path("/proc/driver/nvidia/version").exists():
+        return True
+    return shutil.which("nvidia-smi") is not None
+
+
+def _rosetta() -> bool:
+    if platform.system() != "Darwin" or platform.machine() != "x86_64":
+        return False
+    try:
+        out = subprocess.run(
+            ["sysctl", "-n", "sysctl.proc_translated"], capture_output=True, text=True, timeout=2
+        )
+        return out.stdout.strip() == "1"
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def installed_as_uv_tool() -> bool:
+    return f"uv{os.sep}tools{os.sep}logogram" in sys.prefix.replace("/", os.sep)
+
+
+def reinstall_command(backend: str) -> str:
+    if installed_as_uv_tool():
+        return f"uv tool install --reinstall --torch-backend={backend} logogram"
+    return f"uv pip install --reinstall torch --torch-backend={backend}"
+
+
+def hf_cache_dir() -> Path:
+    from huggingface_hub import constants
+
+    return Path(constants.HF_HUB_CACHE)
+
+
+def system_report() -> SystemReport:
+    import torch
+
+    vm = psutil.virtual_memory()
+    gpus: list[GPU] = []
+    backend = "cpu"
+    if torch.cuda.is_available():
+        backend = "cuda"
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            free = None
+            with contextlib.suppress(Exception):
+                free, _ = torch.cuda.mem_get_info(i)
+            gpus.append(
+                GPU(
+                    name=props.name,
+                    memory_total=int(props.total_memory),
+                    memory_free=int(free) if free is not None else None,
+                    bf16=bool(torch.cuda.is_bf16_supported()),
+                )
+            )
+    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        backend = "mps"
+        gpus.append(
+            GPU(name="Apple GPU (Metal)", memory_total=int(vm.total), memory_free=None, bf16=True)
+        )
+
+    issues: list[Issue] = []
+    torch_cuda = torch.version.cuda
+    if backend != "cuda" and nvidia_driver_present():
+        if torch_cuda is None:
+            issues.append(
+                Issue(
+                    "warning",
+                    "Your NVIDIA GPU isn't being used",
+                    "This PyTorch build is CPU-only, so experiments run on the CPU. Reinstall "
+                    "PyTorch with CUDA support to use the GPU.",
+                    reinstall_command("auto"),
+                )
+            )
+        else:
+            issues.append(
+                Issue(
+                    "warning",
+                    "PyTorch can't reach your NVIDIA GPU",
+                    f"PyTorch was built for CUDA {torch_cuda}, but the NVIDIA driver can't run it. "
+                    "Update the NVIDIA driver, or install a PyTorch build for an older CUDA "
+                    "version (for example cu126).",
+                    reinstall_command("cu126"),
+                )
+            )
+    if platform.system() == "Darwin" and backend == "cpu":
+        if _rosetta():
+            issues.append(
+                Issue(
+                    "warning",
+                    "Python is running under Rosetta",
+                    "This is an Intel (x86_64) Python on Apple Silicon, so the GPU (MPS) can't be "
+                    "used. Install a native arm64 Python and reinstall Logogram with it.",
+                    "uv python install 3.12 && uv tool install --reinstall --python 3.12 logogram",
+                )
+            )
+        elif platform.machine() == "arm64":
+            issues.append(
+                Issue(
+                    "warning",
+                    "Apple GPU (MPS) unavailable",
+                    "PyTorch can't use Metal on this Mac. MPS needs macOS 12.3 or later.",
+                    "Update macOS, then reinstall PyTorch.",
+                )
+            )
+    if vm.total < 8 * GB:
+        issues.append(
+            Issue(
+                "warning",
+                "Less than 8 GB of memory",
+                "GPT-2 small fits, but larger models may not. Close other applications before "
+                "running sweeps.",
+            )
+        )
+    cache_free = None
+    try:
+        cache = hf_cache_dir()
+        probe = cache if cache.exists() else Path.home()
+        cache_free = shutil.disk_usage(probe).free
+        if cache_free < 2 * GB:
+            issues.append(
+                Issue(
+                    "warning",
+                    "Low disk space for model downloads",
+                    f"Only {cache_free / GB:.1f} GB is free where Hugging Face stores models. "
+                    "GPT-2 small needs about 0.6 GB.",
+                    "Free some space, or set HF_HOME to a disk with more room.",
+                )
+            )
+    except Exception:  # noqa: BLE001
+        pass
+
+    bf16 = any(g.bf16 for g in gpus)
+    note = "float32 keeps results exact; use it whenever the model fits."
+    if backend == "cuda" and bf16:
+        note += " bfloat16 is available on this GPU for models that don't fit in float32."
+    elif backend == "cpu":
+        note += " 16-bit types are slow on most CPUs."
+
+    try:
+        cores = psutil.cpu_count(logical=False)
+    except Exception:  # noqa: BLE001
+        cores = None
+    return SystemReport(
+        os=_os_label(),
+        machine=platform.machine(),
+        python=platform.python_version(),
+        cpu=cpu_name(),
+        cpu_cores=cores,
+        cpu_threads=psutil.cpu_count(logical=True),
+        memory_total=int(vm.total),
+        memory_available=int(vm.available),
+        gpus=gpus,
+        backend=backend,
+        torch=torch.__version__,
+        torch_cuda=torch_cuda,
+        transformer_lens=_pkg("transformer-lens"),
+        transformers=_pkg("transformers"),
+        recommended_dtype="float32",
+        precision_note=note,
+        hf_cache_free=cache_free,
+        issues=issues,
+    )
+
+
+# -- memory estimate ---------------------------------------------------------------------------
+
+DTYPE_BYTES = {"float32": 4, "float16": 2, "bfloat16": 2}
+
+
+@dataclass
+class MemoryEstimate:
+    device: str
+    dtype: str
+    n_params: int
+    weights: int
+    activations: int
+    margin: int
+    total: int
+    available: int
+    verdict: str  # fits | tight | wont_fit
+    explanation: str
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def available_memory(device: str) -> int:
+    import torch
+
+    if device == "cuda" and torch.cuda.is_available():
+        free, _ = torch.cuda.mem_get_info()
+        return int(free)
+    if device == "mps":
+        try:
+            return int(torch.mps.recommended_max_memory()) - int(
+                torch.mps.current_allocated_memory()
+            )
+        except Exception:  # noqa: BLE001
+            pass
+    return int(psutil.virtual_memory().available)
+
+
+def estimate_memory(
+    *,
+    n_params: int,
+    n_layers: int,
+    n_heads: int,
+    d_model: int,
+    d_mlp: int,
+    d_vocab: int,
+    dtype: str,
+    device: str,
+    batch_size: int = 64,
+    seq_len: int = 32,
+    loaded_bytes: int = 0,
+) -> MemoryEstimate:
+    """Weights + activations of one batch + a safety margin, against free memory.
+
+    ``loaded_bytes`` is memory held by a model that loading this one would replace.
+    """
+    b = DTYPE_BYTES[dtype]
+    weights = n_params * b
+    tokens = batch_size * seq_len
+    # One forward pass: residual-sized tensors per block, MLP hidden states, attention scores and
+    # patterns, final logits for the last position, plus one layer of cached source activations.
+    per_layer = tokens * (8 * d_model + 2 * d_mlp) * b + 3 * batch_size * n_heads * seq_len**2 * b
+    activations = per_layer * 2 + batch_size * d_vocab * 4 + 3 * tokens * d_model * b
+    margin = int(0.15 * (weights + activations)) + (
+        512 * 1024**2 if device == "cuda" else 256 * 1024**2
+    )
+    total = weights + activations + margin
+    available = available_memory(device) + loaded_bytes
+    if total <= 0.8 * available:
+        verdict = "fits"
+    elif total <= available:
+        verdict = "tight"
+    else:
+        verdict = "wont_fit"
+    where = {"cuda": "GPU memory", "mps": "unified memory", "cpu": "RAM"}.get(device, "memory")
+    explanation = (
+        f"{n_params / 1e6:,.0f}M stored values × {b} bytes ({dtype}) for weights; activations for a "
+        f"batch of {batch_size} prompts of {seq_len} tokens; a margin of 15% plus runtime "
+        f"overhead. Compared with free {where}."
+    )
+    return MemoryEstimate(
+        device=device,
+        dtype=dtype,
+        n_params=n_params,
+        weights=weights,
+        activations=activations,
+        margin=margin,
+        total=total,
+        available=available,
+        verdict=verdict,
+        explanation=explanation,
+    )
+
+
+def format_bytes(n: int | float | None) -> str:
+    if n is None:
+        return "—"
+    if n >= GB:
+        return f"{n / GB:.1f} GB"
+    return f"{n / 1024**2:.0f} MB"