initium-deeplearning 0.1.0__py3-none-any.whl

initium/memory_manipulation/dynamic_memory_resize.py

@@ -0,0 +1,157 @@
+"""
+file: dynamic_memory_resize.py
+
+Static Option 2: rebuild a dnc.DNC's Memory submodule at a new nr_cells,
+transplanting every N-independent learned sublayer so no trained weight is
+lost. Safe because Memory's per-episode state (memory/link_matrix/
+precedence/usage_vector/read_weights/write_weights) is rebuilt from scratch
+by Memory.reset() every forward call regardless -- there is no cross-episode
+memory *content* to preserve, only trained *weights*, and none of those are
+sized by nr_cells.
+
+Also transplants Option 1's link-matrix patch (class + mode) onto the
+rebuilt Memory when the old one was patched -- see the isinstance check
+below and link_matrix_ablation.py's module docstring for why this is
+needed: without it, a resize event mid-run would silently revert an active
+link-matrix ablation/sparsification back to dense behavior.
+
+--- optimizer resync (added for Dynamic-N) --------------------------------
+`model.memories[layer] = new_memory` swaps the module tree, but the
+training-loop optimizer (`torch.optim.Adam(list(rnn.parameters()) + ...)`)
+was built ONCE, before this function was ever called, from the Parameter
+OBJECTS live at that time. load_state_dict() above only copies tensor
+*values* into new_memory's freshly-constructed Parameters -- it never makes
+optimizer.param_groups aware of them. So without the `optimizer=` arg
+below, every Memory sublayer except write_vector_transform (moved over as
+the same object, so already tracked) silently stops receiving Adam updates
+for the rest of the run after ANY live mid-run call to this function:
+gradients still compute (they're in the forward graph), but
+optimizer.step() only updates params actually present in a param_group,
+and the newly-constructed ones never are.
+
+This was already true for static Option 2's curriculum-indexed resize
+(TraversalCurriculum.maybe_advance() calls this function without
+`optimizer=`, ~13 times over a 120k-step run) -- it just wasn't obvious in
+aggregate metrics from a handful of one-time freezes near a lesson
+boundary. It's NOT a problem on a checkpoint *resume*, because
+`optimizer.load_state_dict(ckpt["optimizer_state_dict"])` is called right
+after resize_memory() there and rebuilds `state` keyed to the CURRENT
+param objects, matched positionally against the checkpoint's saved
+param_groups (same architecture => same order => this just works). It's
+specifically live, same-process, mid-run resize calls that orphan params --
+which is exactly what Dynamic-N's whole premise is: resizing far more
+often, live, mid-run. So this fix ships as part of Dynamic-N; the existing
+static-Option-2 call site can opt in by passing `optimizer=optimizer` too
+(not changed automatically here, since that changes reproducibility of
+already-running static-Option-2 experiments).
+"""
+
+import torch
+import torch.nn as nn
+from dnc.memory import Memory
+
+from initium.memory_manipulation.link_matrix_ablation import AblatableSparseLinkMemory
+
+
+def _resync_optimizer_after_resize(optimizer, old_named_params: dict, new_memory) -> None:
+    """Repoint optimizer.param_groups at new_memory's Parameter objects
+    wherever a same-named old parameter existed, and migrate that
+    parameter's Adam state (exp_avg, exp_avg_sq, step) onto the new object
+    so momentum isn't reset by a resize. Safe to do unconditionally here
+    because no Memory Linear sublayer's shape depends on nr_cells (see
+    module docstring) -- every migrated tensor is the same shape before and
+    after.
+
+    write_vector_transform is a no-op under this function: `new_p is
+    old_p` for its parameters (the module object itself was moved, not
+    rebuilt -- see resize_memory() below), so it's skipped by the `is not`
+    check and stays exactly as-is in both param_groups and optimizer.state.
+    """
+    new_named_params = dict(new_memory.named_parameters())
+    replace_map = {}  # old Parameter object -> new Parameter object
+    for name, old_p in old_named_params.items():
+        new_p = new_named_params.get(name)
+        if new_p is not None and new_p is not old_p:
+            replace_map[old_p] = new_p
+
+    for old_p, new_p in replace_map.items():
+        if old_p in optimizer.state:
+            optimizer.state[new_p] = optimizer.state.pop(old_p)
+
+    replace_ids = {id(old_p): new_p for old_p, new_p in replace_map.items()}
+    for group in optimizer.param_groups:
+        group["params"] = [replace_ids.get(id(p), p) for p in group["params"]]
+
+
+def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimizer=None):
+    if device is not None:
+        device = torch.device(device)
+    old_memory: Memory = model.memories[layer]  # VERIFY attribute name
+    if old_memory.nr_cells == new_nr_cells:  # VERIFY kwarg/attr name
+        return old_memory
+
+    # Captured before write_vector_transform is swapped to a placeholder
+    # below, so this reflects old_memory's real (pre-mutation) structure --
+    # used by the optimizer resync at the very end of this function.
+    old_named_params = dict(old_memory.named_parameters())
+
+    new_memory = Memory(
+        input_size=old_memory.input_size,
+        nr_cells=new_nr_cells,  # VERIFY kwarg name
+        cell_size=old_memory.cell_size,
+        read_heads=old_memory.read_heads,
+        independent_linears=True,
+        device=device,  # VERIFY kwarg name -- see note below
+    ).to(device)
+
+    # Transplant every N-independent learned sublayer except
+    # write_vector_transform (that one is a StochasticWriteHead, not a
+    # plain Linear, on old_memory -- move the module object itself instead
+    # of trying to load_state_dict it).
+    old_write_head = old_memory.write_vector_transform
+    old_memory.write_vector_transform = nn.Identity()  # placeholder so it's excluded below
+    new_memory.load_state_dict(old_memory.state_dict(), strict=False)
+    new_memory.write_vector_transform = old_write_head  # move the actual module
+
+    # Static Option 1 interaction: if old_memory was patched by
+    # link_matrix_ablation.patch_link_matrix (its __class__ reassigned to
+    # AblatableSparseLinkMemory, carrying link_matrix_mode / optionally
+    # link_matrix_topk as plain instance attributes), new_memory above is a
+    # *different* freshly-constructed plain dnc.memory.Memory object --
+    # load_state_dict only copies tensor weights, never the Python class or
+    # plain instance attributes, so the patch would silently be dropped
+    # right here, reverting new_memory to baseline dense link-matrix
+    # behavior with no error or warning. Re-apply it explicitly so a resize
+    # event (Option 2) can never silently undo an active link-matrix
+    # ablation/sparsification (Option 1) mid-run.
+    if isinstance(old_memory, AblatableSparseLinkMemory):
+        new_memory.__class__ = AblatableSparseLinkMemory
+        new_memory.link_matrix_mode = old_memory.link_matrix_mode
+        if hasattr(old_memory, "link_matrix_topk"):
+            new_memory.link_matrix_topk = old_memory.link_matrix_topk
+
+    model.memories[layer] = new_memory
+    # model.memories is a plain python list, not an nn.ModuleList -- the
+    # line above rebinds the list slot but never touches whatever attribute
+    # PyTorch's own submodule registry (_modules) actually points at
+    # (rnn_layer_memory_shared, or rnn_layer_memory_<layer> when memories
+    # aren't shared -- set once via setattr() at model construction, see
+    # mamba_controller.py / dnc.DNC.__init__). Without this, that attribute
+    # stays pointed at old_memory forever -- the same object whose
+    # write_vector_transform was just swapped to nn.Identity() above as a
+    # transplant placeholder -- so rnn.state_dict()/rnn.parameters()/
+    # rnn.load_state_dict() silently keep walking a stale, Identity-headed
+    # module while forward() (which reads self.memories[layer] directly)
+    # correctly uses new_memory. Found via a resume-time crash: dynamic-N's
+    # floor-start construction forces a resize on resume that a static-only
+    # workflow never used to trigger, which is what surfaced this -- but
+    # the same corruption happens on any LIVE resize, including every
+    # static-Option-2 lesson advance that ever fired.
+    for attr_name, submodule in list(model._modules.items()):
+        if submodule is old_memory:
+            setattr(model, attr_name, new_memory)
+
+    if optimizer is not None:
+        _resync_optimizer_after_resize(optimizer, old_named_params, new_memory)
+
+    return new_memory

initium/memory_manipulation/dynamic_n_controller.py

@@ -0,0 +1,152 @@
+"""
+file: dynamic_n_controller.py
+
+Dynamic-N, macro-scale (between episodes): grow nr_cells when the model's
+own usage_vector says it needs more capacity, instead of reading a fixed
+per-lesson lookup table (LESSON_NR_CELLS / static Option 2). Per the
+strategy doc ("Whether Option 1 can be made dynamic the same way"), the
+macro-scale case is mechanically just static Option 2's resize mechanism
+(dynamic_memory_resize.resize_memory) with a different *trigger* -- no new
+tensor-resizing logic is needed here, only the decision of *when* to call
+it. This module owns only that decision.
+
+--- Why this is split into record() / should_grow() / mark_grown() --------
+The eventual mid-episode (micro-scale) variant needs the exact same
+"sustained saturation, not a single-step spike, with a cooldown after each
+growth event" decision logic (see the strategy doc's point 1 and point 3),
+just fed from a per-timestep usage reading instead of a per-episode one,
+and executing a live in-place pad instead of resize_memory()'s
+rebuild-and-swap. Keeping the decision core (this class) ignorant of *when*
+it's called and *how* growth is physically applied means that future
+variant only has to replace the call site (what feeds record(), and what
+apply_growth does), not re-derive the EMA/cooldown/ceiling logic:
+
+  - record(frac_saturated): update the state used by should_grow(). Caller
+    decides the cadence (once per episode here; once per timestep for the
+    micro-scale variant later).
+  - should_grow(): pure decision, no side effects, no knowledge of *how*
+    growth is applied. Returns the next nr_cells or None.
+  - mark_grown(step, old_n, new_n): call AFTER growth has actually been
+    applied (by whatever mechanism), to reset the cooldown/EMA and append
+    to growth_history. Decoupled from resize_memory() specifically, so a
+    future in-place-pad apply function can call it too.
+
+The actual resize call (dynamic_memory_resize.resize_memory) happens at the
+call site in Alter_PHASE3_mamba.py's training loop, not in this file --
+same "this module doesn't do the tensor work" split link_matrix_ablation.py
+uses relative to dynamic_memory_resize.py.
+"""
+
+
+class DynamicNController:
+    """Usage-triggered nr_cells growth, macro-scale (between episodes).
+
+    Trigger: an EMA of "fraction of memory cells with usage_vector above
+    usage_high", computed fresh each time record() is called. EMA (not the
+    raw per-step value) is the sustaining filter the strategy doc's point 1
+    calls for -- a single saturated episode doesn't trigger growth, only a
+    persistently saturated one does, which also naturally damps against
+    growing too eagerly (the strategy doc's other flagged failure mode:
+    re-importing Option 2's original "excess capacity is an active cost"
+    problem, Concept 17).
+
+    Growth: multiplicative (nr_cells *= growth_factor each event, clamped
+    to ceiling), matching the strategy doc's own 128->256->512 example.
+    Floor is enforced only at construction (the starting nr_cells passed to
+    the model) -- this controller only ever grows, never shrinks; shrinking
+    an active run is a different, harder problem (would need to reconcile
+    with live content) and isn't part of this task.
+
+    Cooldown: after every growth event, should_grow() returns None for
+    cooldown_steps worth of record() calls, regardless of EMA, so the
+    controller doesn't fire again while the model's addressing policy is
+    still re-settling from the last resize (strategy doc point 3,
+    distribution-shift-at-growth risk). The EMA is also reset to 0 on
+    growth so the post-growth cooldown window starts from a clean read
+    rather than the already-saturated value that just triggered it.
+    """
+
+    def __init__(
+        self,
+        floor: int,
+        ceiling: int,
+        growth_factor: float = 2.0,
+        trigger_frac: float = 0.75,
+        ema_decay: float = 0.98,
+        cooldown_steps: int = 2000,
+    ):
+        if floor > ceiling:
+            raise ValueError(f"DynamicNController: floor ({floor}) > ceiling ({ceiling})")
+        self.floor = floor
+        self.ceiling = ceiling
+        self.growth_factor = growth_factor
+        self.trigger_frac = trigger_frac
+        self.ema_decay = ema_decay
+        self.cooldown_steps = cooldown_steps
+
+        self.ema: float = 0.0
+        self.cooldown_remaining: int = 0
+        # (step, old_n, new_n) for every growth event this run has applied
+        # (restored from checkpoint on resume -- see state_dict/load_state_dict).
+        self.growth_history: list[tuple[int, int, int]] = []
+
+    def record(self, frac_saturated: float) -> None:
+        """Update the sustained-saturation EMA. Call once per decision
+        cadence (once per episode/training-step for macro-scale) with the
+        current episode's fraction of memory cells above the usage-high
+        threshold -- see the training-loop call site for how that fraction
+        is computed from usage_vector.
+        """
+        self.ema = self.ema_decay * self.ema + (1.0 - self.ema_decay) * frac_saturated
+        if self.cooldown_remaining > 0:
+            self.cooldown_remaining -= 1
+
+    def should_grow(self, current_nr_cells: int) -> int | None:
+        """Returns the next nr_cells to resize to, or None if no growth
+        should happen right now. Pure decision -- does not itself resize
+        anything or mutate controller state (mark_grown does that, and only
+        after the caller has actually applied the growth).
+        """
+        if self.cooldown_remaining > 0:
+            return None
+        if current_nr_cells >= self.ceiling:
+            return None
+        if self.ema < self.trigger_frac:
+            return None
+        proposed = int(min(current_nr_cells * self.growth_factor, self.ceiling))
+        if proposed <= current_nr_cells:
+            return None
+        return proposed
+
+    def mark_grown(self, step: int, old_n: int, new_n: int) -> None:
+        """Call immediately after growth has actually been applied (by
+        resize_memory() at the macro-scale call site, or by whatever
+        mid-episode mechanism the micro-scale variant eventually uses).
+        Resets the EMA and starts the post-growth cooldown.
+        """
+        self.growth_history.append((step, old_n, new_n))
+        self.ema = 0.0
+        self.cooldown_remaining = self.cooldown_steps
+
+    def state_dict(self) -> dict:
+        """For checkpointing -- see save_checkpoint()'s dynamic_n_state key.
+        Deliberately does NOT include floor/ceiling/growth_factor/
+        trigger_frac/ema_decay/cooldown_steps: those are run configuration
+        (re-supplied from CLI args / module constants on resume, same
+        convention as link_matrix_mode/topk are re-passed rather than
+        read back out of the checkpoint), not run *state*. Current
+        nr_cells is also deliberately NOT stored here -- it's already
+        self-describing via model_config["nr_cells"] (read live from
+        rnn.memories[0].nr_cells), the same single-source-of-truth
+        convention static Option 2's checkpoint format already uses.
+        """
+        return {
+            "ema": self.ema,
+            "cooldown_remaining": self.cooldown_remaining,
+            "growth_history": list(self.growth_history),
+        }
+
+    def load_state_dict(self, state: dict) -> None:
+        self.ema = state.get("ema", 0.0)
+        self.cooldown_remaining = state.get("cooldown_remaining", 0)
+        self.growth_history = list(state.get("growth_history", []))

initium/memory_manipulation/link_matrix_ablation.py

@@ -0,0 +1,130 @@
+"""
+file: link_matrix_ablation.py
+
+Static Option 1: ablate or sparsify the DNC's temporal link matrix at fixed
+N=256, per Session 009's synthesis finding that temporal-linkage addressing
+is near-unused in QA-style (single-hop, content-addressed) tasks. Graph
+traversal chains multiple hops per episode, which is exactly the case that
+finding doesn't directly cover -- see the hop-breakdown logging in
+Alter_PHASE3_mamba.py (TraversalCurriculum.maybe_advance, plus the periodic
+and terminal OOD evaluate_traversal calls) for the diagnostic that watches
+for a hop-count-correlated accuracy collapse this synthesis finding
+wouldn't predict.
+
+This does NOT touch nr_cells (that's Option 2 / dynamic_memory_resize.py)
+and does NOT resize any tensor -- it only changes what get_link_matrix()
+computes at whatever N the Memory was already built at. Applied by
+reassigning __class__ on an already-constructed dnc.memory.Memory instance
+(see patch_link_matrix() below), not by editing the installed dnc package --
+same external-library-as-black-box precedent as dynamic_memory_resize.py's
+weight-transplant approach, just simpler here since no weights move.
+
+INTERACTION WITH OPTION 2: dynamic_memory_resize.resize_memory() builds a
+brand-new plain dnc.memory.Memory instance and swaps it into
+model.memories[layer]. Its own weight-transplant (load_state_dict) has no
+notion of link-matrix mode, so without extra handling a resize event would
+silently drop back to "dense" mid-run. resize_memory() now checks
+isinstance(old_memory, AblatableSparseLinkMemory) and re-applies the class
+and link_matrix_mode/link_matrix_topk onto the freshly-built new_memory --
+see that file. This module doesn't need to do anything for that to work;
+it's noted here so the coupling isn't a surprise when reading either file
+in isolation.
+"""
+
+import torch
+from dnc.memory import Memory
+
+
+class AblatableSparseLinkMemory(Memory):
+    """Drop-in replacement for get_link_matrix() only. Every other Memory
+    method (write, read, content_weightings, directional_weightings, ...)
+    is inherited unchanged, and this class adds no new nn.Module state --
+    so it's safe to attach to an already-constructed Memory instance via
+    __class__ reassignment (see patch_link_matrix) instead of
+    reconstructing/reloading weights.
+
+    Mode is read from instance attributes (set by patch_link_matrix, not
+    passed as get_link_matrix() args), so this overrides cleanly against
+    the existing call signature in Memory.write() -- no caller changes
+    needed anywhere else in dnc.py or memory.py.
+
+    link_matrix_mode:
+      "dense"   -- unchanged Memory behavior (baseline; same as not
+                   patching at all).
+      "ablated" -- link matrix is never updated; it stays at its initial
+                   zero value (from Memory.new()/reset()) for the whole
+                   episode. directional_weightings() (unchanged, inherited)
+                   then always returns zero forward/backward weightings, so
+                   read_weightings' content_mode term is the only nonzero
+                   contribution -- content-based addressing continues to
+                   work exactly as before, temporal addressing contributes
+                   nothing. This is the "ablate" half of the roadmap's
+                   "ablate or aggressively sparsify."
+      "sparse_topk" -- link matrix is updated as normal every write() call,
+                   but immediately after each update, every row keeps only
+                   its link_matrix_topk largest-magnitude entries and zeros
+                   the rest. This is the "sparsify" half -- reduces the
+                   O(N^2) footprint the roadmap's citation flags, while
+                   still letting the model use *some* temporal signal,
+                   unlike full ablation. Recomputed fresh every step (not a
+                   one-time structural prune), so it tracks whichever links
+                   the model is actually relying on at that point in the
+                   episode rather than freezing an initial guess.
+    """
+
+    def get_link_matrix(
+        self, link_matrix: torch.Tensor, write_weights: torch.Tensor, precedence: torch.Tensor
+    ) -> torch.Tensor:
+        mode = getattr(self, "link_matrix_mode", "dense")
+
+        if mode == "dense":
+            return super().get_link_matrix(link_matrix, write_weights, precedence)
+
+        if mode == "ablated":
+            return link_matrix  # never updated -- stays at reset()'s zero value
+
+        if mode == "sparse_topk":
+            updated = super().get_link_matrix(link_matrix, write_weights, precedence)
+            k = getattr(self, "link_matrix_topk", None)
+            if k is None or k >= self.nr_cells:
+                return updated  # no-op if topk isn't set or wouldn't actually sparsify anything
+            # updated: (batch, 1, N, N). Keep each row's top-k
+            # largest-magnitude entries, zero the rest.
+            magnitudes = updated.abs()
+            _, topk_idx = magnitudes.topk(k, dim=-1)
+            mask = torch.zeros_like(updated)
+            mask.scatter_(-1, topk_idx, 1.0)
+            return updated * mask
+
+        raise ValueError(f"Unknown link_matrix_mode: {mode!r}")
+
+
+def patch_link_matrix(model, mode: str, topk: int | None = None, layer: int | None = None):
+    """Apply Static Option 1 to an already-constructed dnc.DNC (or MambaDNC,
+    which subclasses DNC and builds model.memories identically -- see
+    mamba_controller.py) instance, right after model construction.
+
+    mode: "dense" | "ablated" | "sparse_topk" -- see AblatableSparseLinkMemory
+        docstring. "dense" is accepted (as a no-op) so a run can be flipped
+        back to baseline behavior via the same config/CLI path, for the
+        roadmap's required baseline-vs-isolated-ablation comparison,
+        without touching any other run() plumbing.
+    topk: required (and only used) when mode == "sparse_topk".
+    layer: which model.memories[i] to patch. None (default) patches every
+        entry in model.memories -- correct for share_memory_between_layers
+        =True (this project's setup: a single shared Memory, so memories
+        has exactly one entry) and still correct if share_memory_between_
+        layers is ever False (patches every per-layer Memory identically).
+        Pass an int to patch only one layer's Memory when memories are NOT
+        shared and you want just one layer's link matrix affected.
+    """
+    if mode == "sparse_topk" and topk is None:
+        raise ValueError("patch_link_matrix: topk is required when mode='sparse_topk'")
+
+    targets = model.memories if layer is None else [model.memories[layer]]
+    for memory in targets:
+        memory.__class__ = AblatableSparseLinkMemory
+        memory.link_matrix_mode = mode
+        if topk is not None:
+            memory.link_matrix_topk = topk
+    return model

initium/memory_manipulation/nvrtc_compat.py

@@ -0,0 +1,143 @@
+"""
+file: memory_manipulation/nvrtc_compat.py
+
+Environment-compatibility shim, NOT a modeling change. Works around a
+CUDA/NVRTC library mismatch observed in this environment ("nvrtc: error:
+failed to open libnvrtc-builtins.so.13.0" - the installed torch build's
+CUDA minor version has no matching nvrtc-builtins shared library present,
+a known failure mode after a partial torch/CUDA reinstall or a fresh
+Colab runtime whose nvidia-cuda-nvrtc-cuXX package doesn't match the
+torch wheel's CUDA build).
+
+Root cause: on CUDA, `torch.prod(x, dim=...)` and `torch.cumprod(x,
+dim=...)` (reductions/scans along a non-trivial dim) dispatch through
+PyTorch's Jiterator -- a kernel compiled at PROCESS RUNTIME via NVRTC --
+rather than a statically pre-compiled kernel. Any environment where NVRTC
+can't find its builtins library crashes the FIRST time either op runs,
+regardless of model architecture, controller type, or MoE config.
+
+`dnc.memory.Memory` (third-party, installed via pip -- not a project
+file) calls both of these on the hot path of every single write():
+  - get_usage_vector: torch.prod(1 - free_gates*read_weights, dim=1) and
+    torch.prod(1 - write_weights, dim=1)
+  - allocation weighting: torch.cumprod(sorted_usage, dim=1) (Graves et
+    al. 2016's free-list product term, `prod_{i<j} u[phi[i]]`)
+
+Rather than reimplementing each of dnc.memory's internal methods by hand
+(fragile: depends on exactly matching an undocumented third-party
+formula, and any future call site is invisible until it crashes), this
+module patches `torch.prod` / `torch.cumprod` / `Tensor.prod` /
+`Tensor.cumprod` THEMSELVES. Every call inside dnc/memory.py - known or
+not yet discovered - is transparently redirected, with zero changes
+needed to the third-party package or to any call site in this project.
+
+Fix mechanism: for a CUDA tensor with an explicit `dim`, replace the
+product/cumulative-product with the equivalent log-space reduction
+exp(sum(log(clamp(x, min=EPS)), dim)) or exp(cumsum(log(clamp(x)), dim)).
+log, sum/cumsum, and exp are ordinary elementwise/reduction ops with
+statically pre-compiled CUDA kernels for every dtype PyTorch ships - no
+Jiterator involved - so this sidesteps the missing-nvrtc-builtins
+failure without requiring a CUDA toolkit reinstall inside the notebook.
+
+Safety net: the log-space substitution is only mathematically valid for
+NON-NEGATIVE inputs (log of a negative number is NaN). Every known call
+site in dnc.memory operates on gates/usages/weights, which are always in
+[0,1] - but to stay correct even for an unknown future call site (some
+other library, or an edit to dnc.memory) with genuinely negative values,
+the patched functions check `torch.isnan(result).any()` after the
+log-space computation and, on a hit, recompute via the ORIGINAL op on a
+CPU copy of the input (never re-invoking the broken CUDA op) before
+moving the result back to the original device. This costs a device sync
++ small copy only in that (expected-never, for this project) fallback
+path - the common case never sees it.
+
+Call patch_prod_jiterator() once, before running any forward pass -
+idempotent, so importing this from multiple entry points (core_training.py,
+run_inference.py) is safe. If the underlying environment issue is fixed
+instead (matching nvidia-cuda-nvrtc-cuXX package reinstalled), this patch
+is a harmless no-op difference in float rounding at the ~1e-6 level, not
+a correctness regression - there is no reason to remove it once applied.
+"""
+
+from __future__ import annotations
+
+import torch
+
+_EPS = 1e-6
+_PATCHED_ATTR = "_nvrtc_compat_patched"
+
+_orig_prod = torch.prod
+_orig_cumprod = torch.cumprod
+_orig_tensor_prod = torch.Tensor.prod
+_orig_tensor_cumprod = torch.Tensor.cumprod
+
+
+def _log_space(input: torch.Tensor, dim, keepdim: bool, cumulative: bool) -> torch.Tensor:
+    orig_dtype = input.dtype
+    work = input if input.dtype.is_floating_point else input.float()
+    logged = torch.log(work.clamp(min=_EPS))
+    reduced = (
+        torch.cumsum(logged, dim=dim) if cumulative else torch.sum(logged, dim=dim, keepdim=keepdim)
+    )
+    return torch.exp(reduced).to(orig_dtype)
+
+
+def _patched_prod(input, dim=None, keepdim=False, *, dtype=None):
+    if isinstance(input, torch.Tensor) and input.is_cuda and dim is not None:
+        result = _log_space(input, dim, keepdim, cumulative=False)
+        if torch.isnan(result).any():
+            # Negative-valued input -- log-space isn't valid here. Fall back
+            # to the ORIGINAL op on CPU (never re-run the broken CUDA path).
+            cpu_out = _orig_prod(
+                input.detach().cpu(),
+                dim,
+                keepdim=keepdim,
+                **({"dtype": dtype} if dtype is not None else {}),
+            )
+            return cpu_out.to(input.device)
+        return result.to(dtype) if dtype is not None else result
+    if dim is None:
+        return _orig_prod(input, **({"dtype": dtype} if dtype is not None else {}))
+    return _orig_prod(
+        input, dim, keepdim=keepdim, **({"dtype": dtype} if dtype is not None else {})
+    )
+
+
+def _patched_cumprod(input, dim, *, dtype=None):
+    if isinstance(input, torch.Tensor) and input.is_cuda:
+        result = _log_space(input, dim, keepdim=False, cumulative=True)
+        if torch.isnan(result).any():
+            cpu_out = _orig_cumprod(
+                input.detach().cpu(), dim, **({"dtype": dtype} if dtype is not None else {})
+            )
+            return cpu_out.to(input.device)
+        return result.to(dtype) if dtype is not None else result
+    return _orig_cumprod(input, dim, **({"dtype": dtype} if dtype is not None else {}))
+
+
+def _patched_tensor_prod(self, dim=None, keepdim=False, *, dtype=None):
+    return _patched_prod(self, dim, keepdim, dtype=dtype)
+
+
+def _patched_tensor_cumprod(self, dim, *, dtype=None):
+    return _patched_cumprod(self, dim, dtype=dtype)
+
+
+def patch_prod_jiterator() -> None:
+    """Idempotent: safe to call from every entry point (training, inference,
+    a notebook cell re-running setup) without double-patching."""
+    if getattr(torch, _PATCHED_ATTR, False):
+        return
+    torch.prod = _patched_prod
+    torch.cumprod = _patched_cumprod
+    torch.Tensor.prod = _patched_tensor_prod
+    torch.Tensor.cumprod = _patched_tensor_cumprod
+    setattr(torch, _PATCHED_ATTR, True)
+    print(
+        "[nvrtc_compat] patched torch.prod/cumprod (module- and Tensor-level) "
+        "to avoid PyTorch's Jiterator on CUDA, working around this "
+        "environment's missing/mismatched libnvrtc-builtins. Every "
+        "dnc.memory.Memory call site (get_usage_vector, allocation "
+        "weighting, and any other prod/cumprod use) is covered "
+        "automatically - no per-method patching needed."
+    )