fpwap 0.1.0__py3-none-any.whl

fpwap/__init__.py

@@ -0,0 +1,43 @@
+from fpwap.callbacks.base import Callback
+from fpwap.engine import (
+    PreloopTiming,
+    ProfileReport,
+    Result,
+    SetupTiming,
+    Sweep,
+    TeardownTiming,
+    estimate_max_microbatch,
+)
+from fpwap.extractor import Extractor
+from fpwap.preflight import PreflightReport
+from fpwap.types import (
+    Artifact,
+    ArtifactKey,
+    Context,
+    Emit,
+    LayerArtifact,
+    RaggedTensor,
+    ResultArtifact,
+    WriteBack,
+)
+
+__all__ = [
+    "Artifact",
+    "ArtifactKey",
+    "Callback",
+    "Context",
+    "Emit",
+    "Extractor",
+    "LayerArtifact",
+    "PreflightReport",
+    "PreloopTiming",
+    "ProfileReport",
+    "RaggedTensor",
+    "Result",
+    "ResultArtifact",
+    "SetupTiming",
+    "Sweep",
+    "TeardownTiming",
+    "WriteBack",
+    "estimate_max_microbatch",
+]

fpwap/buffer.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import Tensor
+
+
+class ResidualBuffer:
+    """Inter-layer transport for the fpwap loop.
+
+    Two modes:
+    - In-memory (path=None): pinned torch tensor. Fast async D2H via the CUDA
+      copy engine. Default for workloads that fit in host RAM.
+    - Disk-backed (path=<file>): numpy memmap. The OS page cache manages
+      residency; the full [N, seq, H] corpus never needs to fit in RAM at once.
+      bf16 is stored as uint16 bit-patterns (numpy has no bf16 dtype).
+    """
+
+    def __init__(
+        self,
+        n_samples: int,
+        seq_len: int,
+        hidden: int,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | str = "cpu",
+        path: Path | None = None,
+    ) -> None:
+        self.n_samples = n_samples
+        self.seq_len = seq_len
+        self.hidden = hidden
+        self.dtype = dtype
+        self.device = torch.device(device)
+        self.path = path
+        self._shape = (n_samples, seq_len, hidden)
+
+        if path is not None:
+            self._bf16_as_u16 = dtype == torch.bfloat16
+            np_dtype = _torch_to_numpy(dtype)
+            self._np_dtype: type | None = np_dtype
+            self._mm: np.memmap | None = np.memmap(
+                path, dtype=np_dtype, mode="w+", shape=self._shape
+            )
+            if hasattr(os, "posix_madvise"):
+                try:
+                    os.posix_madvise(
+                        self._mm.ctypes.data, self._mm.nbytes,
+                        os.POSIX_MADV_SEQUENTIAL,  # type: ignore[attr-defined]
+                    )
+                except OSError:
+                    pass
+            self._staging: Tensor | None = None
+            self._data: Tensor | None = None
+        else:
+            self._bf16_as_u16 = False
+            self._np_dtype = None
+            self._mm = None
+            self._staging = None
+            pin = self.device.type == "cpu" and torch.cuda.is_available()
+            self._data = torch.zeros(
+                self._shape, dtype=dtype, device=self.device, pin_memory=pin,
+            )
+
+    def _mm_to_tensor(self, arr: np.ndarray) -> Tensor:
+        t = torch.from_numpy(arr)
+        if self._bf16_as_u16:
+            t = t.view(torch.bfloat16)
+        return t
+
+    def _tensor_to_np(self, values: Tensor) -> np.ndarray:
+        host = values.detach().to(device="cpu", dtype=self.dtype)
+        if self._bf16_as_u16:
+            host = host.view(torch.uint16)
+        return host.numpy()
+
+    def __getitem__(self, sample_ids: Tensor) -> Tensor:
+        if self._data is not None:
+            return self._data[sample_ids]
+        assert self._mm is not None
+        ids_np = sample_ids.detach().to(device="cpu", dtype=torch.int64).numpy()
+        return self._mm_to_tensor(np.asarray(self._mm[ids_np]).copy())
+
+    def __setitem__(self, sample_ids: Tensor, values: Tensor) -> None:
+        if self._data is not None:
+            self._data[sample_ids] = values.to(dtype=self.dtype, device=self.device)
+            return
+        assert self._mm is not None
+        ids_np = sample_ids.detach().to(device="cpu", dtype=torch.int64).numpy()
+        if values.device.type == "cuda":
+            staging = self._ensure_staging(values.shape)
+            if values.dtype != self.dtype:
+                values = values.to(dtype=self.dtype)
+            staging.copy_(values, non_blocking=True)
+            torch.cuda.synchronize()
+            host = staging
+        else:
+            host = values.detach().to(device="cpu", dtype=self.dtype)
+        if self._bf16_as_u16:
+            host = host.view(torch.uint16)
+        self._mm[ids_np] = host.numpy()
+
+    def read_slice(self, start: int, stop: int) -> Tensor:
+        if self._data is not None:
+            return self._data[start:stop]
+        assert self._mm is not None
+        return self._mm_to_tensor(np.asarray(self._mm[start:stop]).copy())
+
+    def _ensure_staging(self, shape: tuple[int, ...]) -> Tensor:
+        if self._staging is not None and self._staging.shape == shape:
+            return self._staging
+        self._staging = torch.zeros(shape, dtype=self.dtype, pin_memory=True)
+        return self._staging
+
+    def write_slice(self, start: int, stop: int, values: Tensor) -> None:
+        if self._data is not None:
+            if values.dtype != self.dtype:
+                values = values.to(dtype=self.dtype)
+            self._data[start:stop].copy_(values, non_blocking=True)
+            return
+        assert self._mm is not None
+        if values.device.type == "cuda":
+            staging = self._ensure_staging(values.shape)
+            if values.dtype != self.dtype:
+                values = values.to(dtype=self.dtype)
+            staging.copy_(values, non_blocking=True)
+            torch.cuda.synchronize()
+            host = staging
+        else:
+            host = values.detach().to(device="cpu", dtype=self.dtype)
+        if self._bf16_as_u16:
+            host = host.view(torch.uint16)
+        self._mm[start:stop] = host.numpy()
+
+    def flush(self) -> None:
+        if self._mm is not None:
+            self._mm.flush()
+
+    def close(self) -> None:
+        if self._mm is not None:
+            self._mm.flush()
+            del self._mm
+            self._mm = None
+
+
+_TORCH_TO_NUMPY = {
+    torch.float32: np.float32,
+    torch.float64: np.float64,
+    torch.float16: np.float16,
+    torch.bfloat16: np.uint16,
+    torch.int32: np.int32,
+    torch.int64: np.int64,
+    torch.int16: np.int16,
+    torch.int8: np.int8,
+    torch.uint8: np.uint8,
+}
+
+
+def _torch_to_numpy(dtype: torch.dtype) -> type:
+    np_dtype = _TORCH_TO_NUMPY.get(dtype)
+    if np_dtype is None:
+        raise ValueError(f"unsupported dtype for memmap: {dtype}")
+    return np_dtype

fpwap/callbacks/__init__.py

@@ -0,0 +1,3 @@
+from fpwap.callbacks.base import Callback
+
+__all__ = ["Callback"]

fpwap/callbacks/base.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from fpwap.types import (
+    Artifact,
+    BatchResult,
+    Context,
+    HookName,
+    LayerArtifact,
+    Phase,
+)
+
+
+class Callback:
+    target_layers: Sequence[int] | Literal["all"] = "all"
+    target_hooks: Sequence[HookName] = ("residual_post",)
+    phase: Phase = "read"
+    needs_grad: bool = False
+    accum_dtype: torch.dtype = torch.float32
+
+    def on_sweep_start(self, ctx: Context) -> None:
+        return None
+
+    def on_layer_start(self, layer_idx: int) -> None:
+        return None
+
+    def on_batch(
+        self,
+        layer_idx: int,
+        hook: HookName,
+        acts: Tensor,
+        sample_ids: Tensor,
+    ) -> BatchResult:
+        return None
+
+    def on_layer_end(self, layer_idx: int) -> LayerArtifact | None:
+        return None
+
+    def on_sweep_end(self) -> Artifact | None:
+        return None
+
+    def checkpoint_state(self) -> bytes:
+        return b""
+
+    def restore_state(self, state: bytes) -> None:
+        return None

fpwap/callbacks/common.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from fpwap.callbacks.base import Callback
+from fpwap.types import (
+    Artifact,
+    BatchResult,
+    HookName,
+    LayerArtifact,
+)
+
+
+class RawActivations(Callback):
+    """Persist per-sample activations, pooled (`last_token_only=True`) by default.
+
+    Returns an `Emit` each microbatch; the engine routes these into the run's
+    result so `result.activations(layer, hook)` can return them concatenated
+    in sample order. For datasets too large to hold in memory, swap in a
+    disk-backed StorageBackend.
+    """
+
+    phase = "read"
+
+    def __init__(
+        self,
+        layers: Sequence[int] | Literal["all"] = "all",
+        hook: HookName = "residual_post",
+        last_token_only: bool = True,
+        out_dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        self.target_layers = layers
+        self.target_hooks = (hook,)
+        self.last_token_only = last_token_only
+        self.out_dtype = out_dtype
+
+    def on_batch(
+        self,
+        layer_idx: int,
+        hook: HookName,
+        acts: Tensor,
+        sample_ids: Tensor,
+    ) -> BatchResult:
+        from fpwap.types import Emit
+
+        pooled = acts[:, -1, :] if self.last_token_only else acts
+        return Emit(pooled.to(self.out_dtype), dtype=self.out_dtype)
+
+
+class IncrementalPCA(Callback):
+    """Streaming PCA over the dataset. One pass, O(H²) memory per layer.
+
+    Accumulates the running mean and `X^T X` in fp32 across microbatches;
+    at `on_layer_end` computes `cov = E[XX^T] - mean mean^T` and returns
+    the top-k eigenvectors as a `LayerArtifact(kind="pca_basis", ...)`.
+    The engine routes it into `result.artifact("pca_basis", layer=i)`.
+
+    Pooling: `last_token_only=True` by default (matches RawActivations).
+    Any 3D `[N, S, H]` activation is pooled to `[N, H]` before accumulation;
+    users needing other pooling should pre-pool in an upstream callback or
+    subclass. Accumulators run on the execution device to avoid per-batch
+    H2D; the final basis is moved to CPU in the artifact payload.
+    """
+
+    phase = "read"
+
+    def __init__(
+        self,
+        layers: Sequence[int] | Literal["all"] = "all",
+        n_components: int = 64,
+        hook: HookName = "residual_post",
+        last_token_only: bool = True,
+    ) -> None:
+        self.target_layers = layers
+        self.target_hooks = (hook,)
+        self.n_components = n_components
+        self.last_token_only = last_token_only
+        self._sums: dict[int, Tensor] = {}
+        self._sumxx: dict[int, Tensor] = {}
+        self._counts: dict[int, int] = {}
+
+    def on_batch(
+        self,
+        layer_idx: int,
+        hook: HookName,
+        acts: Tensor,
+        sample_ids: Tensor,
+    ) -> BatchResult:
+        x = acts[:, -1, :] if (self.last_token_only and acts.dim() == 3) else acts
+        if x.dim() != 2:
+            raise ValueError(
+                f"IncrementalPCA expects 2D activations [N, H] after pooling, "
+                f"got {tuple(x.shape)}"
+            )
+        x = x.to(torch.float32)
+        hdim = x.shape[-1]
+        if layer_idx not in self._sums:
+            self._sums[layer_idx] = torch.zeros(hdim, dtype=torch.float32, device=x.device)
+            self._sumxx[layer_idx] = torch.zeros(
+                hdim, hdim, dtype=torch.float32, device=x.device
+            )
+            self._counts[layer_idx] = 0
+        self._sums[layer_idx] += x.sum(dim=0)
+        self._sumxx[layer_idx] += x.T @ x
+        self._counts[layer_idx] += int(x.shape[0])
+        return None
+
+    def on_layer_end(self, layer_idx: int) -> LayerArtifact | None:
+        if layer_idx not in self._sums:
+            return None
+        n = self._counts[layer_idx]
+        if n == 0:
+            return None
+        mean = self._sums[layer_idx] / n
+        cov = self._sumxx[layer_idx] / n - torch.outer(mean, mean)
+        # eigh returns ascending eigenvalues; reverse for PCA convention.
+        eigvals, eigvecs = torch.linalg.eigh(cov)
+        order = torch.argsort(eigvals, descending=True)
+        eigvals = eigvals[order]
+        eigvecs = eigvecs[:, order]
+        k = min(self.n_components, eigvecs.shape[1])
+        basis = eigvecs[:, :k].contiguous()
+        # Drop per-layer accumulators once artifact is materialized.
+        del self._sums[layer_idx]
+        del self._sumxx[layer_idx]
+        del self._counts[layer_idx]
+        return LayerArtifact(
+            kind="pca_basis",
+            payload={
+                "basis": basis.cpu(),
+                "mean": mean.cpu(),
+                "explained_variance": eigvals[:k].cpu(),
+            },
+        )
+
+
+class DiffOfMeans(Callback):
+    """Per-class activation means for binary-labeled data.
+
+    User-supplied `labels` is a 1D int tensor indexed by sample_id — this
+    sidesteps the dataset-items-in-callback API question. Accumulates
+    streaming sums per class (in fp32, on the execution device), and at
+    `on_layer_end` returns a `LayerArtifact(kind="diff_of_means", payload)`
+    carrying `(mean_1 - mean_0)`, both class means, and their counts. The
+    returned direction is a common probing target ("does the model encode
+    label L in this layer's residual stream?").
+
+    Pooling: `last_token_only=True` by default, matching RawActivations.
+    """
+
+    phase = "read"
+
+    def __init__(
+        self,
+        labels: Tensor,
+        layers: Sequence[int] | Literal["all"] = "all",
+        hook: HookName = "residual_post",
+        last_token_only: bool = True,
+    ) -> None:
+        self.target_layers = layers
+        self.target_hooks = (hook,)
+        self.labels = labels.to(torch.int64)
+        self.last_token_only = last_token_only
+        self._sums: dict[int, dict[int, Tensor]] = {}
+        self._counts: dict[int, dict[int, int]] = {}
+
+    def on_batch(
+        self,
+        layer_idx: int,
+        hook: HookName,
+        acts: Tensor,
+        sample_ids: Tensor,
+    ) -> BatchResult:
+        x = acts[:, -1, :] if (self.last_token_only and acts.dim() == 3) else acts
+        if x.dim() != 2:
+            raise ValueError(
+                f"DiffOfMeans expects 2D activations [N, H] after pooling, "
+                f"got {tuple(x.shape)}"
+            )
+        x = x.to(torch.float32)
+        # sample_ids → labels (cpu lookup into user-provided tensor, then push
+        # back to acts' device for the masked-sum below).
+        labels = self.labels[sample_ids.detach().cpu()].to(x.device)
+        sums = self._sums.setdefault(layer_idx, {})
+        counts = self._counts.setdefault(layer_idx, {})
+        for cls in torch.unique(labels).tolist():
+            mask = labels == cls
+            cls_x = x[mask]
+            if cls not in sums:
+                sums[cls] = torch.zeros(x.shape[-1], dtype=torch.float32, device=x.device)
+                counts[cls] = 0
+            sums[cls] += cls_x.sum(dim=0)
+            counts[cls] += int(cls_x.shape[0])
+        return None
+
+    def on_layer_end(self, layer_idx: int) -> LayerArtifact | None:
+        if layer_idx not in self._sums:
+            return None
+        sums = self._sums[layer_idx]
+        counts = self._counts[layer_idx]
+        means = {cls: sums[cls] / counts[cls] for cls in sums}
+        # Drop accumulators.
+        del self._sums[layer_idx]
+        del self._counts[layer_idx]
+        # Binary-labeled convention: require at least labels 0 and 1 present.
+        if 0 in means and 1 in means:
+            direction = means[1] - means[0]
+        else:
+            direction = None
+        return LayerArtifact(
+            kind="diff_of_means",
+            payload={
+                "direction": direction.cpu() if direction is not None else None,
+                "means": {int(c): v.cpu() for c, v in means.items()},
+                "counts": {int(c): counts[c] for c in counts},
+            },
+        )
+
+
+class SteerInBasis(Callback):
+    """Additive intervention in a pre-computed basis.
+
+    `acts + alpha * basis[:, direction_idx]`, broadcast across the batch and
+    sequence dims. `basis_artifact.payload["basis"]` is expected to be
+    `[H, n_components]` (the shape produced by `IncrementalPCA`); any
+    Artifact with a compatible payload works. Lives in `phase="write"`, so
+    the returned `WriteBack` replaces the residual that feeds the next
+    layer (or the buffer, if targeted at `residual_post`).
+    """
+
+    phase = "write"
+
+    def __init__(
+        self,
+        basis_artifact: Artifact,
+        direction_idx: int,
+        alpha: float,
+        layers: Sequence[int] | Literal["all"] = "all",
+        hook: HookName = "residual_post",
+    ) -> None:
+        self.target_layers = layers
+        self.target_hooks = (hook,)
+        self.basis = basis_artifact
+        self.direction_idx = direction_idx
+        self.alpha = alpha
+
+    def on_batch(
+        self,
+        layer_idx: int,
+        hook: HookName,
+        acts: Tensor,
+        sample_ids: Tensor,
+    ) -> BatchResult:
+        from fpwap.types import WriteBack
+
+        payload = self.basis.payload
+        if isinstance(payload, dict):
+            # Prefer "basis" [H, K]; fall back to "direction" [H] (what
+            # DiffOfMeans returns). direction_idx is ignored for 1D payloads.
+            if "basis" in payload:
+                basis = payload["basis"]
+                direction = basis[:, self.direction_idx]
+            elif "direction" in payload and payload["direction"] is not None:
+                direction = payload["direction"]
+            else:
+                raise KeyError(
+                    "basis_artifact payload has no 'basis' or 'direction' key"
+                )
+        else:
+            direction = (
+                payload[:, self.direction_idx] if payload.dim() == 2 else payload
+            )
+        direction = direction.to(device=acts.device, dtype=acts.dtype)
+        return WriteBack(acts + self.alpha * direction)

fpwap/cost_model.py

@@ -0,0 +1,108 @@
+"""Preflight cost model: predict per-layer latency and recommend config.
+
+Pure arithmetic — no GPU, no model weights, no torch. CI-safe.
+
+The cost model per layer:
+  compute = fwd_per_microbatch_s × ceil(n_samples / microbatch_size)
+  load    = weight_load_s
+
+Without prefetch: per_layer = load + compute
+With prefetch:    per_layer = max(load, compute)
+
+Total wall = embed_s + per_layer × n_layers
+"""
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class CostModelInput:
+    n_layers: int
+    n_samples: int
+    seq_len: int
+    microbatch_size: int
+    weight_load_s: float
+    fwd_per_microbatch_s: float
+    embed_s: float
+    layer_weight_bytes: int
+
+
+@dataclass(frozen=True)
+class CostModelPrediction:
+    per_layer_s: float
+    total_wall_s: float
+    throughput_tok_s: float
+    bottleneck: str  # "load", "compute", or "balanced"
+    load_pct: float
+    compute_pct: float
+    weight_io_gb: float
+    prefetch: bool
+
+
+def predict(inp: CostModelInput, *, prefetch: bool) -> CostModelPrediction:
+    n_microbatches = math.ceil(inp.n_samples / inp.microbatch_size)
+    compute_s = inp.fwd_per_microbatch_s * n_microbatches
+    load_s = inp.weight_load_s
+
+    if prefetch:
+        per_layer_s = max(load_s, compute_s)
+    else:
+        per_layer_s = load_s + compute_s
+
+    total_wall_s = inp.embed_s + per_layer_s * inp.n_layers
+    total_tokens = inp.n_samples * inp.seq_len
+    throughput = total_tokens / total_wall_s if total_wall_s > 0 else 0.0
+
+    weight_io_gb = inp.layer_weight_bytes * inp.n_layers / 1e9
+
+    if per_layer_s > 0:
+        load_pct = load_s / per_layer_s
+        compute_pct = compute_s / per_layer_s
+    else:
+        load_pct = 0.0
+        compute_pct = 0.0
+
+    ratio = load_s / compute_s if compute_s > 0 else float("inf")
+    if ratio > 1.2:
+        bottleneck = "load"
+    elif ratio < 1 / 1.2:
+        bottleneck = "compute"
+    else:
+        bottleneck = "balanced"
+
+    return CostModelPrediction(
+        per_layer_s=per_layer_s,
+        total_wall_s=total_wall_s,
+        throughput_tok_s=throughput,
+        bottleneck=bottleneck,
+        load_pct=load_pct,
+        compute_pct=compute_pct,
+        weight_io_gb=weight_io_gb,
+        prefetch=prefetch,
+    )
+
+
+@dataclass(frozen=True)
+class Recommendation:
+    input: CostModelInput
+    prefetch: bool
+    prediction: CostModelPrediction
+
+
+def recommend(
+    candidates: Sequence[tuple[CostModelInput, bool]],
+) -> Recommendation:
+    if not candidates:
+        raise ValueError("candidates must be non-empty")
+
+    best: Recommendation | None = None
+    for inp, pf in candidates:
+        pred = predict(inp, prefetch=pf)
+        if best is None or pred.throughput_tok_s > best.prediction.throughput_tok_s:
+            best = Recommendation(input=inp, prefetch=pf, prediction=pred)
+
+    assert best is not None
+    return best