initium-deeplearning 0.1.0__py3-none-any.whl

initium/inference/cache/cache_config.py

@@ -0,0 +1,21 @@
+"""
+file: inference/cache/cache_config.py
+
+Cache-side constants. Pure constants, imported explicitly (same convention as
+inference_config.py).
+"""
+
+# "off" | "prefix" | "result" | "prefix,result" | "all". Default off: the default episodes
+# reshuffle the edge list every time (nothing to reuse), and persistent-mode experiments
+# should not change behaviour silently. Set to "all" to make caching on-by-default.
+DEFAULT_CACHE = "off"
+
+DEFAULT_CACHE_RAM_MB = 512
+DEFAULT_CACHE_DISK_MB = 2048
+
+# Number of cache hits that are re-computed WITHOUT the cache and compared. On a
+# mismatch the cache is disabled for the rest of the run (correctness over speed).
+DEFAULT_VERIFY_HITS = 2
+DEFAULT_VERIFY_TOL = 1e-3  # max abs diff on scored outputs and on carried state
+
+ROOT_CHAIN = "root"  # history-chain seed (persistent-mode result keys)

initium/inference/cache/cache_registry.py

@@ -0,0 +1,99 @@
+"""
+file: inference/cache/cache_registry.py
+
+Single entry point used by run_inference.py (counterpart of data/dataset_registry.py):
+    caches, notes, fingerprint = setup_caches(spec, ...)
+Adding a cache type = subclass BaseCache + one line in CACHE_CLASSES.
+"""
+
+from __future__ import annotations
+
+import os
+
+from initium.inference.cache.base_cache import BaseCache
+from initium.inference.cache.keys import model_fingerprint
+from initium.inference.cache.prefix_cache import PrefixStateCache
+from initium.inference.cache.result_cache import ResultCache
+from initium.inference.cache.stores import DiskStore, MemoryStore, TieredStore
+
+CACHE_CLASSES = {"prefix": PrefixStateCache, "result": ResultCache}
+_OFF = {"", "off", "none", "false", "0", "no"}
+
+
+def parse_cache_spec(spec: str) -> list[str]:
+    s = (spec or "off").strip().lower()
+    if s in _OFF:
+        return []
+    if s in ("all", "on"):
+        return list(CACHE_CLASSES)
+    names = []
+    for part in s.replace("+", ",").split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if part not in CACHE_CLASSES:
+            raise ValueError(
+                f"unknown cache type {part!r}; expected off, all, or a comma list of "
+                f"{sorted(CACHE_CLASSES)}"
+            )
+        if part not in names:
+            names.append(part)
+    return names
+
+
+def setup_caches(
+    spec,
+    *,
+    ckpt,
+    device,
+    sampled_writes,
+    deterministic_write,
+    ablate_memory,
+    combiner_skip_stages=None,
+    cache_dir=None,
+    ram_mb=512.0,
+    disk_mb=2048.0,
+    clear=False,
+    allow_stochastic=False,
+):
+    """-> (caches, notes, model_fingerprint|None). Returns no caches (with a note) when
+    the requested caching cannot be done safely."""
+    notes: list[str] = []
+    names = parse_cache_spec(spec)
+    if not names:
+        return [], notes, None
+
+    if sampled_writes and not allow_stochastic:
+        notes.append(
+            "caching disabled: this model samples its write vectors (beta>0 without "
+            "--deterministic-write), so every forward pass is random and cached "
+            "results/states would not be reproducible. Use --deterministic-write, or "
+            "--cache-allow-stochastic to override."
+        )
+        return [], notes, None
+    if sampled_writes:
+        notes.append(
+            "WARNING: caching a stochastic model -- cached entries are ONE random draw "
+            "reused for every hit (and across runs if --cache-dir is set); "
+            "verification is skipped."
+        )
+
+    print("[cache] fingerprinting model weights (once) ...")
+    fp = model_fingerprint(
+        ckpt,
+        {
+            "deterministic_write": bool(deterministic_write),
+            "sampled_writes": bool(sampled_writes),
+            "ablate_memory": bool(ablate_memory),
+            "combiner_skip_stages": sorted(combiner_skip_stages) if combiner_skip_stages else None,
+        },
+    )
+    disk = None
+    if cache_dir:
+        disk = DiskStore(os.path.join(cache_dir, fp), int(disk_mb * 1e6))
+        if clear:
+            disk.clear()
+            notes.append(f"disk cache cleared: {disk.root}")
+    store = TieredStore(MemoryStore(int(ram_mb * 1e6)), disk)
+    caches: list[BaseCache] = [CACHE_CLASSES[n](store, fp, device) for n in names]
+    return caches, notes, fp

initium/inference/cache/cached_engine.py

@@ -0,0 +1,264 @@
+"""
+file: inference/cache/cached_engine.py
+
+InferenceEngine + caches. Only used when at least one cache is active, so with caching
+off the plain engine (and its behaviour) is untouched.
+
+Per episode:
+  1. episode-level lookups   (ResultCache: skip the forward pass entirely)
+  2. else compute, resuming from the longest cached prefix snapshot (PrefixStateCache)
+     and snapshotting every not-yet-cached boundary on the way, in segments
+  3. store the episode result in the episode-level caches
+
+Correctness net: for the first `verify_hits` cache hits the episode is ALSO recomputed
+without any cache and compared (scored outputs + carried state, tolerance verify_tol).
+A mismatch (or NaN) prints a warning, uses the recomputed values, and disables all
+caches for the remainder of the run.
+
+Segment execution:
+  MambaDNC / stock DNC : model(x[start:end], hidden, reset_experience=False)  (state passes through)
+  SplitGraphDNC        : model(x[:end], hidden, ..., start_step=start)  -- the parallel
+                         backbone is recomputed over the whole prefix (cheap, stateless
+                         across calls, keeps outputs exact); only the sequential memory
+                         loop is skipped for t < start.
+"""
+
+from __future__ import annotations
+
+import inspect
+import time
+from typing import Any
+
+import torch
+
+from initium.inference.cache.base_cache import EpisodeCtx
+from initium.inference.cache.cache_config import DEFAULT_VERIFY_HITS, DEFAULT_VERIFY_TOL, ROOT_CHAIN
+from initium.inference.cache.keys import chain_next, hash_tensor
+from initium.inference.cache.state_utils import clone_tree, tree_max_abs_diff
+from initium.inference.engine import FRESH_HIDDEN, InferenceEngine
+from initium.inference.metrics import EpisodeResult
+
+
+def model_supports_resume(model) -> bool:
+    if hasattr(model, "backbone"):  # SplitGraphDNC
+        return "start_step" in inspect.signature(model.forward).parameters
+    return True  # MambaDNC / stock dnc: hx passes through
+
+
+class CachedInferenceEngine(InferenceEngine):
+    def __init__(
+        self,
+        model,
+        output_proj,
+        task,
+        device,
+        caches,
+        ablate_memory: bool = False,
+        combiner_skip_stages: frozenset | None = None,
+        verify_hits: int = DEFAULT_VERIFY_HITS,
+        verify_tol: float = DEFAULT_VERIFY_TOL,
+    ):
+        super().__init__(model, output_proj, task, device, ablate_memory, combiner_skip_stages)
+        self.caches = list(caches)
+        self.resumable = model_supports_resume(model)
+        self.verify_hits, self.verify_tol = verify_hits, verify_tol
+        self.verify = {
+            "checked": 0,
+            "failures": 0,
+            "max_output_diff": 0.0,
+            "max_state_diff": 0.0,
+            "disabled_after_failure": False,
+        }
+        self.notes: list[str] = []
+        self._disabled = False
+        self._warned_bounds = False
+
+    # ---- small utils -----------------------------------------------------------
+    def _note(self, msg: str) -> None:
+        if msg not in self.notes:
+            self.notes.append(msg)
+            print(f"[cache] {msg}")
+
+    def _sync(self) -> None:
+        if self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
+
+    def cache_report(self, reset: bool = False) -> dict:
+        return {c.name: c.report(reset=reset) for c in self.caches}
+
+    def _select_active(self, reset_experience: bool):
+        active = []
+        for c in self.caches:
+            ok, why = c.applicable(reset_experience=reset_experience, resumable=self.resumable)
+            if ok:
+                active.append(c)
+            else:
+                self._note(f"{c.name} cache inactive: {why}")
+        return active
+
+    def _valid_boundaries(self, ep) -> list[int]:
+        T = ep.input_seq.shape[0]
+        bs = sorted({int(b) for b in ep.cache_boundaries})
+        ok = [b for b in bs if 0 < b < T and float(ep.mask[:b].sum()) == 0.0]
+        if len(ok) != len(bs) and not self._warned_bounds:
+            self._warned_bounds = True
+            self._note(
+                "some declared cache boundaries were dropped (out of range, or a scored "
+                "step lies before them)"
+            )
+        return ok
+
+    # ---- model execution ---------------------------------------------------------
+    @torch.no_grad()
+    def _forward_segment(
+        self, x_full: torch.Tensor, start: int, end: int, hidden, reset_experience: bool
+    ):
+        kw: dict[str, Any] = dict(
+            reset_experience=reset_experience, pass_through_memory=not self.ablate_memory
+        )
+        if self.combiner_skip_stages:
+            kw["combiner_skip_stages"] = self.combiner_skip_stages
+        if hasattr(self.model, "backbone"):  # SplitGraphDNC
+            x = x_full[:end].unsqueeze(0).to(self.device)
+            if start > 0:
+                kw["start_step"] = start
+        else:
+            x = x_full[start:end].unsqueeze(0).to(self.device)
+        output, new_hidden = self.model(x, hidden, **kw)
+        output = output.transpose(0, 1).contiguous().squeeze(0)
+        output = self.output_proj(output)
+        return output.float().cpu(), new_hidden
+
+    def _compute_full(self, ctx: EpisodeCtx, hidden):
+        """Uncached reference: one forward over the whole episode."""
+        x = ctx.episode.input_seq
+        t0 = time.perf_counter()
+        out, hid = self._forward_segment(x, 0, x.shape[0], hidden, ctx.reset_experience)
+        self._sync()
+        return out, hid, (time.perf_counter() - t0) * 1000.0
+
+    def _compute(self, ctx: EpisodeCtx, hidden, active):
+        x = ctx.episode.input_seq
+        T = x.shape[0]
+        use_bounds = ctx.reset_experience and self.resumable and bool(ctx.boundaries)
+        resume = None
+        if use_bounds:
+            for c in active:
+                r = c.find_resume(ctx)
+                if r is not None and (resume is None or r.position > resume.position):
+                    resume = r
+        start = resume.position if resume is not None else 0
+        hid = resume.hidden if resume is not None else hidden
+        seg_reset = ctx.reset_experience and resume is None
+        cuts = [b for b in ctx.boundaries if b > start] if use_bounds else []
+        cuts.append(T)
+
+        cold_ms = resume.saved_ms if resume is not None else 0.0  # cost as if computed from scratch
+        segs, pos = [], start
+        for cut in cuts:
+            t0 = time.perf_counter()
+            seg, hid = self._forward_segment(x, pos, cut, hid, seg_reset)
+            self._sync()
+            cold_ms += (time.perf_counter() - t0) * 1000.0
+            segs.append(seg)
+            if cut < T:  # snapshot BEFORE the next segment mutates state
+                for c in active:
+                    c.store_boundary(ctx, cut, hid, cold_ms)
+            pos, seg_reset = cut, False
+        out = torch.cat(segs, dim=0)
+        if start > 0:  # prefix rows are never scored (validated)
+            out = torch.cat([torch.zeros(start, out.shape[1]), out], dim=0)
+        return out, hid, resume, cold_ms
+
+    # ---- verification -------------------------------------------------------------
+    def _maybe_verify(self, ctx, incoming_hidden, out, new_hidden, event, state_valid):
+        v = self.verify
+        if self._disabled or v["checked"] >= self.verify_hits:
+            return out, new_hidden, event
+        ref_out, ref_hidden, _ = self._compute_full(ctx, clone_tree(incoming_hidden))
+        mask = ctx.episode.mask == 1
+        d_out = (out[mask] - ref_out[mask]).abs().max().item() if bool(mask.any()) else 0.0
+        d_state = tree_max_abs_diff(new_hidden, ref_hidden) if state_valid else 0.0
+        v["checked"] += 1
+        v["max_output_diff"] = max(v["max_output_diff"], d_out)
+        v["max_state_diff"] = max(v["max_state_diff"], d_state)
+        if d_out <= self.verify_tol and d_state <= self.verify_tol:  # NaN-safe form
+            return out, new_hidden, event
+        v["failures"] += 1
+        v["disabled_after_failure"] = True
+        self._disabled = True
+        self._note(
+            f"VERIFICATION FAILED at episode {ctx.index} ({event}): max|d output|={d_out:.3e}, "
+            f"max|d state|={d_state:.3e} > tol {self.verify_tol:g}. Caches disabled for the "
+            f"rest of this run; using the recomputed result."
+        )
+        return ref_out, ref_hidden, event + ":VERIFY-FAILED"
+
+    # ---- one episode --------------------------------------------------------------
+    def _run_one(self, ctx: EpisodeCtx, hidden, active):
+        for c in active:
+            hit = c.lookup_episode(ctx)
+            if hit is not None:
+                new_hidden = FRESH_HIDDEN if ctx.reset_experience else hit.hidden
+                return self._maybe_verify(
+                    ctx,
+                    hidden,
+                    hit.output,
+                    new_hidden,
+                    f"hit:{c.name}",
+                    state_valid=not ctx.reset_experience,
+                )
+        out, new_hidden, resume, cold_ms = self._compute(ctx, hidden, active)
+        event = "miss"
+        if resume is not None:
+            event = f"hit:{resume.cache_name}@{resume.position}"
+            out, new_hidden, event = self._maybe_verify(
+                ctx, hidden, out, new_hidden, event, state_valid=True
+            )
+        if not self._disabled:
+            for c in active:
+                c.store_episode(ctx, out, new_hidden, cold_ms)
+        return out, new_hidden, event
+
+    def run(
+        self,
+        episodes,
+        reset_experience: bool,
+        verbose_n: int = 0,
+        progress_every: int = 0,
+        label: str = "",
+    ) -> list[EpisodeResult]:
+        active = self._select_active(reset_experience)
+        want_bounds = any(c.wants_boundaries for c in active)
+        hidden, chain = FRESH_HIDDEN, ROOT_CHAIN
+        results: list[EpisodeResult] = []
+        run_items = run_correct = 0
+        for i, ep in enumerate(episodes):
+            if reset_experience:
+                hidden = FRESH_HIDDEN
+            t0 = time.perf_counter()
+            act = [] if self._disabled else active
+            ctx = EpisodeCtx(
+                index=i,
+                episode=ep,
+                reset_experience=reset_experience,
+                in_hash=hash_tensor(ep.input_seq) if act else "",
+                chain=chain,
+                boundaries=self._valid_boundaries(ep) if (want_bounds and act) else [],
+            )
+            out, hidden, event = self._run_one(ctx, hidden, act)
+            self._sync()
+            elapsed_ms = (time.perf_counter() - t0) * 1000.0
+            if not reset_experience:  # fresh mode: chain stays "root" (state is always fresh)
+                chain = chain_next(chain, ctx.in_hash)
+
+            score = self.task.score_episode(out, ep, verbose=(i < verbose_n))
+            results.append(EpisodeResult(i, score, elapsed_ms, reset_experience, cache_event=event))
+            run_items += score.n_items
+            run_correct += score.n_correct
+            if progress_every and ((i + 1) % progress_every == 0 or i + 1 == len(episodes)):
+                print(
+                    f"[engine{label}] episode {i + 1}/{len(episodes)} | "
+                    f"running item acc {100.0 * run_correct / max(run_items, 1):.2f}%"
+                )
+        return results

initium/inference/cache/keys.py

@@ -0,0 +1,77 @@
+"""
+file: inference/cache/keys.py
+
+Hashing for cache keys. blake2b (fast, stdlib). A cache key always contains the
+MODEL FINGERPRINT (hash of every weight + config + the flags that change the
+forward pass), so entries can never leak between different checkpoints.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any
+
+import torch
+
+
+def new_hasher():
+    return hashlib.blake2b(digest_size=16)
+
+
+def update_with_tensor(h, t: torch.Tensor) -> None:
+    t = t.detach().cpu().contiguous()
+    h.update(str(t.dtype).encode())
+    h.update(str(tuple(t.shape)).encode())
+    if t.numel():
+        h.update(t.reshape(-1).view(torch.uint8).numpy())
+
+
+def hash_tensor(t: torch.Tensor) -> str:
+    h = new_hasher()
+    update_with_tensor(h, t)
+    return h.hexdigest()
+
+
+def digest(*parts: Any) -> str:
+    h = new_hasher()
+    for p in parts:
+        h.update(str(p).encode())
+        h.update(b"\x1f")
+    return h.hexdigest()
+
+
+def chain_next(chain: str, in_hash: str) -> str:
+    """Persistent mode: the state after episode k is a function of (initial state,
+    inputs 0..k). Hashing the inputs into a chain therefore identifies the state
+    exactly, without hashing the (large) memory tensors every episode."""
+    return digest(chain, in_hash)
+
+
+def _update_obj(h, obj: Any) -> None:
+    if isinstance(obj, torch.Tensor):
+        update_with_tensor(h, obj)
+    elif isinstance(obj, dict):
+        for k in sorted(obj, key=str):
+            h.update(str(k).encode())
+            _update_obj(h, obj[k])
+    elif isinstance(obj, list | tuple):
+        for v in obj:
+            _update_obj(h, v)
+    else:
+        h.update(repr(obj).encode())
+
+
+def model_fingerprint(ckpt: dict, extras: dict) -> str:
+    """Hash of weights (rnn + output_proj + prior buffers) + model_config + `extras`
+    (flags that alter the forward pass: deterministic_write, ablate_memory, ...).
+    Takes a couple of seconds for a large model; only computed when caching is on.
+    Deliberately NOT included: device / torch version (tiny numeric differences are
+    what --cache-verify is for)."""
+    h = new_hasher()
+    _update_obj(h, ckpt["rnn_state_dict"])
+    _update_obj(h, ckpt["output_proj_state_dict"])
+    _update_obj(h, ckpt.get("prior_state"))
+    h.update(json.dumps(ckpt["model_config"], sort_keys=True, default=str).encode())
+    h.update(json.dumps(extras, sort_keys=True, default=str).encode())
+    return h.hexdigest()

initium/inference/cache/prefix_cache.py

@@ -0,0 +1,59 @@
+"""
+file: inference/cache/prefix_cache.py
+
+Prefix STATE cache (the DNC analogue of a KV / context cache). After the static
+prefix of an episode (e.g. the whole edge listing) has been processed, the complete
+state (controller state, Memory dict incl. link matrix / usage / weights, last read
+vector) is snapshotted, keyed by (model fingerprint, hash of the prefix tokens).
+Later episodes with the SAME prefix restore it and process only their suffix.
+
+Valid because the model is causal and deterministic: state after step p depends only
+on inputs [0, p). Requires reset_experience=True (identical start state) and a task
+that declares cache_boundaries. If several boundaries are declared, the LONGEST cached
+one is used (longest-prefix match) and the remaining ones are stored on the way.
+"""
+
+from __future__ import annotations
+
+from initium.inference.cache.base_cache import BaseCache, EpisodeCtx, ResumePoint
+from initium.inference.cache.state_utils import restore, snapshot
+
+
+class PrefixStateCache(BaseCache):
+    name = "prefix"
+    wants_boundaries = True
+
+    def applicable(self, *, reset_experience: bool, resumable: bool) -> tuple[bool, str]:
+        if not reset_experience:
+            return False, (
+                "needs reset_experience=True: in persistent mode every episode starts "
+                "from a different state, so a prefix snapshot is never reusable"
+            )
+        if not resumable:
+            return False, (
+                "this model cannot resume mid-sequence (SplitGraphDNC needs the "
+                "start_step edit in split_graph_dnc.py)"
+            )
+        return True, ""
+
+    def _key_for(self, ctx: EpisodeCtx, boundary: int) -> str:
+        return self._key(boundary, ctx.prefix_hash(boundary))
+
+    def find_resume(self, ctx: EpisodeCtx) -> ResumePoint | None:
+        if not ctx.boundaries:
+            return None
+        self.stats.lookups += 1
+        for b in reversed(ctx.boundaries):  # longest prefix first
+            got = self._get(self._key_for(ctx, b))
+            if got is None:
+                continue
+            entry, tier = got
+            self.stats.hit(tier, entry["compute_ms"])
+            return ResumePoint(
+                b, restore(entry["payload"], self.device), entry["compute_ms"], self.name
+            )
+        self.stats.misses += 1
+        return None
+
+    def store_boundary(self, ctx: EpisodeCtx, boundary: int, hidden, compute_ms: float) -> None:
+        self._put(self._key_for(ctx, boundary), snapshot(hidden), compute_ms)

initium/inference/cache/result_cache.py

@@ -0,0 +1,57 @@
+"""
+file: inference/cache/result_cache.py
+
+Result / lookup cache (GATI-style): memoizes the model output of a whole episode.
+
+  fresh mode      : key = (model fp, "F", hash of the episode inputs). A hit skips the
+                    forward pass entirely.
+  persistent mode : the output also depends on the carried-over state, so the key adds
+                    a history CHAIN hash (see keys.chain_next) and the entry also stores
+                    the resulting state, so the run can continue after a hit.
+
+Exact-match only, deterministic models only (the setup gate refuses stochastic write
+heads unless overridden).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from initium.inference.cache.base_cache import BaseCache, EpisodeCtx, EpisodeHit
+from initium.inference.cache.state_utils import restore, snapshot
+
+
+class ResultCache(BaseCache):
+    name = "result"
+
+    def applicable(self, *, reset_experience: bool, resumable: bool) -> tuple[bool, str]:
+        return True, ""
+
+    def _key_for(self, ctx: EpisodeCtx) -> str:
+        return self._key("F" if ctx.reset_experience else "P", ctx.chain, ctx.in_hash)
+
+    def lookup_episode(self, ctx: EpisodeCtx) -> EpisodeHit | None:
+        self.stats.lookups += 1
+        got = self._get(self._key_for(ctx))
+        if got is None:
+            self.stats.misses += 1
+            return None
+        entry, tier = got
+        payload = entry["payload"]
+        if not ctx.reset_experience and payload.get("hidden") is None:
+            self.stats.misses += 1
+            return None
+        self.stats.hit(tier, entry["compute_ms"])
+        hidden = (
+            restore(payload["hidden"], self.device) if payload.get("hidden") is not None else None
+        )
+        return EpisodeHit(payload["output"].clone(), hidden, entry["compute_ms"])
+
+    def store_episode(
+        self, ctx: EpisodeCtx, output: torch.Tensor, hidden_after, compute_ms: float
+    ) -> None:
+        payload = {
+            "output": output.detach().cpu().clone(),
+            "hidden": None if ctx.reset_experience else snapshot(hidden_after),
+        }
+        self._put(self._key_for(ctx), payload, compute_ms)

initium/inference/cache/state_utils.py

@@ -0,0 +1,74 @@
+"""
+file: inference/cache/state_utils.py
+
+Helpers for the nested model state (chx, mhx, last_read): controller states
+(tuples/lists), the Memory dict, tensors, None. Snapshots are stored on CPU and
+ALWAYS deep-copied on both sides: Memory.reset(..., hidden) rebinds entries of the
+dict it is given and the write path rebinds them again, so a cached dict passed
+straight into forward() would silently change under us.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+
+def tree_map(fn: Callable[[torch.Tensor], Any], obj: Any) -> Any:
+    if isinstance(obj, torch.Tensor):
+        return fn(obj)
+    if isinstance(obj, dict):
+        return {k: tree_map(fn, v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [tree_map(fn, v) for v in obj]
+    if isinstance(obj, tuple):
+        items = [tree_map(fn, v) for v in obj]
+        return type(obj)(*items) if hasattr(obj, "_fields") else tuple(items)
+    return obj  # None, numbers, strings
+
+
+def snapshot(tree: Any) -> Any:  # -> independent CPU copy
+    return tree_map(lambda t: t.detach().to("cpu", copy=True), tree)
+
+
+def restore(tree: Any, device: torch.device) -> Any:  # -> independent copy on `device`
+    return tree_map(lambda t: t.detach().to(device, copy=True), tree)
+
+
+def clone_tree(tree: Any) -> Any:
+    return tree_map(lambda t: t.detach().clone(), tree)
+
+
+def tree_nbytes(obj: Any) -> int:
+    if isinstance(obj, torch.Tensor):
+        return obj.numel() * obj.element_size()
+    if isinstance(obj, dict):
+        return sum(tree_nbytes(v) for v in obj.values())
+    if isinstance(obj, list | tuple):
+        return sum(tree_nbytes(v) for v in obj)
+    return 0
+
+
+def tree_max_abs_diff(a: Any, b: Any) -> float:
+    """Max |a-b| over all tensors; inf on any structural mismatch; NaN propagates
+    (callers must test `diff <= tol`, never `diff > tol`)."""
+    if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
+        if a.shape != b.shape:
+            return math.inf
+        if a.numel() == 0:
+            return 0.0
+        return (a.detach().float() - b.detach().float().to(a.device)).abs().max().item()
+    if isinstance(a, dict) and isinstance(b, dict):
+        if a.keys() != b.keys():
+            return math.inf
+        return max([0.0] + [tree_max_abs_diff(a[k], b[k]) for k in a])
+    if isinstance(a, list | tuple) and isinstance(b, list | tuple):
+        if len(a) != len(b):
+            return math.inf
+        return max([0.0] + [tree_max_abs_diff(x, y) for x, y in zip(a, b)])
+    if a is None and b is None:
+        return 0.0
+    return 0.0 if a == b else math.inf