sparse-delta-memory 2.0.0__py3-none-any.whl

sparse_delta_memory/memory_stats.py

@@ -0,0 +1,221 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""Opt-in statistics for SparseDeltaMemory (memory-access / activation / gradient stats).
+
+Everything here is logging only, gated by ``SparseDeltaMemoryArgs.log_memory_access_stats``
+and ``log_memory_norms``, and factored out of ``layer.py`` so that file holds the model
+math. ``SparseDeltaMemory`` mixes in ``SDMStatsMixin`` and calls the ``_record_*`` hooks only
+in eval mode outside CUDA-graph capture: they read tensor values on the host (``.item()``).
+The buffers are registered on the layer (``self``) so they follow device/FSDP moves.
+"""
+import torch
+
+
+class SDMStatsMixin:
+    # ------------------------------------------------------------------ setup
+    def _init_stats(self, args, num_heads: int) -> None:
+        """Allocate the stats buffers/accumulators (called from the layer __init__)."""
+        if args.log_memory_access_stats:
+            stats_slots = num_heads * self.slots_per_head
+            # Unique-index access counts for read and write (across all heads).
+            for name in ("_read_idx_counts", "_write_idx_counts"):
+                self.register_buffer(name, torch.zeros(stats_slots, dtype=torch.int64), persistent=False)
+            # Activation statistics accumulators (for stability debugging).
+            self._activation_stats: dict[str, list[float]] = {
+                "read_weight_max": [],
+                "write_strength_max": [],
+                "write_strength_mean": [],
+                "output_norm": [],
+                "forget_gate_min": [],
+                "forget_gate_mean": [],
+                "forget_gate_max": [],
+                "input_gate_min": [],
+                "input_gate_mean": [],
+                "input_gate_max": [],
+            }
+            if args.log_memory_norms:
+                self._activation_stats.update({
+                    "memory_norm": [],
+                    "memory_max": [],
+                    "updated_memory_norm": [],
+                    "updated_memory_max": [],
+                    "slot_norm_mean": [],
+                    "slot_norm_max": [],
+                    "updated_slot_norm_mean": [],
+                    "updated_slot_norm_max": [],
+                })
+            # Full slot access histograms (accumulated, converted to entropy at log time).
+            for name in ("_read_slot_access_weights", "_write_slot_access_weights"):
+                self.register_buffer(name, torch.zeros(stats_slots, dtype=torch.float32), persistent=False)
+            # Gradient statistics, collected via collect_grad_stats() after backward.
+            self._grad_stats: dict[str, list[float]] = {
+                "memory_grad_norm": [],
+                "qk_proj_grad_norm": [],
+                "output_proj_grad_norm": [],
+            }
+        else:
+            self._read_idx_counts = None
+            self._write_idx_counts = None
+            self._read_slot_access_weights = None
+            self._write_slot_access_weights = None
+            self._activation_stats = None
+            self._grad_stats = None
+
+    # ------------------------------------------------------------ forward hooks
+    def _record_gate_stats(self, g: torch.Tensor, beta: torch.Tensor) -> None:
+        if self._activation_stats is None:
+            return
+        with torch.no_grad():
+            decay = torch.exp(g)
+            s = self._activation_stats
+            s["forget_gate_min"].append(decay.min().item())
+            s["forget_gate_mean"].append(decay.mean().item())
+            s["forget_gate_max"].append(decay.max().item())
+            s["input_gate_min"].append(beta.min().item())
+            s["input_gate_mean"].append(beta.mean().item())
+            s["input_gate_max"].append(beta.max().item())
+
+    def _record_read_stats(self, q_idx, q_val, BH, H, sph) -> None:
+        if self._read_idx_counts is None:
+            return
+        with torch.no_grad():
+            # Indices are [BH, T, num_reads] in [0, sph). Add head offset for [0, H*sph).
+            head_offset = (torch.arange(BH, device=q_idx.device) % H * sph).view(BH, 1, 1)
+            q_idx_flat = (q_idx + head_offset).view(-1)
+            self._read_idx_counts.index_add_(
+                0, q_idx_flat.long(), torch.ones_like(q_idx_flat, dtype=torch.int64)
+            )
+            if self._activation_stats is not None:
+                self._activation_stats["read_weight_max"].append(q_val.max().item())
+            if self._read_slot_access_weights is not None:
+                # index_put_ with accumulate=True sums duplicates in sorted order on CUDA
+                # (deterministic); a float index_add_ uses atomics.
+                q_val_flat = q_val.reshape(-1)
+                self._read_slot_access_weights.index_put_(
+                    (q_idx_flat.long(),), q_val_flat.to(self._read_slot_access_weights.dtype),
+                    accumulate=True)
+
+    def _record_memory_norms_pre(self, memory: torch.Tensor) -> None:
+        if not (self.args.log_memory_norms and self._activation_stats is not None):
+            return
+        with torch.no_grad():
+            s = self._activation_stats
+            s["memory_norm"].append(memory.norm().item())
+            s["memory_max"].append(memory.abs().max().item())
+            slot_norms = memory.norm(dim=-1)
+            s["slot_norm_mean"].append(slot_norms.mean().item())
+            s["slot_norm_max"].append(slot_norms.max().item())
+
+    def _record_write_stats(self, k_idx, k_val, memory, output, BH, H, sph) -> None:
+        if self._write_idx_counts is None:
+            return
+        with torch.no_grad():
+            # k_idx has batch offsets; mod by sph then add head offset for [0, H*sph).
+            head_offset_w = (torch.arange(BH, device=k_idx.device) % H * sph).view(BH, 1, 1)
+            write_idx_flat = (k_idx % sph + head_offset_w).view(-1)
+            self._write_idx_counts.index_add_(
+                0, write_idx_flat.long(), torch.ones_like(write_idx_flat, dtype=torch.int64)
+            )
+            if self._activation_stats is not None:
+                s = self._activation_stats
+                s["write_strength_max"].append(k_val.max().item())
+                s["write_strength_mean"].append(k_val.mean().item())
+                if self.args.log_memory_norms and memory is not None:
+                    s["updated_memory_norm"].append(memory.norm().item())
+                    s["updated_memory_max"].append(memory.abs().max().item())
+                    updated_slot_norms = memory.norm(dim=-1)
+                    s["updated_slot_norm_mean"].append(updated_slot_norms.mean().item())
+                    s["updated_slot_norm_max"].append(updated_slot_norms.max().item())
+                s["output_norm"].append(output.norm(dim=-1).mean().item())
+            if self._write_slot_access_weights is not None:
+                write_strength_flat = k_val.reshape(-1)
+                self._write_slot_access_weights.index_put_(
+                    (write_idx_flat.long(),),
+                    write_strength_flat.to(self._write_slot_access_weights.dtype),
+                    accumulate=True)
+
+    # ------------------------------------------------------------ public API
+    def collect_grad_stats(self) -> None:
+        """Collect gradient statistics after backward pass.
+
+        Call this after optimizer.step() to record gradient norms.
+        Hooks don't work reliably with torch.compile, so we check .grad directly.
+        """
+        if self._grad_stats is None:
+            return
+
+        # Memory parameter gradient
+        if self.memory is not None and self.memory.grad is not None:
+            self._grad_stats["memory_grad_norm"].append(self.memory.grad.norm().item())
+
+        # Q/K projection gradients
+        qk_grad_norms = []
+        for proj in [self.Wq_read, self.Wk_write]:
+            if proj is not None:
+                for param in proj.parameters():
+                    if param.grad is not None:
+                        qk_grad_norms.append(param.grad.norm().item())
+                        break
+        if qk_grad_norms:
+            self._grad_stats["qk_proj_grad_norm"].append(sum(qk_grad_norms) / len(qk_grad_norms))
+
+        # Output projection gradient
+        if self.Wo is not None:
+            for param in self.Wo.parameters():
+                if param.grad is not None:
+                    self._grad_stats["output_proj_grad_norm"].append(param.grad.norm().item())
+                    break
+
+    def get_memory_access_stats(self, reset: bool = True) -> "dict[str, float] | None":
+        """Get memory access statistics and optionally reset counters.
+
+        Returns a dict of unique-access percentages, read/write slot-access entropy, and the
+        averaged activation + gradient statistics since the last reset. Returns None if
+        ``log_memory_access_stats`` is disabled.
+        """
+        if not self.args.log_memory_access_stats:
+            return None
+
+        if self._read_idx_counts is None or self._write_idx_counts is None:
+            return None
+
+        num_slots = self._read_idx_counts.numel()
+
+        # Compute stats - only count unique slots
+        read_unique = (self._read_idx_counts > 0).sum().item()
+        write_unique = (self._write_idx_counts > 0).sum().item()
+
+        stats = {
+            "read_unique_pct": 100.0 * read_unique / num_slots if num_slots > 0 else 0.0,
+            "write_unique_pct": 100.0 * write_unique / num_slots if num_slots > 0 else 0.0,
+        }
+
+        # Compute entropy from full slot access histograms
+        eps = 1e-10
+        for kind in ("read", "write"):
+            weights = getattr(self, f"_{kind}_slot_access_weights")
+            if weights is not None:
+                total_weight = weights.sum()
+                if total_weight > eps:
+                    probs = weights / total_weight
+                    max_entropy = torch.log(torch.tensor(num_slots, dtype=torch.float32))
+                    entropy = -(probs * probs.clamp(min=eps).log()).sum()
+                    stats[f"{kind}_slot_entropy_normalized"] = (entropy / max_entropy).item()
+                else:
+                    stats[f"{kind}_slot_entropy_normalized"] = 0.0
+
+        # Add activation statistics (averaged over forward passes)
+        # and gradient statistics (averaged over backward passes)
+        for acc in (self._activation_stats, self._grad_stats):
+            for key, values in (acc or {}).items():
+                stats[key] = sum(values) / len(values) if values else 0.0
+
+        if reset:
+            for buf in (self._read_idx_counts, self._write_idx_counts, self._read_slot_access_weights,
+                        self._write_slot_access_weights):
+                if buf is not None:
+                    buf.zero_()
+            for acc in (self._activation_stats, self._grad_stats):
+                for key in acc or ():
+                    acc[key] = []
+
+        return stats

sparse_delta_memory/movers.py

@@ -0,0 +1,198 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""The decay-and-write (DW) and read-side backward (GDS) step launchers of the op.
+
+``make_mover_dw`` and ``make_mover_gds`` run the CUDA movers (cuda/mover_fx_cuda.py and cuda/mover_da2_cuda.py for the
+DW; cuda/gx_gds_cuda.py and cuda/mover_da2_cuda.py for the GDS) wherever ``dw_admits`` / ``gds_admits`` hold. There is
+no other route (2.0.0): memory_ops.refuse_non_native refuses every call they do not take before any work, so a launcher
+asked for one raises (NO_NATIVE_MOVER_MSG, an internal error).
+Phase 1's records (krec, qrec / qrec8, tile bounds) are passed by the caller; without them the CUDA movers build their
+own (the same values). Every mover gives each bank (or bank-gradient) row one writer per launch and sums a slot
+group's entries in a fixed order, so its results are bitwise reproducible.
+"""
+
+import torch
+
+from .cuda import mover_da2_cuda as da2
+from .cuda._device import backend
+
+STATS = {"da2_dw": 0, "shipped_dw": 0, "da2_gds": 0, "shipped_gds": 0, "p1_krec": 0, "p1_qrec": 0}
+
+
+# Tail classes (shapes outside the 16-byte-slice predicate): native only when EVERY kernel's hook admits them.
+_TAIL_HOOKS = (("mover_fx_cuda", "tail_supported"), ("gather_cuda", "tail_supported"),
+               ("gather_ns_cuda", "tail_supported"), ("gx_pa2_cuda", "tail_supported"),
+               ("gx_gds_cuda", "tail_supported"), ("pa_dc_cuda", "gwr_tail_supported"))
+
+
+NO_NATIVE_MOVER_MSG = ("sparse_delta_memory: internal error: no native {kind} mover for bank {dtype} {shape}, chunk_size {C}, "
+                       "snapshot_quant {quant} (2.0.0 removed the Triton movers; memory_ops.refuse_non_native should have "
+                       "refused this call); please report it")
+
+
+def tail_admits(bank, chunk_size, quant="none") -> bool:
+    """The ONE tail predicate: every kernel of the call admits this tail class (the AND of _TAIL_HOOKS); forward and
+    backward admissions read it."""
+    import importlib
+    return all(getattr(importlib.import_module(f".cuda.{mod}", __package__), fn)(bank, chunk_size, quant)
+               for mod, fn in _TAIL_HOOKS)
+
+
+def dw_admits(bank, chunk_size, quant="none") -> bool:
+    """The CUDA DW takes calls on ``bank`` with chunks of ``chunk_size`` tokens: mover_da2_cuda's admission on the
+    modern backend, mover_fx_cuda's on a100 / rtx; on every backend also tail_admits."""
+    if bank.is_cuda and backend(bank.device) != "modern":
+        from .cuda import mover_fx_cuda
+        # fp8 / int4 rows above the old 16384 head bound: not admitted (QT class T0)
+        fx_ok = mover_fx_cuda.admits(bank, chunk_size) and (
+            quant == "none" or bank.shape[-1] <= mover_fx_cuda._QUANT_MAX_HEAD_DIM)
+        return fx_ok or tail_admits(bank, chunk_size, quant)
+    return (da2.supported(bank, chunk_size) and da2.fits(bank, chunk_size, "dw")) or tail_admits(bank, chunk_size, quant)
+
+
+def gds_admits(bank, grad_dtype, chunk_size, quant="none") -> bool:
+    """The CUDA GDS takes calls on ``bank`` with a ``grad_dtype`` gradient bank (the bank's, or float32 for a bf16 bank)
+    and chunks of ``chunk_size`` tokens: gx_gds_cuda's GDS (modern: mover_da2_cuda's for the shapes it declines); tail_admits."""
+    grad_ok = grad_dtype == bank.dtype or (bank.dtype == torch.bfloat16 and grad_dtype == torch.float32)
+    if not grad_ok or not bank.is_cuda:
+        return False
+    if backend(bank.device) != "modern":
+        from .cuda import gx_gds_cuda
+        return gx_gds_cuda.admits(bank, grad_dtype, chunk_size, 1, quant) or tail_admits(bank, chunk_size, quant)
+    return ((da2.supported(bank, chunk_size) and da2.fits(bank, chunk_size, "gds", grad_dtype))
+            or tail_admits(bank, chunk_size, quant))
+
+
+def read_records(p1):
+    """Phase 1's read records for the GDS: the 16-byte qrec, or, when the forward asked Phase 1 for the compact form
+    only (DET_QREC8: int32 [N, C * R, 2] = qrec's words 1-2, {absolute slot, scale bits}), qrec8, which
+    gx_gds_cuda's GDS reads with spay_q (same bits); None when Phase 1 emitted neither. make_mover_gds builds 16-byte
+    records itself whenever gx_gds_cuda's GDS declines a call that got qrec8 (only speed lost)."""
+    q = p1.qrec
+    if q is not None and q.numel() > 0:
+        return q
+    q8 = p1.qrec8
+    return q8 if q8 is not None and q8.numel() > 0 else None
+
+
+def make_mover_dw(bank, snapshot, delta_v, slot, cpe, k_val, rex, payload=None, scale=None, *, chunk_size,
+                  num_parallel, valid_tokens, snapshot_quant="none", sorted_k=None, krec=None, tiles_k=None,
+                  segb_k=None):
+    """The DW step launcher (the CUDA DW; a call dw_admits declines raises NO_NATIVE_MOVER_MSG: the op refuses it up
+    front). krec, tiles_k, segb_k: Phase 1's records and tile bounds for these entries (``phase1_forward(..., beta=...)``),
+    or None."""
+    # krec: Phase 1's write records, [N, C*W, 4] int32; when given they replace the mover's own record build.
+    # Quantized snapshots (fp8 / int4) stay on the CUDA DW, which also quantizes them; the bank update is unchanged.
+    # Tails run only on mover_fx_cuda's bank-direct DW (layouts 2 / 4); on layout 0 they raise below.
+    if snapshot_quant == "bf16":   # bf16 rows: the forward passes "none" with its bf16 compact buffer (no other route)
+        raise RuntimeError("make_mover_dw: snapshot_quant bf16 runs as the unquantized DW with bf16 compact rows")
+    tail_off = (sorted_k is not None and not da2._shape_supported(bank, chunk_size)
+                and getattr(snapshot, "_dc_snapc_layout", 0) not in (2, 4))
+    if sorted_k is None or tail_off or not dw_admits(bank, chunk_size, snapshot_quant):
+        # 2.0.0 (A4): phase2_kernels' Triton DW is gone; memory_ops.refuse_non_native refuses such calls up front.
+        raise RuntimeError(NO_NATIVE_MOVER_MSG.format(kind="DW", shape=tuple(bank.shape), dtype=bank.dtype, C=chunk_size,
+                                                      quant=snapshot_quant))
+    STATS["da2_dw"] += 1
+    n = sorted_k.shape[-1]
+    sorted_k = sorted_k.reshape(-1, n)
+    w = slot.shape[-1]
+    steps = slot.numel() // (num_parallel * chunk_size * w)
+    if krec is not None and krec.numel() == sorted_k.numel() * 4 and krec.is_contiguous() \
+            and krec.data_ptr() % 16 == 0 and krec.dtype == torch.int32:   # read as uint4: 16-byte aligned
+        STATS["p1_krec"] += 1
+        krec = krec.reshape(-1, n, 4)
+    else:
+        krec = da2.build_krec(sorted_k, slot, k_val, rex, cpe)
+    # (a misaligned snapshot only disables the compact-snapshot read inside make_dw: it is never copied)
+    return _with_aligned(
+        lambda bk, dv: da2.make_dw(bk, dv, sorted_k, krec, chunk_size=chunk_size, num_parallel=num_parallel,
+                                   steps=steps, snapshot=snapshot, tiles_k=tiles_k, segb_k=segb_k,
+                                   snapshot_quant=snapshot_quant, payload=payload, scale=scale,
+                                   valid_tokens=valid_tokens),
+        (bank, delta_v), outs={0})
+
+
+class _StepRecords:
+    """Record provider of gx_gds_cuda's GDS over px_qrec_cuda's output: the 8-byte read records of every chunk and
+    one CUDA event per step (or None)."""
+
+    def __init__(self, rec8, events):
+        self.rec8, self.events = rec8, events
+
+    def get(self, ti):
+        return self.rec8, (self.events[ti] if self.events is not None else None)
+
+
+def _px_qrec_provider(sorted_q, slot, q_val, log_q, num_parallel, steps):
+    """px_qrec_cuda's per-step side-stream builder of the 8-byte read records: build_qrec8_steps(spay_q, q_idx, q_val,
+    log_q, *, P, Nt) -> (int32 [Nt * P, C * R, 2], list of Nt events or None); GDS(ti) waits on events[ti]. Same bits
+    as Phase 1's qrec words 1-2."""
+    from .cuda import px_qrec_cuda
+    rec8, events = px_qrec_cuda.build_qrec8_steps(sorted_q, slot, q_val, log_q, P=num_parallel, Nt=steps)
+    STATS["px_qrec8"] = STATS.get("px_qrec8", 0) + 1
+    return _StepRecords(rec8, events)
+
+
+def make_mover_gds(bank, grad_bank, grad_readings, slot, q_val, log_q, d_q_w, *, chunk_size, num_parallel,
+                   valid_tokens, sorted_q=None, qrec=None, tiles_q=None, snapshot_quant="none"):
+    """The GDS step launcher: gx_gds_cuda's mover (one sorted walk per tile), which delegates to mover_da2_cuda's (the
+    same bits) outside its scope; a call gds_admits declines raises NO_NATIVE_MOVER_MSG (the op refuses it up front). qrec, tiles_q: Phase 1's read
+    records (``read_records``) and tile bounds for these entries, or None. snapshot_quant: the op's mode (the shared
+    tail predicate admits tail classes per mode: fp8 / int4 tails are refused up front)."""
+    if sorted_q is None or not gds_admits(bank, grad_bank.dtype, chunk_size, snapshot_quant):
+        # 2.0.0 (A4): phase2_kernels' Triton GDS is gone; memory_ops.refuse_non_native refuses such calls up front.
+        raise RuntimeError(NO_NATIVE_MOVER_MSG.format(kind="GDS", shape=tuple(bank.shape), dtype=bank.dtype, C=chunk_size,
+                                                      quant=snapshot_quant))
+    STATS["da2_gds"] += 1
+    n = sorted_q.shape[-1]
+    sorted_q = sorted_q.reshape(-1, n)
+    r = slot.shape[-1]
+    steps = slot.numel() // (num_parallel * chunk_size * r)
+    from .cuda import gx_gds_cuda as gxm
+    if gxm.supported(bank, grad_bank, chunk_size, r):
+        # gx_gds_cuda's mover reads Phase 1's qrec (16 B) or qrec8 (8 B and spay_q) as given; with neither,
+        # px_qrec_cuda's per-step side-stream builder (records of all steps and one event per step).
+        if gxm.records_ok(qrec, sorted_q):
+            STATS["gx_qrec_in"] = STATS.get("gx_qrec_in", 0) + 1
+        else:
+            qrec = _px_qrec_provider(sorted_q, slot, q_val, log_q, num_parallel, steps)
+    elif qrec is not None and torch.is_tensor(qrec) and qrec.numel() == sorted_q.numel() * 4 and \
+            qrec.is_contiguous() and qrec.data_ptr() % 16 == 0 and qrec.dtype == torch.int32:   # Phase 1's records
+        STATS["p1_qrec"] += 1
+        qrec = qrec.reshape(-1, n, 4)
+    else:
+        qrec = da2.build_qrec(sorted_q, slot, q_val, log_q)
+    return _with_aligned(
+        lambda bk, gb, gr_: gxm.make_gds(bk, gb, gr_, sorted_q, qrec, d_q_w, chunk_size=chunk_size,
+                                         num_parallel=num_parallel, steps=steps, q_idx=slot,
+                                         valid_tokens=valid_tokens, tiles_q=tiles_q),
+        (bank, grad_bank, grad_readings), outs={1})
+
+
+def _aligned16(*ts):
+    return all(t is None or t.numel() == 0 or t.data_ptr() % 16 == 0 for t in ts)
+
+
+def _with_aligned(build, tensors, outs):
+    """Dispatch never depends on memory layout: when one of `tensors` is not 16-byte aligned (a
+    view), the CUDA launcher runs on aligned scratch copies instead. Before every launch the scratch is refreshed from
+    the caller's tensor (other kernels may change it between steps); after it, the tensors named in `outs` are copied
+    back. build(*aligned) returns the launcher. Same arithmetic, so the same bits as with aligned inputs."""
+    if _aligned16(*tensors):
+        return build(*tensors)
+    STATS["aligned_scratch"] = STATS.get("aligned_scratch", 0) + 1
+    scratch = [t if (t is None or t.numel() == 0 or t.data_ptr() % 16 == 0) else t.clone(memory_format=torch.contiguous_format)
+               for t in tensors]
+    launch = build(*scratch)
+    pairs = [(t, sc) for t, sc in zip(tensors, scratch) if sc is not t]
+    back = [(t, sc) for i, (t, sc) in enumerate(zip(tensors, scratch)) if sc is not t and i in outs]
+
+    def run(ti):
+        for t, sc in pairs:
+            sc.copy_(t)
+        launch(ti)
+        for t, sc in back:
+            t.copy_(sc)
+    for attr in ("reads_snapc", "writes_snapc", "geom"):
+        if hasattr(launch, attr):
+            setattr(run, attr, getattr(launch, attr))
+    return run