From 55e33860d3666b307671925df9f4434b6591b265 Mon Sep 17 00:00:00 2001 From: Samuel Date: Fri, 9 Oct 2026 12:43:37 +0200 Subject: [PATCH 01/14] Add MoE workspace and diagnostics Adds an optional shared workspace for split-graph MoE models, including competitive slot writes and broadcast readback to specialists. This also adds per-bank MoE diagnostics, segment-level activation/gradient logging, paired memory-on/off eval replay, and stricter finite-gradient/NaN safeguards for recurrent training. The change keeps existing behavior unchanged unless the new --workspace and related flags are enabled. --- README.md | 28 ++ .../LNN_controller/cfc_backbone_parallel.py | 15 + src/initium/MoE/moe_layer.py | 312 ++++++++++++---- src/initium/config/controller_config.py | 7 + src/initium/config/train_config.py | 2 +- src/initium/core_training.py | 344 ++++++++++++++++-- src/initium/diagnostics/__init__.py | 2 + src/initium/diagnostics/segment_logging.py | 244 +++++++++++++ src/initium/inference/model_loader.py | 9 + src/initium/inference/run_inference.py | 7 +- .../mamba_backbone_parallel.py | 20 + .../mamba_controller/split_graph_dnc.py | 110 +++++- src/initium/workspace/__init__.py | 5 + src/initium/workspace/broadcast.py | 171 +++++++++ 14 files changed, 1164 insertions(+), 112 deletions(-) create mode 100644 src/initium/diagnostics/__init__.py create mode 100644 src/initium/diagnostics/segment_logging.py create mode 100644 src/initium/workspace/__init__.py create mode 100644 src/initium/workspace/broadcast.py diff --git a/README.md b/README.md index f009511..f51682e 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,34 @@ python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 - python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant mamba+cfc ``` +### Phase 0 shared workspace (MoE broadcast) + +The optional `--workspace` path adds a persistent shared slot matrix to a split-graph model with MoE enabled. The final parallel-backbone MoE bank supplies the specialist candidates and its existing noisy top-k routes. At each sequence timestep, selected expert outputs compete for slot writes; a learned RMC-style gate blends proposed slot content with the previous matrix. Every expert in that bank then attends to the updated matrix, and the existing sparse MoE routes combine the context-enriched expert outputs. The workspace resets at the beginning of each episode and persists across its timesteps. This is a shared-object broadcast implementation; it does not implement per-recipient translation. + +Use the roadmap's proposed split-graph/CfC-combiner configuration with MoE and the workspace enabled as follows: + +```bash +python -m src.initium.core_training 0.0 --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --workspace +``` + +`--workspace` requires both `--split-graph` and `--moe`. The default is 5 slots, based on Goyal et al.'s shared-workspace reference result; it is a starting configuration, not a claimed optimum for graph traversal. `--workspace-train-mode coadapted` trains the full model. Choose `frozen` to train only workspace parameters while freezing the backbone, MoE/combiner, DNC, stochastic write heads, and task output projection. `--workspace-no-selection` replaces competitive selected writes with random unweighted writes for the no-Selection control. `--workspace-reset-each-step` resets slots before each timestep to test persistence. `--workspace-num-slots 0` bypasses workspace writes and readback, providing a plumbing comparison against the existing MoE model. Run IDs encode these conditions to avoid log/checkpoint collisions. + +The console prints a startup line that reports whether the workspace is active or bypassed, its capacity, write-selection mode, persistence setting, and train mode. At the normal training log cadence, `workspace[...]` reports selected write mass and active-write fraction, the learned write-gate mean, slot-state RMS, readback-to-expert-output RMS ratio, and read-gate value. These values show whether the path is active and the relative scale of its write and broadcast signals; they do not by themselves establish a task benefit. + +The shared object operates on the final backbone MoE bank; earlier backbone MoE layers retain their existing behavior. This is the Phase 0 plumbing target, not evidence that Broadcast improves task performance. The first short run should compare the zero-slot bypass against the same MoE configuration without `--workspace`; only after this plumbing gate should the frozen/co-adapted and no-Selection comparisons begin. Record seed, full command, commit, hardware/software environment, curriculum completion, ID/OOD metrics, convergence steps, and any failures in `07-Experiment/Broadcasting/Progress/Phase-0-Implementation.md`. All current evidence is graph-traversal-specific; transfer across task families remains untested. + +### MoE numerical and memory-use diagnostics + +The SwitchMoE router projection and expert MLPs run in FP32 even when the surrounding model uses CUDA autocast; low-precision inputs receive FP32 MoE outputs so large finite expert values are not converted into fp16 infinities at the residual boundary. Ordinary routing dispatches activations into fixed-capacity per-expert buffers and applies each expert's original Linear parameters directly. It does not expand or gather full expert weight matrices per token. This matters for recurrent controllers: the intermediate batched-weight route materialized a 512×1536 projection matrix for every route and autograd retained those copies across timesteps. With 32 source tokens, one such projection copy was 96 MiB; input and output projections were both affected. Each expert now evaluates its capacity buffer, including any unused capacity rows. Route-expert and route-slot indices are both shaped `(tokens, top_k)` before gathering; this avoids advanced-index broadcasting across unrelated tokens. The workspace's specialist-candidate path intentionally evaluates every expert because broadcast requires all candidates, but it shares the expert parameters instead of copying them per token. Training logs include the fraction of selected routes dropped by capacity. Router and per-source diagnostics stay on-device during recurrent steps and synchronize only when needed for log output. The load-balancing loss is averaged across repeated calls to a recurrent MoE layer (for example, one call per CfC timestep), then summed across distinct MoE layers; its weight no longer scales with episode length. + +Periodic memory-dependency checks replay the exact random episode batch used by the corresponding memory-on ID evaluation, then restore the RNG streams. This makes `memory on - memory off` a paired functional ablation rather than a difference between two independently sampled batches. The training loop also clips all trainable optimizer parameters, including the task output projection, stops before backward on non-finite forward outputs/losses, reports the first parameter with a non-finite unscaled gradient, logs skipped non-finite-gradient updates, and aborts after eight consecutive such steps instead of continuing a run that is no longer updating safely. + +The supplied Phase 1/2 traces are pre-fix evidence. One Phase 2 trace first reports a non-finite task loss at step 11,665, recurrent outputs become non-finite at step 11,666, and the process continues to step 14,106 with an AMP scale of zero before manual interruption. Memory-on/off checks in those logs used separate random episode batches, so their small positive or negative differences are not reliable paired evidence. See `07-Experiment/Broadcasting/Progress/Phase-1-MoE-Failure-Analysis.md` for the causal breakdown and re-run gate. Re-run the Phase 1 selection-only condition before Phase 2; keep the proposed 8-expert, top-1 configuration and matched seed/settings fixed while checking finite gradients, per-bank routing, and paired memory dependency. + +Training console logs now show **one MoE-bank record per 100-step logging window**: route and router-importance fractions by expert, router entropy and maximum probability, input/expert/MoE-delta activation RMS and maximum magnitude, capacity-drop fraction, and accepted gate mass. The `moe_residual:` segment records the complete MoE block output after its skip connection, separately from `moe_delta_rms/max`. A separate `[segment-diag:train]` line reports each recognized module's call count, input/output RMS and max, non-finite fraction, boundary input/output gradient RMS and non-finite fraction, and hierarchical parameter-gradient norms. Boundary gradients are collected during training backward and during failure traces; evaluation remains forward-only. The collector recognizes Mamba and CfC backbones/blocks, Mamba/CfC controller wrappers/blocks/cells, MoE blocks, chained controllers, DNC memory, workspace, and task output; names include their module paths, so the same logger works across supported controller combinations. + +Evaluation emits separate channel lines for ID, OOD, memory-off, combiner-stage ablations, and final ID/OOD checks. These summarize the evaluations independently from training windows. If a forward value or gradient becomes non-finite, the current segment window and MoE-bank window are printed immediately before training stops or skips the unsafe update. The summaries aggregate detached device-side scalars and synchronize at log/evaluation boundaries; they do not retain activation graphs. + Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 1.0` to use the cosine-anchored Track A reader. Lower scales such as `0.5` are available for follow-up comparisons. Use `--dnc-write-variant kanerva-closed-form` for Track B. Add a reader and writer variant together only for the combined cell after each track is characterized. Dense temporal links remain the default. The short-run baseline reproduction gate and the Q99 paired-order evaluation are deliberate user-run experiment steps; no run results are assumed by this repository change. ### Controllers diff --git a/src/initium/LNN_controller/cfc_backbone_parallel.py b/src/initium/LNN_controller/cfc_backbone_parallel.py index a6f7035..d00239b 100644 --- a/src/initium/LNN_controller/cfc_backbone_parallel.py +++ b/src/initium/LNN_controller/cfc_backbone_parallel.py @@ -160,6 +160,21 @@ def forward(self, x): # (B, L, in_dim) -> (B, L, d_model), called ONCE per trai h = self.moe_blocks[i](h) return h + def forward_with_specialists(self, x: torch.Tensor): + """Return final MoE candidates and routes for shared-workspace use.""" + if not self.moe_enabled or self.moe_blocks is None: + raise RuntimeError("workspace broadcast requires an enabled MoE expert pool") + h = self.in_adapter(x) + final_index = len(self.blocks) - 1 + for i, block in enumerate(self.blocks): + h = block(h) + if i == final_index: + base = h + _, experts, indices, weights = self.moe_blocks[i](h, return_details=True) + return base, experts, indices, weights + h = self.moe_blocks[i](h) + raise RuntimeError("CfC parallel backbone has no blocks") + def build_parallel_backbone( in_dim, diff --git a/src/initium/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index fd06de4..0d4165b 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -81,8 +81,6 @@ from __future__ import annotations -from typing import cast - import torch import torch.nn as nn import torch.nn.functional as F @@ -108,7 +106,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class SwitchMoE(nn.Module): - """Switch-style (top-1) sparsely-gated MoE feed-forward sublayer. + """Switch-style top-k sparsely-gated MoE feed-forward sublayer. Controller-agnostic: operates on the last dimension of whatever shape is passed in (..., d_model), and returns the same shape -- it does NOT apply its own residual connection (see MoEBlock below for that), so it @@ -150,11 +148,13 @@ def __init__( self._aux_losses: list[torch.Tensor] = [] self._last_diag: dict = {} + self._diag_accum: dict[str, torch.Tensor] = {} + self._diag_calls = 0 self._last_topk_idx: torch.Tensor | None = ( None # (T, k), detached -- for per-source breakdowns ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: torch.Tensor, *, return_details: bool = False): orig_shape = x.shape d_model = orig_shape[-1] assert d_model == self.d_model, ( @@ -163,9 +163,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: flat = x.reshape(-1, d_model) # (num_tokens, d_model) num_tokens = flat.shape[0] - logits32 = self.router( - flat - ).float() # router kept in fp32 (Switch's own selective-precision fix) + # ``.float()`` after Linear is too late under autocast: CUDA AMP may + # execute the router matmul in fp16 and only cast its already-rounded + # result back to fp32. Switch's selective-precision rule requires the + # projection itself to run in fp32. + with torch.autocast(device_type=x.device.type, enabled=False): + flat32 = flat.float() + logits32 = F.linear( + flat32, + self.router.weight.float(), + None if self.router.bias is None else self.router.bias.float(), + ) if self.training and self.router_noise_eps > 0: # Switch's own exploration mechanism (Appendix C): multiplicative # jitter noise. We apply it directly to the router logits here @@ -181,12 +189,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # own k=1 special case: each token selects its top_k highest-prob # experts, gate weights renormalized to sum to 1 across just those # k slots -- top_k=1 reduces to the original Switch gate exactly. - # Every expert is computed ONCE for the WHOLE token batch (dense, - # single pair of batched einsums below) regardless of top_k, so the - # k selected experts for a given token are genuinely active - # SIMULTANEOUSLY (one fused kernel launch covers all E experts x - # all T tokens), not run as k sequential passes -- this is what - # "multiple experts active at once" means at this project's scale. + # Normal training dispatches selected, capacity-accepted routes into + # bounded per-expert buffers. The optional all-expert result is + # reserved for the workspace specialist path, which explicitly needs + # every specialist candidate. This keeps top-k routing sparse in + # compute and avoids expanding expert weights per route. top_k = self.top_k topk_prob, topk_idx = probs.topk(top_k, dim=-1) # (T, k) each gate_weights = topk_prob / topk_prob.sum(dim=-1, keepdim=True).clamp(min=1e-9) @@ -206,37 +213,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self._aux_losses.append(aux_loss) with torch.no_grad(): - self._last_diag = { - "cv_importance": (P_i.std() / P_i.mean().clamp(min=1e-8)).item(), - "cv_load": (f_i.std() / f_i.mean().clamp(min=1e-8)).item(), - "max_load_frac": f_i.max().item(), - } self._last_topk_idx = topk_idx.detach() # Expert capacity, scaled by top_k (each slot competes for the same # per-expert buffer). capacity = max(1, int((num_tokens * top_k / self.num_experts) * self.capacity_factor)) - w_in = torch.stack( - [cast(nn.Linear, e.w_in).weight for e in self.experts], dim=0 - ) # (E, expert_dim, d_model) - b_in = torch.stack([cast(nn.Linear, e.w_in).bias for e in self.experts], dim=0) - w_out = torch.stack( - [cast(nn.Linear, e.w_out).weight for e in self.experts], dim=0 - ) # (E, d_model, expert_dim) - b_out = torch.stack([cast(nn.Linear, e.w_out).bias for e in self.experts], dim=0) - - # Dense pass over EVERY expert for EVERY token -- the single fused - # computation that makes the top-k selected experts per token - # simultaneous rather than sequential. - hidden = torch.einsum("td,exd->tex", flat, w_in) + b_in # (T, E, expert_dim) - hidden = F.relu(hidden) - expert_out_all = torch.einsum("tex,edx->ted", hidden, w_out) + b_out # (T, E, d_model) - - gathered = torch.gather( - expert_out_all, 1, topk_idx.unsqueeze(-1).expand(-1, -1, d_model) - ) # (T, k, d_model) - # Capacity mask: keep only the first `capacity` (token,slot) pairs, # in original order, routed to each expert -- same drop semantics # as the prior top-1 implementation, generalized over the @@ -245,11 +227,133 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: one_hot_flat = F.one_hot(flat_idx, num_classes=self.num_experts) rank_in_expert = one_hot_flat.cumsum(dim=0).gather(1, flat_idx.unsqueeze(1)).squeeze(1) - 1 keep = (rank_in_expert < capacity).reshape(num_tokens, top_k) + accepted_gates = gate_weights * keep.to(gate_weights.dtype) + accepted_gates = accepted_gates / accepted_gates.sum( + dim=-1, keepdim=True + ).clamp_min(1e-9) + + with torch.autocast(device_type=x.device.type, enabled=False): + # Dispatch activations into fixed-capacity expert buffers. The + # previous batched-weight implementation expanded each expert's + # (D, hidden) matrices once per token/route; autograd retained + # those copies for every recurrent call (96 MiB for just one + # 32-token, 512x1536 projection). Running each expert's Linear on + # its capacity buffer shares the original Parameter storage and + # limits saved activations to dispatched tokens. + route_slots = rank_in_expert.clamp(max=capacity - 1) + route_tokens = torch.arange(num_tokens, device=x.device).repeat_interleave(top_k) + route_keep = keep.reshape(-1) + route_inputs = flat32.index_select(0, route_tokens) + route_inputs = route_inputs * route_keep.unsqueeze(-1).to(route_inputs.dtype) + dispatch_indices = flat_idx * capacity + route_slots + dispatch = flat32.new_zeros((self.num_experts * capacity, d_model)) + dispatch = dispatch.index_add(0, dispatch_indices, route_inputs) + dispatch = dispatch.reshape(self.num_experts, capacity, d_model) + + expert_outputs = [] + for expert_index, expert in enumerate(self.experts): + expert_input = dispatch[expert_index] + expert_hidden = F.linear( + expert_input, expert.w_in.weight.float(), expert.w_in.bias.float() + ).relu() + expert_outputs.append( + F.linear( + expert_hidden, + expert.w_out.weight.float(), + expert.w_out.bias.float(), + ) + ) + expert_outputs = torch.stack(expert_outputs, dim=0) + + # Both advanced indices must have identical (T, top_k) shapes. + # Leaving flat_idx flat (T*top_k,) makes PyTorch broadcast it + # against the (T, top_k) slot index. For top_k=1 this silently + # creates a (T, T, D) selection, mixing unrelated tokens and + # summing T expert outputs per token; the resulting oversized + # MoE delta caused the immediate forward-scale and NaN-gradient + # failure in the Phase 1 trace. + route_experts = flat_idx.reshape(num_tokens, top_k) + route_slots_2d = route_slots.reshape(num_tokens, top_k) + selected_outputs = expert_outputs[route_experts, route_slots_2d] + output32 = (selected_outputs * accepted_gates.unsqueeze(-1)).sum(dim=1) + accepted_outputs = selected_outputs * keep.unsqueeze(-1).to(selected_outputs.dtype) + + expert_out_all = None + if return_details: + # Workspace broadcasting is the one caller that needs all + # specialist candidates, including experts not selected by + # the sparse route. Keep this expensive activation path + # explicit, while still using each expert's shared weights. + expert_out_all = torch.stack( + [ + F.linear( + F.relu( + F.linear( + flat32, + expert.w_in.weight.float(), + expert.w_in.bias.float(), + ) + ), + expert.w_out.weight.float(), + expert.w_out.bias.float(), + ) + for expert in self.experts + ], + dim=1, + ) - gate = (gate_weights * keep.to(gate_weights.dtype)).unsqueeze(-1).to(gathered.dtype) - output = (gathered * gate).sum(dim=1).to(flat.dtype) # (T, d_model) + # Keep the result in fp32 when the input was autocast to fp16/bf16. + # Casting a large but finite fp32 expert result back to fp16 can + # create an Inf before the residual or DNC path has a chance to + # stabilize it. The surrounding residual naturally promotes to fp32. + output_dtype = ( + torch.float32 + if flat.dtype in (torch.float16, torch.bfloat16) + else flat.dtype + ) + output = output32.to(output_dtype) - return output.reshape(orig_shape) + with torch.no_grad(): + self._last_diag = { + "cv_importance": (P_i.std() / P_i.mean().clamp(min=1e-8)).detach(), + "cv_load": (f_i.std() / f_i.mean().clamp(min=1e-8)).detach(), + "max_load_frac": f_i.max().detach(), + "expert_load_frac": f_i.detach(), + "router_importance_frac": P_i.detach(), + "router_entropy": ( + -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1).mean() + ).detach(), + "router_max_prob": probs.max(dim=-1).values.mean().detach(), + "topk_gate_mean": topk_prob.mean().detach(), + "input_rms": flat32.square().mean().sqrt().detach(), + "input_absmax": flat32.abs().amax().detach(), + "expert_output_rms": accepted_outputs.square().mean().sqrt().detach(), + "expert_output_absmax": accepted_outputs.abs().amax().detach(), + "routed_output_rms": output32.square().mean().sqrt().detach(), + "routed_output_absmax": output32.abs().amax().detach(), + "capacity_drop_frac": (1.0 - keep.float().mean()).detach(), + "accepted_gate_mass": accepted_gates.sum(dim=-1).mean().detach(), + } + if self.training: + for key, value in self._last_diag.items(): + self._diag_accum[key] = ( + self._diag_accum[key] + value + if key in self._diag_accum + else value.clone() + ) + self._diag_calls += 1 + + output = output.reshape(orig_shape) + if return_details: + expert_shape = (*orig_shape[:-1], self.num_experts, d_model) + route_shape = (*orig_shape[:-1], top_k) + return ( + output, + expert_out_all.to(output_dtype).reshape(expert_shape), + topk_idx.reshape(route_shape), + accepted_gates.reshape(route_shape), + ) + return output def pop_aux_loss(self) -> torch.Tensor: """Consume and clear the accumulated load-balancing loss (mirrors @@ -259,7 +363,11 @@ def pop_aux_loss(self) -> torch.Tensor: across both mechanisms).""" if not self._aux_losses: return torch.zeros((), device=self.router.weight.device) - total = torch.stack(self._aux_losses).sum() + # A recurrent controller invokes this layer once per timestep. The + # Switch objective is per routing batch, so sum() would silently + # multiply its strength by episode length. Average calls within this + # layer; pop_total_moe_aux_loss still sums distinct model layers. + total = torch.stack(self._aux_losses).mean() self._aux_losses = [] return total @@ -273,6 +381,15 @@ def last_diagnostics(self) -> dict: not a routing-internal metric.""" return dict(self._last_diag) + def pop_diagnostics(self) -> dict: + """Consume mean training diagnostics accumulated since the last pop.""" + if not self._diag_calls: + return {} + result = {key: value / self._diag_calls for key, value in self._diag_accum.items()} + self._diag_accum = {} + self._diag_calls = 0 + return result + def last_routing(self) -> torch.Tensor | None: """(num_tokens, top_k) expert indices chosen on the most recent forward() call, or None before any call. Callers that know the @@ -323,8 +440,11 @@ def __init__( if device is not None and getattr(device, "type", None) == "cuda": self.to(device) - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x + self.moe(self.norm(x)) + def forward(self, x: torch.Tensor, *, return_details: bool = False): + if not return_details: + return x + self.moe(self.norm(x)) + routed, experts, indices, weights = self.moe(self.norm(x), return_details=True) + return x + routed, experts, indices, weights def pop_aux_loss(self) -> torch.Tensor: return self.moe.pop_aux_loss() @@ -332,6 +452,9 @@ def pop_aux_loss(self) -> torch.Tensor: def last_diagnostics(self) -> dict: return self.moe.last_diagnostics() + def pop_diagnostics(self) -> dict: + return self.moe.pop_diagnostics() + class SourceEmbedding(nn.Module): """Learned per-source bias added to a token's representation before @@ -358,8 +481,8 @@ class MultiSourceMoEBlock(nn.Module): """CfC-oriented MoE sublayer: accepts a VARIABLE-length list of per-source token tensors -- each (B, d_model) -- and returns the same number of outputs, one per source, each individually routed and gated - through ONE shared bank of Top-K experts (SwitchMoE.forward -- every - expert still runs exactly once per call, dense and simultaneous). + through ONE shared bank of Top-K experts (SwitchMoE.forward computes + only selected, capacity-accepted token/expert pairs). This is what lets a CfC controller/combiner keep its "many inputs in, many outputs out, with per-input specialization" property once MoE @@ -432,25 +555,18 @@ def forward(self, sources: list[torch.Tensor]) -> list[torch.Tensor]: if topk_idx is not None: num_experts = self.moe.num_experts if self._cumulative_source_counts is None: - self._cumulative_source_counts = [torch.zeros(num_experts) for _ in batch_sizes] - source_diag, offset = [], 0 + self._cumulative_source_counts = [ + torch.zeros(num_experts, device=topk_idx.device) for _ in batch_sizes + ] + offset = 0 for i, b in enumerate(batch_sizes): - idx_slice = topk_idx[offset : offset + b].reshape(-1).cpu() + idx_slice = topk_idx[offset : offset + b].reshape(-1) counts = torch.bincount(idx_slice, minlength=num_experts).float() - self._cumulative_source_counts[i] += counts # never reset -- spans the whole run - cum = self._cumulative_source_counts[i] - frac = (cum / cum.sum().clamp(min=1)).tolist() - top_expert = int(cum.argmax().item()) - source_diag.append( - { - "expert_frac": frac, - "top_expert": top_expert, - "top_expert_frac": frac[top_expert], - "n_calls_accumulated": int(cum.sum().item()), - } - ) + # Keep routing counters on device and defer CPU synchronization + # until inference explicitly requests a report. This block runs + # once per CfC timestep during training. + self._cumulative_source_counts[i] += counts offset += b - self._last_source_diag = source_diag outs, offset = [], 0 for i, b in enumerate(batch_sizes): @@ -464,6 +580,9 @@ def pop_aux_loss(self) -> torch.Tensor: def last_diagnostics(self) -> dict: return self.moe.last_diagnostics() + def pop_diagnostics(self) -> dict: + return self.moe.pop_diagnostics() + def last_source_diagnostics(self) -> list[dict] | None: """Per-source expert-usage breakdown from the most recent forward() call: one dict per source (same order as the `sources` list passed @@ -472,7 +591,22 @@ def last_source_diagnostics(self) -> list[dict] | None: Comparing this ACROSS sources is the concrete test of whether the router is genuinely specializing by input identity rather than routing every source near-identically.""" - return getattr(self, "_last_source_diag", None) + if self._cumulative_source_counts is None: + return None + source_diag = [] + for counts in self._cumulative_source_counts: + counts_cpu = counts.detach().float().cpu() + fractions = counts_cpu / counts_cpu.sum().clamp_min(1.0) + top_expert = int(counts_cpu.argmax().item()) + source_diag.append( + { + "expert_frac": fractions.tolist(), + "top_expert": top_expert, + "top_expert_frac": float(fractions[top_expert]), + "n_calls_accumulated": int(counts_cpu.sum().item()), + } + ) + return source_diag class MoERNNWrapper(nn.Module): @@ -526,27 +660,67 @@ def forward(self, input: torch.Tensor, hx): return out, new_hx -def pop_total_moe_aux_loss(moe_layers: list) -> tuple: +def pop_total_moe_aux_loss(moe_layers: list, *, include_diagnostics: bool = True) -> tuple: """Sum pop_aux_loss() across all installed MoEBlock/SwitchMoE layers and merge their diagnostics (mean across layers for CV metrics, max for the worst-case load fraction), mirroring stochastic_write_head_v2's pop_total_kl() pattern exactly, so the training loop combines this auxiliary loss the same way it already combines the KL loss.""" total = None - cv_imp, cv_load, max_load = [], [], [] + cv_imp, cv_load, max_load, capacity_drop = [], [], [], [] + layer_reports = [] for layer in moe_layers: loss = layer.pop_aux_loss() total = loss if total is None else total + loss - diag = layer.last_diagnostics() + if include_diagnostics: + pop_diagnostics = getattr(layer, "pop_diagnostics", None) + diag = pop_diagnostics() if pop_diagnostics is not None else layer.last_diagnostics() + else: + diag = {} if diag: - cv_imp.append(diag["cv_importance"]) - cv_load.append(diag["cv_load"]) - max_load.append(diag["max_load_frac"]) + cv_imp.append(float(diag["cv_importance"])) + cv_load.append(float(diag["cv_load"])) + max_load.append(float(diag["max_load_frac"])) + capacity_drop.append(float(diag.get("capacity_drop_frac", 0.0))) + layer_reports.append( + { + "name": getattr(layer, "diagnostic_name", type(layer).__name__), + **{ + key: float(diag[key]) + for key in ( + "cv_importance", + "cv_load", + "max_load_frac", + "router_entropy", + "router_max_prob", + "topk_gate_mean", + "input_rms", + "input_absmax", + "expert_output_rms", + "expert_output_absmax", + "routed_output_rms", + "routed_output_absmax", + "capacity_drop_frac", + "accepted_gate_mass", + ) + if key in diag + }, + **{ + key: diag[key].detach().float().cpu().tolist() + for key in ("expert_load_frac", "router_importance_frac") + if key in diag + }, + } + ) if total is None: total = torch.zeros(()) merged = { "moe_cv_importance": sum(cv_imp) / len(cv_imp) if cv_imp else 0.0, "moe_cv_load": sum(cv_load) / len(cv_load) if cv_load else 0.0, "moe_max_load_frac": max(max_load) if max_load else 0.0, + "moe_capacity_drop_frac": ( + sum(capacity_drop) / len(capacity_drop) if capacity_drop else 0.0 + ), + "moe_layers": layer_reports, } return total, merged diff --git a/src/initium/config/controller_config.py b/src/initium/config/controller_config.py index 5da51c8..ef1bf12 100644 --- a/src/initium/config/controller_config.py +++ b/src/initium/config/controller_config.py @@ -67,6 +67,13 @@ SPLIT_GRAPH_COMBINER_VARIANT = "mamba1" SPLIT_GRAPH_COMBINER_NUM_BLOCKS = 1 +# Phase 0 shared global workspace (requires split-graph + MoE). +WORKSPACE_ENABLED = False +WORKSPACE_NUM_SLOTS = 5 +WORKSPACE_NO_SELECTION = False # random/unweighted write control +WORKSPACE_RESET_EACH_STEP = False # persistence ablation +WORKSPACE_TRAIN_MODE = "coadapted" # "coadapted" | "frozen" + # DNC memory mechanism variants (only applied by the split-graph assembly). DNC_READ_VARIANT = "cosine" # "cosine" | "relational-mlp" | "relational-residual" DNC_READ_RESIDUAL_SCALE = 1.0 # Gain applied before tanh for relational-residual. diff --git a/src/initium/config/train_config.py b/src/initium/config/train_config.py index 825f77e..db07c96 100644 --- a/src/initium/config/train_config.py +++ b/src/initium/config/train_config.py @@ -12,7 +12,7 @@ BATCH_SIZE = 16 # bucketed, so padding waste stays low even >1 TOTAL_STEPS = 120000 LOG_EVERY = 100 -EVAL_EVERY = 1000 # curriculum-advance eval cadence +EVAL_EVERY = 500 # curriculum-advance eval cadence LR = 3e-4 LR_MIN = 3e-5 LR_DECAY_STEPS = TOTAL_STEPS # cosine anneal length. Derived from TOTAL_STEPS here on diff --git a/src/initium/core_training.py b/src/initium/core_training.py index e25209a..13c5820 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -105,6 +105,11 @@ SPLIT_GRAPH_MAMBA_HEADDIM, SPLIT_GRAPH_MAMBA_VARIANT, SPLIT_GRAPH_NUM_BLOCKS, + WORKSPACE_ENABLED, + WORKSPACE_NO_SELECTION, + WORKSPACE_NUM_SLOTS, + WORKSPACE_RESET_EACH_STEP, + WORKSPACE_TRAIN_MODE, ) from initium.config.train_config import ( AMP_INIT_SCALE, @@ -154,6 +159,7 @@ score_probe_cases, ) from initium.data.prepared_dataset import load_prepared_dataset, save_prepared_dataset +from initium.diagnostics.segment_logging import SegmentDiagnostics from initium.mamba_controller.mamba_controller import MambaDNC from initium.mamba_controller.split_graph_dnc import SplitGraphDNC from initium.memory_manipulation.dynamic_memory_resize import resize_memory @@ -204,6 +210,73 @@ def _first_nonfinite_report(tensor, name, step): return ok +def _capture_rng_state(): + """Capture all RNG streams used by task generation and model evaluation. + + Periodic dependency evaluations replay this state so the memory-on and + memory-off scores use the same episodes, then restore the stream so the + paired check does not perturb subsequent evaluation/training randomness. + """ + return { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch_cpu": torch.random.get_rng_state(), + "torch_cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, + } + + +def _restore_rng_state(state): + random.setstate(state["python"]) + np.random.set_state(state["numpy"]) + torch.random.set_rng_state(state["torch_cpu"]) + if state["torch_cuda"] is not None: + torch.cuda.set_rng_state_all(state["torch_cuda"]) + + +def _report_first_nonfinite_gradient(named_parameters, step): + """Identify the first parameter whose unscaled gradient is non-finite.""" + for name, parameter in named_parameters: + grad = parameter.grad + if grad is not None and not torch.isfinite(grad).all(): + bad_fraction = (~torch.isfinite(grad)).float().mean().item() + print( + f"[NaN-TRACE] step {step}: first non-finite gradient in " + f"'{name}' ({bad_fraction * 100:.2f}% of elements)" + ) + return name + print(f"[NaN-TRACE] step {step}: gradient norm is non-finite but no bad gradient tensor was found") + return None + + +def _format_moe_layer_reports(diagnostics: dict) -> list[str]: + """Format per-bank routing/activation summaries for console output.""" + reports = [] + for layer_diag in diagnostics.get("moe_layers", []): + load = ",".join(f"{value:.2f}" for value in layer_diag.get("expert_load_frac", [])) + importance = ",".join( + f"{value:.2f}" for value in layer_diag.get("router_importance_frac", []) + ) + reports.append( + f"moe-bank[{layer_diag['name']} " + f"load=[{load}] importance=[{importance}] " + f"entropy={layer_diag.get('router_entropy', 0.0):.2f} " + f"pmax={layer_diag.get('router_max_prob', 0.0):.2f} " + f"cv_load={layer_diag.get('cv_load', 0.0):.2f} " + f"cv_imp={layer_diag.get('cv_importance', 0.0):.2f} " + f"max_load={layer_diag.get('max_load_frac', 0.0):.2f} " + f"topk_gate={layer_diag.get('topk_gate_mean', 0.0):.2f} " + f"in_rms={layer_diag.get('input_rms', 0.0):.3g} " + f"in_max={layer_diag.get('input_absmax', 0.0):.3g} " + f"expert_rms/max={layer_diag.get('expert_output_rms', 0.0):.3g}/" + f"{layer_diag.get('expert_output_absmax', 0.0):.3g} " + f"moe_delta_rms/max={layer_diag.get('routed_output_rms', 0.0):.3g}/" + f"{layer_diag.get('routed_output_absmax', 0.0):.3g} " + f"drop={layer_diag.get('capacity_drop_frac', 0.0):.1%} " + f"gate_mass={layer_diag.get('accepted_gate_mass', 0.0):.2f}]" + ) + return reports + + def _ogs_components(id_acc, ood_acc, ood_perfect, ood_memory_off, lesson): depth = max(int(lesson), 1) offset_normalized = max(0.0, 1.0 - abs(float(id_acc) - float(ood_acc)) / (100.0 * depth)) @@ -580,6 +653,11 @@ def run( split_graph_combiner_mode: str = SPLIT_GRAPH_COMBINER_MODE, split_graph_combiner_variant: str = SPLIT_GRAPH_COMBINER_VARIANT, split_graph_combiner_num_blocks: int = SPLIT_GRAPH_COMBINER_NUM_BLOCKS, + workspace_enabled: bool = WORKSPACE_ENABLED, + workspace_num_slots: int = WORKSPACE_NUM_SLOTS, + workspace_no_selection: bool = WORKSPACE_NO_SELECTION, + workspace_reset_each_step: bool = WORKSPACE_RESET_EACH_STEP, + workspace_train_mode: str = WORKSPACE_TRAIN_MODE, dnc_read_variant: str = DNC_READ_VARIANT, dnc_read_residual_scale: float = DNC_READ_RESIDUAL_SCALE, dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, @@ -645,6 +723,7 @@ def run( "moe_cv_importance", "moe_cv_load", "moe_max_load_frac", + "moe_capacity_drop_frac", ] ) @@ -1147,6 +1226,11 @@ def run( dnc_read_residual_max_ratio=dnc_read_residual_max_ratio, dnc_write_variant=dnc_write_variant, dnc_write_observation_variance=dnc_write_observation_variance, + workspace_enabled=workspace_enabled, + workspace_num_slots=workspace_num_slots, + workspace_no_selection=workspace_no_selection, + workspace_reset_each_step=workspace_reset_each_step, + workspace_train_mode=workspace_train_mode, device=device, ).to(device) else: @@ -1172,6 +1256,23 @@ def run( **mamba_kwargs, ).to(device) + if split_graph_enabled: + if workspace_enabled: + workspace_active = isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None + write_mode = "random-control" if workspace_no_selection else "MoE-top-k" + persistence = "off (reset each timestep)" if workspace_reset_each_step else "on" + print( + f"[{run_id}] WORKSPACE {'active' if workspace_active else 'bypassed'} | " + f"slots={workspace_num_slots} | write_selection={write_mode} | " + f"persistence={persistence} | train_mode={workspace_train_mode} | " + "scope=final-backbone-MoE-bank; each specialist reads shared slots" + ) + else: + print( + f"[{run_id}] WORKSPACE disabled | Selection-only split-graph run; " + "no shared-slot writes or broadcast readback" + ) + combiner_stage_kinds = getattr( getattr(rnn, "combiner_wrapper", None), "stage_kinds", None ) # set only for a hybrid ("+") split-graph combiner; None otherwise @@ -1213,6 +1314,11 @@ def run( "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, "dnc_write_variant": dnc_write_variant, "dnc_write_observation_variance": dnc_write_observation_variance, + "workspace_enabled": workspace_enabled, + "workspace_num_slots": workspace_num_slots, + "workspace_no_selection": workspace_no_selection, + "workspace_reset_each_step": workspace_reset_each_step, + "workspace_train_mode": workspace_train_mode, } if link_matrix_mode != "dense": @@ -1225,6 +1331,22 @@ def run( output_proj = nn.Linear(INPUT_DIM, TRIPLE_DIM).to(device) dataset.set_output_proj(output_proj) + segment_diagnostics = SegmentDiagnostics(rnn) + segment_diagnostics.add_module(output_proj, "task_output") + moe_layer_ids = {id(layer) for layer in getattr(rnn, "moe_layers", [])} + for module_name, module in rnn.named_modules(): + if id(module) in moe_layer_ids: + module.diagnostic_name = module_name or type(module).__name__ + + def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: + segment_diagnostics.emit(channel, include_gradients=include_gradients) + if moe_enabled: + _, failure_moe_diag = pop_total_moe_aux_loss( + rnn.moe_layers, include_diagnostics=True + ) + reports = _format_moe_layer_reports(failure_moe_diag) + if reports: + print(f"[moe-diag:{channel}] " + " | ".join(reports)) # Install the stochastic write head(s) BEFORE building the # optimizer, so their parameters (mu_transform, logvar_transform) are @@ -1243,7 +1365,22 @@ def run( rnn, device=device, sample=(beta_target > 0.0) ) - optimizer = torch.optim.Adam(list(rnn.parameters()) + list(output_proj.parameters()), lr=LR) + if workspace_enabled and workspace_train_mode == "frozen": + for parameter in rnn.parameters(): + parameter.requires_grad_(False) + assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None + for parameter in rnn.workspace.parameters(): + parameter.requires_grad_(True) + for parameter in output_proj.parameters(): + parameter.requires_grad_(False) + trainable_parameters = [ + parameter + for parameter in list(rnn.parameters()) + list(output_proj.parameters()) + if parameter.requires_grad + ] + if not trainable_parameters: + raise ValueError("training configuration has no trainable parameters") + optimizer = torch.optim.Adam(trainable_parameters, lr=LR) # Option 3 (Concept 24/LB-17): multi-axis efficiency accounting. # param_count is fixed for the life of this run -- static Option 2's # nr_cells resize never changes it (every Memory sublayer is sized by @@ -1481,6 +1618,8 @@ def set_lr(step_now): # whatever LR was saved inside optimizer_state_dict at checkpoint time # (potentially still reflecting the old buggy drifted value). running_task_loss, running_kl_loss, running_div, running_grad_norm = 0.0, 0.0, 0, 0.0 + running_finite_grad_steps = 0 + running_nonfinite_grad_steps = 0 # Dynamic-beta controller state (no-op / unused when beta_mode == "static"). # Initialized post-resume so a resumed dynamic run starts its health check @@ -1488,6 +1627,7 @@ def set_lr(step_now): beta_constraint_ema = 0.0 scale_at_last_eval = scaler.get_scale() grad_was_finite_since_eval = True + consecutive_nonfinite_grad_steps = 0 t0 = time.time() @@ -1508,36 +1648,82 @@ def set_lr(step_now): answer_mask = answer_mask.to(device, non_blocking=True) check_finite_this_step = True if check_finite_this_step: - _first_nonfinite_report(input_seq, "input_seq", step) + if not _first_nonfinite_report(input_seq, "input_seq", step): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite training input at step {step}; stopping." + ) hidden = (None, None, None) optimizer.zero_grad(set_to_none=True) + segment_diagnostics.set_channel("train") with torch.amp.autocast("cuda", enabled=amp_enabled): output, hidden = rnn(input_seq, hidden, reset_experience=True) - if check_finite_this_step: - _first_nonfinite_report(output, "rnn_output", step) + if check_finite_this_step and not _first_nonfinite_report( + output, "rnn_output", step + ): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite recurrent output at step {step}; " + "stopping before backward/optimizer update to protect the checkpoint." + ) output = output.transpose(0, 1) # (B, T, 92); Linear accepts non-contiguous input. output = output_proj(output) # (B, T, 90) + if check_finite_this_step and not _first_nonfinite_report( + output, "output_logits", step + ): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite output logits at step {step}; " + "stopping before backward/optimizer update to protect the checkpoint." + ) task_loss = dataset.loss(output, target_digits, answer_mask) - if check_finite_this_step: - _first_nonfinite_report(task_loss, "task_loss", step) + if check_finite_this_step and not _first_nonfinite_report( + task_loss, "task_loss", step + ): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite task loss at step {step}; " + "stopping before backward/optimizer update to protect the checkpoint." + ) kl_loss, kl_diag = pop_total_kl(stochastic_heads, free_bits=FREE_BITS) # Phase 1 - if check_finite_this_step: - _first_nonfinite_report(kl_loss, "kl_loss", step) + if check_finite_this_step and not _first_nonfinite_report( + kl_loss, "kl_loss", step + ): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite KL loss at step {step}; stopping before backward." + ) moe_aux_loss, moe_diag = ( - pop_total_moe_aux_loss(rnn.moe_layers) + pop_total_moe_aux_loss( + rnn.moe_layers, + include_diagnostics=(step + 1) % LOG_EVERY == 0, + ) if moe_enabled else (torch.zeros((), device=device), {}) ) # v8 (Option 4): Switch-style load-balancing loss, already scaled # by moe_load_balance_alpha inside SwitchMoE.pop_aux_loss() -- no # extra weighting applied here, unlike beta_eff*kl_loss below. + if check_finite_this_step and not _first_nonfinite_report( + moe_aux_loss, "moe_aux_loss", step + ): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite MoE auxiliary loss at step {step}; " + "stopping before backward." + ) global_ramp = min(1.0, (step - anneal_start_step) / max(1, KL_ANNEAL_STEPS)) lesson_dip_ramp = min(1.0, (step - lesson_dip_start_step) / max(1, LESSON_KL_DIP_STEPS)) beta_eff = beta_target * global_ramp * lesson_dip_ramp loss = ( task_loss + beta_eff * kl_loss + moe_aux_loss ) # Phase 1 + Option 4: L = L_task + beta*L_KL + L_moe_aux + if check_finite_this_step and not _first_nonfinite_report(loss, "total_loss", step): + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] non-finite total loss at step {step}; stopping before backward." + ) scaler.scale(loss).backward() scaler.unscale_(optimizer) # Log addition #2: capture the pre-clip gradient norm instead of @@ -1545,21 +1731,48 @@ def set_lr(step_now): # gradient-scale trajectory to compare a future β>0 switch-in # against, useful for telling a genuine instability apart from the # expected one-off "shock" when sampling noise is switched on. - grad_norm = torch.nn.utils.clip_grad_norm_(rnn.parameters(), max_norm=10.0) + # Clip every parameter in this optimizer, including output_proj. + # Clipping only rnn.parameters() left the trainable prediction head + # unbounded while its gradients were still being applied by Adam. + try: + grad_norm = torch.nn.utils.clip_grad_norm_( + trainable_parameters, max_norm=10.0, error_if_nonfinite=True + ) + except RuntimeError as error: + if "non-finite" not in str(error): + raise + # error_if_nonfinite raises before applying the clip coefficient, + # preserving the original bad gradients for the diagnostic below. + grad_norm = torch.tensor(float("nan"), device=device) if not math.isfinite(float(grad_norm)): grad_was_finite_since_eval = False - # clip_grad_norm_ computes clip_coef = max_norm/(total_norm+eps) - # and multiplies EVERY gradient by it. A NaN total_norm therefore - # NaNs out every gradient in the model, including ones that were - # finite. GradScaler's found_inf was recorded during unscale_, - # BEFORE this clip, so it can be clean while the post-clip grads - # are all NaN -- scaler.step() then applies them and the weights - # are permanently poisoned, which is exactly the one-way collapse - # at step 23700. Zero the grads and skip the step instead: this - # costs one wasted batch and is fully recoverable. + consecutive_nonfinite_grad_steps += 1 + running_nonfinite_grad_steps += 1 + _report_first_nonfinite_gradient( + [ + (f"rnn.{name}", parameter) + for name, parameter in rnn.named_parameters() + ] + + [ + (f"output_proj.{name}", parameter) + for name, parameter in output_proj.named_parameters() + ], + step, + ) + emit_failure_diagnostics("train", include_gradients=True) + # Skip this optimizer update. GradScaler's found_inf scan happens + # during unscale_; the explicit finite-norm check also covers + # non-finite gradients that arise after that scan. optimizer.zero_grad(set_to_none=True) scaler.update() + if consecutive_nonfinite_grad_steps >= 8: + raise FloatingPointError( + f"[{run_id}] {consecutive_nonfinite_grad_steps} consecutive optimizer " + f"steps had non-finite gradients (AMP scale={scaler.get_scale():.4g}); " + "stopping rather than silently training without parameter updates." + ) else: + consecutive_nonfinite_grad_steps = 0 scaler.step(optimizer) scaler.update() # Log addition #3: AMP loss-scale value. A collapsing/repeatedly @@ -1571,7 +1784,9 @@ def set_lr(step_now): running_task_loss += task_loss.detach() running_kl_loss += kl_loss.detach() running_div += dataset.diversity(output.detach(), target_digits, answer_mask) - running_grad_norm += float(grad_norm) + if math.isfinite(float(grad_norm)): + running_grad_norm += float(grad_norm) + running_finite_grad_steps += 1 step += 1 current_lr = set_lr(step) # Phase 1 FIX: clamped manual LR, replaces scheduler.step() @@ -1647,7 +1862,7 @@ def set_lr(step_now): running_kl_loss.item() if torch.is_tensor(running_kl_loss) else running_kl_loss ) / LOG_EVERY avg_div = running_div / LOG_EVERY - avg_grad_norm = running_grad_norm / LOG_EVERY + avg_grad_norm = running_grad_norm / max(running_finite_grad_steps, 1) kl_contrib = ( beta_eff * avg_kl ) # actual beta*L_KL added to total loss, vs. avg_kl (pre-beta, raw clamped sum) @@ -1680,6 +1895,7 @@ def set_lr(step_now): f"{LOG_EVERY / elapsed:.2f} steps/s", f"LR {current_lr:.6f}", f"grad_norm {avg_grad_norm:.4f}", + f"nonfinite_grad_steps {running_nonfinite_grad_steps}", ] if amp_enabled: parts.append(f"amp_scale {amp_scale:.1f}") @@ -1689,10 +1905,28 @@ def set_lr(step_now): f"moe_aux {float(moe_aux_loss.detach()):.4f}", f"moe_cv_load {moe_diag.get('moe_cv_load', 0.0):.4f}", f"moe_max_load_frac {moe_diag.get('moe_max_load_frac', 0.0):.4f}", + f"moe_capacity_drop {moe_diag.get('moe_capacity_drop_frac', 0.0):.1%}", ] + parts.extend(_format_moe_layer_reports(moe_diag)) + if workspace_enabled and isinstance(rnn, SplitGraphDNC): + if rnn.workspace is None: + parts.append("workspace[bypassed: slots=0]") + else: + workspace_diag = rnn.workspace.last_diagnostics() + if workspace_diag: + parts.append( + "workspace[" + f"write_mass {workspace_diag['write_mass'].item():.3f} " + f"active_writes {workspace_diag['active_write_frac'].item():.1%} " + f"write_gate {workspace_diag['write_gate_mean'].item():.3f} " + f"slot_rms {workspace_diag['slot_rms'].item():.3f} " + f"broadcast/specialist {workspace_diag['broadcast_to_specialist_ratio'].item():.3f} " + f"read_gate {workspace_diag['read_gate'].item():.3f}]" + ) if dyn_n_str: parts.append(dyn_n_str.strip(" |")) print(" | ".join(parts)) + segment_diagnostics.emit("train", include_gradients=True) if ( track_b_writer is not None @@ -1752,12 +1986,15 @@ def set_lr(step_now): moe_diag.get("moe_cv_importance", 0.0), moe_diag.get("moe_cv_load", 0.0), moe_diag.get("moe_max_load_frac", 0.0), + moe_diag.get("moe_capacity_drop_frac", 0.0), ] ) if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats(device) # so next window's peak isn't cumulative log_file.flush() running_task_loss, running_kl_loss, running_div, running_grad_norm = 0.0, 0.0, 0, 0.0 + running_finite_grad_steps = 0 + running_nonfinite_grad_steps = 0 if kl_on and step % PRIOR_SNAPSHOT_EVERY == 0: # v5 (Phase 2): refit (mu_g, Sigma_g) from the writes @@ -1807,9 +2044,12 @@ def set_lr(step_now): id_stage_memory = rnn.memories[0] if hasattr(id_stage_memory, "begin_stage_diagnostics"): id_stage_memory.begin_stage_diagnostics() + id_eval_rng_state = _capture_rng_state() + segment_diagnostics.set_channel("eval-id") _, id_triple_acc, id_perfect_frac = curriculum.maybe_advance( rnn, device, step=step, optimizer=optimizer ) + segment_diagnostics.emit("eval-id") id_stage_diagnostics = ( id_stage_memory.finish_stage_diagnostics() if hasattr(id_stage_memory, "finish_stage_diagnostics") @@ -1896,6 +2136,7 @@ def set_lr(step_now): ood_stage_memory = rnn.memories[0] if hasattr(ood_stage_memory, "begin_stage_diagnostics"): ood_stage_memory.begin_stage_diagnostics() + segment_diagnostics.set_channel("eval-ood") ood_triple_acc, ood_perfect_frac, ood_hop_breakdown = dataset.evaluate_ood( rnn, device, @@ -1903,6 +2144,7 @@ def set_lr(step_now): rng=ood_rng, field_log=ood_field_log, ) + segment_diagnostics.emit("eval-ood") ood_stage_diagnostics = ( ood_stage_memory.finish_stage_diagnostics() if hasattr(ood_stage_memory, "finish_stage_diagnostics") @@ -1912,10 +2154,12 @@ def set_lr(step_now): if hasattr(dataset, "evaluate_ood_ablated"): ood_ablated_rng = random.Random() ood_ablated_rng.setstate(ood_rng_state) + segment_diagnostics.set_channel("eval-ood-memory-off") with _replay_torch_rng(ood_torch_rng): ood_ablated_acc, _ood_ablated_perfect = dataset.evaluate_ood_ablated( rnn, device, periodic_ood_episodes, ood_ablated_rng ) + segment_diagnostics.emit("eval-ood-memory-off") ood_log_writer.writerow( [ @@ -2068,9 +2312,22 @@ def set_lr(step_now): # Functional-usage check: same lesson distribution the ID eval # above just used (pre_advance_lesson, not curriculum.lesson -- # this call may have just advanced it). - ablated_triple_acc, ablated_perfect_frac = dataset.evaluate_id_ablated( - rnn, device, curriculum, pre_advance_lesson - ) + # Replay the exact ID evaluation episodes with memory bypassed. + # Comparing two independently sampled episode sets made this + # purported causal dependency metric too noisy to diagnose + # whether the model actually used memory. Restore the post-ID + # stream afterward so the paired evaluation consumes no extra + # randomness from the run. + rng_state_after_other_evaluations = _capture_rng_state() + _restore_rng_state(id_eval_rng_state) + segment_diagnostics.set_channel("eval-id-memory-off") + try: + ablated_triple_acc, ablated_perfect_frac = dataset.evaluate_id_ablated( + rnn, device, curriculum, pre_advance_lesson + ) + finally: + _restore_rng_state(rng_state_after_other_evaluations) + segment_diagnostics.emit("eval-id-memory-off") mem_check_writer.writerow( [ @@ -2098,11 +2355,13 @@ def set_lr(step_now): # contribution to accuracy is measured directly instead of # inferred from loss curves / plateau shape. for stage_idx, stage_kind in enumerate(combiner_stage_kinds): + segment_diagnostics.set_channel(f"eval-combiner-{stage_idx}") stage_triple_acc, stage_perfect_frac = ( dataset.evaluate_id_combiner_stage_ablated( rnn, device, curriculum, pre_advance_lesson, skip_stages={stage_idx} ) ) + segment_diagnostics.emit(f"eval-combiner-{stage_idx}") combiner_stage_log_writer.writerow( [ step, @@ -2188,9 +2447,11 @@ def set_lr(step_now): final_id_stage_memory = rnn.memories[0] if hasattr(final_id_stage_memory, "begin_stage_diagnostics"): final_id_stage_memory.begin_stage_diagnostics() + segment_diagnostics.set_channel("eval-final-id") _, id_triple_acc, id_perfect_frac = curriculum.maybe_advance( rnn, device, step=step, optimizer=optimizer ) + segment_diagnostics.emit("eval-final-id") final_id_stage_diagnostics = ( final_id_stage_memory.finish_stage_diagnostics() if hasattr(final_id_stage_memory, "finish_stage_diagnostics") @@ -2222,6 +2483,7 @@ def set_lr(step_now): final_ood_stage_memory = rnn.memories[0] if hasattr(final_ood_stage_memory, "begin_stage_diagnostics"): final_ood_stage_memory.begin_stage_diagnostics() + segment_diagnostics.set_channel("eval-final-ood") ood_triple_acc, ood_perfect_frac, ood_hop_breakdown = dataset.evaluate_ood( rnn, device, @@ -2230,6 +2492,7 @@ def set_lr(step_now): verbose_n=10, field_log=final_ood_field_log, ) + segment_diagnostics.emit("eval-final-ood") final_ood_stage_diagnostics = ( final_ood_stage_memory.finish_stage_diagnostics() if hasattr(final_ood_stage_memory, "finish_stage_diagnostics") @@ -2339,10 +2602,12 @@ def _write_diagnostic_values(diagnostics): if hasattr(dataset, "evaluate_ood_ablated"): final_ood_ablated_rng = random.Random() final_ood_ablated_rng.setstate(final_ood_rng_state) + segment_diagnostics.set_channel("eval-final-ood-memory-off") with _replay_torch_rng(final_ood_torch_rng): final_ood_ablated_acc, _final_ood_ablated_perfect = dataset.evaluate_ood_ablated( rnn, device, final_ood_episodes, final_ood_ablated_rng ) + segment_diagnostics.emit("eval-final-ood-memory-off") q99_metrics = None if split_graph_enabled and dnc_write_variant == "kanerva-closed-form": @@ -2462,6 +2727,7 @@ def _write_diagnostic_values(diagnostics): if dynamic_n_ctrl is not None: dynamic_n_log_file.close() + segment_diagnostics.close() return summary @@ -2638,6 +2904,19 @@ def _write_diagnostic_values(diagnostics): "SwitchMoE blocks between Mamba controller blocks (mamba_controller.py " "/ moe_layer.py).", ) + parser.add_argument( + "--workspace", action="store_true", + help="Phase 0 shared workspace: competitively write selected MoE expert outputs to persistent slots, then broadcast slots back to every specialist. Requires --split-graph and --moe.", + ) + parser.add_argument("--workspace-num-slots", type=int, default=WORKSPACE_NUM_SLOTS, + help="Persistent workspace slot count (0 is the plumbing/bypass control).") + parser.add_argument("--workspace-no-selection", action="store_true", + help="Control: use random, unweighted specialist writes instead of MoE-selected writes.") + parser.add_argument("--workspace-reset-each-step", action="store_true", + help="Persistence ablation: reset workspace before each timestep's write/read cycle.") + parser.add_argument("--workspace-train-mode", choices=["coadapted", "frozen"], + default=WORKSPACE_TRAIN_MODE, + help="coadapted trains the full model; frozen trains workspace parameters and freezes substrate/task-head parameters.") parser.add_argument( "--moe-num-experts", type=int, @@ -2891,6 +3170,12 @@ def _write_diagnostic_values(diagnostics): if args.link_matrix_mode == "sparse_topk" and args.link_matrix_topk is None: raise SystemExit("--link-matrix-mode=sparse_topk requires --link-matrix-topk") + if args.workspace and not args.split_graph: + parser.error("--workspace requires --split-graph") + if args.workspace and not args.moe: + parser.error("--workspace requires --moe") + if args.workspace_num_slots < 0: + parser.error("--workspace-num-slots must be >= 0") if args.beta is not None: betas_to_run = [args.beta] @@ -2930,6 +3215,12 @@ def _write_diagnostic_values(diagnostics): run_id = f"{run_id}_dynN" if args.moe: run_id = f"{run_id}_moe{args.moe_num_experts}e" + if args.workspace: + run_id += f"_workspace{args.workspace_num_slots}s_{args.workspace_train_mode}" + if args.workspace_no_selection: + run_id += "_randomwrite" + if args.workspace_reset_each_step: + run_id += "_resetstep" if args.split_graph: run_id = f"{run_id}_splitgraph_{args.split_graph_variant.replace('+', '')}" if args.dnc_read_variant != "cosine" or args.dnc_write_variant != "learned": @@ -2978,6 +3269,11 @@ def _write_diagnostic_values(diagnostics): split_graph_combiner_mode=args.split_graph_combiner_mode, split_graph_combiner_variant=args.split_graph_combiner_variant, split_graph_combiner_num_blocks=args.split_graph_combiner_num_blocks, + workspace_enabled=args.workspace, + workspace_num_slots=args.workspace_num_slots, + workspace_no_selection=args.workspace_no_selection, + workspace_reset_each_step=args.workspace_reset_each_step, + workspace_train_mode=args.workspace_train_mode, dnc_read_variant=args.dnc_read_variant, dnc_read_residual_scale=args.dnc_read_residual_scale, dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, diff --git a/src/initium/diagnostics/__init__.py b/src/initium/diagnostics/__init__.py new file mode 100644 index 0000000..fb48747 --- /dev/null +++ b/src/initium/diagnostics/__init__.py @@ -0,0 +1,2 @@ +"""Reusable training diagnostics.""" + diff --git a/src/initium/diagnostics/segment_logging.py b/src/initium/diagnostics/segment_logging.py new file mode 100644 index 0000000..a902b31 --- /dev/null +++ b/src/initium/diagnostics/segment_logging.py @@ -0,0 +1,244 @@ +"""Low-synchronization activation and gradient summaries for model segments. + +Hooks are selected by module type rather than a particular model topology, +so the same collector covers split-graph, standalone Mamba/CfC, chained +Mamba+CfC, recurrent DNC, and workspace models. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Any + +import torch +from torch import nn + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if torch.is_tensor(value): + return value + if isinstance(value, (tuple, list)): + for item in value: + found = _first_tensor(item) + if found is not None: + return found + if isinstance(value, dict): + for item in value.values(): + found = _first_tensor(item) + if found is not None: + return found + return None + + +def _segment_label(path: str, module: nn.Module) -> str | None: + class_name = type(module).__name__ + lower_name = class_name.lower() + leaf = path or "model" + if class_name in {"MambaBackboneParallel", "CfCBackboneParallel"}: + return f"backbone.{class_name.removesuffix('BackboneParallel').lower()}:{leaf}" + if class_name in {"_ParallelBlock", "_CfCParallelBlock"}: + kind = "mamba" if class_name == "_ParallelBlock" else "cfc" + return f"backbone.{kind}_block:{leaf}" + if class_name in {"MoEBlock", "MultiSourceMoEBlock"}: + return f"moe_residual:{leaf}" + if class_name == "SplitGraphDNC": + return f"split_graph:{leaf}" + if class_name == "ChainedControllerWrapper": + return f"controller_chain:{leaf}" + if class_name == "CfCControllerWrapper": + return f"controller.cfc:{leaf}" + if class_name == "CfCControllerBlock": + return f"controller.cfc_block:{leaf}" + if class_name in {"MambaControllerWrapper", "Mamba2ControllerWrapper", "Mamba3ControllerWrapper"}: + return f"controller.{class_name.removesuffix('ControllerWrapper').lower()}:{leaf}" + if class_name in { + "MambaControllerBlock", + "Mamba2ControllerBlock", + "Mamba3ControllerBlock", + }: + return f"controller.mamba_block:{leaf}" + if class_name in {"MambaControllerCell", "Mamba2ControllerCell", "Mamba3ControllerCell"}: + return f"controller.mamba_cell:{leaf}" + if class_name == "WorkspaceBroadcast": + return f"workspace:{leaf}" + if class_name == "Memory" or "memory" in path.lower() and "memory" in lower_name: + return f"dnc_memory:{leaf}" + if isinstance(module, nn.LSTM): + return f"controller.lstm:{leaf}" + if isinstance(module, nn.GRU): + return f"controller.gru:{leaf}" + if isinstance(module, nn.RNN): + return f"controller.rnn:{leaf}" + return None + + +class SegmentDiagnostics: + """Aggregate segment input/output health and hierarchical gradient norms. + + Scalar reductions remain detached on their source device. The only + device-to-host transfer occurs when ``emit`` is called at a log/eval + boundary or when a numerical failure needs an immediate trace. + """ + + def __init__(self, model: nn.Module): + self.model = model + self.channel: str | None = "train" + self._records: dict[str, OrderedDict[str, dict[str, Any]]] = {} + self._segments: list[tuple[str, str]] = [] + self._external_modules: list[tuple[str, nn.Module]] = [] + self._handles = [] + for path, module in model.named_modules(): + label = _segment_label(path, module) + if label is not None: + self._register(module, path, label) + + def _register(self, module: nn.Module, path: str, label: str) -> None: + self._segments.append((path, label)) + + def hook(_module, inputs, output): + if self.channel is None: + return + channel = self.channel + input_tensor = _first_tensor(inputs) + output_tensor = _first_tensor(output) + record = self._records.setdefault(channel, OrderedDict()).setdefault( + label, {"path": path, "calls": 0} + ) + record["calls"] += 1 + if input_tensor is not None: + self._accumulate_tensor(record, "in", input_tensor) + if input_tensor.requires_grad: + input_tensor.register_hook( + lambda grad, target=record: self._accumulate_tensor( + target, "grad_in", grad + ) + ) + if output_tensor is not None: + self._accumulate_tensor(record, "out", output_tensor) + if output_tensor.requires_grad: + output_tensor.register_hook( + lambda grad, target=record: self._accumulate_tensor( + target, "grad_out", grad + ) + ) + + self._handles.append(module.register_forward_hook(hook)) + + @staticmethod + def _accumulate_tensor(record: dict[str, Any], prefix: str, tensor: torch.Tensor) -> None: + value = tensor.detach().float().reshape(-1) + finite = torch.isfinite(value) + safe = torch.where(finite, value, torch.zeros_like(value)) + values = { + f"{prefix}_sum_sq": safe.square().sum(), + f"{prefix}_max": safe.abs().amax() if safe.numel() else safe.new_zeros(()), + f"{prefix}_bad": (~finite).sum().float(), + } + for key, item in values.items(): + record[key] = record[key] + item if key in record else item + record[f"{prefix}_numel"] = record.get(f"{prefix}_numel", 0) + value.numel() + + def add_module(self, module: nn.Module, label: str) -> None: + """Add a caller-named module such as the task output projection.""" + self._register(module, label, label) + self._external_modules.append((label, module)) + + def set_channel(self, channel: str | None) -> None: + self.channel = channel + + def _gradient_records(self) -> OrderedDict[str, tuple[torch.Tensor, int]]: + records: OrderedDict[str, list[tuple[torch.Tensor, int]]] = OrderedDict() + unassigned: list[tuple[torch.Tensor, int]] = [] + for name, parameter in self.model.named_parameters(): + grad = parameter.grad + if grad is None: + continue + value = grad.detach().float() + contribution = value.square().sum() + matched = False + for path, label in self._segments: + if path and (name == path or name.startswith(path + ".")): + records.setdefault(label, []).append((contribution, parameter.numel())) + matched = True + if not matched: + unassigned.append((contribution, parameter.numel())) + for label, module in self._external_modules: + for parameter in module.parameters(): + grad = parameter.grad + if grad is not None: + records.setdefault(label, []).append( + (grad.detach().float().square().sum(), parameter.numel()) + ) + if unassigned: + records["other"] = unassigned + result = OrderedDict() + for label, tensors in records.items(): + result[label] = ( + torch.stack([value for value, _ in tensors]).sum().sqrt(), + sum(numel for _, numel in tensors), + ) + return result + + def emit(self, channel: str, *, include_gradients: bool = False) -> None: + records = self._records.pop(channel, None) + gradient_records = self._gradient_records() if include_gradients else OrderedDict() + if not records and not gradient_records: + return + + scalar_tensors: list[torch.Tensor] = [] + record_shapes: list[tuple[str, int, int, int, int, int]] = [] + for label, record in (records or {}).items(): + record_shapes.append( + ( + label, + record["calls"], + record.get("in_numel", 0), + record.get("out_numel", 0), + record.get("grad_in_numel", 0), + record.get("grad_out_numel", 0), + ) + ) + for prefix in ("in", "out", "grad_in", "grad_out"): + for suffix in ("sum_sq", "max", "bad"): + item = record.get(f"{prefix}_{suffix}") + if item is None: + item = torch.zeros((), device=next(self.model.parameters()).device) + scalar_tensors.append(item.detach().float().reshape(())) + gradient_shapes = list(gradient_records.items()) + scalar_tensors.extend(value.detach().float().reshape(()) for value, _ in gradient_records.values()) + if scalar_tensors: + values = torch.stack(scalar_tensors).cpu().tolist() + else: + values = [] + + parts = [] + index = 0 + for label, calls, in_numel, out_numel, grad_in_numel, grad_out_numel in record_shapes: + summaries = [] + for prefix, numel, display in ( + ("in", in_numel, "in"), + ("out", out_numel, "out"), + ("grad_in", grad_in_numel, "g_in"), + ("grad_out", grad_out_numel, "g_out"), + ): + sum_sq, maximum, bad = values[index : index + 3] + index += 3 + if numel: + summaries.append( + f"{display}:rms={(sum_sq / numel) ** 0.5:.3g}," + f"max={maximum:.3g},nonfinite={bad / numel:.2%}" + ) + parts.append(f"{label}[calls={calls} {' '.join(summaries)}]") + if gradient_shapes: + gradients = [] + for label, (_, parameter_count) in gradient_shapes: + norm = values[index] + index += 1 + gradients.append(f"{label}={norm:.3g}/n{parameter_count}") + parts.append("grad_norm[" + "; ".join(gradients) + "]") + print(f"[segment-diag:{channel}] " + " | ".join(parts)) + + def close(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 0b68989..a7f94e2 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -198,6 +198,15 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea "split_graph_combiner_cfc_multi_source_moe", cc.SPLIT_GRAPH_COMBINER_CFC_MULTI_SOURCE_MOE, ), + workspace_enabled=c.get("workspace_enabled", cc.WORKSPACE_ENABLED), + workspace_num_slots=c.get("workspace_num_slots", cc.WORKSPACE_NUM_SLOTS), + workspace_no_selection=c.get( + "workspace_no_selection", cc.WORKSPACE_NO_SELECTION + ), + workspace_reset_each_step=c.get( + "workspace_reset_each_step", cc.WORKSPACE_RESET_EACH_STEP + ), + workspace_train_mode=c.get("workspace_train_mode", cc.WORKSPACE_TRAIN_MODE), device=device, ).to(device) diff --git a/src/initium/inference/run_inference.py b/src/initium/inference/run_inference.py index 0eaa6b6..d8905e7 100644 --- a/src/initium/inference/run_inference.py +++ b/src/initium/inference/run_inference.py @@ -558,9 +558,10 @@ def main(argv=None) -> int: " (multi-source combiner)" if hasattr(layer, "last_source_diagnostics") else "" ) print( - f" [{i}]{tag} cv_load={diag['cv_load']:.4f} " - f"cv_importance={diag['cv_importance']:.4f} " - f"max_load_frac={diag['max_load_frac']:.4f}" + f" [{i}]{tag} cv_load={float(diag['cv_load']):.4f} " + f"cv_importance={float(diag['cv_importance']):.4f} " + f"max_load_frac={float(diag['max_load_frac']):.4f} " + f"capacity_drop_frac={float(diag.get('capacity_drop_frac', 0.0)):.4f}" ) src_diag = ( layer.last_source_diagnostics() diff --git a/src/initium/mamba_controller/mamba_backbone_parallel.py b/src/initium/mamba_controller/mamba_backbone_parallel.py index d109ffd..f9455d5 100644 --- a/src/initium/mamba_controller/mamba_backbone_parallel.py +++ b/src/initium/mamba_controller/mamba_backbone_parallel.py @@ -289,3 +289,23 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: assert self.moe_blocks is not None h = self.moe_blocks[i](h) return h + + def forward_with_specialists(self, x: torch.Tensor): + """Return final MoE candidates and sparse routes for workspace use. + + Earlier backbone/MoE blocks remain unchanged. The final MoE block + exposes expert deltas before routing so the workspace can broadcast + context to every specialist, then reuse the original sparse route. + """ + if not self.moe_enabled or self.moe_blocks is None: + raise RuntimeError("workspace broadcast requires an enabled MoE expert pool") + h = self.in_adapter(x) + final_index = len(self.blocks) - 1 + for i, block in enumerate(self.blocks): + h = block(h) + if i == final_index: + base = h + _, experts, indices, weights = self.moe_blocks[i](h, return_details=True) + return base, experts, indices, weights + h = self.moe_blocks[i](h) + raise RuntimeError("parallel backbone has no blocks") diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 11e3af8..63052d2 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -42,10 +42,10 @@ this project has built, so this doesn't add a new bottleneck, it just stops making the (much more expensive) backbone share it. -Requirement from the roadmap's own methodological note on this option: "do -not layer them all at once... Concept 16 / Suspected Pattern SP-10". This -module therefore deliberately does NOT wire in Option 4's MoE (moe_enabled -raises if set) -- Option 5 must be run and ablated on its own first. +The split graph can compose with the existing external MoE blocks. The +optional shared workspace is a separate Phase 0 feature and requires those +MoE specialists; it operates on the final parallel-backbone MoE bank while +leaving earlier MoE blocks unchanged. --- Built-in ablation for the roadmap's own required verification -------- "Verify [the design] preserves the exact 'read now, decide next hop @@ -128,6 +128,7 @@ from initium.mamba_controller.mamba3_controller import Mamba3ControllerWrapper from initium.mamba_controller.mamba_controller import MambaControllerWrapper from initium.memory_manipulation.dnc_mechanisms import build_memory +from initium.workspace.broadcast import WorkspaceBroadcast class SplitGraphDNC(nn.Module): @@ -200,6 +201,11 @@ def __init__( dnc_read_residual_max_ratio: float = 0.5, dnc_write_variant: str = "learned", dnc_write_observation_variance: float = 1.0, + workspace_enabled: bool = False, + workspace_num_slots: int = 5, + workspace_no_selection: bool = False, + workspace_reset_each_step: bool = False, + workspace_train_mode: str = "coadapted", ): super().__init__() if not independent_linears: @@ -228,6 +234,17 @@ def __init__( self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio self.dnc_write_variant = dnc_write_variant self.dnc_write_observation_variance = dnc_write_observation_variance + if workspace_enabled and not moe_enabled: + raise ValueError("--workspace requires --moe so it can broadcast to MoE specialists") + if workspace_train_mode not in ("coadapted", "frozen"): + raise ValueError("workspace_train_mode must be 'coadapted' or 'frozen'") + if workspace_enabled and workspace_train_mode == "frozen" and workspace_num_slots == 0: + raise ValueError("frozen workspace training requires workspace_num_slots > 0") + self.workspace_enabled = workspace_enabled + self.workspace_num_slots = workspace_num_slots + self.workspace_no_selection = workspace_no_selection + self.workspace_reset_each_step = workspace_reset_each_step + self.workspace_train_mode = workspace_train_mode # Mamba bookkeeping, for checkpoint self-description -- same # rationale/convention as MambaDNC (mamba_controller.py). @@ -467,15 +484,26 @@ def __init__( self.output = nn.Linear(self.nn_output_size, self.input_size, device=device) torch.nn.init.kaiming_uniform_(self.output.weight) + if workspace_num_slots < 0: + raise ValueError("workspace_num_slots must be >= 0") + # Build the optional module last so enabling it leaves the existing + # backbone, combiner, DNC, and output-projection initializations and + # subsequent global RNG stream intact for matched-seed controls. + rng_state = ( + torch.random.get_rng_state() + if workspace_enabled and workspace_num_slots > 0 + else None + ) + self.workspace = ( + WorkspaceBroadcast(hidden_size, moe_num_experts, workspace_num_slots) + if workspace_enabled and workspace_num_slots > 0 + else None + ) + if rng_state is not None: + torch.random.set_rng_state(rng_state) - # Option 4 (MoE) interoperability: kept OFF and deliberately NOT - # wired into the backbone blocks here (unlike mamba_controller.py's - # MambaControllerWrapper) -- Concept 16/SP-10 (non-additive - # combination) is why Option 5 is being built and ablated on its - # own first. moe_layers stays an empty list so - # pop_total_moe_aux_loss(...) (moe_layer.py) is still safe to call - # unconditionally from the training script regardless of which - # controller is active this run. + # Keep a flat list of every active MoE block so the training loop can + # collect auxiliary losses and routing diagnostics uniformly. self.moe_enabled = moe_enabled self.moe_cfc_multi_source = moe_cfc_multi_source # Flat list every pop_total_moe_aux_loss() call sums over: the @@ -525,6 +553,20 @@ def forward( """ chx, mhx, last_read = hx B, T, _ = input.shape + combiner_hx_in = chx + workspace_state = None + if self.workspace is not None: + if isinstance(chx, dict) and "_initium_workspace_state" in chx: + workspace_state = chx["_initium_workspace_state"] + combiner_hx_in = chx.get("_initium_combiner_state") + if start_step > 0 and workspace_state is None: + raise ValueError( + "workspace-enabled start_step>0 requires workspace state in hx[0]" + ) + if workspace_state is None or reset_experience: + workspace_state = self.workspace.init_state( + B, device=input.device, dtype=input.dtype + ) if start_step: # Backbone still runs over the FULL input (cheap, parallel, stateless across # calls); only the sequential memory loop is skipped for t < start_step. @@ -539,7 +581,35 @@ def forward( # Single call, whole sequence, zero dependency on memory state. # This is the entire "removed from the sequential critical path" # half of Option 5. - H = self.backbone(input) # (B, T, hidden_size) + if self.workspace is not None: + backbone_stages = ( + list(self.backbone) + if isinstance(self.backbone, nn.Sequential) + else [self.backbone] + ) + stage_input = input + for stage in backbone_stages[:-1]: + stage_input = stage(stage_input) + final_stage = backbone_stages[-1] + if not hasattr(final_stage, "forward_with_specialists"): + raise RuntimeError( + "workspace requires the final parallel-backbone stage to expose its MoE specialists" + ) + base, specialists, route_indices, route_weights = ( + final_stage.forward_with_specialists(stage_input) + ) + broadcast_delta, workspace_state = self.workspace( + specialists[:, start_step:], + route_indices[:, start_step:], + route_weights[:, start_step:], + workspace_state, + no_selection=self.workspace_no_selection, + reset_each_step=self.workspace_reset_each_step, + ) + H = base.clone() + H[:, start_step:] = base[:, start_step:] + broadcast_delta + else: + H = self.backbone(input) # (B, T, hidden_size) # ---- memory hidden-state init (byte-identical convention to # dnc.DNC._init_hidden / MambaDNC._init_hidden's memory branch) --- @@ -576,8 +646,8 @@ def forward( if self.combiner_mode == "controller": assert self.combiner_wrapper is not None combiner_hx = ( - chx - if chx is not None + combiner_hx_in + if combiner_hx_in is not None else self.combiner_wrapper.init_state(B, device=device, dtype=H.dtype) ) @@ -632,6 +702,16 @@ def forward( chx = combiner_hx # v11: expose the controller combiner's final state, same # convention as mem_state -> mhx just above + if self.workspace is not None: + chx = { + "_initium_workspace_state": workspace_state, + "_initium_combiner_state": ( + combiner_hx + if self.combiner_mode == "controller" + else combiner_hx_in + ), + } + mhx = [mem_state] output = torch.stack(outputs, dim=0) # (T, B, input_size) diff --git a/src/initium/workspace/__init__.py b/src/initium/workspace/__init__.py new file mode 100644 index 0000000..d5a04e1 --- /dev/null +++ b/src/initium/workspace/__init__.py @@ -0,0 +1,5 @@ +"""Shared workspace selection and broadcast components.""" + +from initium.workspace.broadcast import WorkspaceBroadcast + +__all__ = ["WorkspaceBroadcast"] diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py new file mode 100644 index 0000000..e86e466 --- /dev/null +++ b/src/initium/workspace/broadcast.py @@ -0,0 +1,171 @@ +"""Persistent shared-object write and broadcast read for MoE specialists. + +The module owns no episode state. Callers pass the slot matrix between steps, +which keeps resets explicit and makes continuation/checkpointing possible. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + + +class WorkspaceBroadcast(nn.Module): + """Competitive MoE-selected writes and all-expert shared-slot readback. + + Inputs are specialist deltas ``(B, E, D)`` and sparse MoE routes + ``(B, K)``/weights ``(B, K)``. The router's selection determines which + specialists write. Every specialist independently queries all slots and + receives the resulting context before the selected deltas are combined. + """ + + def __init__(self, d_model: int, num_experts: int, num_slots: int = 5): + super().__init__() + if d_model < 1 or num_experts < 1: + raise ValueError("d_model and num_experts must be positive") + if num_slots < 0: + raise ValueError("num_slots must be >= 0") + self.d_model = d_model + self.num_experts = num_experts + self.num_slots = num_slots + self.query = nn.Linear(d_model, d_model, bias=False) + self.key = nn.Linear(d_model, d_model, bias=False) + self.value = nn.Linear(d_model, d_model, bias=False) + self.write_gate = nn.Linear(2 * d_model, d_model) + self.read_gate = nn.Parameter(torch.zeros(())) + self.write_gate_closed = False + self._last_step_diag: dict[str, torch.Tensor] = {} + self._last_diag: dict[str, torch.Tensor] = {} + nn.init.zeros_(self.write_gate.weight) + nn.init.zeros_(self.write_gate.bias) + + def init_state(self, batch: int, *, device, dtype) -> torch.Tensor: + return torch.zeros(batch, self.num_slots, self.d_model, device=device, dtype=dtype) + + def forward_step( + self, + specialist_deltas: torch.Tensor, + route_indices: torch.Tensor, + route_weights: torch.Tensor, + state: torch.Tensor, + *, + no_selection: bool = False, + reset_each_step: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.num_slots == 0 or self.write_gate_closed: + selected = specialist_deltas.gather( + 1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model) + ) + zero = specialist_deltas.new_zeros((), dtype=torch.float32) + self._last_step_diag = { + "write_mass": zero, + "active_write_frac": zero, + "write_gate_mean": zero, + "slot_rms": zero, + "broadcast_rms": zero, + "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + } + return (selected * route_weights.unsqueeze(-1)).sum(dim=1), state + + batch, experts, width = specialist_deltas.shape + if (experts, width) != (self.num_experts, self.d_model): + raise ValueError("specialist delta shape does not match workspace configuration") + if state.shape != (batch, self.num_slots, self.d_model): + raise ValueError("workspace state has an incompatible batch/slot/feature shape") + + if reset_each_step: + state = torch.zeros_like(state) + if no_selection: + write_weights = torch.rand( + batch, experts, device=specialist_deltas.device, dtype=torch.float32 + ) + write_weights = write_weights / write_weights.sum(-1, keepdim=True).clamp_min(1e-8) + else: + write_weights = torch.zeros( + batch, experts, device=specialist_deltas.device, dtype=torch.float32 + ) + write_weights.scatter_add_(1, route_indices, route_weights.float()) + + # Candidate keys compete over slots using content-based attention; + # slot writes are weighted by the existing sparse MoE selection. + keys = self.key(specialist_deltas) + scores = torch.einsum("bed,bsd->bes", keys.float(), self.query(state).float()) + scores = scores / math.sqrt(self.d_model) + slot_attention = scores.softmax(dim=-1).to(specialist_deltas.dtype) + values = self.value(specialist_deltas) + slot_proposal = torch.einsum( + "be,bes,bed->bsd", write_weights.to(values.dtype), slot_attention, values + ) + mass = torch.einsum("be,bes->bs", write_weights.to(values.dtype), slot_attention) + slot_proposal = slot_proposal / mass.unsqueeze(-1).clamp_min(1e-6) + slot_inputs = slot_proposal + old_inputs = state + update = torch.sigmoid(self.write_gate(torch.cat([old_inputs, slot_inputs], dim=-1))) + has_writer = write_weights.sum(dim=-1, keepdim=True).unsqueeze(-1) > 0 + update = update * has_writer.to(update.dtype) + next_state = update.mul(slot_inputs).add((1.0 - update).mul(old_inputs)) + + # All experts query the shared object, including those not selected + # for this step's write. + read_queries = self.query(specialist_deltas) + read_scores = torch.einsum( + "bed,bsd->bes", read_queries.float(), self.key(next_state).float() + ) + read_scores = read_scores / math.sqrt(self.d_model) + read_weights = read_scores.softmax(dim=-1).to(specialist_deltas.dtype) + read_values = self.value(next_state) + broadcast = torch.einsum("bes,bsd->bed", read_weights, read_values) + read_scale = torch.sigmoid(self.read_gate).to(specialist_deltas.dtype) + enriched = specialist_deltas + read_scale * broadcast + selected = enriched.gather(1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model)) + self._last_step_diag = { + "write_mass": write_weights.detach().sum(dim=-1).mean(), + "active_write_frac": (write_weights.detach().sum(dim=-1) > 0).float().mean(), + "write_gate_mean": update.detach().float().mean(), + "slot_rms": next_state.detach().float().square().mean().sqrt(), + "broadcast_rms": ( + read_scale.detach().float() * broadcast.detach().float() + ).square().mean().sqrt(), + "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + } + return (selected * route_weights.unsqueeze(-1)).sum(dim=1), next_state + + def forward( + self, + specialist_deltas: torch.Tensor, + route_indices: torch.Tensor, + route_weights: torch.Tensor, + state: torch.Tensor, + *, + no_selection: bool = False, + reset_each_step: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Process time-major specialists ``(B,T,E,D)`` sequentially.""" + outputs = [] + step_diags = [] + for step in range(specialist_deltas.shape[1]): + out, state = self.forward_step( + specialist_deltas[:, step], route_indices[:, step], route_weights[:, step], + state, + no_selection=no_selection, + reset_each_step=reset_each_step, + ) + outputs.append(out) + step_diags.append(self._last_step_diag) + if step_diags: + self._last_diag = { + key: torch.stack([diag[key] for diag in step_diags]).mean().detach() + for key in step_diags[0] + } + self._last_diag["read_gate"] = torch.sigmoid(self.read_gate.detach()).float() + self._last_diag["broadcast_to_specialist_ratio"] = ( + self._last_diag["broadcast_rms"] + / self._last_diag["specialist_rms"].clamp_min(1e-8) + ) + return torch.stack(outputs, dim=1), state + + def last_diagnostics(self) -> dict[str, torch.Tensor]: + """Detached activity metrics for the most recent sequence pass.""" + return dict(self._last_diag) From ee021163e62b28ffd29bcd0fb71c735cc1d6c5bc Mon Sep 17 00:00:00 2001 From: Samuel Date: Fri, 9 Oct 2026 21:23:49 +0200 Subject: [PATCH 02/14] Improve split-graph MoE routing and diagnostics This patch keeps MoE additive around the original controller signal instead of replacing it at initialization, zero-initializes expert output projections, and fixes top-1 routing so the router still receives useful task gradients. Capacity is now assigned by route confidence rather than flattened token order, which avoids first-source bias during multi-source CfC routing and adds accepted/dropped source diagnostics for training logs and inference reports. It also adds split-graph MoE placement controls (`all`, `backbone`, `combiner`) so backbone and combiner MoE effects can be isolated cleanly, and updates the README with the placement behavior and a short diagnostic smoke-test matrix for comparing MoE effects on memory-learning transitions. The change keeps the baseline CfC fusion path intact while letting source-aware experts provide learned residual corrections. --- README.md | 15 ++ src/initium/LNN_controller/cfc_controller.py | 16 ++- src/initium/MoE/moe_layer.py | 132 +++++++++++++++--- src/initium/core_training.py | 43 +++++- src/initium/inference/run_inference.py | 4 +- .../mamba_controller/split_graph_dnc.py | 32 +++-- 6 files changed, 205 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index f51682e..a71cfc1 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,8 @@ The shared object operates on the final backbone MoE bank; earlier backbone MoE ### MoE numerical and memory-use diagnostics +MoE experts now initialize as a zero-output residual branch: each expert's final projection starts at zero, so enabling MoE preserves the original controller signal at step 0 while the output projections learn. Top-1 routing retains the selected expert's probability from the full expert softmax; normalizing a one-expert selected set would turn every gate into 1 and eliminate task-loss gradients to the router. Top-k greater than one continues to normalize over selected experts. The CfC multi-source combiner likewise retains its original affine projection over the concatenated backbone and previous-read inputs; its source-aware experts contribute an additive correction around that baseline. Capacity slots go to the highest-confidence routes rather than the earliest rows in the flattened token batch, avoiding a systematic advantage for the first source when source batches share an expert's capacity. CfC source diagnostics report accepted and dropped route fractions separately for each source at training-log cadence. + The SwitchMoE router projection and expert MLPs run in FP32 even when the surrounding model uses CUDA autocast; low-precision inputs receive FP32 MoE outputs so large finite expert values are not converted into fp16 infinities at the residual boundary. Ordinary routing dispatches activations into fixed-capacity per-expert buffers and applies each expert's original Linear parameters directly. It does not expand or gather full expert weight matrices per token. This matters for recurrent controllers: the intermediate batched-weight route materialized a 512×1536 projection matrix for every route and autograd retained those copies across timesteps. With 32 source tokens, one such projection copy was 96 MiB; input and output projections were both affected. Each expert now evaluates its capacity buffer, including any unused capacity rows. Route-expert and route-slot indices are both shaped `(tokens, top_k)` before gathering; this avoids advanced-index broadcasting across unrelated tokens. The workspace's specialist-candidate path intentionally evaluates every expert because broadcast requires all candidates, but it shares the expert parameters instead of copying them per token. Training logs include the fraction of selected routes dropped by capacity. Router and per-source diagnostics stay on-device during recurrent steps and synchronize only when needed for log output. The load-balancing loss is averaged across repeated calls to a recurrent MoE layer (for example, one call per CfC timestep), then summed across distinct MoE layers; its weight no longer scales with episode length. Periodic memory-dependency checks replay the exact random episode batch used by the corresponding memory-on ID evaluation, then restore the RNG streams. This makes `memory on - memory off` a paired functional ablation rather than a difference between two independently sampled batches. The training loop also clips all trainable optimizer parameters, including the task output projection, stops before backward on non-finite forward outputs/losses, reports the first parameter with a non-finite unscaled gradient, logs skipped non-finite-gradient updates, and aborts after eight consecutive such steps instead of continuing a run that is no longer updating safely. @@ -132,6 +134,19 @@ Training console logs now show **one MoE-bank record per 100-step logging window Evaluation emits separate channel lines for ID, OOD, memory-off, combiner-stage ablations, and final ID/OOD checks. These summarize the evaluations independently from training windows. If a forward value or gradient becomes non-finite, the current segment window and MoE-bank window are printed immediately before training stops or skips the unsafe update. The summaries aggregate detached device-side scalars and synchronize at log/evaluation boundaries; they do not retain activation graphs. +To isolate where split-graph MoE changes the memory-learning transition, `--moe-placement` selects which component receives expert banks while keeping the model, seed, and router settings matched. `all` (the default) preserves the existing behavior; `backbone` enables MoE only in the parallel backbone; `combiner` enables MoE only in the sequential controller combiner. For the CfC combiner, the latter is its source-aware multi-source MoE bank. Component placement requires `--moe --split-graph`; combiner-only additionally requires controller combiner mode. The workspace requires backbone MoE and therefore cannot be combined with combiner-only placement. + +Run these short, seed-matched diagnostics on the same source revision. Each is capped at 5,000 steps so the usual early memory-use transition can be compared before committing to the full Phase 1 horizon: + +```bash +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-smoke-no-moe --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --total-steps 5000 +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-smoke-moe-backbone --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-smoke-moe-combiner --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement combiner --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-smoke-moe-all --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +``` + +Compare the paired memory-on/off dependency, ID/OOD metrics, curriculum progress, per-component segment diagnostics, and per-bank routing/activation summaries. The placement runs distinguish backbone-MoE effects from the CfC multi-source path; the `all` condition checks their interaction. These are diagnostic smoke runs, not evidence of a completed Phase 1 result. + Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 1.0` to use the cosine-anchored Track A reader. Lower scales such as `0.5` are available for follow-up comparisons. Use `--dnc-write-variant kanerva-closed-form` for Track B. Add a reader and writer variant together only for the combined cell after each track is characterized. Dense temporal links remain the default. The short-run baseline reproduction gate and the Q99 paired-order evaluation are deliberate user-run experiment steps; no run results are assumed by this repository change. ### Controllers diff --git a/src/initium/LNN_controller/cfc_controller.py b/src/initium/LNN_controller/cfc_controller.py index 012f84b..3edf67a 100644 --- a/src/initium/LNN_controller/cfc_controller.py +++ b/src/initium/LNN_controller/cfc_controller.py @@ -283,9 +283,21 @@ def forward_multi_source(self, sources: list[torch.Tensor], hx): if self.moe_source_dims is None: raise RuntimeError("forward_multi_source requires moe_source_dims to have been set") assert self.source_in_adapters is not None and self.moe_blocks is not None + # Preserve the stock CfC input-fusion path exactly: the baseline + # receives one affine projection of concat([backbone, previous + # read]). The source-aware MoE is an additive residual around that + # fused representation. This keeps the memory signal and controller + # initialization intact while still allowing the experts to learn + # source-specific corrections. + concatenated = torch.cat(sources, dim=-1) + x = self.in_adapter(concatenated) projected = [adapter(s) for adapter, s in zip(self.source_in_adapters, sources)] - fused_per_source = self.moe_blocks[0](projected) # list[Tensor], one per source - x = torch.stack(fused_per_source, dim=0).sum(dim=0) # combine into one fused input + fused_per_source = self.moe_blocks[0](projected) # source + routed delta + moe_delta = torch.stack( + [fused - source for fused, source in zip(fused_per_source, projected)], + dim=0, + ).sum(dim=0) + x = x + moe_delta if hx is None: hx = self.init_state(x.size(0), device=x.device) new_hx = [] diff --git a/src/initium/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index 0d4165b..cb10b8e 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -100,6 +100,14 @@ def __init__(self, d_model: int, expert_dim: int): super().__init__() self.w_in = nn.Linear(d_model, expert_dim) self.w_out = nn.Linear(expert_dim, d_model) + # Make inserting an MoE residual a no-op at initialization. The + # output projections still receive task gradients on the first + # update; their hidden projections begin learning as soon as those + # output weights move away from zero. This avoids a randomly + # initialized expert bank changing the controller's signal scale + # before routing/expert specialization has learned anything. + nn.init.zeros_(self.w_out.weight) + nn.init.zeros_(self.w_out.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w_out(F.relu(self.w_in(x))) @@ -153,6 +161,7 @@ def __init__( self._last_topk_idx: torch.Tensor | None = ( None # (T, k), detached -- for per-source breakdowns ) + self._last_accepted: torch.Tensor | None = None # (T, k), detached def forward(self, x: torch.Tensor, *, return_details: bool = False): orig_shape = x.shape @@ -186,9 +195,9 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): probs = torch.softmax(logits32, dim=-1) # (num_tokens, num_experts) # Top-K routing (Shazeer et al. 2017 / GShard), generalizing Switch's - # own k=1 special case: each token selects its top_k highest-prob - # experts, gate weights renormalized to sum to 1 across just those - # k slots -- top_k=1 reduces to the original Switch gate exactly. + # own k=1 special case: top-1 retains its selected probability from + # the full expert softmax (so the task objective can train the + # router); top-k>1 normalizes among the selected experts. # Normal training dispatches selected, capacity-accepted routes into # bounded per-expert buffers. The optional all-expert result is # reserved for the workspace specialist path, which explicitly needs @@ -196,7 +205,15 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): # compute and avoids expanding expert weights per route. top_k = self.top_k topk_prob, topk_idx = probs.topk(top_k, dim=-1) # (T, k) each - gate_weights = topk_prob / topk_prob.sum(dim=-1, keepdim=True).clamp(min=1e-9) + if top_k == 1: + # Switch's top-1 gate multiplies the expert output by its + # probability in the full expert softmax. Renormalizing a + # one-element selected set would make every gate exactly 1 and + # cancel the softmax derivative, leaving the task loss unable + # to train the router at all. + gate_weights = topk_prob + else: + gate_weights = topk_prob / topk_prob.sum(dim=-1, keepdim=True).clamp(min=1e-9) # f_i/P_i are cheap to compute regardless of mode; only the aux LOSS # accumulation (needed for backward()) is training-gated below. @@ -219,18 +236,30 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): # per-expert buffer). capacity = max(1, int((num_tokens * top_k / self.num_experts) * self.capacity_factor)) - # Capacity mask: keep only the first `capacity` (token,slot) pairs, - # in original order, routed to each expert -- same drop semantics - # as the prior top-1 implementation, generalized over the - # flattened (T*k) slot order. + # Capacity is assigned by router confidence, not flattened token + # order. In recurrent multi-source calls the source batches are + # concatenated (backbone first, memory read second); first-come + # capacity therefore silently favored the first source whenever an + # expert overflowed. Confidence-priority acceptance is independent + # of source ordering and preserves the strongest routes. flat_idx = topk_idx.reshape(-1) - one_hot_flat = F.one_hot(flat_idx, num_classes=self.num_experts) - rank_in_expert = one_hot_flat.cumsum(dim=0).gather(1, flat_idx.unsqueeze(1)).squeeze(1) - 1 + flat_priority = topk_prob.reshape(-1) + priority_order = torch.argsort(flat_priority, descending=True, stable=True) + expert_in_priority_order = flat_idx.index_select(0, priority_order) + one_hot_priority = F.one_hot( + expert_in_priority_order, num_classes=self.num_experts + ) + ranks_in_priority_order = one_hot_priority.cumsum(dim=0).gather( + 1, expert_in_priority_order.unsqueeze(1) + ).squeeze(1) - 1 + rank_in_expert = torch.empty_like(ranks_in_priority_order) + rank_in_expert.scatter_(0, priority_order, ranks_in_priority_order) keep = (rank_in_expert < capacity).reshape(num_tokens, top_k) accepted_gates = gate_weights * keep.to(gate_weights.dtype) - accepted_gates = accepted_gates / accepted_gates.sum( - dim=-1, keepdim=True - ).clamp_min(1e-9) + if top_k > 1: + accepted_gates = accepted_gates / accepted_gates.sum( + dim=-1, keepdim=True + ).clamp_min(1e-9) with torch.autocast(device_type=x.device.type, enabled=False): # Dispatch activations into fixed-capacity expert buffers. The @@ -314,6 +343,7 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): output = output32.to(output_dtype) with torch.no_grad(): + self._last_accepted = keep.detach() self._last_diag = { "cv_importance": (P_i.std() / P_i.mean().clamp(min=1e-8)).detach(), "cv_load": (f_i.std() / f_i.mean().clamp(min=1e-8)).detach(), @@ -399,6 +429,11 @@ def last_routing(self) -> torch.Tensor | None: class has no notion of on its own.""" return self._last_topk_idx + def last_accepted_routing(self) -> torch.Tensor | None: + """Capacity acceptance mask ``(num_tokens, top_k)`` for the most + recent call, or ``None`` before the first call.""" + return self._last_accepted + class MoEBlock(nn.Module): """Pre-norm residual wrapper around a SwitchMoE sublayer -- the same @@ -499,6 +534,8 @@ class MultiSourceMoEBlock(nn.Module): """ _cumulative_source_counts: list[torch.Tensor] | None + _cumulative_source_accepted: list[torch.Tensor] | None + _cumulative_source_routes: list[torch.Tensor] | None def __init__( self, @@ -529,6 +566,8 @@ def __init__( load_balance_alpha=load_balance_alpha, ) self._cumulative_source_counts = None + self._cumulative_source_accepted = None + self._cumulative_source_routes = None if device is not None and getattr(device, "type", None) == "cuda": self.to(device) @@ -552,20 +591,37 @@ def forward(self, sources: list[torch.Tensor]) -> list[torch.Tensor]: # concentrates on experts {2,5}, source 1 (read vector) on {1,7}" # -- rather than every source routing near-identically by chance. topk_idx = self.moe.last_routing() # (sum(batch_sizes), top_k) or None - if topk_idx is not None: + accepted = self.moe.last_accepted_routing() + if topk_idx is not None and accepted is not None: num_experts = self.moe.num_experts if self._cumulative_source_counts is None: self._cumulative_source_counts = [ torch.zeros(num_experts, device=topk_idx.device) for _ in batch_sizes ] + self._cumulative_source_accepted = [ + torch.zeros(num_experts, device=topk_idx.device) for _ in batch_sizes + ] + self._cumulative_source_routes = [ + torch.zeros((), device=topk_idx.device) for _ in batch_sizes + ] offset = 0 for i, b in enumerate(batch_sizes): idx_slice = topk_idx[offset : offset + b].reshape(-1) counts = torch.bincount(idx_slice, minlength=num_experts).float() + accepted_slice = accepted[offset : offset + b].reshape(-1) + accepted_counts = torch.bincount( + idx_slice, + weights=accepted_slice.to(dtype=torch.float32), + minlength=num_experts, + ) # Keep routing counters on device and defer CPU synchronization # until inference explicitly requests a report. This block runs # once per CfC timestep during training. self._cumulative_source_counts[i] += counts + assert self._cumulative_source_accepted is not None + assert self._cumulative_source_routes is not None + self._cumulative_source_accepted[i] += accepted_counts + self._cumulative_source_routes[i] += accepted_slice.sum() offset += b outs, offset = [], 0 @@ -594,20 +650,58 @@ def last_source_diagnostics(self) -> list[dict] | None: if self._cumulative_source_counts is None: return None source_diag = [] - for counts in self._cumulative_source_counts: + assert self._cumulative_source_accepted is not None + assert self._cumulative_source_routes is not None + for counts, accepted_counts, accepted_routes in zip( + self._cumulative_source_counts, + self._cumulative_source_accepted, + self._cumulative_source_routes, + ): counts_cpu = counts.detach().float().cpu() + accepted_cpu = accepted_counts.detach().float().cpu() fractions = counts_cpu / counts_cpu.sum().clamp_min(1.0) top_expert = int(counts_cpu.argmax().item()) + total_routes = float(counts_cpu.sum().item()) + accepted_total = float(accepted_cpu.sum().item()) source_diag.append( { "expert_frac": fractions.tolist(), "top_expert": top_expert, "top_expert_frac": float(fractions[top_expert]), - "n_calls_accumulated": int(counts_cpu.sum().item()), + "accepted_expert_frac": ( + accepted_cpu / accepted_cpu.sum().clamp_min(1.0) + ).tolist(), + "accepted_route_frac": accepted_total / max(total_routes, 1.0), + "capacity_drop_frac": 1.0 - accepted_total / max(total_routes, 1.0), + "n_routes_accumulated": int(total_routes), + "n_accepted_accumulated": int(accepted_routes.detach().item()), } ) return source_diag + def pop_source_diagnostics(self) -> list[dict] | None: + """Return and reset per-source routing/capacity counters. + + This is intended for periodic training logs, where source-level + overflow must be compared over the same window as the pooled MoE + diagnostics. ``last_source_diagnostics`` remains cumulative for + inference callers. + """ + report = self.last_source_diagnostics() + if report is not None: + assert self._cumulative_source_counts is not None + assert self._cumulative_source_accepted is not None + assert self._cumulative_source_routes is not None + for counts, accepted, routes in zip( + self._cumulative_source_counts, + self._cumulative_source_accepted, + self._cumulative_source_routes, + ): + counts.zero_() + accepted.zero_() + routes.zero_() + return report + class MoERNNWrapper(nn.Module): """Minimal external-MoE add-on for the stock nn.LSTM/GRU/RNN controller @@ -678,6 +772,11 @@ def pop_total_moe_aux_loss(moe_layers: list, *, include_diagnostics: bool = True else: diag = {} if diag: + source_report = None + if include_diagnostics: + pop_source_diagnostics = getattr(layer, "pop_source_diagnostics", None) + if pop_source_diagnostics is not None: + source_report = pop_source_diagnostics() cv_imp.append(float(diag["cv_importance"])) cv_load.append(float(diag["cv_load"])) max_load.append(float(diag["max_load_frac"])) @@ -710,6 +809,7 @@ def pop_total_moe_aux_loss(moe_layers: list, *, include_diagnostics: bool = True for key in ("expert_load_frac", "router_importance_frac") if key in diag }, + **({"source_routes": source_report} if source_report is not None else {}), } ) if total is None: diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 13c5820..8d6c8ae 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -256,6 +256,12 @@ def _format_moe_layer_reports(diagnostics: dict) -> list[str]: importance = ",".join( f"{value:.2f}" for value in layer_diag.get("router_importance_frac", []) ) + source_routes = layer_diag.get("source_routes") or [] + source_summary = " ".join( + f"src{i}:accepted={source.get('accepted_route_frac', 0.0):.1%}," + f"drop={source.get('capacity_drop_frac', 0.0):.1%}" + for i, source in enumerate(source_routes) + ) reports.append( f"moe-bank[{layer_diag['name']} " f"load=[{load}] importance=[{importance}] " @@ -272,7 +278,8 @@ def _format_moe_layer_reports(diagnostics: dict) -> list[str]: f"moe_delta_rms/max={layer_diag.get('routed_output_rms', 0.0):.3g}/" f"{layer_diag.get('routed_output_absmax', 0.0):.3g} " f"drop={layer_diag.get('capacity_drop_frac', 0.0):.1%} " - f"gate_mass={layer_diag.get('accepted_gate_mass', 0.0):.2f}]" + f"gate_mass={layer_diag.get('accepted_gate_mass', 0.0):.2f}" + f"{(' ' + source_summary) if source_summary else ''}]" ) return reports @@ -639,6 +646,7 @@ def run( dynamic_n_trigger_frac: float = DYNAMIC_N_TRIGGER_FRAC, dynamic_n_cooldown_steps: int = DYNAMIC_N_COOLDOWN_STEPS, moe_enabled: bool = MOE_ENABLED, + split_graph_moe_placement: str = "all", moe_num_experts: int = MOE_NUM_EXPERTS, moe_expert_dim: int | None = MOE_EXPERT_DIM, moe_capacity_factor: float = MOE_CAPACITY_FACTOR, @@ -706,12 +714,14 @@ def run( "beta_effective", "task_loss", "kl_loss", + "kl_contrib", "total_loss", "digit_diversity", "kl_mean", "kl_max", "kl_min", "kl_std", + "clamp_frac", "lr", "grad_norm", "amp_scale", @@ -1168,12 +1178,13 @@ def run( if split_graph_enabled: if moe_enabled: print( - f"[{run_id}] moe_enabled=True with split_graph_enabled=True: MoE is now " - f"wired into both the parallel backbone and the sequential combiner (see " - f"split_graph_dnc.py). A combiner_variant='cfc' additionally routes its two " + f"[{run_id}] split-graph MoE placement={split_graph_moe_placement}: " + f"backbone={'on' if split_graph_moe_placement in ('all', 'backbone') else 'off'}, " + f"combiner={'on' if split_graph_moe_placement in ('all', 'combiner') else 'off'} " + f"(see split_graph_dnc.py). A combiner_variant='cfc' additionally routes its two " f"natural input sources (backbone output, previous read vector) through " - f"per-source specialized experts when " - f"split_graph_combiner_cfc_multi_source_moe={split_graph_combiner_cfc_multi_source_moe} " + f"per-source specialized experts when active=" + f"{split_graph_moe_placement in ('all', 'combiner') and split_graph_combiner_cfc_multi_source_moe} " f"(see moe_layer.MultiSourceMoEBlock)." ) @@ -1215,6 +1226,7 @@ def run( combiner_num_blocks=split_graph_combiner_num_blocks, independent_linears=True, moe_enabled=moe_enabled, + moe_placement=split_graph_moe_placement, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -1309,6 +1321,7 @@ def run( else None ), "split_graph_combiner_cfc_multi_source_moe": split_graph_combiner_cfc_multi_source_moe, + "split_graph_moe_placement": split_graph_moe_placement, "dnc_read_variant": dnc_read_variant, "dnc_read_residual_scale": dnc_read_residual_scale, "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, @@ -2904,6 +2917,13 @@ def _write_diagnostic_values(diagnostics): "SwitchMoE blocks between Mamba controller blocks (mamba_controller.py " "/ moe_layer.py).", ) + parser.add_argument( + "--moe-placement", + choices=("all", "backbone", "combiner"), + default="all", + help="For split-graph MoE, enable experts in both components (all), " + "only the parallel backbone, or only the sequential combiner. Default: all.", + ) parser.add_argument( "--workspace", action="store_true", help="Phase 0 shared workspace: competitively write selected MoE expert outputs to persistent slots, then broadcast slots back to every specialist. Requires --split-graph and --moe.", @@ -3174,6 +3194,14 @@ def _write_diagnostic_values(diagnostics): parser.error("--workspace requires --split-graph") if args.workspace and not args.moe: parser.error("--workspace requires --moe") + if args.moe_placement != "all" and not args.moe: + parser.error("--moe-placement backbone/combiner requires --moe") + if args.moe_placement != "all" and not args.split_graph: + parser.error("--moe-placement backbone/combiner requires --split-graph") + if args.moe_placement == "combiner" and args.split_graph_combiner_mode != "controller": + parser.error("--moe-placement combiner requires --split-graph-combiner-mode controller") + if args.workspace and args.moe_placement == "combiner": + parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: parser.error("--workspace-num-slots must be >= 0") @@ -3215,6 +3243,8 @@ def _write_diagnostic_values(diagnostics): run_id = f"{run_id}_dynN" if args.moe: run_id = f"{run_id}_moe{args.moe_num_experts}e" + if args.split_graph and args.moe_placement != "all": + run_id += f"_{args.moe_placement}only" if args.workspace: run_id += f"_workspace{args.workspace_num_slots}s_{args.workspace_train_mode}" if args.workspace_no_selection: @@ -3256,6 +3286,7 @@ def _write_diagnostic_values(diagnostics): dynamic_n_trigger_frac=args.dynamic_n_trigger_frac, dynamic_n_cooldown_steps=args.dynamic_n_cooldown_steps, moe_enabled=args.moe, + split_graph_moe_placement=args.moe_placement, moe_num_experts=args.moe_num_experts, moe_expert_dim=args.moe_expert_dim, moe_capacity_factor=args.moe_capacity_factor, diff --git a/src/initium/inference/run_inference.py b/src/initium/inference/run_inference.py index d8905e7..7ab96e0 100644 --- a/src/initium/inference/run_inference.py +++ b/src/initium/inference/run_inference.py @@ -575,7 +575,9 @@ def main(argv=None) -> int: ) print( f" source {s_idx}: top_expert={s['top_expert']} " - f"({s['top_expert_frac']:.2f}) | dist [{frac_str}]" + f"({s['top_expert_frac']:.2f}) | dist [{frac_str}] " + f"accepted={s.get('accepted_route_frac', 0.0):.1%} " + f"drop={s.get('capacity_drop_frac', 0.0):.1%}" ) rl.print_summary(summary) diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 63052d2..e9a5b1c 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -181,6 +181,7 @@ def __init__( independent_linears: bool = True, device: torch.device | None = None, moe_enabled: bool = False, + moe_placement: str = "all", moe_num_experts: int = 8, moe_expert_dim: int | None = None, moe_top_k: int = 1, @@ -234,8 +235,14 @@ def __init__( self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio self.dnc_write_variant = dnc_write_variant self.dnc_write_observation_variance = dnc_write_observation_variance - if workspace_enabled and not moe_enabled: - raise ValueError("--workspace requires --moe so it can broadcast to MoE specialists") + if moe_placement not in ("all", "backbone", "combiner"): + raise ValueError("moe_placement must be 'all', 'backbone', or 'combiner'") + backbone_moe_enabled = moe_enabled and moe_placement in ("all", "backbone") + combiner_moe_enabled = moe_enabled and moe_placement in ("all", "combiner") + if workspace_enabled and not backbone_moe_enabled: + raise ValueError( + "--workspace requires backbone MoE because it broadcasts to backbone specialists" + ) if workspace_train_mode not in ("coadapted", "frozen"): raise ValueError("workspace_train_mode must be 'coadapted' or 'frozen'") if workspace_enabled and workspace_train_mode == "frozen" and workspace_num_slots == 0: @@ -263,7 +270,7 @@ def __init__( expand=mamba_expand, headdim=mamba_headdim, cfc_kwargs=cfc_kwargs, - moe_enabled=moe_enabled, + moe_enabled=backbone_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -335,7 +342,7 @@ def __init__( expand=combiner_expand, headdim=combiner_headdim, ngroups=combiner_ngroups, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -351,7 +358,7 @@ def __init__( d_state=_d_state, expand=combiner_expand, headdim=combiner_headdim, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -374,7 +381,7 @@ def __init__( d_state=16, d_conv=combiner_d_conv, expand=combiner_expand, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -387,7 +394,7 @@ def __init__( expand=combiner_expand, headdim=combiner_headdim, ngroups=combiner_ngroups, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -398,7 +405,7 @@ def __init__( d_state=64, expand=combiner_expand, headdim=combiner_headdim, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -414,7 +421,7 @@ def __init__( in_dim=combiner_in_dim, d_model=hidden_size, num_blocks=combiner_num_blocks, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -428,7 +435,7 @@ def __init__( # concat -- see moe_layer.MultiSourceMoEBlock. moe_source_dims=( [hidden_size, self.read_vectors_size] - if (moe_enabled and moe_cfc_multi_source) + if (combiner_moe_enabled and moe_cfc_multi_source) else None ), device=device, @@ -441,7 +448,7 @@ def __init__( d_state=_d_state, d_conv=combiner_d_conv, expand=combiner_expand, - moe_enabled=moe_enabled, + moe_enabled=combiner_moe_enabled, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -504,7 +511,8 @@ def __init__( # Keep a flat list of every active MoE block so the training loop can # collect auxiliary losses and routing diagnostics uniformly. - self.moe_enabled = moe_enabled + self.moe_enabled = backbone_moe_enabled or combiner_moe_enabled + self.moe_placement = moe_placement self.moe_cfc_multi_source = moe_cfc_multi_source # Flat list every pop_total_moe_aux_loss() call sums over: the # backbone's own per-block MoE sublayers, PLUS the combiner's (a From f6047b09bd851d8f3a9f1b35874a42a66974d9cb Mon Sep 17 00:00:00 2001 From: Samuel Date: Sat, 10 Oct 2026 09:15:02 +0200 Subject: [PATCH 03/14] Add backbone MoE isolation diagnostics Adds optional split-graph backbone MoE isolation per zero-based block and wires it through the training/run, loader, and model configs. This enables backbone-only MoE experiments, validates incompatible options (workspace/bad ranges), and emits valid-vs-padded routing diagnostics for graph padding without masking the actual execution path. --- README.md | 9 ++ .../LNN_controller/cfc_backbone_parallel.py | 60 +++++++++---- src/initium/MoE/moe_layer.py | 82 +++++++++++++++-- src/initium/core_training.py | 87 ++++++++++++++++++- src/initium/inference/model_loader.py | 4 + .../mamba_backbone_parallel.py | 53 +++++++---- .../mamba_controller/split_graph_dnc.py | 26 +++++- 7 files changed, 279 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index a71cfc1..a7f28fd 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,15 @@ python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-smoke-mo Compare the paired memory-on/off dependency, ID/OOD metrics, curriculum progress, per-component segment diagnostics, and per-bank routing/activation summaries. The placement runs distinguish backbone-MoE effects from the CfC multi-source path; the `all` condition checks their interaction. These are diagnostic smoke runs, not evidence of a completed Phase 1 result. +For a backbone-only run, `--moe-backbone-block-isolation INDEX` installs an MoE bank after only the selected zero-based parallel-backbone block. The other backbone blocks still run their Mamba/CfC computation but have no MoE; the sequential combiner must also be excluded with `--moe-placement backbone`. The flag is range-checked against `--split-graph-num-blocks`, included in run IDs and checkpoint metadata, and cannot be combined with the workspace broadcast path. For a two-block Mamba backbone, run one isolated condition per block: + +```bash +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-backbone-block0 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 0 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-backbone-block1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +``` + +Each `moe-bank[backbone.moe_blocks.INDEX ...]` line reports router-selected expert load, router importance, entropy, maximum probability, and capacity-drop rate separately for valid graph timesteps and all-zero padded timesteps. The valid-token mask is derived from the graph input encoding and is diagnostic-only in this experiment: routing, shared capacity, balancing loss, and outputs still include padding, so the run measures the current behavior without applying a padding fix. The per-block isolation changes which banks are active but leaves the routing policy otherwise unchanged. Compare the two blocks' valid/padded distributions and memory-dependency curves against the non-isolated backbone-only trace before deciding whether to test mask-aware routing. + Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 1.0` to use the cosine-anchored Track A reader. Lower scales such as `0.5` are available for follow-up comparisons. Use `--dnc-write-variant kanerva-closed-form` for Track B. Add a reader and writer variant together only for the combined cell after each track is characterized. Dense temporal links remain the default. The short-run baseline reproduction gate and the Q99 paired-order evaluation are deliberate user-run experiment steps; no run results are assumed by this repository change. ### Controllers diff --git a/src/initium/LNN_controller/cfc_backbone_parallel.py b/src/initium/LNN_controller/cfc_backbone_parallel.py index d00239b..1b39114 100644 --- a/src/initium/LNN_controller/cfc_backbone_parallel.py +++ b/src/initium/LNN_controller/cfc_backbone_parallel.py @@ -100,6 +100,7 @@ def __init__( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_block_isolation: int | None = None, device=None, dtype=None, ): @@ -134,45 +135,65 @@ def __init__( # token batch, so this costs nothing extra to support here). self.moe_enabled = moe_enabled self.moe_blocks: nn.ModuleList | None = None + if moe_block_isolation is not None and not 0 <= moe_block_isolation < num_blocks: + raise ValueError( + f"moe_block_isolation={moe_block_isolation} is outside " + f"the configured backbone block range [0, {num_blocks - 1}]" + ) if moe_enabled: self.moe_blocks = nn.ModuleList( [ - MoEBlock( - d_model, - num_experts=moe_num_experts, - expert_dim=moe_expert_dim, - top_k=moe_top_k, - capacity_factor=moe_capacity_factor, - load_balance_alpha=moe_load_balance_alpha, - device=device, - dtype=dtype, + ( + MoEBlock( + d_model, + num_experts=moe_num_experts, + expert_dim=moe_expert_dim, + top_k=moe_top_k, + capacity_factor=moe_capacity_factor, + load_balance_alpha=moe_load_balance_alpha, + device=device, + dtype=dtype, + ) + if moe_block_isolation is None or i == moe_block_isolation + else nn.Identity() ) - for _ in range(num_blocks) + for i in range(num_blocks) ] ) - def forward(self, x): # (B, L, in_dim) -> (B, L, d_model), called ONCE per training step + def forward( + self, x, *, token_mask: torch.Tensor | None = None + ): # (B, L, in_dim) -> (B, L, d_model), called ONCE per training step h = self.in_adapter(x) for i, block in enumerate(self.blocks): h = block(h) - if self.moe_enabled: + if self.moe_enabled and isinstance(self.moe_blocks[i], MoEBlock): assert self.moe_blocks is not None - h = self.moe_blocks[i](h) + h = self.moe_blocks[i](h, token_mask=token_mask) return h - def forward_with_specialists(self, x: torch.Tensor): + def forward_with_specialists( + self, x: torch.Tensor, *, token_mask: torch.Tensor | None = None + ): """Return final MoE candidates and routes for shared-workspace use.""" if not self.moe_enabled or self.moe_blocks is None: raise RuntimeError("workspace broadcast requires an enabled MoE expert pool") h = self.in_adapter(x) final_index = len(self.blocks) - 1 + if not isinstance(self.moe_blocks[final_index], MoEBlock): + raise RuntimeError( + "workspace broadcast requires the final backbone block to have an active MoE bank" + ) for i, block in enumerate(self.blocks): h = block(h) if i == final_index: base = h - _, experts, indices, weights = self.moe_blocks[i](h, return_details=True) + _, experts, indices, weights = self.moe_blocks[i]( + h, return_details=True, token_mask=token_mask + ) return base, experts, indices, weights - h = self.moe_blocks[i](h) + if isinstance(self.moe_blocks[i], MoEBlock): + h = self.moe_blocks[i](h, token_mask=token_mask) raise RuntimeError("CfC parallel backbone has no blocks") @@ -192,6 +213,7 @@ def build_parallel_backbone( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_block_isolation: int | None = None, device=None, dtype=None, ): @@ -219,6 +241,7 @@ def build_parallel_backbone( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_block_isolation=moe_block_isolation, **(cfc_kwargs or {}), ) ) @@ -241,6 +264,7 @@ def build_parallel_backbone( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_block_isolation=moe_block_isolation, ) ) cur = d_model @@ -261,7 +285,9 @@ def collect_backbone_moe_layers(backbone: nn.Module) -> list: raise TypeError( "MoE-enabled backbone stages must expose a ModuleList of moe_blocks" ) - moe_layers.extend(blocks) + moe_layers.extend( + block for block in blocks if callable(getattr(block, "pop_aux_loss", None)) + ) return moe_layers diff --git a/src/initium/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index cb10b8e..143bd4c 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -163,7 +163,13 @@ def __init__( ) self._last_accepted: torch.Tensor | None = None # (T, k), detached - def forward(self, x: torch.Tensor, *, return_details: bool = False): + def forward( + self, + x: torch.Tensor, + *, + return_details: bool = False, + token_mask: torch.Tensor | None = None, + ): orig_shape = x.shape d_model = orig_shape[-1] assert d_model == self.d_model, ( @@ -171,6 +177,15 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): ) flat = x.reshape(-1, d_model) # (num_tokens, d_model) num_tokens = flat.shape[0] + if token_mask is not None: + if tuple(token_mask.shape) != tuple(orig_shape[:-1]): + raise ValueError( + "SwitchMoE: token_mask must match the leading input dimensions " + f"{tuple(orig_shape[:-1])}, got {tuple(token_mask.shape)}" + ) + token_mask_flat = token_mask.reshape(-1).to(device=x.device, dtype=torch.bool) + else: + token_mask_flat = None # ``.float()`` after Linear is too late under autocast: CUDA AMP may # execute the router matmul in fp16 and only cast its already-rounded @@ -364,6 +379,40 @@ def forward(self, x: torch.Tensor, *, return_details: bool = False): "capacity_drop_frac": (1.0 - keep.float().mean()).detach(), "accepted_gate_mass": accepted_gates.sum(dim=-1).mean().detach(), } + if token_mask_flat is not None: + valid = token_mask_flat + padded = ~valid + per_token_entropy = -( + probs * probs.clamp_min(1e-9).log() + ).sum(dim=-1) + per_token_max_prob = probs.max(dim=-1).values + per_token_drop = 1.0 - keep.float().mean(dim=-1) + + def subset_routing_stats(subset: torch.Tensor, prefix: str) -> None: + subset_weight = subset.to(probs.dtype) + count = subset_weight.sum() + denom = count.clamp_min(1.0) + subset_load = ( + one_hot_k + * subset_weight[:, None, None] + ).sum(dim=(0, 1)) / (denom * top_k) + subset_importance = ( + probs * subset_weight[:, None] + ).sum(dim=0) / denom + entropy = (per_token_entropy * subset_weight).sum() / denom + max_prob = (per_token_max_prob * subset_weight).sum() / denom + drop_frac = (per_token_drop * subset_weight).sum() / denom + self._last_diag[f"{prefix}_token_count"] = count.detach() + self._last_diag[f"{prefix}_expert_load_frac"] = subset_load.detach() + self._last_diag[f"{prefix}_router_importance_frac"] = ( + subset_importance.detach() + ) + self._last_diag[f"{prefix}_router_entropy"] = entropy.detach() + self._last_diag[f"{prefix}_router_max_prob"] = max_prob.detach() + self._last_diag[f"{prefix}_capacity_drop_frac"] = drop_frac.detach() + + subset_routing_stats(valid, "valid") + subset_routing_stats(padded, "padded") if self.training: for key, value in self._last_diag.items(): self._diag_accum[key] = ( @@ -475,10 +524,18 @@ def __init__( if device is not None and getattr(device, "type", None) == "cuda": self.to(device) - def forward(self, x: torch.Tensor, *, return_details: bool = False): + def forward( + self, + x: torch.Tensor, + *, + return_details: bool = False, + token_mask: torch.Tensor | None = None, + ): if not return_details: - return x + self.moe(self.norm(x)) - routed, experts, indices, weights = self.moe(self.norm(x), return_details=True) + return x + self.moe(self.norm(x), token_mask=token_mask) + routed, experts, indices, weights = self.moe( + self.norm(x), return_details=True, token_mask=token_mask + ) return x + routed, experts, indices, weights def pop_aux_loss(self) -> torch.Tensor: @@ -801,12 +858,27 @@ def pop_total_moe_aux_loss(moe_layers: list, *, include_diagnostics: bool = True "routed_output_absmax", "capacity_drop_frac", "accepted_gate_mass", + "valid_token_count", + "padded_token_count", + "valid_router_entropy", + "valid_router_max_prob", + "valid_capacity_drop_frac", + "padded_router_entropy", + "padded_router_max_prob", + "padded_capacity_drop_frac", ) if key in diag }, **{ key: diag[key].detach().float().cpu().tolist() - for key in ("expert_load_frac", "router_importance_frac") + for key in ( + "expert_load_frac", + "router_importance_frac", + "valid_expert_load_frac", + "valid_router_importance_frac", + "padded_expert_load_frac", + "padded_router_importance_frac", + ) if key in diag }, **({"source_routes": source_report} if source_report is not None else {}), diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 8d6c8ae..466db6d 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -256,12 +256,38 @@ def _format_moe_layer_reports(diagnostics: dict) -> list[str]: importance = ",".join( f"{value:.2f}" for value in layer_diag.get("router_importance_frac", []) ) + valid_load = ",".join( + f"{value:.2f}" for value in layer_diag.get("valid_expert_load_frac", []) + ) + padded_load = ",".join( + f"{value:.2f}" for value in layer_diag.get("padded_expert_load_frac", []) + ) + valid_importance = ",".join( + f"{value:.2f}" for value in layer_diag.get("valid_router_importance_frac", []) + ) + padded_importance = ",".join( + f"{value:.2f}" for value in layer_diag.get("padded_router_importance_frac", []) + ) source_routes = layer_diag.get("source_routes") or [] source_summary = " ".join( f"src{i}:accepted={source.get('accepted_route_frac', 0.0):.1%}," f"drop={source.get('capacity_drop_frac', 0.0):.1%}" for i, source in enumerate(source_routes) ) + split_summary = "" + if "valid_token_count" in layer_diag: + valid_count = layer_diag.get("valid_token_count", 0.0) + padded_count = layer_diag.get("padded_token_count", 0.0) + split_summary = ( + f" valid[n={valid_count:.1f} load=[{valid_load}] imp=[{valid_importance}] " + f"H={layer_diag.get('valid_router_entropy', 0.0):.2f} " + f"pmax={layer_diag.get('valid_router_max_prob', 0.0):.2f} " + f"drop={layer_diag.get('valid_capacity_drop_frac', 0.0):.1%}]" + f" padded[n={padded_count:.1f} load=[{padded_load}] imp=[{padded_importance}] " + f"H={layer_diag.get('padded_router_entropy', 0.0):.2f} " + f"pmax={layer_diag.get('padded_router_max_prob', 0.0):.2f} " + f"drop={layer_diag.get('padded_capacity_drop_frac', 0.0):.1%}]" + ) reports.append( f"moe-bank[{layer_diag['name']} " f"load=[{load}] importance=[{importance}] " @@ -279,6 +305,7 @@ def _format_moe_layer_reports(diagnostics: dict) -> list[str]: f"{layer_diag.get('routed_output_absmax', 0.0):.3g} " f"drop={layer_diag.get('capacity_drop_frac', 0.0):.1%} " f"gate_mass={layer_diag.get('accepted_gate_mass', 0.0):.2f}" + f"{split_summary}" f"{(' ' + source_summary) if source_summary else ''}]" ) return reports @@ -647,6 +674,7 @@ def run( dynamic_n_cooldown_steps: int = DYNAMIC_N_COOLDOWN_STEPS, moe_enabled: bool = MOE_ENABLED, split_graph_moe_placement: str = "all", + split_graph_moe_backbone_block_isolation: int | None = None, moe_num_experts: int = MOE_NUM_EXPERTS, moe_expert_dim: int | None = MOE_EXPERT_DIM, moe_capacity_factor: float = MOE_CAPACITY_FACTOR, @@ -1187,6 +1215,17 @@ def run( f"{split_graph_moe_placement in ('all', 'combiner') and split_graph_combiner_cfc_multi_source_moe} " f"(see moe_layer.MultiSourceMoEBlock)." ) + if split_graph_moe_backbone_block_isolation is not None: + routing_note = ( + "Graph routing diagnostics split selected routes into valid and zero-padded timesteps." + if dataset_type == "graph" + else "Token-subset routing diagnostics are unavailable for this dataset." + ) + print( + f"[{run_id}] backbone MoE block isolation={split_graph_moe_backbone_block_isolation} " + f"(zero-based per backbone stage; all other backbone blocks run without MoE). " + f"{routing_note}" + ) # v10: variant-matched hyperparameter defaults instead of always # reusing the Mamba-1 constants here -- harmless previously (any @@ -1227,6 +1266,7 @@ def run( independent_linears=True, moe_enabled=moe_enabled, moe_placement=split_graph_moe_placement, + moe_backbone_block_isolation=split_graph_moe_backbone_block_isolation, moe_num_experts=moe_num_experts, moe_expert_dim=moe_expert_dim, moe_top_k=moe_top_k, @@ -1322,6 +1362,7 @@ def run( ), "split_graph_combiner_cfc_multi_source_moe": split_graph_combiner_cfc_multi_source_moe, "split_graph_moe_placement": split_graph_moe_placement, + "split_graph_moe_backbone_block_isolation": split_graph_moe_backbone_block_isolation, "dnc_read_variant": dnc_read_variant, "dnc_read_residual_scale": dnc_read_residual_scale, "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, @@ -1671,7 +1712,23 @@ def set_lr(step_now): segment_diagnostics.set_channel("train") with torch.amp.autocast("cuda", enabled=amp_enabled): - output, hidden = rnn(input_seq, hidden, reset_experience=True) + if split_graph_enabled: + # Graph episodes are variable length and collate_fn pads them + # with all-zero feature rows. This mask is diagnostic-only for + # MoE routing: routing/capacity behavior remains unchanged. + sequence_valid_mask = ( + input_seq.ne(0).any(dim=-1) + if moe_enabled and dataset_type == "graph" + else None + ) + output, hidden = rnn( + input_seq, + hidden, + reset_experience=True, + sequence_valid_mask=sequence_valid_mask, + ) + else: + output, hidden = rnn(input_seq, hidden, reset_experience=True) if check_finite_this_step and not _first_nonfinite_report( output, "rnn_output", step ): @@ -2924,6 +2981,16 @@ def _write_diagnostic_values(diagnostics): help="For split-graph MoE, enable experts in both components (all), " "only the parallel backbone, or only the sequential combiner. Default: all.", ) + parser.add_argument( + "--moe-backbone-block-isolation", + type=int, + default=None, + metavar="INDEX", + help="For split-graph backbone MoE, install a bank only after the selected " + "zero-based backbone block; other blocks run without MoE. Use one run per " + "block index in every backbone stage to isolate routing (requires --moe and " + "--moe-placement backbone).", + ) parser.add_argument( "--workspace", action="store_true", help="Phase 0 shared workspace: competitively write selected MoE expert outputs to persistent slots, then broadcast slots back to every specialist. Requires --split-graph and --moe.", @@ -3200,6 +3267,21 @@ def _write_diagnostic_values(diagnostics): parser.error("--moe-placement backbone/combiner requires --split-graph") if args.moe_placement == "combiner" and args.split_graph_combiner_mode != "controller": parser.error("--moe-placement combiner requires --split-graph-combiner-mode controller") + if args.moe_backbone_block_isolation is not None: + if not args.moe or not args.split_graph: + parser.error("--moe-backbone-block-isolation requires --moe and --split-graph") + if args.moe_placement != "backbone": + parser.error( + "--moe-backbone-block-isolation requires --moe-placement backbone " + "to keep combiner MoE out of the isolation run" + ) + if args.workspace: + parser.error("--moe-backbone-block-isolation cannot be combined with --workspace") + if not 0 <= args.moe_backbone_block_isolation < args.split_graph_num_blocks: + parser.error( + "--moe-backbone-block-isolation must be in [0, " + f"{args.split_graph_num_blocks - 1}] for the configured number of blocks" + ) if args.workspace and args.moe_placement == "combiner": parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: @@ -3245,6 +3327,8 @@ def _write_diagnostic_values(diagnostics): run_id = f"{run_id}_moe{args.moe_num_experts}e" if args.split_graph and args.moe_placement != "all": run_id += f"_{args.moe_placement}only" + if args.moe_backbone_block_isolation is not None: + run_id += f"_backboneblock{args.moe_backbone_block_isolation}only" if args.workspace: run_id += f"_workspace{args.workspace_num_slots}s_{args.workspace_train_mode}" if args.workspace_no_selection: @@ -3287,6 +3371,7 @@ def _write_diagnostic_values(diagnostics): dynamic_n_cooldown_steps=args.dynamic_n_cooldown_steps, moe_enabled=args.moe, split_graph_moe_placement=args.moe_placement, + split_graph_moe_backbone_block_isolation=args.moe_backbone_block_isolation, moe_num_experts=args.moe_num_experts, moe_expert_dim=args.moe_expert_dim, moe_capacity_factor=args.moe_capacity_factor, diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index a7f94e2..3defe11 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -189,6 +189,10 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea ), independent_linears=True, moe_enabled=_moe["moe_enabled"], + moe_placement=c.get("split_graph_moe_placement", "all"), + moe_backbone_block_isolation=c.get( + "split_graph_moe_backbone_block_isolation", None + ), moe_num_experts=_moe["moe_num_experts"], moe_expert_dim=_moe["moe_expert_dim"], moe_top_k=_moe["moe_top_k"], diff --git a/src/initium/mamba_controller/mamba_backbone_parallel.py b/src/initium/mamba_controller/mamba_backbone_parallel.py index f9455d5..e15bcab 100644 --- a/src/initium/mamba_controller/mamba_backbone_parallel.py +++ b/src/initium/mamba_controller/mamba_backbone_parallel.py @@ -224,6 +224,7 @@ def __init__( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_block_isolation: int | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, ): @@ -233,6 +234,11 @@ def __init__( self.d_model = d_model self.moe_enabled = moe_enabled self.moe_blocks: nn.ModuleList | None = None + if moe_block_isolation is not None and not 0 <= moe_block_isolation < num_blocks: + raise ValueError( + f"moe_block_isolation={moe_block_isolation} is outside " + f"the configured backbone block range [0, {num_blocks - 1}]" + ) # Same role as MambaControllerWrapper's in_adapter: DNC's raw # per-timestep input (input_size) generally differs from @@ -263,21 +269,27 @@ def __init__( if moe_enabled: self.moe_blocks = nn.ModuleList( [ - MoEBlock( - d_model, - num_experts=moe_num_experts, - expert_dim=moe_expert_dim, - top_k=moe_top_k, - capacity_factor=moe_capacity_factor, - load_balance_alpha=moe_load_balance_alpha, - device=device, - dtype=dtype, + ( + MoEBlock( + d_model, + num_experts=moe_num_experts, + expert_dim=moe_expert_dim, + top_k=moe_top_k, + capacity_factor=moe_capacity_factor, + load_balance_alpha=moe_load_balance_alpha, + device=device, + dtype=dtype, + ) + if moe_block_isolation is None or i == moe_block_isolation + else nn.Identity() ) - for _ in range(num_blocks) + for i in range(num_blocks) ] ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, x: torch.Tensor, *, token_mask: torch.Tensor | None = None + ) -> torch.Tensor: # x: (B, L, in_dim) -> (B, L, d_model). Called EXACTLY ONCE per # training step (not once per DNC timestep, unlike # MambaControllerWrapper) -- this single call is the entire @@ -285,12 +297,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.in_adapter(x) for i, block in enumerate(self.blocks): h = block(h) - if self.moe_enabled: + if self.moe_enabled and isinstance(self.moe_blocks[i], MoEBlock): assert self.moe_blocks is not None - h = self.moe_blocks[i](h) + h = self.moe_blocks[i](h, token_mask=token_mask) return h - def forward_with_specialists(self, x: torch.Tensor): + def forward_with_specialists( + self, x: torch.Tensor, *, token_mask: torch.Tensor | None = None + ): """Return final MoE candidates and sparse routes for workspace use. Earlier backbone/MoE blocks remain unchanged. The final MoE block @@ -301,11 +315,18 @@ def forward_with_specialists(self, x: torch.Tensor): raise RuntimeError("workspace broadcast requires an enabled MoE expert pool") h = self.in_adapter(x) final_index = len(self.blocks) - 1 + if not isinstance(self.moe_blocks[final_index], MoEBlock): + raise RuntimeError( + "workspace broadcast requires the final backbone block to have an active MoE bank" + ) for i, block in enumerate(self.blocks): h = block(h) if i == final_index: base = h - _, experts, indices, weights = self.moe_blocks[i](h, return_details=True) + _, experts, indices, weights = self.moe_blocks[i]( + h, return_details=True, token_mask=token_mask + ) return base, experts, indices, weights - h = self.moe_blocks[i](h) + if isinstance(self.moe_blocks[i], MoEBlock): + h = self.moe_blocks[i](h, token_mask=token_mask) raise RuntimeError("parallel backbone has no blocks") diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index e9a5b1c..108cae1 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -182,6 +182,7 @@ def __init__( device: torch.device | None = None, moe_enabled: bool = False, moe_placement: str = "all", + moe_backbone_block_isolation: int | None = None, moe_num_experts: int = 8, moe_expert_dim: int | None = None, moe_top_k: int = 1, @@ -239,6 +240,11 @@ def __init__( raise ValueError("moe_placement must be 'all', 'backbone', or 'combiner'") backbone_moe_enabled = moe_enabled and moe_placement in ("all", "backbone") combiner_moe_enabled = moe_enabled and moe_placement in ("all", "combiner") + if moe_backbone_block_isolation is not None and not backbone_moe_enabled: + raise ValueError( + "moe_backbone_block_isolation requires active backbone MoE " + "(moe_enabled=True and moe_placement='backbone')" + ) if workspace_enabled and not backbone_moe_enabled: raise ValueError( "--workspace requires backbone MoE because it broadcasts to backbone specialists" @@ -276,6 +282,7 @@ def __init__( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_block_isolation=moe_backbone_block_isolation, device=device, ) @@ -537,6 +544,7 @@ def forward( pass_through_memory: bool = True, combiner_skip_stages=None, start_step: int = 0, # v20: resume the sequential loop mid-sequence (inference caching) + sequence_valid_mask: torch.Tensor | None = None, ): """ input: (B, T, input_size) -- batch-first, the whole padded episode @@ -554,6 +562,9 @@ def forward( the functional-usage ablation check (evaluate_traversal(..., ablate_memory=True)) already wired into this project's training loop; see module docstring. + sequence_valid_mask: optional (B, T) mask forwarded to backbone MoE + banks for valid-versus-padding routing diagnostics. It does not + mask backbone computation, MoE dispatch, or memory execution. Returns (output, (chx, mhx, last_read)) with output shaped (T, B, input_size) -- see module docstring for why this shape, @@ -597,14 +608,16 @@ def forward( ) stage_input = input for stage in backbone_stages[:-1]: - stage_input = stage(stage_input) + stage_input = stage(stage_input, token_mask=sequence_valid_mask) final_stage = backbone_stages[-1] if not hasattr(final_stage, "forward_with_specialists"): raise RuntimeError( "workspace requires the final parallel-backbone stage to expose its MoE specialists" ) base, specialists, route_indices, route_weights = ( - final_stage.forward_with_specialists(stage_input) + final_stage.forward_with_specialists( + stage_input, token_mask=sequence_valid_mask + ) ) broadcast_delta, workspace_state = self.workspace( specialists[:, start_step:], @@ -617,7 +630,14 @@ def forward( H = base.clone() H[:, start_step:] = base[:, start_step:] + broadcast_delta else: - H = self.backbone(input) # (B, T, hidden_size) + backbone_stages = ( + list(self.backbone) + if isinstance(self.backbone, nn.Sequential) + else [self.backbone] + ) + H = input + for stage in backbone_stages: + H = stage(H, token_mask=sequence_valid_mask) # ---- memory hidden-state init (byte-identical convention to # dnc.DNC._init_hidden / MambaDNC._init_hidden's memory branch) --- From 00b2bac96ecb5b5970e181ce3d392f4b701ae633 Mon Sep 17 00:00:00 2001 From: Samuel Date: Sat, 10 Oct 2026 15:05:34 +0200 Subject: [PATCH 04/14] Add backbone MoE residual scaling control This change adds a configurable backbone MoE residual scale for split-graph runs, defaulting to 1.0 to preserve existing behavior. The scale is validated as finite and non-negative, only allowed with backbone/all placement, rejected with the workspace broadcast path, and included in run IDs and checkpoint metadata. The residual is applied before the backbone residual addition as x + scale * moe_delta, while leaving routing and auxiliary load-balancing behavior active. The new option is wired through the MoE layer, Mamba and CfC parallel backbone builders, and inference config loading. The README also documents the experimental usage and follow-up commands for conservative backbone-only scale testing. --- README.md | 12 +++++++ .../LNN_controller/cfc_backbone_parallel.py | 5 +++ src/initium/MoE/moe_layer.py | 13 +++++-- src/initium/core_training.py | 35 +++++++++++++++++++ src/initium/inference/model_loader.py | 3 ++ .../mamba_backbone_parallel.py | 2 ++ .../mamba_controller/split_graph_dnc.py | 8 +++++ 7 files changed, 76 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a7f28fd..90b8480 100644 --- a/README.md +++ b/README.md @@ -154,6 +154,18 @@ python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-back python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-backbone-block1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 ``` +`--moe-backbone-residual-scale SCALE` multiplies the active parallel-backbone MoE delta before its residual addition (`backbone_output + SCALE * moe_delta`). Its default is `1.0`, preserving prior behavior. The scale must be finite and non-negative, requires split-graph MoE with `--moe-placement backbone` or `all`, is included in run IDs/checkpoint metadata, and is not supported with the workspace broadcast path. With `all`, this changes only the parallel-backbone MoE residuals; combiner MoE residuals remain at 1.0. It scales the task-loss path through the backbone MoE branch while leaving routing and the load-balancing auxiliary loss active. In an isolated-block run, only that block has an MoE residual to scale. For the current block-1 follow-up, use a deliberately conservative 0.1 scale: + +```bash +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-block1-residual-scale-0p1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 1 --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +``` + +To test backbone and combiner MoE together after the backbone-only scale control, omit block isolation and select `--moe-placement all`. The residual scale still applies only to the backbone; the combiner MoE keeps its default residual scale of 1.0: + +```bash +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-backbone-combiner-backbone-scale-0p1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 +``` + Each `moe-bank[backbone.moe_blocks.INDEX ...]` line reports router-selected expert load, router importance, entropy, maximum probability, and capacity-drop rate separately for valid graph timesteps and all-zero padded timesteps. The valid-token mask is derived from the graph input encoding and is diagnostic-only in this experiment: routing, shared capacity, balancing loss, and outputs still include padding, so the run measures the current behavior without applying a padding fix. The per-block isolation changes which banks are active but leaves the routing policy otherwise unchanged. Compare the two blocks' valid/padded distributions and memory-dependency curves against the non-isolated backbone-only trace before deciding whether to test mask-aware routing. Add `--dnc-read-variant relational-mlp` to reproduce the original Track A reader, or `--dnc-read-variant relational-residual --dnc-read-residual-scale 1.0` to use the cosine-anchored Track A reader. Lower scales such as `0.5` are available for follow-up comparisons. Use `--dnc-write-variant kanerva-closed-form` for Track B. Add a reader and writer variant together only for the combined cell after each track is characterized. Dense temporal links remain the default. The short-run baseline reproduction gate and the Q99 paired-order evaluation are deliberate user-run experiment steps; no run results are assumed by this repository change. diff --git a/src/initium/LNN_controller/cfc_backbone_parallel.py b/src/initium/LNN_controller/cfc_backbone_parallel.py index 1b39114..24b9297 100644 --- a/src/initium/LNN_controller/cfc_backbone_parallel.py +++ b/src/initium/LNN_controller/cfc_backbone_parallel.py @@ -100,6 +100,7 @@ def __init__( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_residual_scale: float = 1.0, moe_block_isolation: int | None = None, device=None, dtype=None, @@ -151,6 +152,7 @@ def __init__( top_k=moe_top_k, capacity_factor=moe_capacity_factor, load_balance_alpha=moe_load_balance_alpha, + residual_scale=moe_residual_scale, device=device, dtype=dtype, ) @@ -213,6 +215,7 @@ def build_parallel_backbone( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_residual_scale: float = 1.0, moe_block_isolation: int | None = None, device=None, dtype=None, @@ -241,6 +244,7 @@ def build_parallel_backbone( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_residual_scale=moe_residual_scale, moe_block_isolation=moe_block_isolation, **(cfc_kwargs or {}), ) @@ -264,6 +268,7 @@ def build_parallel_backbone( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_residual_scale=moe_residual_scale, moe_block_isolation=moe_block_isolation, ) ) diff --git a/src/initium/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index 143bd4c..2cfa205 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -81,6 +81,8 @@ from __future__ import annotations +import math + import torch import torch.nn as nn import torch.nn.functional as F @@ -504,10 +506,16 @@ def __init__( router_noise_eps: float = 1e-2, load_balance_alpha: float = 0.01, top_k: int = 1, + residual_scale: float = 1.0, device: torch.device | None = None, dtype: torch.dtype | None = None, ): super().__init__() + if not math.isfinite(residual_scale) or residual_scale < 0.0: + raise ValueError( + f"MoEBlock residual_scale must be finite and >= 0, got {residual_scale!r}" + ) + self.residual_scale = float(residual_scale) self.norm = nn.LayerNorm(d_model, device=device, dtype=dtype) self.moe = SwitchMoE( d_model, @@ -532,11 +540,12 @@ def forward( token_mask: torch.Tensor | None = None, ): if not return_details: - return x + self.moe(self.norm(x), token_mask=token_mask) + routed = self.moe(self.norm(x), token_mask=token_mask) + return x + self.residual_scale * routed routed, experts, indices, weights = self.moe( self.norm(x), return_details=True, token_mask=token_mask ) - return x + routed, experts, indices, weights + return x + self.residual_scale * routed, experts, indices, weights def pop_aux_loss(self) -> torch.Tensor: return self.moe.pop_aux_loss() diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 466db6d..397ac8d 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -675,6 +675,7 @@ def run( moe_enabled: bool = MOE_ENABLED, split_graph_moe_placement: str = "all", split_graph_moe_backbone_block_isolation: int | None = None, + split_graph_moe_backbone_residual_scale: float = 1.0, moe_num_experts: int = MOE_NUM_EXPERTS, moe_expert_dim: int | None = MOE_EXPERT_DIM, moe_capacity_factor: float = MOE_CAPACITY_FACTOR, @@ -1226,6 +1227,13 @@ def run( f"(zero-based per backbone stage; all other backbone blocks run without MoE). " f"{routing_note}" ) + if split_graph_moe_backbone_residual_scale != 1.0: + print( + f"[{run_id}] backbone MoE residual scale=" + f"{split_graph_moe_backbone_residual_scale:g} | " + "MoE delta is scaled before its residual addition; router, " + "experts, and auxiliary routing loss remain active." + ) # v10: variant-matched hyperparameter defaults instead of always # reusing the Mamba-1 constants here -- harmless previously (any @@ -1272,6 +1280,7 @@ def run( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_backbone_residual_scale=split_graph_moe_backbone_residual_scale, moe_cfc_multi_source=split_graph_combiner_cfc_multi_source_moe, dnc_read_variant=dnc_read_variant, dnc_read_residual_scale=dnc_read_residual_scale, @@ -1363,6 +1372,7 @@ def run( "split_graph_combiner_cfc_multi_source_moe": split_graph_combiner_cfc_multi_source_moe, "split_graph_moe_placement": split_graph_moe_placement, "split_graph_moe_backbone_block_isolation": split_graph_moe_backbone_block_isolation, + "split_graph_moe_backbone_residual_scale": split_graph_moe_backbone_residual_scale, "dnc_read_variant": dnc_read_variant, "dnc_read_residual_scale": dnc_read_residual_scale, "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, @@ -2991,6 +3001,17 @@ def _write_diagnostic_values(diagnostics): "block index in every backbone stage to isolate routing (requires --moe and " "--moe-placement backbone).", ) + parser.add_argument( + "--moe-backbone-residual-scale", + type=float, + default=1.0, + metavar="SCALE", + help="Multiply each active split-graph backbone MoE delta by SCALE " + "before adding it to the backbone residual (default: 1.0, unchanged). " + "Use with --moe-placement backbone or all; with all, combiner MoE " + "residuals remain unscaled. In an isolation run, only the selected " + "block's MoE residual is scaled.", + ) parser.add_argument( "--workspace", action="store_true", help="Phase 0 shared workspace: competitively write selected MoE expert outputs to persistent slots, then broadcast slots back to every specialist. Requires --split-graph and --moe.", @@ -3282,6 +3303,17 @@ def _write_diagnostic_values(diagnostics): "--moe-backbone-block-isolation must be in [0, " f"{args.split_graph_num_blocks - 1}] for the configured number of blocks" ) + if not math.isfinite(args.moe_backbone_residual_scale) or args.moe_backbone_residual_scale < 0.0: + parser.error("--moe-backbone-residual-scale must be finite and >= 0") + if args.moe_backbone_residual_scale != 1.0: + if not args.moe or not args.split_graph: + parser.error("--moe-backbone-residual-scale requires --moe and --split-graph") + if args.moe_placement not in ("backbone", "all"): + parser.error( + "--moe-backbone-residual-scale requires --moe-placement backbone or all" + ) + if args.workspace: + parser.error("--moe-backbone-residual-scale cannot be combined with --workspace") if args.workspace and args.moe_placement == "combiner": parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: @@ -3329,6 +3361,8 @@ def _write_diagnostic_values(diagnostics): run_id += f"_{args.moe_placement}only" if args.moe_backbone_block_isolation is not None: run_id += f"_backboneblock{args.moe_backbone_block_isolation}only" + if args.moe_backbone_residual_scale != 1.0: + run_id += f"_backbonerscale{args.moe_backbone_residual_scale:g}" if args.workspace: run_id += f"_workspace{args.workspace_num_slots}s_{args.workspace_train_mode}" if args.workspace_no_selection: @@ -3372,6 +3406,7 @@ def _write_diagnostic_values(diagnostics): moe_enabled=args.moe, split_graph_moe_placement=args.moe_placement, split_graph_moe_backbone_block_isolation=args.moe_backbone_block_isolation, + split_graph_moe_backbone_residual_scale=args.moe_backbone_residual_scale, moe_num_experts=args.moe_num_experts, moe_expert_dim=args.moe_expert_dim, moe_capacity_factor=args.moe_capacity_factor, diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 3defe11..4276fff 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -193,6 +193,9 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea moe_backbone_block_isolation=c.get( "split_graph_moe_backbone_block_isolation", None ), + moe_backbone_residual_scale=c.get( + "split_graph_moe_backbone_residual_scale", 1.0 + ), moe_num_experts=_moe["moe_num_experts"], moe_expert_dim=_moe["moe_expert_dim"], moe_top_k=_moe["moe_top_k"], diff --git a/src/initium/mamba_controller/mamba_backbone_parallel.py b/src/initium/mamba_controller/mamba_backbone_parallel.py index e15bcab..f885801 100644 --- a/src/initium/mamba_controller/mamba_backbone_parallel.py +++ b/src/initium/mamba_controller/mamba_backbone_parallel.py @@ -224,6 +224,7 @@ def __init__( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_residual_scale: float = 1.0, moe_block_isolation: int | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, @@ -277,6 +278,7 @@ def __init__( top_k=moe_top_k, capacity_factor=moe_capacity_factor, load_balance_alpha=moe_load_balance_alpha, + residual_scale=moe_residual_scale, device=device, dtype=dtype, ) diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 108cae1..3474472 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -188,6 +188,7 @@ def __init__( moe_top_k: int = 1, moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, + moe_backbone_residual_scale: float = 1.0, # When combiner_variant == "cfc" (standalone, not part of a hybrid # combiner chain) and moe_enabled, route the combiner's two natural # inputs -- backbone output h_t and the previous read vector -- @@ -236,6 +237,7 @@ def __init__( self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio self.dnc_write_variant = dnc_write_variant self.dnc_write_observation_variance = dnc_write_observation_variance + self.moe_backbone_residual_scale = moe_backbone_residual_scale if moe_placement not in ("all", "backbone", "combiner"): raise ValueError("moe_placement must be 'all', 'backbone', or 'combiner'") backbone_moe_enabled = moe_enabled and moe_placement in ("all", "backbone") @@ -249,6 +251,11 @@ def __init__( raise ValueError( "--workspace requires backbone MoE because it broadcasts to backbone specialists" ) + if workspace_enabled and moe_backbone_residual_scale != 1.0: + raise ValueError( + "workspace specialist broadcasts do not use the backbone MoE residual; " + "moe_backbone_residual_scale must remain 1.0 with workspace_enabled" + ) if workspace_train_mode not in ("coadapted", "frozen"): raise ValueError("workspace_train_mode must be 'coadapted' or 'frozen'") if workspace_enabled and workspace_train_mode == "frozen" and workspace_num_slots == 0: @@ -282,6 +289,7 @@ def __init__( moe_top_k=moe_top_k, moe_capacity_factor=moe_capacity_factor, moe_load_balance_alpha=moe_load_balance_alpha, + moe_residual_scale=moe_backbone_residual_scale, moe_block_isolation=moe_backbone_block_isolation, device=device, ) From c22613cf65e8f169b3a88021362bed547108046a Mon Sep 17 00:00:00 2001 From: Samuel Date: Sat, 10 Oct 2026 19:15:12 +0200 Subject: [PATCH 05/14] Enable fp32 DNC tracing and diagnostics This commit forces DNC interface, addressing, and recurrent memory math onto fp32 even under CUDA autocast, and records the precision choice in startup logs/checkpoint metadata. It adds --dnc-backward-trace for split-graph runs to watch DNC read/write addressing, temporal links, usage/allocation, prod/cumprod numerics, and the first non-finite backward boundary with batch-row context; collected diagnostics are emitted during training and failure windows and include product-input stats, read-weight mass checks, and memory/read/link magnitudes. --- README.md | 10 + src/initium/core_training.py | 43 +++ .../memory_manipulation/dnc_mechanisms.py | 287 +++++++++++++++++- .../memory_manipulation/nvrtc_compat.py | 67 +++- 4 files changed, 402 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 90b8480..f30b547 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,16 @@ The supplied Phase 1/2 traces are pre-fix evidence. One Phase 2 trace first repo Training console logs now show **one MoE-bank record per 100-step logging window**: route and router-importance fractions by expert, router entropy and maximum probability, input/expert/MoE-delta activation RMS and maximum magnitude, capacity-drop fraction, and accepted gate mass. The `moe_residual:` segment records the complete MoE block output after its skip connection, separately from `moe_delta_rms/max`. A separate `[segment-diag:train]` line reports each recognized module's call count, input/output RMS and max, non-finite fraction, boundary input/output gradient RMS and non-finite fraction, and hierarchical parameter-gradient norms. Boundary gradients are collected during training backward and during failure traces; evaluation remains forward-only. The collector recognizes Mamba and CfC backbones/blocks, Mamba/CfC controller wrappers/blocks/cells, MoE blocks, chained controllers, DNC memory, workspace, and task output; names include their module paths, so the same logger works across supported controller combinations. +Split-graph DNC interface transforms, cosine content addressing, usage/allocation, write/read weighting, temporal links, and recurrent memory updates now run in fp32 even when the outer controller uses CUDA autocast. The stock cosine rule divides by the memory/key norm product plus `1e-6`; the instrumented trace showed finite score-side gradients but a non-finite key-side gradient, so fp32 protects that path from low-precision intermediate-gradient overflow. The startup log and checkpoint metadata identify this precision choice. The equations remain unchanged. + +To inspect DNC internals, add `--dnc-backward-trace` to a split-graph training command. The flag watches gradients at read/write interface values, cosine-addressing scores/logits/softmax, usage, allocation, write weights, temporal links, precedence, recurrent memory, and read vectors, plus the patched CUDA `torch.prod`/`torch.cumprod` outputs and their inputs. It reports every distinct hooked boundary with a non-finite gradient on the first affected optimizer step, rather than stopping after the first hook label, and includes affected batch rows. At the regular logging interval and on failure, `[dnc-numerics:...]` reports memory/read/link magnitudes, read-weight mass violations, and per-product input min/max, fraction below `1e-6`, negative/non-finite fraction, and call count. The flag synchronizes at watched backward boundaries, so use a short diagnostic capture around the known first event rather than a full long run. For the attached trace, which first reports a bad memory input gradient at step 1,120, 1,800 steps covers the onset: + +```bash +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix dnc-backward-trace --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 1800 --dnc-backward-trace +``` + +Keep other flags exactly matched to the run being investigated. The fp32 DNC path is the targeted numerical correction; repeat the short capture to check whether it removes the non-finite key gradients and restores paired memory dependence. If failures remain, the expanded trace still reports all bad stages on their first affected step. + Evaluation emits separate channel lines for ID, OOD, memory-off, combiner-stage ablations, and final ID/OOD checks. These summarize the evaluations independently from training windows. If a forward value or gradient becomes non-finite, the current segment window and MoE-bank window are printed immediately before training stops or skips the unsafe update. The summaries aggregate detached device-side scalars and synchronize at log/evaluation boundaries; they do not retain activation graphs. To isolate where split-graph MoE changes the memory-learning transition, `--moe-placement` selects which component receives expert banks while keeping the model, seed, and router settings matched. `all` (the default) preserves the existing behavior; `backbone` enables MoE only in the parallel backbone; `combiner` enables MoE only in the sequential controller combiner. For the CfC combiner, the latter is its source-aware multi-source MoE bank. Component placement requires `--moe --split-graph`; combiner-only additionally requires controller combiner mode. The workspace requires backbone MoE and therefore cannot be combined with combiner-only placement. diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 397ac8d..d659720 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -700,6 +700,7 @@ def run( dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, dnc_write_variant: str = DNC_WRITE_VARIANT, dnc_write_observation_variance: float = DNC_WRITE_OBSERVATION_VARIANCE, + dnc_backward_trace: bool = False, dataset_type: str = DATASET_TYPE, dataset_link: str | None = DATASET_LINK, test_dataset_link: str | None = None, @@ -1294,6 +1295,17 @@ def run( workspace_train_mode=workspace_train_mode, device=device, ).to(device) + if dnc_backward_trace: + rnn.memories[0].set_backward_trace(True) + print( + f"[{run_id}] DNC backward trace enabled | watches usage, allocation, " + "patched prod/cumprod, write/read addressing, temporal links, " + "precedence, memory state and read vectors | forward equations unchanged" + ) + print( + f"[{run_id}] DNC compute precision=fp32 | memory interface, addressing, " + "and recurrent state updates run outside CUDA autocast" + ) else: rnn = MambaDNC( input_size=INPUT_DIM, @@ -1378,6 +1390,8 @@ def run( "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, "dnc_write_variant": dnc_write_variant, "dnc_write_observation_variance": dnc_write_observation_variance, + "dnc_compute_precision": "fp32", + "dnc_backward_trace": dnc_backward_trace, "workspace_enabled": workspace_enabled, "workspace_num_slots": workspace_num_slots, "workspace_no_selection": workspace_no_selection, @@ -1404,6 +1418,13 @@ def run( def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: segment_diagnostics.emit(channel, include_gradients=include_gradients) + if dnc_backward_trace and split_graph_enabled: + memory_diag = rnn.memories[0].pop_numerical_diagnostics() + if memory_diag: + print( + f"[dnc-numerics:{channel}] " + + " ".join(f"{name}={value:.4g}" for name, value in memory_diag.items()) + ) if moe_enabled: _, failure_moe_diag = pop_total_moe_aux_loss( rnn.moe_layers, include_diagnostics=True @@ -1721,6 +1742,8 @@ def set_lr(step_now): optimizer.zero_grad(set_to_none=True) segment_diagnostics.set_channel("train") + if dnc_backward_trace and split_graph_enabled: + rnn.memories[0].set_backward_trace(True, step=step) with torch.amp.autocast("cuda", enabled=amp_enabled): if split_graph_enabled: # Graph episodes are variable length and collate_fn pads them @@ -2007,6 +2030,15 @@ def set_lr(step_now): parts.append(dyn_n_str.strip(" |")) print(" | ".join(parts)) segment_diagnostics.emit("train", include_gradients=True) + if dnc_backward_trace and split_graph_enabled: + memory_diag = rnn.memories[0].pop_numerical_diagnostics() + if memory_diag: + print( + f"[dnc-numerics:train] step={step} " + + " ".join( + f"{name}={value:.4g}" for name, value in memory_diag.items() + ) + ) if ( track_b_writer is not None @@ -3170,6 +3202,12 @@ def _write_diagnostic_values(diagnostics): default=DNC_WRITE_OBSERVATION_VARIANCE, help="Fixed positive observation variance for the Kanerva-style write.", ) + parser.add_argument( + "--dnc-backward-trace", + action="store_true", + help="Trace the first non-finite gradient through DNC internals and report " + "windowed memory/read/link magnitudes. Diagnostic only; adds synchronization.", + ) parser.add_argument( "--dataset-type", type=str, @@ -3237,6 +3275,8 @@ def _write_diagnostic_values(diagnostics): parser.error("--prepare-only cannot be combined with --load-prepared") if args.dnc_write_observation_variance <= 0: parser.error("--dnc-write-observation-variance must be positive") + if args.dnc_backward_trace and not args.split_graph: + parser.error("--dnc-backward-trace requires --split-graph") if args.dnc_read_residual_scale < 0: parser.error("--dnc-read-residual-scale must be non-negative") if args.dnc_read_residual_max_ratio < 0: @@ -3371,6 +3411,8 @@ def _write_diagnostic_values(diagnostics): run_id += "_resetstep" if args.split_graph: run_id = f"{run_id}_splitgraph_{args.split_graph_variant.replace('+', '')}" + if args.dnc_backward_trace: + run_id += "_dncbacktrace" if args.dnc_read_variant != "cosine" or args.dnc_write_variant != "learned": run_id += f"_read-{args.dnc_read_variant}_write-{args.dnc_write_variant}" if args.dnc_read_variant == "relational-residual": @@ -3430,6 +3472,7 @@ def _write_diagnostic_values(diagnostics): dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, dnc_write_variant=args.dnc_write_variant, dnc_write_observation_variance=args.dnc_write_observation_variance, + dnc_backward_trace=args.dnc_backward_trace, dataset_type=args.dataset_type, dataset_link=args.dataset_link, test_dataset_link=args.test_dataset_link, diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py index 5a7571a..b8fd456 100644 --- a/src/initium/memory_manipulation/dnc_mechanisms.py +++ b/src/initium/memory_manipulation/dnc_mechanisms.py @@ -12,13 +12,23 @@ import torch.nn as nn import torch.nn.functional as F from dnc.memory import Memory +from dnc.util import σ, θ +from initium.memory_manipulation.nvrtc_compat import ( + pop_product_numerical_diagnostics, + set_product_backward_trace, +) READ_VARIANTS = ("cosine", "relational-mlp", "relational-residual") WRITE_VARIANTS = ("learned", "kanerva-closed-form") class MechanismMemory(Memory): - """Stock DNC memory with independently selectable read and write rules.""" + """Stock DNC memory with independently selectable read and write rules. + + The memory interface and state are evaluated in fp32 even when the outer + controller uses CUDA autocast. This protects cosine-addressing backward + from low-precision intermediate gradients near the DNC norm epsilon. + """ def __init__( self, @@ -52,6 +62,18 @@ def __init__( self._stage_diagnostic_counts: dict[str, int] = {} self._stage_diagnostic_min: dict[str, torch.Tensor] = {} self._stage_diagnostic_max: dict[str, torch.Tensor] = {} + self._backward_trace_enabled = False + self._backward_trace_step: int | None = None + self._backward_trace_first_bad_step: int | None = None + self._backward_trace_reported_stages: set[str] = set() + self._activation_peaks: dict[str, torch.Tensor] = {} + self._read_weight_violation_count: torch.Tensor | None = None + self._read_weight_raw_mass_max: torch.Tensor | None = None + self._read_weight_mass_max: torch.Tensor | None = None + self._memory_value_max: torch.Tensor | None = None + self._read_vector_max: torch.Tensor | None = None + self._link_row_mass_max: torch.Tensor | None = None + self._link_column_mass_max: torch.Tensor | None = None self.relational_score: nn.Module | None = None if read_variant in {"relational-mlp", "relational-residual"}: hidden = relational_hidden_size or max(32, self.cell_size) @@ -73,6 +95,24 @@ def __init__( if self.device is not None: self.relational_score.to(self.device) + def forward(self, xi: torch.Tensor, hidden: dict[str, torch.Tensor]): + """Run all DNC interface, addressing, and state math in fp32. + + Controller autocast is useful for the large sequence blocks, but the + DNC cosine address divides by a norm product plus a small epsilon. + Keeping its full forward/backward path in fp32 avoids casting large + intermediate derivatives back to fp16. The memory state is fp32 by + construction; the conversion also makes resumed/custom states safe. + """ + fp32_hidden = { + name: value.float() if value.is_floating_point() else value + for name, value in hidden.items() + } + if xi.is_cuda: + with torch.autocast(device_type="cuda", enabled=False): + return super().forward(xi.float(), fp32_hidden) + return super().forward(xi.float(), fp32_hidden) + def begin_stage_diagnostics(self) -> None: """Collect lightweight read/write statistics until finish is called.""" self._collect_stage_diagnostics = True @@ -112,6 +152,212 @@ def finish_stage_diagnostics(self) -> dict[str, float]: result["diagnostic_steps"] = float(max(self._stage_diagnostic_counts.values(), default=0)) return result + def set_backward_trace(self, enabled: bool, *, step: int | None = None) -> None: + """Watch DNC operation boundaries for the first non-finite gradient. + + This is intended for short diagnostic runs: each watched gradient + boundary synchronizes while checking for a non-finite value. + """ + self._backward_trace_enabled = bool(enabled) + self._backward_trace_step = step + set_product_backward_trace(enabled, self._report_bad_gradient if enabled else None) + + def _report_bad_gradient(self, stage: str, grad: torch.Tensor) -> torch.Tensor: + step = self._backward_trace_step + if self._backward_trace_first_bad_step is not None: + if step != self._backward_trace_first_bad_step: + return grad + if stage in self._backward_trace_reported_stages: + return grad + if bool(torch.isfinite(grad).all()): + return grad + if self._backward_trace_first_bad_step is None: + self._backward_trace_first_bad_step = step + + finite = torch.isfinite(grad) + finite_values = grad.detach().float()[finite] + finite_max = ( + float(finite_values.abs().max().item()) if finite_values.numel() else float("nan") + ) + finite_rms = ( + float(finite_values.square().mean().sqrt().item()) + if finite_values.numel() + else float("nan") + ) + bad_fraction = float((~finite).float().mean().item()) + bad_rows = [] + if grad.ndim >= 2: + per_row_bad_fraction = (~finite).reshape(grad.shape[0], -1).float().mean(dim=1) + bad_rows = [ + int(index) + for index in torch.nonzero(per_row_bad_fraction > 0, as_tuple=False) + .flatten() + .detach() + .cpu() + .tolist() + ] + is_first = not self._backward_trace_reported_stages + self._backward_trace_reported_stages.add(stage) + label = "first" if is_first else "also" + print( + f"[DNC-BWD-TRACE] {label} non-finite gradient through " + f"{stage} at step {step}: bad={bad_fraction:.3%} " + f"finite_rms={finite_rms:.4g} finite_max={finite_max:.4g} " + f"shape={tuple(grad.shape)} bad_batch_rows={bad_rows}" + ) + return grad + + def _watch_backward(self, stage: str, value: torch.Tensor) -> torch.Tensor: + if not self._backward_trace_enabled: + return value + peak = value.detach().float().abs().amax() + previous = self._activation_peaks.get(stage) + self._activation_peaks[stage] = peak if previous is None else torch.maximum(previous, peak) + if not value.requires_grad: + return value + + value.register_hook(lambda grad: self._report_bad_gradient(stage, grad)) + return value + + @staticmethod + def _track_peak(previous: torch.Tensor | None, value: torch.Tensor) -> torch.Tensor: + peak = value.detach().float().abs().amax() + return peak if previous is None else torch.maximum(previous, peak) + + def pop_numerical_diagnostics(self) -> dict[str, float]: + """Return/reset compact DNC invariant and magnitude metrics.""" + tensors = { + "read_weight_raw_mass_abs_max": self._read_weight_raw_mass_max, + "read_weight_mass_abs_max": self._read_weight_mass_max, + "memory_abs_max": self._memory_value_max, + "read_vector_abs_max": self._read_vector_max, + "link_row_mass_max": self._link_row_mass_max, + "link_column_mass_max": self._link_column_mass_max, + } + device_tensors = [value for value in tensors.values() if value is not None] + values = torch.stack(device_tensors).cpu().tolist() if device_tensors else [] + result: dict[str, float] = {} + index = 0 + for name, value in tensors.items(): + if value is not None: + result[name] = float(values[index]) + index += 1 + if self._read_weight_violation_count is not None: + result["read_weight_violation_calls"] = float( + self._read_weight_violation_count.detach().cpu().item() + ) + else: + result["read_weight_violation_calls"] = 0.0 + self._read_weight_violation_count = None + self._read_weight_raw_mass_max = None + self._read_weight_mass_max = None + self._memory_value_max = None + self._read_vector_max = None + self._link_row_mass_max = None + self._link_column_mass_max = None + result.update( + { + f"{stage}_abs_max": float(value.detach().cpu().item()) + for stage, value in self._activation_peaks.items() + } + ) + self._activation_peaks.clear() + result.update(pop_product_numerical_diagnostics()) + return result + + def get_usage_vector(self, usage, free_gates, read_weights, write_weights): + self._watch_backward("usage_input", usage) + self._watch_backward("usage_free_gates", free_gates) + self._watch_backward("usage_read_weights", read_weights) + self._watch_backward("usage_write_weights", write_weights) + updated = super().get_usage_vector(usage, free_gates, read_weights, write_weights) + return self._watch_backward("usage_update", updated) + + def allocate(self, usage, write_gate): + self._watch_backward("allocation_usage_input", usage) + self._watch_backward("allocation_write_gate", write_gate) + allocation, updated_usage = super().allocate(usage, write_gate) + return self._watch_backward("allocation_weights", allocation), self._watch_backward( + "allocation_usage", updated_usage + ) + + def write_weighting(self, memory, write_content_weights, allocation_weights, write_gate, allocation_gate): + self._watch_backward("write_address_content_weights", write_content_weights) + self._watch_backward("write_address_allocation_weights", allocation_weights) + self._watch_backward("write_address_gate", write_gate) + self._watch_backward("write_address_allocation_gate", allocation_gate) + weights = super().write_weighting( + memory, write_content_weights, allocation_weights, write_gate, allocation_gate + ) + return self._watch_backward("write_weights", weights) + + def get_link_matrix(self, link_matrix, write_weights, precedence): + updated = super().get_link_matrix(link_matrix, write_weights, precedence) + if self._backward_trace_enabled: + row_mass = updated.detach().float().clamp_min(0).sum(dim=-1).amax() + column_mass = updated.detach().float().clamp_min(0).sum(dim=-2).amax() + self._link_row_mass_max = ( + row_mass + if self._link_row_mass_max is None + else torch.maximum(self._link_row_mass_max, row_mass) + ) + self._link_column_mass_max = ( + column_mass + if self._link_column_mass_max is None + else torch.maximum(self._link_column_mass_max, column_mass) + ) + return self._watch_backward("temporal_link_matrix", updated) + + def update_precedence(self, precedence, write_weights): + updated = super().update_precedence(precedence, write_weights) + return self._watch_backward("precedence", updated) + + def content_weightings(self, memory, keys, strengths): + self._watch_backward("content_address_memory", memory) + self._watch_backward("content_address_keys", keys) + self._watch_backward("content_address_strengths", strengths) + # Preserve pytorch-dnc's exact operation order while exposing stages + # hidden inside its cosine-addressing helper. + scores = self._watch_backward("content_cosine_scores", θ(memory, keys)) + logits = self._watch_backward( + "content_strength_scaled_scores", scores * strengths.unsqueeze(2) + ) + weights = self._watch_backward("content_softmax_weights", σ(logits, 2)) + return self._watch_backward("content_weights", weights) + + def read_weightings(self, memory, content_weights, link_matrix, read_modes, read_weights): + raw = super().read_weightings(memory, content_weights, link_matrix, read_modes, read_weights) + if self._backward_trace_enabled: + raw_mass = raw.detach().float().sum(dim=-1, keepdim=True) + self._read_weight_raw_mass_max = self._track_peak( + self._read_weight_raw_mass_max, raw_mass + ) + + self._watch_backward("read_weights_raw", raw) + if self._backward_trace_enabled: + self._read_weight_mass_max = self._track_peak( + self._read_weight_mass_max, raw.sum(dim=-1, keepdim=True) + ) + violation = ( + (raw < 0).any() + | (raw.sum(dim=-1) > 1.0 + 1e-5).any() + | (~torch.isfinite(raw).all()) + ).to(dtype=torch.float32) + self._read_weight_violation_count = ( + violation + if self._read_weight_violation_count is None + else self._read_weight_violation_count + violation + ) + return raw + + def read_vectors(self, memory, read_weights): + if self._backward_trace_enabled: + self._memory_value_max = self._track_peak(self._memory_value_max, memory) + vectors = super().read_vectors(memory, read_weights) + if self._backward_trace_enabled: + self._read_vector_max = self._track_peak(self._read_vector_max, vectors) + return self._watch_backward("read_vectors", vectors) + def new(self, batch_size: int = 1): hidden = super().new(batch_size) if self.write_variant == "kanerva-closed-form": @@ -184,10 +430,21 @@ def _dnc_cosine_scores(memory, keys): def read_content_weightings(self, memory, keys, strengths): """Score rows, retaining DNC's key-strength softmax normalization.""" + self._watch_backward("read_content_memory", memory) + self._watch_backward("read_content_keys", keys) + self._watch_backward("read_content_strengths", strengths) scores = self._read_similarity_scores(memory, keys) - return F.softmax(scores * strengths.unsqueeze(-1), dim=-1) + self._watch_backward("read_content_scores", scores) + weights = F.softmax(scores * strengths.unsqueeze(-1), dim=-1) + return self._watch_backward("read_content_weights", weights) def read(self, read_keys, read_strengths, read_modes, hidden): + self._watch_backward("memory_before_read", hidden["memory"]) + self._watch_backward("read_weights_previous", hidden["read_weights"]) + self._watch_backward("link_matrix_before_read", hidden["link_matrix"]) + self._watch_backward("read_keys", read_keys) + self._watch_backward("read_strengths", read_strengths) + self._watch_backward("read_modes", read_modes) if self.read_variant == "cosine": if self._collect_stage_diagnostics: # Use the upstream content-weighting implementation so the @@ -264,8 +521,21 @@ def write( allocation_gate, hidden, ): + self._watch_backward("memory_before_write", hidden["memory"]) + self._watch_backward("usage_before_write", hidden["usage_vector"]) + self._watch_backward("write_weights_previous", hidden["write_weights"]) + self._watch_backward("precedence_before_write", hidden["precedence"]) + self._watch_backward("link_matrix_before_write", hidden["link_matrix"]) + self._watch_backward("write_key", write_key) + self._watch_backward("write_vector", write_vector) + self._watch_backward("erase_vector", erase_vector) + self._watch_backward("free_gates", free_gates) + self._watch_backward("write_read_strengths", read_strengths) + self._watch_backward("write_strength", write_strength) + self._watch_backward("write_gate", write_gate) + self._watch_backward("allocation_gate", allocation_gate) if self.write_variant == "learned": - return super().write( + updated = super().write( write_key, write_vector, erase_vector, @@ -276,6 +546,12 @@ def write( allocation_gate, hidden, ) + self._watch_backward("memory_after_write", updated["memory"]) + self._watch_backward("usage_after_write", updated["usage_vector"]) + self._watch_backward("write_weights_after_write", updated["write_weights"]) + self._watch_backward("precedence_after_write", updated["precedence"]) + self._watch_backward("link_matrix_after_write", updated["link_matrix"]) + return updated # Keep DNC's usage, allocation and content-based write addressing. hidden["usage_vector"] = self.get_usage_vector( @@ -322,6 +598,11 @@ def write( hidden["link_matrix"], hidden["write_weights"], hidden["precedence"] ) hidden["precedence"] = self.update_precedence(hidden["precedence"], hidden["write_weights"]) + self._watch_backward("memory_after_write", hidden["memory"]) + self._watch_backward("usage_after_write", hidden["usage_vector"]) + self._watch_backward("write_weights_after_write", hidden["write_weights"]) + self._watch_backward("precedence_after_write", hidden["precedence"]) + self._watch_backward("link_matrix_after_write", hidden["link_matrix"]) return hidden diff --git a/src/initium/memory_manipulation/nvrtc_compat.py b/src/initium/memory_manipulation/nvrtc_compat.py index 67bf0d2..3469169 100644 --- a/src/initium/memory_manipulation/nvrtc_compat.py +++ b/src/initium/memory_manipulation/nvrtc_compat.py @@ -66,6 +66,8 @@ _EPS = 1e-6 _PATCHED_ATTR = "_nvrtc_compat_patched" _STRICT_NEGATIVE_CHECK = os.environ.get("INITIUM_NVRTC_COMPAT_STRICT", "0") == "1" +_PRODUCT_BACKWARD_TRACE: tuple[bool, object | None] = (False, None) +_PRODUCT_INPUT_STATS: dict[str, dict[str, torch.Tensor | int]] = {} _orig_prod = torch.prod _orig_cumprod = torch.cumprod @@ -84,8 +86,66 @@ def _log_space(input: torch.Tensor, dim, keepdim: bool, cumulative: bool) -> tor return torch.exp(reduced).to(orig_dtype) +def set_product_backward_trace(enabled: bool, reporter=None) -> None: + """Attach an optional diagnostic hook to patched product outputs.""" + global _PRODUCT_BACKWARD_TRACE + _PRODUCT_BACKWARD_TRACE = (bool(enabled), reporter) + + +def _watch_product_output( + result: torch.Tensor, input: torch.Tensor, operation: str +) -> torch.Tensor: + enabled, reporter = _PRODUCT_BACKWARD_TRACE + if enabled and reporter is not None and result.requires_grad: + result.register_hook(lambda grad: reporter(operation, grad)) + if input.requires_grad: + input.register_hook(lambda grad: reporter(f"{operation}_input_gradient", grad)) + return result + + +def _record_product_input(input: torch.Tensor, operation: str) -> None: + enabled, _ = _PRODUCT_BACKWARD_TRACE + if not enabled: + return + values = input.detach().float() + metrics = { + "min": values.amin(), + "max": values.amax(), + "lt_1e-6_fraction": (values.abs() < _EPS).float().mean(), + "negative_fraction": (values < 0).float().mean(), + "nonfinite_fraction": (~torch.isfinite(values)).float().mean(), + } + if operation in _PRODUCT_INPUT_STATS: + previous = _PRODUCT_INPUT_STATS[operation] + previous["min"] = torch.minimum(previous["min"], metrics["min"]) + previous["max"] = torch.maximum(previous["max"], metrics["max"]) + for name in ("lt_1e-6_fraction", "negative_fraction", "nonfinite_fraction"): + previous[name] = previous[name] + metrics[name] + previous["calls"] = int(previous["calls"]) + 1 + else: + metrics["calls"] = 1 + _PRODUCT_INPUT_STATS[operation] = metrics + + +def pop_product_numerical_diagnostics() -> dict[str, float]: + """Return/reset compact input summaries for patched product operations.""" + result: dict[str, float] = {} + for operation, metrics in _PRODUCT_INPUT_STATS.items(): + calls = max(1, int(metrics["calls"])) + result[f"{operation}_input_min"] = float(metrics["min"].detach().cpu().item()) + result[f"{operation}_input_max"] = float(metrics["max"].detach().cpu().item()) + for name in ("lt_1e-6_fraction", "negative_fraction", "nonfinite_fraction"): + result[f"{operation}_input_{name}"] = float( + (metrics[name] / calls).detach().cpu().item() + ) + result[f"{operation}_calls"] = float(calls) + _PRODUCT_INPUT_STATS.clear() + return result + + def _patched_prod(input, dim=None, keepdim=False, *, dtype=None): if isinstance(input, torch.Tensor) and input.is_cuda and dim is not None: + _record_product_input(input, "torch_prod") if not input.dtype.is_floating_point: return _orig_prod(input, dim, keepdim=keepdim, dtype=dtype) if _STRICT_NEGATIVE_CHECK and bool((input < 0).any()): @@ -93,7 +153,8 @@ def _patched_prod(input, dim=None, keepdim=False, *, dtype=None): input.device ) result = _log_space(input, dim, keepdim, cumulative=False) - return result.to(dtype) if dtype is not None else result + result = result.to(dtype) if dtype is not None else result + return _watch_product_output(result, input, "patched_torch_prod") if dim is None: return _orig_prod(input, **({"dtype": dtype} if dtype is not None else {})) return _orig_prod( @@ -103,12 +164,14 @@ def _patched_prod(input, dim=None, keepdim=False, *, dtype=None): def _patched_cumprod(input, dim, *, dtype=None): if isinstance(input, torch.Tensor) and input.is_cuda: + _record_product_input(input, "torch_cumprod") if not input.dtype.is_floating_point: return _orig_cumprod(input, dim, dtype=dtype) if _STRICT_NEGATIVE_CHECK and bool((input < 0).any()): return _orig_cumprod(input.detach().cpu(), dim, dtype=dtype).to(input.device) result = _log_space(input, dim, keepdim=False, cumulative=True) - return result.to(dtype) if dtype is not None else result + result = result.to(dtype) if dtype is not None else result + return _watch_product_output(result, input, "patched_torch_cumprod") return _orig_cumprod(input, dim, **({"dtype": dtype} if dtype is not None else {})) From 4a61e38285213530bcff85a7983fcf7d7db80dde Mon Sep 17 00:00:00 2001 From: Samuel Date: Sat, 10 Oct 2026 22:16:30 +0200 Subject: [PATCH 06/14] Enable residual-scaled workspace MoE broadcasts This change allows the workspace broadcast path to run with the split-graph backbone MoE residual scale instead of rejecting it. The final-bank broadcast delta is scaled with the matched backbone setting, and active workspace attention/slot updates run in fp32 outside CUDA autocast while the zero-slot/closed-gate bypass remains unchanged. Startup reporting and the README were updated to describe the new phase-1 configuration, fp32 precision behavior, and the relevant MoE/DNC validation commands. --- README.md | 14 +++---- src/initium/core_training.py | 9 ++-- .../mamba_backbone_parallel.py | 2 + .../mamba_controller/split_graph_dnc.py | 13 +++--- src/initium/workspace/broadcast.py | 41 +++++++++++++++++++ 5 files changed, 62 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index f30b547..0a5d596 100644 --- a/README.md +++ b/README.md @@ -106,17 +106,17 @@ python -m initium.core_training 0.0 --split-graph --split-graph-variant mamba1 - ### Phase 0 shared workspace (MoE broadcast) -The optional `--workspace` path adds a persistent shared slot matrix to a split-graph model with MoE enabled. The final parallel-backbone MoE bank supplies the specialist candidates and its existing noisy top-k routes. At each sequence timestep, selected expert outputs compete for slot writes; a learned RMC-style gate blends proposed slot content with the previous matrix. Every expert in that bank then attends to the updated matrix, and the existing sparse MoE routes combine the context-enriched expert outputs. The workspace resets at the beginning of each episode and persists across its timesteps. This is a shared-object broadcast implementation; it does not implement per-recipient translation. +The optional `--workspace` path adds a persistent shared slot matrix to a split-graph model with MoE enabled. The final parallel-backbone MoE bank supplies the specialist candidates and its existing noisy top-k routes. At each sequence timestep, selected expert outputs compete for slot writes; a learned RMC-style gate blends proposed slot content with the previous matrix. Every expert in that bank then attends to the updated matrix, and the existing sparse MoE routes combine the context-enriched expert outputs. The workspace resets at the beginning of each episode and persists across its timesteps. The `--moe-backbone-residual-scale` setting scales the complete final-bank delta, including its broadcast contribution, just as it scales the ordinary MoE residuals in the selection-only backbone. The DNC memory uses the same fp32/autocast-disabled path with or without workspace. This is a shared-object broadcast implementation; it does not implement per-recipient translation. -Use the roadmap's proposed split-graph/CfC-combiner configuration with MoE and the workspace enabled as follows: +Use the Phase 1 matched split-graph/CfC-combiner setup with MoE and the workspace enabled as follows: ```bash -python -m src.initium.core_training 0.0 --split-graph --split-graph-variant mamba1 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --workspace +python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase2-workspace-coadapted --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --workspace --workspace-num-slots 5 --workspace-train-mode coadapted --total-steps 30000 ``` `--workspace` requires both `--split-graph` and `--moe`. The default is 5 slots, based on Goyal et al.'s shared-workspace reference result; it is a starting configuration, not a claimed optimum for graph traversal. `--workspace-train-mode coadapted` trains the full model. Choose `frozen` to train only workspace parameters while freezing the backbone, MoE/combiner, DNC, stochastic write heads, and task output projection. `--workspace-no-selection` replaces competitive selected writes with random unweighted writes for the no-Selection control. `--workspace-reset-each-step` resets slots before each timestep to test persistence. `--workspace-num-slots 0` bypasses workspace writes and readback, providing a plumbing comparison against the existing MoE model. Run IDs encode these conditions to avoid log/checkpoint collisions. -The console prints a startup line that reports whether the workspace is active or bypassed, its capacity, write-selection mode, persistence setting, and train mode. At the normal training log cadence, `workspace[...]` reports selected write mass and active-write fraction, the learned write-gate mean, slot-state RMS, readback-to-expert-output RMS ratio, and read-gate value. These values show whether the path is active and the relative scale of its write and broadcast signals; they do not by themselves establish a task benefit. +The console prints a startup line that reports whether the workspace is active or bypassed, its capacity, write-selection mode, persistence setting, train mode, and active compute precision. Active workspace attention and slot-state updates run in fp32 outside CUDA autocast; zero-slot/closed-gate bypass retains the selection-only path. At the normal training log cadence, `workspace[...]` reports selected write mass and active-write fraction, the learned write-gate mean, slot-state RMS, readback-to-expert-output RMS ratio, and read-gate value. These values show whether the path is active and the relative scale of its write and broadcast signals; they do not by themselves establish a task benefit. The shared object operates on the final backbone MoE bank; earlier backbone MoE layers retain their existing behavior. This is the Phase 0 plumbing target, not evidence that Broadcast improves task performance. The first short run should compare the zero-slot bypass against the same MoE configuration without `--workspace`; only after this plumbing gate should the frozen/co-adapted and no-Selection comparisons begin. Record seed, full command, commit, hardware/software environment, curriculum completion, ID/OOD metrics, convergence steps, and any failures in `07-Experiment/Broadcasting/Progress/Phase-0-Implementation.md`. All current evidence is graph-traversal-specific; transfer across task families remains untested. @@ -134,13 +134,13 @@ Training console logs now show **one MoE-bank record per 100-step logging window Split-graph DNC interface transforms, cosine content addressing, usage/allocation, write/read weighting, temporal links, and recurrent memory updates now run in fp32 even when the outer controller uses CUDA autocast. The stock cosine rule divides by the memory/key norm product plus `1e-6`; the instrumented trace showed finite score-side gradients but a non-finite key-side gradient, so fp32 protects that path from low-precision intermediate-gradient overflow. The startup log and checkpoint metadata identify this precision choice. The equations remain unchanged. -To inspect DNC internals, add `--dnc-backward-trace` to a split-graph training command. The flag watches gradients at read/write interface values, cosine-addressing scores/logits/softmax, usage, allocation, write weights, temporal links, precedence, recurrent memory, and read vectors, plus the patched CUDA `torch.prod`/`torch.cumprod` outputs and their inputs. It reports every distinct hooked boundary with a non-finite gradient on the first affected optimizer step, rather than stopping after the first hook label, and includes affected batch rows. At the regular logging interval and on failure, `[dnc-numerics:...]` reports memory/read/link magnitudes, read-weight mass violations, and per-product input min/max, fraction below `1e-6`, negative/non-finite fraction, and call count. The flag synchronizes at watched backward boundaries, so use a short diagnostic capture around the known first event rather than a full long run. For the attached trace, which first reports a bad memory input gradient at step 1,120, 1,800 steps covers the onset: +To inspect DNC internals, add `--dnc-backward-trace` to a split-graph training command. The flag watches gradients at read/write interface values, cosine-addressing scores/logits/softmax, usage, allocation, write weights, temporal links, precedence, recurrent memory, and read vectors, plus the patched CUDA `torch.prod`/`torch.cumprod` outputs and their inputs. It reports every distinct hooked boundary with a non-finite gradient on the first affected optimizer step, rather than stopping after the first hook label, and includes affected batch rows. At the regular logging interval and on failure, `[dnc-numerics:...]` reports memory/read/link magnitudes, read-weight mass violations, and per-product input min/max, fraction below `1e-6`, negative/non-finite fraction, and call count. The flag synchronizes at watched backward boundaries, so use a short diagnostic capture around a known first event. Earlier captures reported their first bad memory-input gradients at steps 1,120 and 626; the matched fp32 validation run below completed 5,000 steps without a bad DNC hook or non-finite-gradient skip. ```bash python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix dnc-backward-trace --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 1800 --dnc-backward-trace ``` -Keep other flags exactly matched to the run being investigated. The fp32 DNC path is the targeted numerical correction; repeat the short capture to check whether it removes the non-finite key gradients and restores paired memory dependence. If failures remain, the expanded trace still reports all bad stages on their first affected step. +Keep other flags exactly matched to the run being investigated. The fp32 DNC path removed the observed bad-key-gradient pattern in the 5,000-step matched validation capture. That run also acquired strong paired ID memory dependence by step 2,500; see `AI-Architecture-Research/07-Experiment/Broadcasting/Progress/Phase-1-MoE-Failure-Analysis.md` for the complete result and remaining generalization limits. Repeat across seeds before treating the numerical and learning outcomes as robust. Evaluation emits separate channel lines for ID, OOD, memory-off, combiner-stage ablations, and final ID/OOD checks. These summarize the evaluations independently from training windows. If a forward value or gradient becomes non-finite, the current segment window and MoE-bank window are printed immediately before training stops or skips the unsafe update. The summaries aggregate detached device-side scalars and synchronize at log/evaluation boundaries; they do not retain activation graphs. @@ -164,7 +164,7 @@ python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-back python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-moe-backbone-block1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 ``` -`--moe-backbone-residual-scale SCALE` multiplies the active parallel-backbone MoE delta before its residual addition (`backbone_output + SCALE * moe_delta`). Its default is `1.0`, preserving prior behavior. The scale must be finite and non-negative, requires split-graph MoE with `--moe-placement backbone` or `all`, is included in run IDs/checkpoint metadata, and is not supported with the workspace broadcast path. With `all`, this changes only the parallel-backbone MoE residuals; combiner MoE residuals remain at 1.0. It scales the task-loss path through the backbone MoE branch while leaving routing and the load-balancing auxiliary loss active. In an isolated-block run, only that block has an MoE residual to scale. For the current block-1 follow-up, use a deliberately conservative 0.1 scale: +`--moe-backbone-residual-scale SCALE` multiplies the active parallel-backbone MoE delta before its residual addition (`backbone_output + SCALE * moe_delta`). Its default is `1.0`, preserving prior behavior. The scale must be finite and non-negative, requires split-graph MoE with `--moe-placement backbone` or `all`, and is included in run IDs/checkpoint metadata. With `all`, this changes only the parallel-backbone MoE residuals; combiner MoE residuals remain at 1.0. It scales the task-loss path through the backbone MoE branch while leaving routing and the load-balancing auxiliary loss active. When workspace broadcast is enabled, the scale applies to the complete final-bank delta after expert readback, including both the selected expert outputs and their broadcast contribution. Workspace slot writes and router/load-balancing calculations remain active at their existing scales. Active workspace attention and slot updates run in fp32 outside CUDA autocast, and split-graph DNC memory math uses its fp32 path in both workspace and selection-only runs. In an isolated-block run, only that block has an MoE residual to scale. For the current block-1 follow-up, use a deliberately conservative 0.1 scale: ```bash python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase1-block1-residual-scale-0p1 --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement backbone --moe-backbone-block-isolation 1 --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --total-steps 5000 diff --git a/src/initium/core_training.py b/src/initium/core_training.py index d659720..3af6d87 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -1232,8 +1232,9 @@ def run( print( f"[{run_id}] backbone MoE residual scale=" f"{split_graph_moe_backbone_residual_scale:g} | " - "MoE delta is scaled before its residual addition; router, " - "experts, and auxiliary routing loss remain active." + "applies to each backbone MoE residual, including the " + "broadcast-enriched final bank when workspace is active; " + "router, experts, and auxiliary routing loss remain active." ) # v10: variant-matched hyperparameter defaults instead of always @@ -1338,6 +1339,7 @@ def run( f"[{run_id}] WORKSPACE {'active' if workspace_active else 'bypassed'} | " f"slots={workspace_num_slots} | write_selection={write_mode} | " f"persistence={persistence} | train_mode={workspace_train_mode} | " + f"compute_precision={'fp32' if workspace_active else 'bypass'} | " "scope=final-backbone-MoE-bank; each specialist reads shared slots" ) else: @@ -3040,6 +3042,7 @@ def _write_diagnostic_values(diagnostics): metavar="SCALE", help="Multiply each active split-graph backbone MoE delta by SCALE " "before adding it to the backbone residual (default: 1.0, unchanged). " + "With --workspace, this also scales the final-bank broadcast contribution. " "Use with --moe-placement backbone or all; with all, combiner MoE " "residuals remain unscaled. In an isolation run, only the selected " "block's MoE residual is scaled.", @@ -3352,8 +3355,6 @@ def _write_diagnostic_values(diagnostics): parser.error( "--moe-backbone-residual-scale requires --moe-placement backbone or all" ) - if args.workspace: - parser.error("--moe-backbone-residual-scale cannot be combined with --workspace") if args.workspace and args.moe_placement == "combiner": parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: diff --git a/src/initium/mamba_controller/mamba_backbone_parallel.py b/src/initium/mamba_controller/mamba_backbone_parallel.py index f885801..2f1c35e 100644 --- a/src/initium/mamba_controller/mamba_backbone_parallel.py +++ b/src/initium/mamba_controller/mamba_backbone_parallel.py @@ -312,6 +312,8 @@ def forward_with_specialists( Earlier backbone/MoE blocks remain unchanged. The final MoE block exposes expert deltas before routing so the workspace can broadcast context to every specialist, then reuse the original sparse route. + The split-graph caller applies this block's configured residual scale + to the combined routed and broadcast delta. """ if not self.moe_enabled or self.moe_blocks is None: raise RuntimeError("workspace broadcast requires an enabled MoE expert pool") diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 3474472..feb4cdf 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -251,11 +251,6 @@ def __init__( raise ValueError( "--workspace requires backbone MoE because it broadcasts to backbone specialists" ) - if workspace_enabled and moe_backbone_residual_scale != 1.0: - raise ValueError( - "workspace specialist broadcasts do not use the backbone MoE residual; " - "moe_backbone_residual_scale must remain 1.0 with workspace_enabled" - ) if workspace_train_mode not in ("coadapted", "frozen"): raise ValueError("workspace_train_mode must be 'coadapted' or 'frozen'") if workspace_enabled and workspace_train_mode == "frozen" and workspace_num_slots == 0: @@ -636,7 +631,13 @@ def forward( reset_each_step=self.workspace_reset_each_step, ) H = base.clone() - H[:, start_step:] = base[:, start_step:] + broadcast_delta + # The workspace returns the selected, broadcast-enriched MoE + # delta. Apply the same residual scale as the ordinary final MoE + # block to the complete delta so workspace and selection-only + # runs share a matched backbone residual setting. + H[:, start_step:] = base[:, start_step:] + ( + self.moe_backbone_residual_scale * broadcast_delta + ) else: backbone_stages = ( list(self.backbone) diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index e86e466..99df1d3 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -53,6 +53,47 @@ def forward_step( *, no_selection: bool = False, reset_each_step: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Keep the zero-slot/closed-gate bypass identical to the selection-only + # MoE path. Active workspace attention and recurrent slot updates use + # fp32 even when the surrounding controller is under CUDA autocast. + if self.num_slots == 0 or self.write_gate_closed: + return self._forward_step( + specialist_deltas, + route_indices, + route_weights, + state, + no_selection=no_selection, + reset_each_step=reset_each_step, + ) + if specialist_deltas.is_cuda: + with torch.autocast(device_type="cuda", enabled=False): + return self._forward_step( + specialist_deltas.float(), + route_indices, + route_weights.float(), + state.float(), + no_selection=no_selection, + reset_each_step=reset_each_step, + ) + return self._forward_step( + specialist_deltas.float(), + route_indices, + route_weights.float(), + state.float(), + no_selection=no_selection, + reset_each_step=reset_each_step, + ) + + def _forward_step( + self, + specialist_deltas: torch.Tensor, + route_indices: torch.Tensor, + route_weights: torch.Tensor, + state: torch.Tensor, + *, + no_selection: bool = False, + reset_each_step: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if self.num_slots == 0 or self.write_gate_closed: selected = specialist_deltas.gather( From 0f48290efc91e94f63a3e1e6333f79d702f360e9 Mon Sep 17 00:00:00 2001 From: Samuel Date: Sun, 11 Oct 2026 12:18:50 +0200 Subject: [PATCH 07/14] Add frozen workspace warm-start and MoE diagnostics This commit adds a warm-start path for frozen workspace experiments: a compatible workspace-disabled checkpoint can be resumed into a new frozen-workspace run without reinitializing backbone, MoE, combiner, DNC, stochastic-write, or task-head weights. It validates model metadata and shared tensor shapes, keeps newly introduced workspace parameters freshly initialized, and resets optimizer, scaler, RNG, curriculum, and step counters to start the workspace from lesson 1. Existing checkpoints with workspace state still follow the normal exact-resume behavior. It also adds optional MoE routing diagnostics control. Detailed routing/activation diagnostics can be disabled to reduce overhead while keeping the auxiliary load-balancing loss and routing logic active. The code resets cached diagnostics cleanly when toggled, and the training loop now conditionally emits the more expensive MoE diagnostics. The README is updated to document the new warm-start workflow and expected behavior. --- README.md | 2 + src/initium/MoE/moe_layer.py | 198 ++++++++++++-------- src/initium/core_training.py | 350 +++++++++++++++++++++++++++++++---- 3 files changed, 440 insertions(+), 110 deletions(-) diff --git a/README.md b/README.md index 0a5d596..ab5f9ba 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,8 @@ python -m src.initium.core_training 0.0 --seed 0 --run-id-suffix phase2-workspac `--workspace` requires both `--split-graph` and `--moe`. The default is 5 slots, based on Goyal et al.'s shared-workspace reference result; it is a starting configuration, not a claimed optimum for graph traversal. `--workspace-train-mode coadapted` trains the full model. Choose `frozen` to train only workspace parameters while freezing the backbone, MoE/combiner, DNC, stochastic write heads, and task output projection. `--workspace-no-selection` replaces competitive selected writes with random unweighted writes for the no-Selection control. `--workspace-reset-each-step` resets slots before each timestep to test persistence. `--workspace-num-slots 0` bypasses workspace writes and readback, providing a plumbing comparison against the existing MoE model. Run IDs encode these conditions to avoid log/checkpoint collisions. +For a frozen-workspace experiment, pass a compatible workspace-disabled baseline checkpoint through `--resume`. Initium detects this case, verifies the saved architecture/settings and shared tensor shapes, loads the matching backbone, MoE, combiner, DNC, stochastic-write, and task-head weights, and keeps the newly added workspace parameters at their fresh initialization. It then resets the optimizer, scaler, RNG streams, training step, and curriculum to step 0 / lesson 1 so the workspace learns across the full curriculum. The checkpoint's final memory size is used to load its weights, then memory is resized to the first lesson's configured size. This warm-start accepts only matching beta/model settings and rejects missing or mismatched non-workspace weights. A checkpoint that already includes workspace state follows the ordinary exact-resume path, including optimizer and progress restoration. Example: add `--resume checkpoints/.pt` to the frozen-workspace command above and use a distinct `--run-id-suffix` for the new experiment. + The console prints a startup line that reports whether the workspace is active or bypassed, its capacity, write-selection mode, persistence setting, train mode, and active compute precision. Active workspace attention and slot-state updates run in fp32 outside CUDA autocast; zero-slot/closed-gate bypass retains the selection-only path. At the normal training log cadence, `workspace[...]` reports selected write mass and active-write fraction, the learned write-gate mean, slot-state RMS, readback-to-expert-output RMS ratio, and read-gate value. These values show whether the path is active and the relative scale of its write and broadcast signals; they do not by themselves establish a task benefit. The shared object operates on the final backbone MoE bank; earlier backbone MoE layers retain their existing behavior. This is the Phase 0 plumbing target, not evidence that Broadcast improves task performance. The first short run should compare the zero-slot bypass against the same MoE configuration without `--workspace`; only after this plumbing gate should the frozen/co-adapted and no-Selection comparisons begin. Record seed, full command, commit, hardware/software environment, curriculum completion, ID/OOD metrics, convergence steps, and any failures in `07-Experiment/Broadcasting/Progress/Phase-0-Implementation.md`. All current evidence is graph-traversal-specific; transfer across task families remains untested. diff --git a/src/initium/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index 2cfa205..d790ea4 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -157,6 +157,11 @@ def __init__( self.experts = nn.ModuleList([Expert(d_model, expert_dim) for _ in range(num_experts)]) self._aux_losses: list[torch.Tensor] = [] + # Routing/activation reports are useful for experiments but can be + # expensive in recurrent controllers, where this module may run once + # per timestep. Keep them independently switchable from the actual + # routing and load-balancing objective. + self.diagnostics_enabled = True self._last_diag: dict = {} self._diag_accum: dict[str, torch.Tensor] = {} self._diag_calls = 0 @@ -232,22 +237,29 @@ def forward( else: gate_weights = topk_prob / topk_prob.sum(dim=-1, keepdim=True).clamp(min=1e-9) - # f_i/P_i are cheap to compute regardless of mode; only the aux LOSS - # accumulation (needed for backward()) is training-gated below. - # Computing diagnostics unconditionally is what lets inference - # (model.eval()) report routing health -- previously last_diagnostics() - # stayed empty forever outside training, which made it impossible to - # verify MoE routing on real held-out data at inference time. - one_hot_k = F.one_hot(topk_idx, num_classes=self.num_experts).float() # (T, k, E) - f_i = one_hot_k.sum(dim=(0, 1)) / max(num_tokens * top_k, 1) - P_i = probs.mean(dim=0) - + collect_diagnostics = self.diagnostics_enabled + # The auxiliary loss needs the non-differentiable assignment fraction + # and differentiable probability mass. Use bincount instead of a + # dense (tokens, top_k, experts) one-hot tensor, which is especially + # costly for recurrent MoE calls. The one-hot form is built only when + # detailed reports are enabled and token-subset statistics need it. if self.training: + f_i = torch.bincount( + topk_idx.reshape(-1), minlength=self.num_experts + ).to(dtype=probs.dtype) / max(num_tokens * top_k, 1) + P_i = probs.mean(dim=0) aux_loss = self.load_balance_alpha * self.num_experts * (f_i * P_i).sum() self._aux_losses.append(aux_loss) + elif collect_diagnostics: + f_i = torch.bincount( + topk_idx.reshape(-1), minlength=self.num_experts + ).to(dtype=probs.dtype) / max(num_tokens * top_k, 1) + P_i = probs.mean(dim=0) + + if collect_diagnostics: + with torch.no_grad(): + self._last_topk_idx = topk_idx.detach() - with torch.no_grad(): - self._last_topk_idx = topk_idx.detach() # Expert capacity, scaled by top_k (each slot competes for the same # per-expert buffer). @@ -322,7 +334,11 @@ def forward( route_slots_2d = route_slots.reshape(num_tokens, top_k) selected_outputs = expert_outputs[route_experts, route_slots_2d] output32 = (selected_outputs * accepted_gates.unsqueeze(-1)).sum(dim=1) - accepted_outputs = selected_outputs * keep.unsqueeze(-1).to(selected_outputs.dtype) + accepted_outputs = ( + selected_outputs * keep.unsqueeze(-1).to(selected_outputs.dtype) + if collect_diagnostics + else None + ) expert_out_all = None if return_details: @@ -359,70 +375,72 @@ def forward( ) output = output32.to(output_dtype) - with torch.no_grad(): - self._last_accepted = keep.detach() - self._last_diag = { - "cv_importance": (P_i.std() / P_i.mean().clamp(min=1e-8)).detach(), - "cv_load": (f_i.std() / f_i.mean().clamp(min=1e-8)).detach(), - "max_load_frac": f_i.max().detach(), - "expert_load_frac": f_i.detach(), - "router_importance_frac": P_i.detach(), - "router_entropy": ( - -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1).mean() - ).detach(), - "router_max_prob": probs.max(dim=-1).values.mean().detach(), - "topk_gate_mean": topk_prob.mean().detach(), - "input_rms": flat32.square().mean().sqrt().detach(), - "input_absmax": flat32.abs().amax().detach(), - "expert_output_rms": accepted_outputs.square().mean().sqrt().detach(), - "expert_output_absmax": accepted_outputs.abs().amax().detach(), - "routed_output_rms": output32.square().mean().sqrt().detach(), - "routed_output_absmax": output32.abs().amax().detach(), - "capacity_drop_frac": (1.0 - keep.float().mean()).detach(), - "accepted_gate_mass": accepted_gates.sum(dim=-1).mean().detach(), - } - if token_mask_flat is not None: - valid = token_mask_flat - padded = ~valid - per_token_entropy = -( - probs * probs.clamp_min(1e-9).log() - ).sum(dim=-1) - per_token_max_prob = probs.max(dim=-1).values - per_token_drop = 1.0 - keep.float().mean(dim=-1) - - def subset_routing_stats(subset: torch.Tensor, prefix: str) -> None: - subset_weight = subset.to(probs.dtype) - count = subset_weight.sum() - denom = count.clamp_min(1.0) - subset_load = ( - one_hot_k - * subset_weight[:, None, None] - ).sum(dim=(0, 1)) / (denom * top_k) - subset_importance = ( - probs * subset_weight[:, None] - ).sum(dim=0) / denom - entropy = (per_token_entropy * subset_weight).sum() / denom - max_prob = (per_token_max_prob * subset_weight).sum() / denom - drop_frac = (per_token_drop * subset_weight).sum() / denom - self._last_diag[f"{prefix}_token_count"] = count.detach() - self._last_diag[f"{prefix}_expert_load_frac"] = subset_load.detach() - self._last_diag[f"{prefix}_router_importance_frac"] = ( - subset_importance.detach() - ) - self._last_diag[f"{prefix}_router_entropy"] = entropy.detach() - self._last_diag[f"{prefix}_router_max_prob"] = max_prob.detach() - self._last_diag[f"{prefix}_capacity_drop_frac"] = drop_frac.detach() - - subset_routing_stats(valid, "valid") - subset_routing_stats(padded, "padded") - if self.training: - for key, value in self._last_diag.items(): - self._diag_accum[key] = ( - self._diag_accum[key] + value - if key in self._diag_accum - else value.clone() - ) - self._diag_calls += 1 + if collect_diagnostics: + with torch.no_grad(): + one_hot_k = F.one_hot(topk_idx, num_classes=self.num_experts).float() + self._last_accepted = keep.detach() + self._last_diag = { + "cv_importance": (P_i.std() / P_i.mean().clamp(min=1e-8)).detach(), + "cv_load": (f_i.std() / f_i.mean().clamp(min=1e-8)).detach(), + "max_load_frac": f_i.max().detach(), + "expert_load_frac": f_i.detach(), + "router_importance_frac": P_i.detach(), + "router_entropy": ( + -(probs * probs.clamp_min(1e-9).log()).sum(dim=-1).mean() + ).detach(), + "router_max_prob": probs.max(dim=-1).values.mean().detach(), + "topk_gate_mean": topk_prob.mean().detach(), + "input_rms": flat32.square().mean().sqrt().detach(), + "input_absmax": flat32.abs().amax().detach(), + "expert_output_rms": accepted_outputs.square().mean().sqrt().detach(), + "expert_output_absmax": accepted_outputs.abs().amax().detach(), + "routed_output_rms": output32.square().mean().sqrt().detach(), + "routed_output_absmax": output32.abs().amax().detach(), + "capacity_drop_frac": (1.0 - keep.float().mean()).detach(), + "accepted_gate_mass": accepted_gates.sum(dim=-1).mean().detach(), + } + if token_mask_flat is not None: + valid = token_mask_flat + padded = ~valid + per_token_entropy = -( + probs * probs.clamp_min(1e-9).log() + ).sum(dim=-1) + per_token_max_prob = probs.max(dim=-1).values + per_token_drop = 1.0 - keep.float().mean(dim=-1) + + def subset_routing_stats(subset: torch.Tensor, prefix: str) -> None: + subset_weight = subset.to(probs.dtype) + count = subset_weight.sum() + denom = count.clamp_min(1.0) + subset_load = ( + one_hot_k + * subset_weight[:, None, None] + ).sum(dim=(0, 1)) / (denom * top_k) + subset_importance = ( + probs * subset_weight[:, None] + ).sum(dim=0) / denom + entropy = (per_token_entropy * subset_weight).sum() / denom + max_prob = (per_token_max_prob * subset_weight).sum() / denom + drop_frac = (per_token_drop * subset_weight).sum() / denom + self._last_diag[f"{prefix}_token_count"] = count.detach() + self._last_diag[f"{prefix}_expert_load_frac"] = subset_load.detach() + self._last_diag[f"{prefix}_router_importance_frac"] = ( + subset_importance.detach() + ) + self._last_diag[f"{prefix}_router_entropy"] = entropy.detach() + self._last_diag[f"{prefix}_router_max_prob"] = max_prob.detach() + self._last_diag[f"{prefix}_capacity_drop_frac"] = drop_frac.detach() + + subset_routing_stats(valid, "valid") + subset_routing_stats(padded, "padded") + if self.training: + for key, value in self._last_diag.items(): + self._diag_accum[key] = ( + self._diag_accum[key] + value + if key in self._diag_accum + else value.clone() + ) + self._diag_calls += 1 output = output.reshape(orig_shape) if return_details: @@ -452,6 +470,19 @@ def pop_aux_loss(self) -> torch.Tensor: self._aux_losses = [] return total + def set_diagnostics_enabled(self, enabled: bool) -> None: + """Enable/disable optional routing and activation diagnostics. + + This does not affect expert routing, outputs, or auxiliary loss. + Clearing cached values also makes runtime toggles well-defined. + """ + self.diagnostics_enabled = bool(enabled) + self._last_diag = {} + self._diag_accum = {} + self._diag_calls = 0 + self._last_topk_idx = None + self._last_accepted = None + def last_diagnostics(self) -> dict: """Read-only routing-health diagnostics from the most recent training forward() call: CV(Importance), CV(Load) (Shazeer et al. @@ -550,6 +581,9 @@ def forward( def pop_aux_loss(self) -> torch.Tensor: return self.moe.pop_aux_loss() + def set_diagnostics_enabled(self, enabled: bool) -> None: + self.moe.set_diagnostics_enabled(enabled) + def last_diagnostics(self) -> dict: return self.moe.last_diagnostics() @@ -656,8 +690,8 @@ def forward(self, sources: list[torch.Tensor]) -> list[torch.Tensor]: # per-source specialization -- e.g. "source 0 (backbone output) # concentrates on experts {2,5}, source 1 (read vector) on {1,7}" # -- rather than every source routing near-identically by chance. - topk_idx = self.moe.last_routing() # (sum(batch_sizes), top_k) or None - accepted = self.moe.last_accepted_routing() + topk_idx = self.moe.last_routing() if self.moe.diagnostics_enabled else None + accepted = self.moe.last_accepted_routing() if self.moe.diagnostics_enabled else None if topk_idx is not None and accepted is not None: num_experts = self.moe.num_experts if self._cumulative_source_counts is None: @@ -699,6 +733,12 @@ def forward(self, sources: list[torch.Tensor]) -> list[torch.Tensor]: def pop_aux_loss(self) -> torch.Tensor: return self.moe.pop_aux_loss() + def set_diagnostics_enabled(self, enabled: bool) -> None: + self.moe.set_diagnostics_enabled(enabled) + self._cumulative_source_counts = None + self._cumulative_source_accepted = None + self._cumulative_source_routes = None + def last_diagnostics(self) -> dict: return self.moe.last_diagnostics() diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 3af6d87..59170b2 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -172,7 +172,7 @@ pop_total_kl, update_all_prior_snapshots, ) -from initium.MoE.moe_layer import pop_total_moe_aux_loss +from initium.MoE.moe_layer import SwitchMoE, pop_total_moe_aux_loss device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") @@ -652,6 +652,118 @@ def restore_rng_state(rng_state): torch.cuda.set_rng_state_all(rng_state["cuda"]) +def load_frozen_workspace_warm_start( + rnn: nn.Module, + output_proj: nn.Module, + checkpoint: dict, + expected_model_config: dict, +) -> list[str]: + """Load a workspace-disabled baseline into a frozen-workspace model. + + Only the newly introduced workspace state may be absent. Every shared + tensor must match by name and shape, and the recorded architecture must + match outside workspace settings and the curriculum-dependent memory + size. Returns the target workspace keys left at their fresh initialization. + """ + checkpoint_config = checkpoint.get("model_config") + if not isinstance(checkpoint_config, dict): + raise ValueError( + "Frozen-workspace warm start requires a checkpoint with model_config metadata." + ) + if checkpoint_config.get("workspace_enabled", False) or any( + key.startswith("workspace.") or ".workspace." in key + for key in checkpoint.get("rnn_state_dict", {}) + ): + raise ValueError( + "Frozen-workspace warm start expects a workspace-disabled baseline checkpoint." + ) + + ignored_config_keys = { + "nr_cells", # resized to the checkpoint size for load, then back to lesson 1 + "dnc_backward_trace", # diagnostic instrumentation, not architecture + "workspace_enabled", + "workspace_num_slots", + "workspace_no_selection", + "workspace_reset_each_step", + "workspace_train_mode", + } + missing_config = sorted( + key + for key in expected_model_config + if key not in ignored_config_keys and key not in checkpoint_config + ) + unexpected_config = sorted( + key + for key in checkpoint_config + if key not in ignored_config_keys and key not in expected_model_config + ) + mismatched_config = sorted( + key + for key, expected in expected_model_config.items() + if key not in ignored_config_keys + and key in checkpoint_config + and checkpoint_config[key] != expected + ) + if missing_config or unexpected_config or mismatched_config: + details = [] + if missing_config: + details.append(f"missing metadata: {', '.join(missing_config)}") + if unexpected_config: + details.append(f"unrecognized metadata: {', '.join(unexpected_config)}") + if mismatched_config: + details.append( + "configuration differs: " + + ", ".join( + f"{key}={checkpoint_config[key]!r} (expected " + f"{expected_model_config[key]!r})" + for key in mismatched_config + ) + ) + raise ValueError( + "Checkpoint is not a compatible frozen-workspace baseline: " + + "; ".join(details) + ) + + checkpoint_state = checkpoint.get("rnn_state_dict") + if not isinstance(checkpoint_state, dict): + raise ValueError("Checkpoint does not contain an rnn_state_dict.") + target_state = rnn.state_dict() + missing_keys = set(target_state) - set(checkpoint_state) + unexpected_keys = set(checkpoint_state) - set(target_state) + + def is_workspace_key(key: str) -> bool: + return key == "workspace" or key.startswith("workspace.") or ".workspace." in key + + missing_workspace_keys = sorted(key for key in missing_keys if is_workspace_key(key)) + missing_non_workspace = sorted(missing_keys - set(missing_workspace_keys)) + if missing_non_workspace or unexpected_keys: + raise ValueError( + "Checkpoint state does not match the frozen-workspace model. " + f"Missing non-workspace keys={missing_non_workspace[:8]}, " + f"unexpected keys={sorted(unexpected_keys)[:8]}" + ) + shape_mismatches = [ + f"{key}: checkpoint{tuple(checkpoint_state[key].shape)} != model{tuple(target_state[key].shape)}" + for key in target_state.keys() & checkpoint_state.keys() + if torch.is_tensor(target_state[key]) + and torch.is_tensor(checkpoint_state[key]) + and target_state[key].shape != checkpoint_state[key].shape + ] + if shape_mismatches: + raise ValueError( + "Checkpoint has incompatible tensor shapes: " + "; ".join(shape_mismatches[:8]) + ) + + load_result = rnn.load_state_dict(checkpoint_state, strict=False) + if load_result.unexpected_keys or sorted(load_result.missing_keys) != missing_workspace_keys: + raise RuntimeError( + "Unexpected state mismatch while loading frozen-workspace baseline: " + f"missing={load_result.missing_keys}, unexpected={load_result.unexpected_keys}" + ) + output_proj.load_state_dict(checkpoint["output_proj_state_dict"]) + return missing_workspace_keys + + # ========================================== # TRAINING LOOP # ========================================== @@ -673,6 +785,7 @@ def run( dynamic_n_trigger_frac: float = DYNAMIC_N_TRIGGER_FRAC, dynamic_n_cooldown_steps: int = DYNAMIC_N_COOLDOWN_STEPS, moe_enabled: bool = MOE_ENABLED, + moe_diagnostics_enabled: bool = True, split_graph_moe_placement: str = "all", split_graph_moe_backbone_block_isolation: int | None = None, split_graph_moe_backbone_residual_scale: float = 1.0, @@ -733,7 +846,36 @@ def run( os.makedirs(LOG_DIR, exist_ok=True) log_path = os.path.join(LOG_DIR, f"run_{run_id}.csv") - resuming = resume_from is not None + resume_requested = resume_from is not None + resume_checkpoint = ( + load_checkpoint_for_resume(resume_from, device) if resume_requested else None + ) + checkpoint_config = ( + (resume_checkpoint.get("model_config") or {}) + if resume_checkpoint is not None + else {} + ) + checkpoint_state = ( + (resume_checkpoint.get("rnn_state_dict") or {}) + if resume_checkpoint is not None + else {} + ) + checkpoint_has_workspace = any( + key.startswith("workspace.") or ".workspace." in key + for key in checkpoint_state + ) + checkpoint_workspace_enabled = bool( + checkpoint_config.get("workspace_enabled", False) or checkpoint_has_workspace + ) + frozen_workspace_warm_start = bool( + resume_requested + and workspace_enabled + and workspace_train_mode == "frozen" + and not checkpoint_workspace_enabled + ) + # A baseline-to-frozen-workspace warm start starts a new experiment: its + # logs must be fresh, and it gets a fresh workspace-only optimizer. + resuming = resume_requested and not frozen_workspace_warm_start log_file = open(log_path, "a" if resuming else "w", newline="") log_writer = csv.writer(log_file) if not resuming: @@ -1091,6 +1233,12 @@ def run( ) + (f" | moe={moe_num_experts}e" if moe_enabled else "") ) + if moe_enabled: + print( + f"[{run_id}] MoE routing diagnostics=" + f"{'on' if moe_diagnostics_enabled else 'off'} | " + "load-balancing auxiliary loss and routing remain active" + ) if isolate_link_ablation: curriculum.lesson_nr_cells = [MODEL_NR_CELLS] * len(curriculum.table) @@ -1400,6 +1548,83 @@ def run( "workspace_reset_each_step": workspace_reset_each_step, "workspace_train_mode": workspace_train_mode, } + controller_model_config = {} + if controller == "mamba": + controller_model_config = { + "mamba_d_state": MAMBA_D_STATE, + "mamba_d_conv": MAMBA_D_CONV, + "mamba_expand": MAMBA_EXPAND, + } + elif controller == "mamba2": + controller_model_config = { + "mamba2_d_state": MAMBA2_D_STATE, + "mamba2_d_conv": MAMBA2_D_CONV, + "mamba2_expand": MAMBA2_EXPAND, + "mamba2_headdim": MAMBA2_HEADDIM, + "mamba2_ngroups": MAMBA2_NGROUPS, + } + elif controller == "mamba3": + controller_model_config = { + "mamba3_d_state": MAMBA3_D_STATE, + "mamba3_expand": MAMBA3_EXPAND, + "mamba3_headdim": MAMBA3_HEADDIM, + "mamba3_rope_fraction": MAMBA3_ROPE_FRACTION, + } + elif controller == "cfc": + controller_model_config = { + "cfc_mode": CFC_MODE, + "cfc_backbone_units": CFC_BACKBONE_UNITS, + "cfc_backbone_layers": CFC_BACKBONE_LAYERS, + "cfc_backbone_dropout": CFC_BACKBONE_DROPOUT, + "cfc_activation": CFC_ACTIVATION, + "cfc_mixed_memory": CFC_MIXED_MEMORY, + "cfc_residual": CFC_RESIDUAL, + } + elif "+" in controller: + controller_model_config = { + "mamba_d_state": MAMBA_D_STATE, + "mamba_d_conv": MAMBA_D_CONV, + "mamba_expand": MAMBA_EXPAND, + "mamba2_d_state": MAMBA2_D_STATE, + "mamba2_d_conv": MAMBA2_D_CONV, + "mamba2_expand": MAMBA2_EXPAND, + "mamba2_headdim": MAMBA2_HEADDIM, + "mamba2_ngroups": MAMBA2_NGROUPS, + "mamba3_d_state": MAMBA3_D_STATE, + "mamba3_expand": MAMBA3_EXPAND, + "mamba3_headdim": MAMBA3_HEADDIM, + "mamba3_rope_fraction": MAMBA3_ROPE_FRACTION, + "cfc_mode": CFC_MODE, + "cfc_backbone_units": CFC_BACKBONE_UNITS, + "cfc_backbone_layers": CFC_BACKBONE_LAYERS, + "cfc_backbone_dropout": CFC_BACKBONE_DROPOUT, + "cfc_activation": CFC_ACTIVATION, + "cfc_mixed_memory": CFC_MIXED_MEMORY, + "cfc_residual": CFC_RESIDUAL, + "hybrid_cfc_num_blocks": HYBRID_CFC_NUM_BLOCKS, + } + expected_resume_model_config = { + "input_size": INPUT_DIM, + "hidden_size": hidden_size, + "num_hidden_layers": getattr(rnn, "num_hidden_layers", None), + "nr_cells": rnn.memories[0].nr_cells, + "cell_size": cell_size, + "read_heads": read_heads, + "controller_type": controller, + "link_matrix_mode": link_matrix_mode, + "link_matrix_topk": link_matrix_topk, + "dynamic_n_mode": dynamic_n_mode, + "moe_enabled": moe_enabled, + "moe_num_experts": moe_num_experts, + "moe_expert_dim": moe_expert_dim, + "moe_capacity_factor": moe_capacity_factor, + "moe_load_balance_alpha": moe_load_balance_alpha, + "moe_top_k": moe_top_k, + "split_graph_enabled": split_graph_enabled, + "split_graph_variant": split_graph_variant, + **controller_model_config, + **extra_model_config, + } if link_matrix_mode != "dense": patch_link_matrix(rnn, mode=link_matrix_mode, topk=link_matrix_topk) @@ -1417,6 +1642,9 @@ def run( for module_name, module in rnn.named_modules(): if id(module) in moe_layer_ids: module.diagnostic_name = module_name or type(module).__name__ + for module in rnn.modules(): + if isinstance(module, SwitchMoE): + module.set_diagnostics_enabled(moe_diagnostics_enabled) def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: segment_diagnostics.emit(channel, include_gradients=include_gradients) @@ -1427,7 +1655,7 @@ def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) - f"[dnc-numerics:{channel}] " + " ".join(f"{name}={value:.4g}" for name, value in memory_diag.items()) ) - if moe_enabled: + if moe_enabled and moe_diagnostics_enabled: _, failure_moe_diag = pop_total_moe_aux_loss( rnn.moe_layers, include_diagnostics=True ) @@ -1568,8 +1796,10 @@ def set_lr(step_now): 10**9 ) # idle at run start (so the dip ramp is already at 1.0 immediately); # reset to `step` on every lesson advance instead - if resuming: - ckpt = load_checkpoint_for_resume(resume_from, device) + if resume_requested: + ckpt = resume_checkpoint + if ckpt is None: + raise RuntimeError("resume_requested was set without a loaded checkpoint") # Static Option 2 resume fix: rnn was just constructed above at # LESSON_NR_CELLS[0] (curriculum.lesson doesn't exist yet at that # point in run()). If this checkpoint is mid-schedule, the live @@ -1595,16 +1825,57 @@ def set_lr(step_now): ) resize_memory(rnn, ckpt_nr_cells, device=device, optimizer=optimizer) - rnn.load_state_dict(ckpt["rnn_state_dict"]) - output_proj.load_state_dict(ckpt["output_proj_state_dict"]) - optimizer.load_state_dict(ckpt["optimizer_state_dict"]) + if frozen_workspace_warm_start: + if ckpt.get("beta_target") != beta_target: + raise ValueError( + "Frozen-workspace warm start requires the baseline checkpoint " + f"beta ({ckpt.get('beta_target')}) to match this run ({beta_target})." + ) + missing_workspace_keys = load_frozen_workspace_warm_start( + rnn, + output_proj, + ckpt, + expected_resume_model_config, + ) + # The baseline was saved at a later curriculum lesson, where + # memory has more cells. Start the workspace experiment at lesson + # 1 with that lesson's configured memory size and a fresh schedule. + resize_memory(rnn, starting_nr_cells, device=device, optimizer=optimizer) + # Memory resize constructs a replacement module with fresh + # Parameter objects; reapply the frozen policy after transplant. + for parameter in rnn.parameters(): + parameter.requires_grad_(False) + assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None + for parameter in rnn.workspace.parameters(): + parameter.requires_grad_(True) + # Model construction and checkpoint loading consume RNG state. + # Start the new workspace experiment from the requested seed. + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + curriculum.lesson = 0 + step = 0 + print( + f"[{run_id}] Warm-started frozen workspace from {resume_from} " + f"(baseline step {ckpt['step']}, lesson " + f"{ckpt.get('curriculum_lesson', 0) + 1}); loaded matching " + f"backbone/combiner/MoE/DNC/task-head weights, kept " + f"{len(missing_workspace_keys)} workspace tensors freshly initialized, " + f"reset optimizer, scaler, RNG, and curriculum to step 0 / lesson 1." + ) + else: + rnn.load_state_dict(ckpt["rnn_state_dict"]) + output_proj.load_state_dict(ckpt["output_proj_state_dict"]) + optimizer.load_state_dict(ckpt["optimizer_state_dict"]) # v3 fix: restore GradScaler state (scale, growth tracker) instead # of leaving `scaler` at its freshly-constructed init_scale -- see # header note. `.get(...)` instead of ["..."] so a pre-v3 checkpoint # (saved before this key existed) still loads, just without this fix. - if "scaler_state_dict" in ckpt: + if not frozen_workspace_warm_start and "scaler_state_dict" in ckpt: scaler.load_state_dict(ckpt["scaler_state_dict"]) - else: + elif not frozen_workspace_warm_start: print( f"[{run_id}] WARNING: checkpoint predates GradScaler-state " f"saving (v3 fix) -- scaler starting fresh at " @@ -1618,7 +1889,8 @@ def set_lr(step_now): # and also means old checkpoints' scheduler_state_dict (which may # already reflect the buggy drifted state) is simply ignored rather # than needing migration. - curriculum.lesson = ckpt["curriculum_lesson"] + if not frozen_workspace_warm_start: + curriculum.lesson = ckpt["curriculum_lesson"] # Dynamic-N (macro-scale): restore the usage-EMA/cooldown/growth # history, so a resumed run doesn't silently reset the trigger's # sustained-saturation window back to 0 (which would just delay the @@ -1628,7 +1900,7 @@ def set_lr(step_now): # saved before dynamic_n_mode existed, or a checkpoint from a # non-dynamic-N run, still loads -- the controller (if this run # even has one) just starts fresh, same as a brand-new run would. - if dynamic_n_ctrl is not None: + if dynamic_n_ctrl is not None and not frozen_workspace_warm_start: if ckpt.get("dynamic_n_state") is not None: dynamic_n_ctrl.load_state_dict(ckpt["dynamic_n_state"]) else: @@ -1649,9 +1921,9 @@ def set_lr(step_now): # instead of ["..."] so a pre-v5 (Phase 1) checkpoint, which has no # prior state to restore at all, still loads -- just starting the # learned prior fresh at N(0,I), same as a brand-new run would. - if "prior_state" in ckpt: + if not frozen_workspace_warm_start and "prior_state" in ckpt: load_prior_state(stochastic_heads, ckpt["prior_state"]) - else: + elif not frozen_workspace_warm_start: print( f"[{run_id}] WARNING: checkpoint predates the learned-prior " f"snapshot (v5/Phase 2) - prior starting fresh at N(0,I) " @@ -1664,18 +1936,19 @@ def set_lr(step_now): # docstring). `.get(...)` so a pre-v6 checkpoint still loads -- # ood_rng just starts fresh in that case, identical to today's # behavior before this fix. - if "ood_rng_state" in ckpt: + if not frozen_workspace_warm_start and "ood_rng_state" in ckpt: ood_rng.setstate(ckpt["ood_rng_state"]) - else: + elif not frozen_workspace_warm_start: print( f"[{run_id}] WARNING: checkpoint predates ood_rng-state " f"saving (v6 fix) -- OOD sampling restarting from its " f"fresh, seed-derived state instead of continuing the " f"pre-resume leg's exact walk sequence." ) - restore_rng_state(ckpt["rng_state"]) - step = ckpt["step"] - if ckpt["beta_target"] != beta_target: + if not frozen_workspace_warm_start: + restore_rng_state(ckpt["rng_state"]) + step = ckpt["step"] + if not frozen_workspace_warm_start and ckpt["beta_target"] != beta_target: anneal_start_step = step # Phase 1: restart the ramp here print( f"[{run_id}] beta_target changed on resume " @@ -1683,11 +1956,12 @@ def set_lr(step_now): f"restarted from step {anneal_start_step}, ramping over " f"the next {KL_ANNEAL_STEPS} steps." ) - print( - f"[{run_id}] Resumed from {resume_from} at step {step} " - f"(lesson {curriculum.lesson + 1}/{len(curriculum.table)}); " - f"continuing to total_steps={total_steps}" - ) + if not frozen_workspace_warm_start: + print( + f"[{run_id}] Resumed from {resume_from} at step {step} " + f"(lesson {curriculum.lesson + 1}/{len(curriculum.table)}); " + f"continuing to total_steps={total_steps}" + ) if step >= total_steps: print( f"[{run_id}] Checkpoint step {step} already >= total_steps " @@ -1697,7 +1971,7 @@ def set_lr(step_now): print( f"\n=== [{run_id}] Training (beta_target={beta_target}, beta_mode={beta_mode}) " - f"{'[resumed]' if resuming else ''} ===" + f"{'[warm-started]' if frozen_workspace_warm_start else '[resumed]' if resuming else ''} ===" ) rnn.train() set_lr(step) # Phase 1 FIX: ensure correct clamped LR from the very first @@ -1803,7 +2077,9 @@ def set_lr(step_now): moe_aux_loss, moe_diag = ( pop_total_moe_aux_loss( rnn.moe_layers, - include_diagnostics=(step + 1) % LOG_EVERY == 0, + include_diagnostics=( + moe_diagnostics_enabled and (step + 1) % LOG_EVERY == 0 + ), ) if moe_enabled else (torch.zeros((), device=device), {}) @@ -2005,9 +2281,9 @@ def set_lr(step_now): if amp_enabled: parts.append(f"amp_scale {amp_scale:.1f}") parts.append(f"gpu_mem_peak_mb {gpu_mem_peak_mb:.1f}") - if moe_enabled: + if moe_enabled and moe_diagnostics_enabled: + parts.append(f"moe_aux {float(moe_aux_loss.detach()):.4f}") parts += [ - f"moe_aux {float(moe_aux_loss.detach()):.4f}", f"moe_cv_load {moe_diag.get('moe_cv_load', 0.0):.4f}", f"moe_max_load_frac {moe_diag.get('moe_max_load_frac', 0.0):.4f}", f"moe_capacity_drop {moe_diag.get('moe_capacity_drop_frac', 0.0):.1%}", @@ -2868,9 +3144,10 @@ def _write_diagnostic_values(diagnostics): "--resume", type=str, default=None, - help="path to a checkpoint .pt to resume from (requires --beta / positional beta " - "to also be given, so run_id/beta_target can be matched against the checkpoint). " - "Training continues from the checkpoint's step up to TOTAL_STEPS.", + help="path to a checkpoint .pt (requires --beta / positional beta). Normally " + "restores model, optimizer, RNG, curriculum, and step. With --workspace " + "--workspace-train-mode=frozen, a compatible workspace-disabled baseline " + "instead warm-starts matching weights and begins a fresh workspace training run.", ) parser.add_argument( "--run-id-suffix", @@ -3018,6 +3295,14 @@ def _write_diagnostic_values(diagnostics): "SwitchMoE blocks between Mamba controller blocks (mamba_controller.py " "/ moe_layer.py).", ) + parser.add_argument( + "--moe-diagnostics", + action=argparse.BooleanOptionalAction, + default=True, + help="Collect and print detailed MoE routing/activation diagnostics " + "(enabled by default; use --no-moe-diagnostics to reduce diagnostic " + "overhead). Expert routing and the load-balancing auxiliary loss stay active.", + ) parser.add_argument( "--moe-placement", choices=("all", "backbone", "combiner"), @@ -3398,6 +3683,8 @@ def _write_diagnostic_values(diagnostics): run_id = f"{run_id}_dynN" if args.moe: run_id = f"{run_id}_moe{args.moe_num_experts}e" + if not args.moe_diagnostics: + run_id += "_moediagoff" if args.split_graph and args.moe_placement != "all": run_id += f"_{args.moe_placement}only" if args.moe_backbone_block_isolation is not None: @@ -3447,6 +3734,7 @@ def _write_diagnostic_values(diagnostics): dynamic_n_trigger_frac=args.dynamic_n_trigger_frac, dynamic_n_cooldown_steps=args.dynamic_n_cooldown_steps, moe_enabled=args.moe, + moe_diagnostics_enabled=args.moe_diagnostics, split_graph_moe_placement=args.moe_placement, split_graph_moe_backbone_block_isolation=args.moe_backbone_block_isolation, split_graph_moe_backbone_residual_scale=args.moe_backbone_residual_scale, From 132e8a1b70fe5363b2540288fa0eade8e064494e Mon Sep 17 00:00:00 2001 From: Samuel Date: Sun, 11 Oct 2026 19:47:31 +0200 Subject: [PATCH 08/14] Add optional segment/workspace diagnostics Adds a segment-diagnostics flag to the training CLI and passes it into SegmentDiagnostics, which now skips hook registration when disabled. WorkspaceBroadcast also gains a diagnostics toggle to avoid expensive per-timestep reductions, and the training entrypoint reports whether MoE and segment diagnostics are enabled. This keeps long recurrent runs lighter while preserving the optional trace output when needed. --- src/initium/core_training.py | 22 +++++- src/initium/diagnostics/segment_logging.py | 9 ++- src/initium/workspace/broadcast.py | 84 +++++++++++++--------- 3 files changed, 79 insertions(+), 36 deletions(-) diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 59170b2..e0d2a0d 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -786,6 +786,7 @@ def run( dynamic_n_cooldown_steps: int = DYNAMIC_N_COOLDOWN_STEPS, moe_enabled: bool = MOE_ENABLED, moe_diagnostics_enabled: bool = True, + segment_diagnostics_enabled: bool = True, split_graph_moe_placement: str = "all", split_graph_moe_backbone_block_isolation: int | None = None, split_graph_moe_backbone_residual_scale: float = 1.0, @@ -1233,6 +1234,10 @@ def run( ) + (f" | moe={moe_num_experts}e" if moe_enabled else "") ) + print( + f"[{run_id}] diagnostics: moe={'on' if moe_diagnostics_enabled else 'off'} " + f"segment={'on' if segment_diagnostics_enabled else 'off'}" + ) if moe_enabled: print( f"[{run_id}] MoE routing diagnostics=" @@ -1636,7 +1641,9 @@ def run( output_proj = nn.Linear(INPUT_DIM, TRIPLE_DIM).to(device) dataset.set_output_proj(output_proj) - segment_diagnostics = SegmentDiagnostics(rnn) + segment_diagnostics = SegmentDiagnostics( + rnn, enabled=segment_diagnostics_enabled + ) segment_diagnostics.add_module(output_proj, "task_output") moe_layer_ids = {id(layer) for layer in getattr(rnn, "moe_layers", [])} for module_name, module in rnn.named_modules(): @@ -1645,6 +1652,8 @@ def run( for module in rnn.modules(): if isinstance(module, SwitchMoE): module.set_diagnostics_enabled(moe_diagnostics_enabled) + if isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None: + rnn.workspace.set_diagnostics_enabled(moe_diagnostics_enabled) def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: segment_diagnostics.emit(channel, include_gradients=include_gradients) @@ -3301,7 +3310,15 @@ def _write_diagnostic_values(diagnostics): default=True, help="Collect and print detailed MoE routing/activation diagnostics " "(enabled by default; use --no-moe-diagnostics to reduce diagnostic " - "overhead). Expert routing and the load-balancing auxiliary loss stay active.", + "overhead, including workspace activity summaries). Expert routing, " + "workspace computation, and the load-balancing auxiliary loss stay active.", + ) + parser.add_argument( + "--segment-diagnostics", + action=argparse.BooleanOptionalAction, + default=True, + help="Collect activation/gradient summaries for model segments (enabled by default; " + "use --no-segment-diagnostics to avoid per-module hooks in long recurrent sequences).", ) parser.add_argument( "--moe-placement", @@ -3735,6 +3752,7 @@ def _write_diagnostic_values(diagnostics): dynamic_n_cooldown_steps=args.dynamic_n_cooldown_steps, moe_enabled=args.moe, moe_diagnostics_enabled=args.moe_diagnostics, + segment_diagnostics_enabled=args.segment_diagnostics, split_graph_moe_placement=args.moe_placement, split_graph_moe_backbone_block_isolation=args.moe_backbone_block_isolation, split_graph_moe_backbone_residual_scale=args.moe_backbone_residual_scale, diff --git a/src/initium/diagnostics/segment_logging.py b/src/initium/diagnostics/segment_logging.py index a902b31..78a308d 100644 --- a/src/initium/diagnostics/segment_logging.py +++ b/src/initium/diagnostics/segment_logging.py @@ -80,13 +80,16 @@ class SegmentDiagnostics: boundary or when a numerical failure needs an immediate trace. """ - def __init__(self, model: nn.Module): + def __init__(self, model: nn.Module, *, enabled: bool = True): self.model = model + self.enabled = bool(enabled) self.channel: str | None = "train" self._records: dict[str, OrderedDict[str, dict[str, Any]]] = {} self._segments: list[tuple[str, str]] = [] self._external_modules: list[tuple[str, nn.Module]] = [] self._handles = [] + if not self.enabled: + return for path, module in model.named_modules(): label = _segment_label(path, module) if label is not None: @@ -140,6 +143,8 @@ def _accumulate_tensor(record: dict[str, Any], prefix: str, tensor: torch.Tensor def add_module(self, module: nn.Module, label: str) -> None: """Add a caller-named module such as the task output projection.""" + if not self.enabled: + return self._register(module, label, label) self._external_modules.append((label, module)) @@ -180,6 +185,8 @@ def _gradient_records(self) -> OrderedDict[str, tuple[torch.Tensor, int]]: return result def emit(self, channel: str, *, include_gradients: bool = False) -> None: + if not self.enabled: + return records = self._records.pop(channel, None) gradient_records = self._gradient_records() if include_gradients else OrderedDict() if not records and not gradient_records: diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index 99df1d3..8df9e50 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -36,6 +36,7 @@ def __init__(self, d_model: int, num_experts: int, num_slots: int = 5): self.write_gate = nn.Linear(2 * d_model, d_model) self.read_gate = nn.Parameter(torch.zeros(())) self.write_gate_closed = False + self.diagnostics_enabled = True self._last_step_diag: dict[str, torch.Tensor] = {} self._last_diag: dict[str, torch.Tensor] = {} nn.init.zeros_(self.write_gate.weight) @@ -99,15 +100,16 @@ def _forward_step( selected = specialist_deltas.gather( 1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model) ) - zero = specialist_deltas.new_zeros((), dtype=torch.float32) - self._last_step_diag = { - "write_mass": zero, - "active_write_frac": zero, - "write_gate_mean": zero, - "slot_rms": zero, - "broadcast_rms": zero, - "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), - } + if self.diagnostics_enabled: + zero = specialist_deltas.new_zeros((), dtype=torch.float32) + self._last_step_diag = { + "write_mass": zero, + "active_write_frac": zero, + "write_gate_mean": zero, + "slot_rms": zero, + "broadcast_rms": zero, + "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + } return (selected * route_weights.unsqueeze(-1)).sum(dim=1), state batch, experts, width = specialist_deltas.shape @@ -118,24 +120,32 @@ def _forward_step( if reset_each_step: state = torch.zeros_like(state) + selected_deltas = specialist_deltas.gather( + 1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model) + ) if no_selection: + # The random-write control intentionally samples across every + # specialist, so it must keep the full expert dimension here. write_weights = torch.rand( batch, experts, device=specialist_deltas.device, dtype=torch.float32 ) write_weights = write_weights / write_weights.sum(-1, keepdim=True).clamp_min(1e-8) + write_candidates = specialist_deltas else: - write_weights = torch.zeros( - batch, experts, device=specialist_deltas.device, dtype=torch.float32 - ) - write_weights.scatter_add_(1, route_indices, route_weights.float()) + # Only routed experts can contribute to the weighted write. + # Keeping the top-k axis avoids projecting and scoring all E + # specialists on every sequential timestep. Duplicate routes, + # if supplied by a caller, retain their original additive effect. + write_weights = route_weights.float() + write_candidates = selected_deltas # Candidate keys compete over slots using content-based attention; # slot writes are weighted by the existing sparse MoE selection. - keys = self.key(specialist_deltas) + keys = self.key(write_candidates) scores = torch.einsum("bed,bsd->bes", keys.float(), self.query(state).float()) scores = scores / math.sqrt(self.d_model) slot_attention = scores.softmax(dim=-1).to(specialist_deltas.dtype) - values = self.value(specialist_deltas) + values = self.value(write_candidates) slot_proposal = torch.einsum( "be,bes,bed->bsd", write_weights.to(values.dtype), slot_attention, values ) @@ -148,9 +158,9 @@ def _forward_step( update = update * has_writer.to(update.dtype) next_state = update.mul(slot_inputs).add((1.0 - update).mul(old_inputs)) - # All experts query the shared object, including those not selected - # for this step's write. - read_queries = self.query(specialist_deltas) + # Only routed experts are combined into the backbone residual, so + # computing readback for unselected experts is dead work. + read_queries = self.query(selected_deltas) read_scores = torch.einsum( "bed,bsd->bes", read_queries.float(), self.key(next_state).float() ) @@ -159,19 +169,19 @@ def _forward_step( read_values = self.value(next_state) broadcast = torch.einsum("bes,bsd->bed", read_weights, read_values) read_scale = torch.sigmoid(self.read_gate).to(specialist_deltas.dtype) - enriched = specialist_deltas + read_scale * broadcast - selected = enriched.gather(1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model)) - self._last_step_diag = { - "write_mass": write_weights.detach().sum(dim=-1).mean(), - "active_write_frac": (write_weights.detach().sum(dim=-1) > 0).float().mean(), - "write_gate_mean": update.detach().float().mean(), - "slot_rms": next_state.detach().float().square().mean().sqrt(), - "broadcast_rms": ( - read_scale.detach().float() * broadcast.detach().float() - ).square().mean().sqrt(), - "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), - } - return (selected * route_weights.unsqueeze(-1)).sum(dim=1), next_state + enriched = selected_deltas + read_scale * broadcast + if self.diagnostics_enabled: + self._last_step_diag = { + "write_mass": write_weights.detach().sum(dim=-1).mean(), + "active_write_frac": (write_weights.detach().sum(dim=-1) > 0).float().mean(), + "write_gate_mean": update.detach().float().mean(), + "slot_rms": next_state.detach().float().square().mean().sqrt(), + "broadcast_rms": ( + read_scale.detach().float() * broadcast.detach().float() + ).square().mean().sqrt(), + "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + } + return (enriched * route_weights.unsqueeze(-1)).sum(dim=1), next_state def forward( self, @@ -185,7 +195,7 @@ def forward( ) -> tuple[torch.Tensor, torch.Tensor]: """Process time-major specialists ``(B,T,E,D)`` sequentially.""" outputs = [] - step_diags = [] + step_diags = [] if self.diagnostics_enabled else None for step in range(specialist_deltas.shape[1]): out, state = self.forward_step( specialist_deltas[:, step], route_indices[:, step], route_weights[:, step], @@ -194,7 +204,8 @@ def forward( reset_each_step=reset_each_step, ) outputs.append(out) - step_diags.append(self._last_step_diag) + if step_diags is not None: + step_diags.append(self._last_step_diag) if step_diags: self._last_diag = { key: torch.stack([diag[key] for diag in step_diags]).mean().detach() @@ -210,3 +221,10 @@ def forward( def last_diagnostics(self) -> dict[str, torch.Tensor]: """Detached activity metrics for the most recent sequence pass.""" return dict(self._last_diag) + + def set_diagnostics_enabled(self, enabled: bool) -> None: + """Enable or skip optional per-timestep workspace health reductions.""" + self.diagnostics_enabled = bool(enabled) + if not self.diagnostics_enabled: + self._last_step_diag = {} + self._last_diag = {} From 03954e194ca8ae5b7b4966502966aaa5e47f49af Mon Sep 17 00:00:00 2001 From: Samuel Date: Mon, 12 Oct 2026 22:41:51 +0200 Subject: [PATCH 09/14] Add eval-only checkpoint evaluation mode Add evaluation-only resume flows that load checkpoint weights without restoring optimizer state or training updates. The new flags support deterministic eval RNGs and workspace counterfactual modes (checkpoint/reset-each-step/bypass), while preventing lesson advancement during evaluation-only passes so checkpoint scoring stays comparable. --- src/initium/core_training.py | 106 ++++++++++++++++-- src/initium/data/common/chain_task.py | 7 +- .../data/graph_traversal/graph_traversal.py | 10 +- .../mamba_controller/split_graph_dnc.py | 5 + src/initium/workspace/broadcast.py | 22 ++++ 5 files changed, 136 insertions(+), 14 deletions(-) diff --git a/src/initium/core_training.py b/src/initium/core_training.py index e0d2a0d..8d5ad80 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -809,6 +809,9 @@ def run( workspace_no_selection: bool = WORKSPACE_NO_SELECTION, workspace_reset_each_step: bool = WORKSPACE_RESET_EACH_STEP, workspace_train_mode: str = WORKSPACE_TRAIN_MODE, + eval_only: bool = False, + eval_workspace_mode: str = "checkpoint", + eval_seed: int | None = None, dnc_read_variant: str = DNC_READ_VARIANT, dnc_read_residual_scale: float = DNC_READ_RESIDUAL_SCALE, dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, @@ -848,6 +851,14 @@ def run( os.makedirs(LOG_DIR, exist_ok=True) log_path = os.path.join(LOG_DIR, f"run_{run_id}.csv") resume_requested = resume_from is not None + if eval_only and not resume_requested: + raise ValueError("eval_only requires --resume with a checkpoint") + if eval_workspace_mode not in ("checkpoint", "reset-each-step", "bypass"): + raise ValueError( + "eval_workspace_mode must be 'checkpoint', 'reset-each-step', or 'bypass'" + ) + if eval_only and eval_workspace_mode != "checkpoint" and not workspace_enabled: + raise ValueError("workspace evaluation modes require --workspace") resume_checkpoint = ( load_checkpoint_for_resume(resume_from, device) if resume_requested else None ) @@ -874,9 +885,17 @@ def run( and workspace_train_mode == "frozen" and not checkpoint_workspace_enabled ) + if eval_only and frozen_workspace_warm_start: + raise ValueError( + "eval_only requires a checkpoint that already contains trained workspace weights; " + "a workspace-disabled baseline would initialize a fresh workspace." + ) # A baseline-to-frozen-workspace warm start starts a new experiment: its # logs must be fresh, and it gets a fresh workspace-only optimizer. - resuming = resume_requested and not frozen_workspace_warm_start + # Evaluation-only runs get independent, fresh logs keyed by their run ID; + # they load checkpoint weights/state but must not append rows to the + # original training logs. + resuming = resume_requested and not frozen_workspace_warm_start and not eval_only log_file = open(log_path, "a" if resuming else "w", newline="") log_writer = csv.writer(log_file) if not resuming: @@ -1877,7 +1896,8 @@ def set_lr(step_now): else: rnn.load_state_dict(ckpt["rnn_state_dict"]) output_proj.load_state_dict(ckpt["output_proj_state_dict"]) - optimizer.load_state_dict(ckpt["optimizer_state_dict"]) + if not eval_only: + optimizer.load_state_dict(ckpt["optimizer_state_dict"]) # v3 fix: restore GradScaler state (scale, growth tracker) instead # of leaving `scaler` at its freshly-constructed init_scale -- see # header note. `.get(...)` instead of ["..."] so a pre-v3 checkpoint @@ -1957,6 +1977,16 @@ def set_lr(step_now): if not frozen_workspace_warm_start: restore_rng_state(ckpt["rng_state"]) step = ckpt["step"] + if eval_only and eval_seed is not None: + # An explicit evaluation seed overrides checkpoint RNG state so + # separate checkpoints/modes can be compared on identical + # generated episodes. Omit it to replay the checkpoint's stream. + random.seed(eval_seed) + np.random.seed(eval_seed) + torch.manual_seed(eval_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(eval_seed) + ood_rng.seed(eval_seed) if not frozen_workspace_warm_start and ckpt["beta_target"] != beta_target: anneal_start_step = step # Phase 1: restart the ramp here print( @@ -1965,25 +1995,45 @@ def set_lr(step_now): f"restarted from step {anneal_start_step}, ramping over " f"the next {KL_ANNEAL_STEPS} steps." ) - if not frozen_workspace_warm_start: + if eval_only: + if not isinstance(rnn, SplitGraphDNC) or rnn.workspace is None: + if eval_workspace_mode != "checkpoint": + raise ValueError( + "workspace evaluation mode requires a checkpoint with an active workspace" + ) + else: + if eval_workspace_mode == "reset-each-step": + rnn.workspace_reset_each_step = True + elif eval_workspace_mode == "bypass": + rnn.workspace_eval_bypass = True + print( + f"[{run_id}] Evaluation-only checkpoint load | step={step} | " + f"workspace_mode={eval_workspace_mode} | " + f"eval_seed={eval_seed if eval_seed is not None else 'checkpoint RNG'} | " + "optimizer state not loaded" + ) + elif not frozen_workspace_warm_start: print( f"[{run_id}] Resumed from {resume_from} at step {step} " f"(lesson {curriculum.lesson + 1}/{len(curriculum.table)}); " f"continuing to total_steps={total_steps}" ) - if step >= total_steps: + if step >= total_steps and not eval_only: print( f"[{run_id}] Checkpoint step {step} already >= total_steps " f"{total_steps} -- nothing to do. Raise --total-steps if you " f"want to extend further." ) + run_mode_label = "Evaluation-only" if eval_only else "Training" print( - f"\n=== [{run_id}] Training (beta_target={beta_target}, beta_mode={beta_mode}) " + f"\n=== [{run_id}] {run_mode_label} (beta_target={beta_target}, beta_mode={beta_mode}) " f"{'[warm-started]' if frozen_workspace_warm_start else '[resumed]' if resuming else ''} ===" ) - rnn.train() - set_lr(step) # Phase 1 FIX: ensure correct clamped LR from the very first + rnn.train(not eval_only) + output_proj.train(not eval_only) + if not eval_only: + set_lr(step) # Phase 1 FIX: ensure correct clamped LR from the very first # iteration -- otherwise a resumed run's first step would briefly use # whatever LR was saved inside optimizer_state_dict at checkpoint time # (potentially still reflecting the old buggy drifted value). @@ -2010,7 +2060,7 @@ def set_lr(step_now): # legs already have their own timestamped log rows if that's needed). t_run_start = time.time() - while step < total_steps: + while step < total_steps and not eval_only: input_seq, target_digits, answer_mask = dataset.sample_batch(curriculum, BATCH_SIZE) input_seq = input_seq.to(device, non_blocking=True) @@ -2842,13 +2892,17 @@ def set_lr(step_now): latest_path = os.path.join(CHECKPOINT_DIR, f"{run_id}_latest.pt") shutil.copyfile(ckpt_path, latest_path) - print(f"\n[{run_id}] Training complete. Final evaluation on training-distribution lesson:") + completion_label = "Evaluation-only pass" if eval_only else "Training complete" + print( + f"\n[{run_id}] {completion_label}. Final evaluation on " + "training-distribution lesson:" + ) final_id_stage_memory = rnn.memories[0] if hasattr(final_id_stage_memory, "begin_stage_diagnostics"): final_id_stage_memory.begin_stage_diagnostics() segment_diagnostics.set_channel("eval-final-id") _, id_triple_acc, id_perfect_frac = curriculum.maybe_advance( - rnn, device, step=step, optimizer=optimizer + rnn, device, step=step, optimizer=optimizer, allow_advance=not eval_only ) segment_diagnostics.emit("eval-final-id") final_id_stage_diagnostics = ( @@ -3158,6 +3212,23 @@ def _write_diagnostic_values(diagnostics): "--workspace-train-mode=frozen, a compatible workspace-disabled baseline " "instead warm-starts matching weights and begins a fresh workspace training run.", ) + parser.add_argument( + "--eval-only", + action="store_true", + help="Load --resume checkpoint and run only the final ID/OOD evaluation; no training updates or optimizer-state restore.", + ) + parser.add_argument( + "--eval-workspace-mode", + choices=("checkpoint", "reset-each-step", "bypass"), + default="checkpoint", + help="Evaluation-only workspace counterfactual: preserve checkpoint behavior, clear slots before every timestep, or bypass workspace readback while retaining the selected MoE residual.", + ) + parser.add_argument( + "--eval-seed", + type=int, + default=None, + help="Evaluation-only RNG seed. Use the same value across checkpoints to evaluate identical generated ID/OOD episodes; omit to replay checkpoint RNG state.", + ) parser.add_argument( "--run-id-suffix", type=str, @@ -3620,6 +3691,18 @@ def _write_diagnostic_values(diagnostics): "--resume requires the beta positional arg too, e.g.:\n" f" python3 {sys.argv[0]} 0.02 --resume checkpoints/beta_0p02_latest.pt" ) + if args.eval_only and args.resume is None: + parser.error("--eval-only requires --resume ") + if not args.eval_only and args.eval_workspace_mode != "checkpoint": + parser.error("--eval-workspace-mode is only valid with --eval-only") + if args.eval_seed is not None and not args.eval_only: + parser.error("--eval-seed is only valid with --eval-only") + if args.eval_only and args.eval_workspace_mode != "checkpoint" and not args.workspace: + parser.error("non-default --eval-workspace-mode requires --workspace") + if args.eval_only and args.run_id_suffix is None: + args.run_id_suffix = f"eval-{args.eval_workspace_mode}" + if args.eval_seed is not None: + args.run_id_suffix += f"-seed{args.eval_seed}" if args.link_matrix_mode == "sparse_topk" and args.link_matrix_topk is None: raise SystemExit("--link-matrix-mode=sparse_topk requires --link-matrix-topk") @@ -3774,6 +3857,9 @@ def _write_diagnostic_values(diagnostics): workspace_no_selection=args.workspace_no_selection, workspace_reset_each_step=args.workspace_reset_each_step, workspace_train_mode=args.workspace_train_mode, + eval_only=args.eval_only, + eval_workspace_mode=args.eval_workspace_mode, + eval_seed=args.eval_seed, dnc_read_variant=args.dnc_read_variant, dnc_read_residual_scale=args.dnc_read_residual_scale, dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, diff --git a/src/initium/data/common/chain_task.py b/src/initium/data/common/chain_task.py index 61b6aab..205c727 100644 --- a/src/initium/data/common/chain_task.py +++ b/src/initium/data/common/chain_task.py @@ -142,7 +142,9 @@ def __init__(self, dataset, table, lesson_nr_cells): self.lesson = 0 self.lesson_nr_cells = list(lesson_nr_cells) - def maybe_advance(self, model, device, step=None, optimizer=None): + def maybe_advance( + self, model, device, step=None, optimizer=None, *, allow_advance=True + ): nf, nq = self.table[self.lesson] combined_acc, perfect_frac, breakdown = self.dataset.evaluate_chain( model, @@ -174,7 +176,8 @@ def maybe_advance(self, model, device, step=None, optimizer=None): ) if ( - combined_acc / 100.0 >= self.dataset.advance_threshold + allow_advance + and combined_acc / 100.0 >= self.dataset.advance_threshold and self.lesson < len(self.table) - 1 ): self.lesson += 1 diff --git a/src/initium/data/graph_traversal/graph_traversal.py b/src/initium/data/graph_traversal/graph_traversal.py index d4d5a8f..1d64fc6 100644 --- a/src/initium/data/graph_traversal/graph_traversal.py +++ b/src/initium/data/graph_traversal/graph_traversal.py @@ -333,7 +333,9 @@ def sample_episode(self): ep = build_traversal_episode(num_nodes, out_degree_range, path_len_range) return ep - def maybe_advance(self, model, device, step=None, optimizer=None): + def maybe_advance( + self, model, device, step=None, optimizer=None, *, allow_advance=True + ): """Returns (lesson, id_triple_acc, id_perfect_frac). Eval samples num_nodes over the lesson's nodes_range (as trained), or walks the fixed custom_graph when one was supplied; the gate advances on @@ -376,7 +378,11 @@ def maybe_advance(self, model, device, step=None, optimizer=None): field_writer, self.field_log_file, step, self.lesson + 1, "id", field_log ) - if triple_acc / 100.0 >= ADVANCE_THRESHOLD and self.lesson < len(self.table) - 1: + if ( + allow_advance + and triple_acc / 100.0 >= ADVANCE_THRESHOLD + and self.lesson < len(self.table) - 1 + ): self.lesson += 1 new_n = self.lesson_nr_cells[self.lesson] if new_n != self.lesson_nr_cells[self.lesson - 1]: diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index feb4cdf..67f75b2 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -260,6 +260,10 @@ def __init__( self.workspace_no_selection = workspace_no_selection self.workspace_reset_each_step = workspace_reset_each_step self.workspace_train_mode = workspace_train_mode + # Runtime-only counterfactual for evaluation. It is deliberately not + # part of checkpoint config/state so a saved model can be scored with + # and without workspace readback using identical learned weights. + self.workspace_eval_bypass = False # Mamba bookkeeping, for checkpoint self-description -- same # rationale/convention as MambaDNC (mamba_controller.py). @@ -629,6 +633,7 @@ def forward( workspace_state, no_selection=self.workspace_no_selection, reset_each_step=self.workspace_reset_each_step, + bypass_broadcast=self.workspace_eval_bypass, ) H = base.clone() # The workspace returns the selected, broadcast-enriched MoE diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index 8df9e50..e7ca473 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -54,6 +54,7 @@ def forward_step( *, no_selection: bool = False, reset_each_step: bool = False, + bypass_broadcast: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # Keep the zero-slot/closed-gate bypass identical to the selection-only # MoE path. Active workspace attention and recurrent slot updates use @@ -158,6 +159,25 @@ def _forward_step( update = update * has_writer.to(update.dtype) next_state = update.mul(slot_inputs).add((1.0 - update).mul(old_inputs)) + if bypass_broadcast: + # Evaluation counterfactual: preserve the selected MoE residual + # while preventing workspace readback from influencing it. The + # workspace write path still executes, so this isolates the + # broadcast contribution without changing routing or the model + # parameters loaded from the checkpoint. + if self.diagnostics_enabled: + self._last_step_diag = { + "write_mass": write_weights.detach().sum(dim=-1).mean(), + "active_write_frac": (write_weights.detach().sum(dim=-1) > 0) + .float() + .mean(), + "write_gate_mean": update.detach().float().mean(), + "slot_rms": next_state.detach().float().square().mean().sqrt(), + "broadcast_rms": next_state.new_zeros((), dtype=torch.float32), + "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + } + return (selected_deltas * route_weights.unsqueeze(-1)).sum(dim=1), next_state + # Only routed experts are combined into the backbone residual, so # computing readback for unselected experts is dead work. read_queries = self.query(selected_deltas) @@ -192,6 +212,7 @@ def forward( *, no_selection: bool = False, reset_each_step: bool = False, + bypass_broadcast: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Process time-major specialists ``(B,T,E,D)`` sequentially.""" outputs = [] @@ -202,6 +223,7 @@ def forward( state, no_selection=no_selection, reset_each_step=reset_each_step, + bypass_broadcast=bypass_broadcast, ) outputs.append(out) if step_diags is not None: From 0b42e7bc4a6e2e404ca64ed0b3e405bfd6635497 Mon Sep 17 00:00:00 2001 From: Samuel Date: Tue, 13 Oct 2026 19:47:57 +0200 Subject: [PATCH 10/14] Add experimental bistable workspace gate Introduce an opt-in recurrent stochastic write gate with checkpointed noise, gate-only warm starts, stability checks, and training/evaluation diagnostics. Document Phase 4 usage and add unit tests; preserve the model's prior training mode during traversal evaluation. --- README.md | 16 + src/initium/core_training.py | 322 +++++++++++++++++- .../data/graph_traversal/graph_traversal.py | 3 +- src/initium/diagnostics/segment_logging.py | 30 +- src/initium/inference/model_loader.py | 11 +- .../mamba_controller/split_graph_dnc.py | 18 +- src/initium/workspace/bistability.py | 182 ++++++++++ src/initium/workspace/broadcast.py | 174 +++++++++- tests/unit/test_bistable_workspace_gate.py | 105 ++++++ 9 files changed, 845 insertions(+), 16 deletions(-) create mode 100644 src/initium/workspace/bistability.py create mode 100644 tests/unit/test_bistable_workspace_gate.py diff --git a/README.md b/README.md index ab5f9ba..76cb799 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,22 @@ The console prints a startup line that reports whether the workspace is active o The shared object operates on the final backbone MoE bank; earlier backbone MoE layers retain their existing behavior. This is the Phase 0 plumbing target, not evidence that Broadcast improves task performance. The first short run should compare the zero-slot bypass against the same MoE configuration without `--workspace`; only after this plumbing gate should the frozen/co-adapted and no-Selection comparisons begin. Record seed, full command, commit, hardware/software environment, curriculum completion, ID/OOD metrics, convergence steps, and any failures in `07-Experiment/Broadcasting/Progress/Phase-0-Implementation.md`. All current evidence is graph-traversal-specific; transfer across task families remains untested. +### Phase 4 experimental bistable write gate + +The default `--workspace-gate-variant continuous` preserves the Phase 0–3 sigmoid update. The opt-in `bistable` variant treats each workspace slot as a competing write candidate. It combines the workspace's content-selection score and the MoE route confidence with independent Gaussian background noise, then iterates recurrent self-excitation and lateral inhibition before making a hard write/no-write decision. A straight-through estimator carries gradients through the pre-commit probabilities. This is an experimental adaptation of GNW ignition, not a mechanism implemented by the source neuroscience papers. + +For the matched gate-only comparison, resume the same five-slot frozen-workspace checkpoint twice. The frozen control updates only the existing continuous `write_gate`; the bistable condition updates only the new recurrent gate parameters. The gate-variant warm start loads shared model/workspace weights, initializes only the new gate, and starts with a fresh optimizer. `--workspace-train-mode coadapted` remains available to train the full model from that same checkpoint as a separate condition. Bistable background noise uses an independent checkpointed RNG stream, so zero-noise and selection-strength ablations do not change the graph episode RNG stream. + +```bash +python -m src.initium.core_training 0.0 --seed 0 --resume --run-id-suffix phase4-continuous-gate-pilot --split-graph --split-graph-variant mamba1 --split-graph-num-blocks 2 --split-graph-combiner-mode controller --split-graph-combiner-variant cfc --moe --moe-placement all --moe-backbone-residual-scale 0.1 --moe-num-experts 8 --moe-top-k 1 --moe-capacity-factor 1.5 --moe-load-balance-alpha 0.01 --workspace --workspace-num-slots 5 --workspace-train-mode frozen --workspace-train-scope gate --workspace-gate-variant continuous --total-steps 8100 +``` + +Use the same command, checkpoint, seed, and step budget for the bistable condition, changing the suffix and adding `--workspace-gate-variant bistable`. Gate-only runs use a safety tripwire that aborts before backward/update if any instrumented segment exceeds `--workspace-stability-max-activation` (default `100000`). Segment diagnostics report activation/gradient RMS and maxima; bistable diagnostics additionally report the pre-commit probability histogram (10 bins), low/middle/high occupancy, hard commit frequency, writer frequencies, slot RMS, and readback scale during training and held-out ID/OOD evaluation. A zero-noise ablation uses `--workspace-bistable-noise-scale 0`; selection-strength sweeps use `--workspace-bistable-selection-scale` while holding `--eval-seed` fixed. The gate still samples noise during evaluation unless `--workspace-bistable-deterministic` is supplied. + +The preregistered operational bimodality check is at least 10% of held-out candidate probabilities below 0.1, at least 10% above 0.9, and at most 20% in [0.1, 0.9], sustained across episodes. Also require finite, nonzero gate gradients, no repeated skipped optimizer updates, and no task or paired-memory-dependency collapse. Run the short stability/pilot budget first; do not start a long Phase 4 run unless the histogram and stability gates pass. Results remain specific to graph traversal and the single available trained model seed. + +After saving a pilot checkpoint, evaluate that same checkpoint and `--eval-seed` under each `--eval-workspace-mode checkpoint`, `reset-each-step`, and `bypass`. For noise and selection ablations, use the same `--eval-seed` and change only the corresponding bistable scale. The gate RNG is separately checkpointed, so these changes keep the episode stream fixed; use distinct run suffixes when recording each condition. + ### MoE numerical and memory-use diagnostics MoE experts now initialize as a zero-output residual branch: each expert's final projection starts at zero, so enabling MoE preserves the original controller signal at step 0 while the output projections learn. Top-1 routing retains the selected expert's probability from the full expert softmax; normalizing a one-expert selected set would turn every gate into 1 and eliminate task-loss gradients to the router. Top-k greater than one continues to normalize over selected experts. The CfC multi-source combiner likewise retains its original affine projection over the concatenated backbone and previous-read inputs; its source-aware experts contribute an additive correction around that baseline. Capacity slots go to the highest-confidence routes rather than the earliest rows in the flattened token batch, avoiding a systematic advantage for the first source when source batches share an expert's capacity. CfC source diagnostics report accepted and dropped route fractions separately for each source at training-log cadence. diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 8d5ad80..a8cb7e5 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -686,6 +686,12 @@ def load_frozen_workspace_warm_start( "workspace_no_selection", "workspace_reset_each_step", "workspace_train_mode", + "workspace_train_scope", + "workspace_gate_variant", + "workspace_bistable_selection_scale", + "workspace_bistable_noise_scale", + "workspace_bistable_deterministic", + "workspace_stability_max_activation", } missing_config = sorted( key @@ -809,6 +815,12 @@ def run( workspace_no_selection: bool = WORKSPACE_NO_SELECTION, workspace_reset_each_step: bool = WORKSPACE_RESET_EACH_STEP, workspace_train_mode: str = WORKSPACE_TRAIN_MODE, + workspace_train_scope: str = "workspace", + workspace_gate_variant: str = "continuous", + workspace_bistable_selection_scale: float = 1.0, + workspace_bistable_noise_scale: float = 0.75, + workspace_bistable_deterministic: bool = False, + workspace_stability_max_activation: float = 100_000.0, eval_only: bool = False, eval_workspace_mode: str = "checkpoint", eval_seed: int | None = None, @@ -879,6 +891,39 @@ def run( checkpoint_workspace_enabled = bool( checkpoint_config.get("workspace_enabled", False) or checkpoint_has_workspace ) + phase4_gate_control_warm_start = bool( + resume_requested + and workspace_enabled + and workspace_train_scope == "gate" + and checkpoint_workspace_enabled + and ( + checkpoint_config.get("workspace_train_scope") != "gate" + or checkpoint_config.get("workspace_train_mode") != workspace_train_mode + or checkpoint_config.get("workspace_gate_variant", "continuous") + != workspace_gate_variant + ) + and not eval_only + ) + workspace_mode_warm_start = bool( + resume_requested + and workspace_enabled + and checkpoint_workspace_enabled + and checkpoint_config.get("workspace_train_mode") != workspace_train_mode + and not eval_only + ) + gate_variant_warm_start = bool( + resume_requested + and workspace_enabled + and workspace_gate_variant == "bistable" + and checkpoint_workspace_enabled + and checkpoint_config.get("workspace_gate_variant", "continuous") == "continuous" + and not eval_only + ) + phase4_gate_warm_start = bool( + phase4_gate_control_warm_start + or workspace_mode_warm_start + or gate_variant_warm_start + ) frozen_workspace_warm_start = bool( resume_requested and workspace_enabled @@ -895,7 +940,12 @@ def run( # Evaluation-only runs get independent, fresh logs keyed by their run ID; # they load checkpoint weights/state but must not append rows to the # original training logs. - resuming = resume_requested and not frozen_workspace_warm_start and not eval_only + resuming = ( + resume_requested + and not frozen_workspace_warm_start + and not phase4_gate_warm_start + and not eval_only + ) log_file = open(log_path, "a" if resuming else "w", newline="") log_writer = csv.writer(log_file) if not resuming: @@ -1466,6 +1516,9 @@ def run( workspace_no_selection=workspace_no_selection, workspace_reset_each_step=workspace_reset_each_step, workspace_train_mode=workspace_train_mode, + workspace_gate_variant=workspace_gate_variant, + workspace_bistable_selection_scale=workspace_bistable_selection_scale, + workspace_bistable_noise_scale=workspace_bistable_noise_scale, device=device, ).to(device) if dnc_backward_trace: @@ -1506,11 +1559,17 @@ def run( if workspace_enabled: workspace_active = isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None write_mode = "random-control" if workspace_no_selection else "MoE-top-k" - persistence = "off (reset each timestep)" if workspace_reset_each_step else "on" + reset_workspace_each_step = workspace_reset_each_step or ( + eval_only and eval_workspace_mode == "reset-each-step" + ) + persistence = ( + "off (reset each timestep)" if reset_workspace_each_step else "on" + ) print( f"[{run_id}] WORKSPACE {'active' if workspace_active else 'bypassed'} | " f"slots={workspace_num_slots} | write_selection={write_mode} | " f"persistence={persistence} | train_mode={workspace_train_mode} | " + f"gate={workspace_gate_variant} | " f"compute_precision={'fp32' if workspace_active else 'bypass'} | " "scope=final-backbone-MoE-bank; each specialist reads shared slots" ) @@ -1571,6 +1630,12 @@ def run( "workspace_no_selection": workspace_no_selection, "workspace_reset_each_step": workspace_reset_each_step, "workspace_train_mode": workspace_train_mode, + "workspace_gate_variant": workspace_gate_variant, + "workspace_train_scope": workspace_train_scope, + "workspace_bistable_selection_scale": workspace_bistable_selection_scale, + "workspace_bistable_noise_scale": workspace_bistable_noise_scale, + "workspace_bistable_deterministic": workspace_bistable_deterministic, + "workspace_stability_max_activation": workspace_stability_max_activation, } controller_model_config = {} if controller == "mamba": @@ -1660,8 +1725,16 @@ def run( output_proj = nn.Linear(INPUT_DIM, TRIPLE_DIM).to(device) dataset.set_output_proj(output_proj) + if workspace_train_scope == "gate" and not segment_diagnostics_enabled: + raise ValueError("Phase 4 gate-only stability checks require segment diagnostics") segment_diagnostics = SegmentDiagnostics( - rnn, enabled=segment_diagnostics_enabled + rnn, + enabled=segment_diagnostics_enabled, + activation_absmax_limit=( + workspace_stability_max_activation + if workspace_train_scope == "gate" + else None + ), ) segment_diagnostics.add_module(output_proj, "task_output") moe_layer_ids = {id(layer) for layer in getattr(rnn, "moe_layers", [])} @@ -1673,6 +1746,7 @@ def run( module.set_diagnostics_enabled(moe_diagnostics_enabled) if isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None: rnn.workspace.set_diagnostics_enabled(moe_diagnostics_enabled) + rnn.workspace_bistable_deterministic = workspace_bistable_deterministic def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: segment_diagnostics.emit(channel, include_gradients=include_gradients) @@ -1691,6 +1765,32 @@ def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) - if reports: print(f"[moe-diag:{channel}] " + " | ".join(reports)) + def report_bistable_evaluation(mode: str) -> None: + if not ( + workspace_gate_variant == "bistable" + and isinstance(rnn, SplitGraphDNC) + and rnn.workspace is not None + ): + return + diagnostics = rnn.workspace.consume_bistable_diagnostics() + if not diagnostics or diagnostics["gate_candidate_count"].item() == 0: + print(f"[{run_id}] [bistable:{mode}] no valid gate candidates were recorded") + return + histogram = ",".join(f"{value:.3f}" for value in diagnostics["gate_histogram"].tolist()) + writers = ",".join(f"{value:.3f}" for value in diagnostics["writer_frequency"].tolist()) + print( + f"[{run_id}] [bistable:{mode}] p_mean=" + f"{diagnostics['gate_probability_mean'].item():.3f} " + f"low/mid/high=" + f"{diagnostics['gate_low_frac'].item():.1%}/" + f"{diagnostics['gate_mid_frac'].item():.1%}/" + f"{diagnostics['gate_high_frac'].item():.1%} " + f"commit_freq={diagnostics['write_frequency'].item():.1%} " + f"selection_peak={diagnostics['selection_strength_mean'].item():.3f} " + f"probability_histogram_10bin={histogram} " + f"writer_frequency={writers}" + ) + # Install the stochastic write head(s) BEFORE building the # optimizer, so their parameters (mu_transform, logvar_transform) are # included in optimizer.parameters(). mu_transform is initialized from @@ -1712,8 +1812,17 @@ def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) - for parameter in rnn.parameters(): parameter.requires_grad_(False) assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None - for parameter in rnn.workspace.parameters(): - parameter.requires_grad_(True) + if workspace_train_scope == "gate": + gate_parameters = ( + rnn.workspace.bistable_gate_parameters() + if workspace_gate_variant == "bistable" + else rnn.workspace.write_gate.parameters() + ) + for parameter in gate_parameters: + parameter.requires_grad_(True) + else: + for parameter in rnn.workspace.parameters(): + parameter.requires_grad_(True) for parameter in output_proj.parameters(): parameter.requires_grad_(False) trainable_parameters = [ @@ -1874,8 +1983,17 @@ def set_lr(step_now): for parameter in rnn.parameters(): parameter.requires_grad_(False) assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None - for parameter in rnn.workspace.parameters(): - parameter.requires_grad_(True) + if workspace_train_scope == "gate": + gate_parameters = ( + rnn.workspace.bistable_gate_parameters() + if workspace_gate_variant == "bistable" + else rnn.workspace.write_gate.parameters() + ) + for parameter in gate_parameters: + parameter.requires_grad_(True) + else: + for parameter in rnn.workspace.parameters(): + parameter.requires_grad_(True) # Model construction and checkpoint loading consume RNG state. # Start the new workspace experiment from the requested seed. random.seed(seed) @@ -1893,6 +2011,32 @@ def set_lr(step_now): f"{len(missing_workspace_keys)} workspace tensors freshly initialized, " f"reset optimizer, scaler, RNG, and curriculum to step 0 / lesson 1." ) + elif gate_variant_warm_start: + incompatible = rnn.load_state_dict(ckpt["rnn_state_dict"], strict=False) + expected_missing = { + key for key in rnn.state_dict() if key.startswith("workspace.bistable_gate.") + } + if set(incompatible.missing_keys) != expected_missing or incompatible.unexpected_keys: + raise ValueError( + "Bistable gate warm start only permits newly initialized " + "workspace.bistable_gate parameters; checkpoint state differs " + f"(missing={incompatible.missing_keys[:8]}, " + f"unexpected={incompatible.unexpected_keys[:8]})." + ) + output_proj.load_state_dict(ckpt["output_proj_state_dict"]) + print( + f"[{run_id}] Initialized bistable gate from continuous-workspace " + f"checkpoint {resume_from} at step {step}; all shared model and " + "workspace tensors loaded, gate parameters initialized fresh, " + "optimizer reset." + ) + elif phase4_gate_warm_start: + rnn.load_state_dict(ckpt["rnn_state_dict"]) + output_proj.load_state_dict(ckpt["output_proj_state_dict"]) + print( + f"[{run_id}] Started a fresh Phase 4 condition from {resume_from} " + f"at step {step}; workspace/model tensors loaded, optimizer reset." + ) else: rnn.load_state_dict(ckpt["rnn_state_dict"]) output_proj.load_state_dict(ckpt["output_proj_state_dict"]) @@ -1902,9 +2046,9 @@ def set_lr(step_now): # of leaving `scaler` at its freshly-constructed init_scale -- see # header note. `.get(...)` instead of ["..."] so a pre-v3 checkpoint # (saved before this key existed) still loads, just without this fix. - if not frozen_workspace_warm_start and "scaler_state_dict" in ckpt: + if not frozen_workspace_warm_start and not phase4_gate_warm_start and "scaler_state_dict" in ckpt: scaler.load_state_dict(ckpt["scaler_state_dict"]) - elif not frozen_workspace_warm_start: + elif not frozen_workspace_warm_start and not phase4_gate_warm_start: print( f"[{run_id}] WARNING: checkpoint predates GradScaler-state " f"saving (v3 fix) -- scaler starting fresh at " @@ -2097,6 +2241,18 @@ def set_lr(step_now): ) else: output, hidden = rnn(input_seq, hidden, reset_experience=True) + if workspace_train_scope == "gate": + exceeded, maximum, segment = segment_diagnostics.activation_absmax_exceeded( + "train" + ) + if exceeded: + emit_failure_diagnostics("train") + raise FloatingPointError( + f"[{run_id}] Phase 4 activation safety limit " + f"{workspace_stability_max_activation:.6g} exceeded in " + f"{segment} (observed absmax={maximum:.6g}) at step {step}; " + "stopping before backward/optimizer update." + ) if check_finite_this_step and not _first_nonfinite_report( output, "rnn_output", step ): @@ -2166,6 +2322,31 @@ def set_lr(step_now): ) scaler.scale(loss).backward() scaler.unscale_(optimizer) + bistable_gate_grad_norm = 0.0 + if workspace_gate_variant == "bistable": + assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None + gate_gradients = [ + parameter.grad.detach().float() + for parameter in rnn.workspace.bistable_gate_parameters() + if parameter.grad is not None + ] + if not gate_gradients: + emit_failure_diagnostics("train", include_gradients=True) + raise FloatingPointError( + f"[{run_id}] bistable write gate received no gradients at step {step}." + ) + bistable_gate_grad_norm = float( + torch.stack([gradient.square().sum() for gradient in gate_gradients]) + .sum() + .sqrt() + .cpu() + ) + if not math.isfinite(bistable_gate_grad_norm) or bistable_gate_grad_norm == 0.0: + emit_failure_diagnostics("train", include_gradients=True) + raise FloatingPointError( + f"[{run_id}] bistable write gate gradient norm is " + f"{bistable_gate_grad_norm} at step {step}." + ) # Log addition #2: capture the pre-clip gradient norm instead of # discarding clip_grad_norm_'s return value. Gives a β=0 baseline # gradient-scale trajectory to compare a future β>0 switch-in @@ -2363,6 +2544,31 @@ def set_lr(step_now): f"broadcast/specialist {workspace_diag['broadcast_to_specialist_ratio'].item():.3f} " f"read_gate {workspace_diag['read_gate'].item():.3f}]" ) + if "gate_low_frac" in workspace_diag: + print( + f"[{run_id}] [bistable] p_mean=" + f"{workspace_diag['gate_probability_mean'].item():.3f} " + f"low/mid/high=" + f"{workspace_diag['gate_low_frac'].item():.1%}/" + f"{workspace_diag['gate_mid_frac'].item():.1%}/" + f"{workspace_diag['gate_high_frac'].item():.1%} " + f"commit_freq={workspace_diag['write_frequency'].item():.1%} " + f"selection_peak=" + f"{workspace_diag['selection_strength_mean'].item():.3f} " + f"gate_grad_norm={bistable_gate_grad_norm:.3g}" + ) + if "gate_histogram" in workspace_diag: + histogram = ",".join( + f"{value:.2f}" + for value in workspace_diag["gate_histogram"].tolist() + ) + print(f"[{run_id}] [bistable] probability_histogram_10bin={histogram}") + if "writer_frequency" in workspace_diag: + writer_frequency = ",".join( + f"{value:.3f}" + for value in workspace_diag["writer_frequency"].tolist() + ) + print(f"[{run_id}] [workspace] writer_frequency={writer_frequency}") if dyn_n_str: parts.append(dyn_n_str.strip(" |")) print(" | ".join(parts)) @@ -2893,11 +3099,19 @@ def set_lr(step_now): shutil.copyfile(ckpt_path, latest_path) completion_label = "Evaluation-only pass" if eval_only else "Training complete" + if eval_only: + # No optimization step ran, so there is no effective training beta or + # KL measurement to attribute in the shared final-reporting path. + beta_eff = 0.0 + kl_loss = None + kl_diag = {} print( f"\n[{run_id}] {completion_label}. Final evaluation on " "training-distribution lesson:" ) final_id_stage_memory = rnn.memories[0] + if isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None: + rnn.workspace.reset_bistable_diagnostics() if hasattr(final_id_stage_memory, "begin_stage_diagnostics"): final_id_stage_memory.begin_stage_diagnostics() segment_diagnostics.set_channel("eval-final-id") @@ -2905,6 +3119,7 @@ def set_lr(step_now): rnn, device, step=step, optimizer=optimizer, allow_advance=not eval_only ) segment_diagnostics.emit("eval-final-id") + report_bistable_evaluation("ID") final_id_stage_diagnostics = ( final_id_stage_memory.finish_stage_diagnostics() if hasattr(final_id_stage_memory, "finish_stage_diagnostics") @@ -2934,6 +3149,8 @@ def set_lr(step_now): final_ood_torch_rng = _capture_torch_rng() final_ood_field_log: dict[tuple[int, int], list[int]] = {} final_ood_stage_memory = rnn.memories[0] + if isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None: + rnn.workspace.reset_bistable_diagnostics() if hasattr(final_ood_stage_memory, "begin_stage_diagnostics"): final_ood_stage_memory.begin_stage_diagnostics() segment_diagnostics.set_channel("eval-final-ood") @@ -2946,6 +3163,7 @@ def set_lr(step_now): field_log=final_ood_field_log, ) segment_diagnostics.emit("eval-final-ood") + report_bistable_evaluation("OOD") final_ood_stage_diagnostics = ( final_ood_stage_memory.finish_stage_diagnostics() if hasattr(final_ood_stage_memory, "finish_stage_diagnostics") @@ -2985,6 +3203,14 @@ def set_lr(step_now): ] ) track_a_file.flush() + residual_ratio = final_id_stage_diagnostics.get( + "relational_read_residual_to_cosine_rms_mean" + ) + residual_ratio_text = ( + f"{residual_ratio:.3f}" + if residual_ratio is not None + else f"n/a (not defined for {dnc_read_variant} read)" + ) print( f"[{run_id}] [Track A / final] stress cosine/read/held-out=" f"{cosine_top1_accuracy(relational_probe_cases):.3f}/" @@ -2995,8 +3221,7 @@ def set_lr(step_now): f"read entropy/max=" f"{final_id_stage_diagnostics.get('read_content_entropy_mean', float('nan')):.3f}/" f"{final_id_stage_diagnostics.get('read_content_max_weight_mean', float('nan')):.3f} | " - f"residual/cosine rms=" - f"{final_id_stage_diagnostics.get('relational_read_residual_to_cosine_rms_mean', float('nan')):.3f}" + f"residual/cosine rms={residual_ratio_text}" ) if track_b_writer is not None and track_b_file is not None: @@ -3433,6 +3658,41 @@ def _write_diagnostic_values(diagnostics): parser.add_argument("--workspace-train-mode", choices=["coadapted", "frozen"], default=WORKSPACE_TRAIN_MODE, help="coadapted trains the full model; frozen trains workspace parameters and freezes substrate/task-head parameters.") + parser.add_argument( + "--workspace-train-scope", + choices=("workspace", "gate"), + default="workspace", + help="With frozen mode, train all workspace parameters or only the active write gate. Gate-only supports matched continuous/bistable Phase 4 controls.", + ) + parser.add_argument( + "--workspace-gate-variant", + choices=("continuous", "bistable"), + default="continuous", + help="Use the existing continuous update gate or the experimental recurrent stochastic bistable write gate.", + ) + parser.add_argument( + "--workspace-bistable-selection-scale", + type=float, + default=1.0, + help="Scale the content-selection drive in the bistable gate; vary this for the preregistered selection-strength ablation.", + ) + parser.add_argument( + "--workspace-bistable-noise-scale", + type=float, + default=0.75, + help="Independent background-noise scale for bistable ignition; set to 0 for the noise ablation.", + ) + parser.add_argument( + "--workspace-bistable-deterministic", + action="store_true", + help="Disable independent gate noise for a deterministic diagnostic; normal bistable evaluation remains stochastic and seed-reproducible.", + ) + parser.add_argument( + "--workspace-stability-max-activation", + type=float, + default=100_000.0, + help="Phase 4 absolute activation maximum; abort before backward/update if any instrumented segment exceeds this value.", + ) parser.add_argument( "--moe-num-experts", type=int, @@ -3744,6 +4004,31 @@ def _write_diagnostic_values(diagnostics): parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: parser.error("--workspace-num-slots must be >= 0") + if args.workspace_train_scope == "gate" and not args.workspace: + parser.error("--workspace-train-scope=gate requires --workspace") + if args.workspace_gate_variant == "bistable" and ( + not args.workspace or args.workspace_num_slots < 2 + ): + parser.error( + "--workspace-gate-variant=bistable requires --workspace and at least two slots" + ) + if args.workspace_gate_variant == "bistable" and not args.moe_diagnostics: + parser.error("bistable gate histogram/diagnostics require --moe-diagnostics") + if ( + not math.isfinite(args.workspace_bistable_selection_scale) + or args.workspace_bistable_selection_scale < 0 + ): + parser.error("--workspace-bistable-selection-scale must be finite and >= 0") + if ( + not math.isfinite(args.workspace_bistable_noise_scale) + or args.workspace_bistable_noise_scale < 0 + ): + parser.error("--workspace-bistable-noise-scale must be finite and >= 0") + if ( + not math.isfinite(args.workspace_stability_max_activation) + or args.workspace_stability_max_activation <= 0 + ): + parser.error("--workspace-stability-max-activation must be finite and > 0") if args.beta is not None: betas_to_run = [args.beta] @@ -3793,6 +4078,15 @@ def _write_diagnostic_values(diagnostics): run_id += f"_backbonerscale{args.moe_backbone_residual_scale:g}" if args.workspace: run_id += f"_workspace{args.workspace_num_slots}s_{args.workspace_train_mode}" + if args.workspace_gate_variant != "continuous" or args.workspace_train_scope == "gate": + run_id += f"_gate{args.workspace_gate_variant}_{args.workspace_train_scope}" + if args.workspace_gate_variant == "bistable": + run_id += ( + f"_sel{args.workspace_bistable_selection_scale:g}" + f"_noise{args.workspace_bistable_noise_scale:g}" + ) + if args.workspace_bistable_deterministic: + run_id += "_detgate" if args.workspace_no_selection: run_id += "_randomwrite" if args.workspace_reset_each_step: @@ -3857,6 +4151,12 @@ def _write_diagnostic_values(diagnostics): workspace_no_selection=args.workspace_no_selection, workspace_reset_each_step=args.workspace_reset_each_step, workspace_train_mode=args.workspace_train_mode, + workspace_train_scope=args.workspace_train_scope, + workspace_gate_variant=args.workspace_gate_variant, + workspace_bistable_selection_scale=args.workspace_bistable_selection_scale, + workspace_bistable_noise_scale=args.workspace_bistable_noise_scale, + workspace_bistable_deterministic=args.workspace_bistable_deterministic, + workspace_stability_max_activation=args.workspace_stability_max_activation, eval_only=args.eval_only, eval_workspace_mode=args.eval_workspace_mode, eval_seed=args.eval_seed, diff --git a/src/initium/data/graph_traversal/graph_traversal.py b/src/initium/data/graph_traversal/graph_traversal.py index 1d64fc6..796e136 100644 --- a/src/initium/data/graph_traversal/graph_traversal.py +++ b/src/initium/data/graph_traversal/graph_traversal.py @@ -445,6 +445,7 @@ def evaluate_traversal( if num_nodes is not None and nodes_range is not None: raise ValueError("evaluate_traversal: pass num_nodes or nodes_range, not both") + was_training = model.training model.eval() total_triples, correct_triples = 0, 0 perfect_episodes = 0 @@ -531,7 +532,7 @@ def evaluate_traversal( f"Eval: triple-level acc {triple_acc:.2f}% | " f"perfect-traversal fraction {perfect_frac:.2f}% ({tested} episodes)" ) - model.train() + model.train(was_training) if hop_breakdown: breakdown = { diff --git a/src/initium/diagnostics/segment_logging.py b/src/initium/diagnostics/segment_logging.py index 78a308d..97cf528 100644 --- a/src/initium/diagnostics/segment_logging.py +++ b/src/initium/diagnostics/segment_logging.py @@ -61,6 +61,8 @@ def _segment_label(path: str, module: nn.Module) -> str | None: return f"controller.mamba_cell:{leaf}" if class_name == "WorkspaceBroadcast": return f"workspace:{leaf}" + if class_name == "BistableWriteGate": + return f"workspace_bistable_gate:{leaf}" if class_name == "Memory" or "memory" in path.lower() and "memory" in lower_name: return f"dnc_memory:{leaf}" if isinstance(module, nn.LSTM): @@ -80,9 +82,16 @@ class SegmentDiagnostics: boundary or when a numerical failure needs an immediate trace. """ - def __init__(self, model: nn.Module, *, enabled: bool = True): + def __init__( + self, + model: nn.Module, + *, + enabled: bool = True, + activation_absmax_limit: float | None = None, + ): self.model = model self.enabled = bool(enabled) + self.activation_absmax_limit = activation_absmax_limit self.channel: str | None = "train" self._records: dict[str, OrderedDict[str, dict[str, Any]]] = {} self._segments: list[tuple[str, str]] = [] @@ -245,6 +254,25 @@ def emit(self, channel: str, *, include_gradients: bool = False) -> None: parts.append("grad_norm[" + "; ".join(gradients) + "]") print(f"[segment-diag:{channel}] " + " | ".join(parts)) + def activation_absmax_exceeded(self, channel: str) -> tuple[bool, float, str | None]: + """Check captured forward maxima against the optional safety limit.""" + if not self.enabled or self.activation_absmax_limit is None: + return False, 0.0, None + records = self._records.get(channel, {}) + candidates: list[tuple[str, torch.Tensor]] = [] + for label, record in records.items(): + for key in ("in_max", "out_max"): + value = record.get(key) + if torch.is_tensor(value): + candidates.append((label, value)) + if not candidates: + return False, 0.0, None + values = torch.stack([value.detach().float() for _, value in candidates]) + maximum, index = values.max(dim=0) + maximum_value = float(maximum.cpu()) + label = candidates[int(index.cpu())][0] + return maximum_value > self.activation_absmax_limit, maximum_value, label + def close(self) -> None: for handle in self._handles: handle.remove() diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 4276fff..d511f49 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -157,7 +157,7 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea ), ) _moe = _moe_kwargs(c) - return SplitGraphDNC( + model = SplitGraphDNC( input_size=input_dim, hidden_size=hidden, nr_cells=nr_cells, @@ -214,8 +214,17 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea "workspace_reset_each_step", cc.WORKSPACE_RESET_EACH_STEP ), workspace_train_mode=c.get("workspace_train_mode", cc.WORKSPACE_TRAIN_MODE), + workspace_gate_variant=c.get("workspace_gate_variant", "continuous"), + workspace_bistable_selection_scale=c.get( + "workspace_bistable_selection_scale", 1.0 + ), + workspace_bistable_noise_scale=c.get("workspace_bistable_noise_scale", 0.75), device=device, ).to(device) + model.workspace_bistable_deterministic = c.get( + "workspace_bistable_deterministic", False + ) + return model def load_model(ckpt: dict, device: torch.device, deterministic_write: bool = False) -> LoadedModel: diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 67f75b2..ac0dffd 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -209,6 +209,9 @@ def __init__( workspace_no_selection: bool = False, workspace_reset_each_step: bool = False, workspace_train_mode: str = "coadapted", + workspace_gate_variant: str = "continuous", + workspace_bistable_selection_scale: float = 1.0, + workspace_bistable_noise_scale: float = 0.75, ): super().__init__() if not independent_linears: @@ -260,6 +263,9 @@ def __init__( self.workspace_no_selection = workspace_no_selection self.workspace_reset_each_step = workspace_reset_each_step self.workspace_train_mode = workspace_train_mode + self.workspace_bistable_noise_scale = float(workspace_bistable_noise_scale) + self.workspace_bistable_selection_scale = float(workspace_bistable_selection_scale) + self.workspace_bistable_deterministic = False # Runtime-only counterfactual for evaluation. It is deliberately not # part of checkpoint config/state so a saved model can be scored with # and without workspace readback using identical learned weights. @@ -516,7 +522,14 @@ def __init__( else None ) self.workspace = ( - WorkspaceBroadcast(hidden_size, moe_num_experts, workspace_num_slots) + WorkspaceBroadcast( + hidden_size, + moe_num_experts, + workspace_num_slots, + gate_variant=workspace_gate_variant, + bistable_selection_scale=workspace_bistable_selection_scale, + bistable_noise_scale=workspace_bistable_noise_scale, + ) if workspace_enabled and workspace_num_slots > 0 else None ) @@ -634,6 +647,9 @@ def forward( no_selection=self.workspace_no_selection, reset_each_step=self.workspace_reset_each_step, bypass_broadcast=self.workspace_eval_bypass, + bistable_noise_scale=self.workspace_bistable_noise_scale, + bistable_selection_scale=self.workspace_bistable_selection_scale, + deterministic_gate=self.workspace_bistable_deterministic, ) H = base.clone() # The workspace returns the selected, broadcast-enriched MoE diff --git a/src/initium/workspace/bistability.py b/src/initium/workspace/bistability.py new file mode 100644 index 0000000..e463322 --- /dev/null +++ b/src/initium/workspace/bistability.py @@ -0,0 +1,182 @@ +"""Stochastic competitive write ignition for shared workspaces. + +The gate is an experimental, opt-in approximation of population ignition: +slot candidates receive a content-selection drive and independent background +noise, then settle through recurrent self-excitation and lateral inhibition. +Forward writes are binary; a straight-through estimator carries gradients +through the settled probabilities during training. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class BistableWriteGate(nn.Module): + """Resolve slot-write candidates with recurrent competitive dynamics.""" + + def __init__( + self, + *, + selection_scale: float = 1.0, + noise_scale: float = 0.75, + iterations: int = 6, + temperature: float = 0.4, + ) -> None: + super().__init__() + if selection_scale < 0 or noise_scale < 0: + raise ValueError("selection_scale and noise_scale must be >= 0") + if iterations < 1: + raise ValueError("iterations must be >= 1") + if temperature <= 0: + raise ValueError("temperature must be > 0") + self.selection_scale = float(selection_scale) + self.noise_scale = float(noise_scale) + self.iterations = int(iterations) + self.temperature = float(temperature) + self.bias = nn.Parameter(torch.tensor(1.1)) + self.self_excitation_logit = nn.Parameter(torch.tensor(1.1)) + self.inhibition_logit = nn.Parameter(torch.tensor(0.5)) + self._last_diagnostics: dict[str, torch.Tensor] = {} + self._noise_seed = int(torch.randint(0, 2**31 - 1, ()).item()) + self._noise_generators: dict[str, torch.Generator] = {} + self._pending_generator_states: dict[str, torch.Tensor] = {} + + def get_extra_state(self) -> dict: + """Persist the independent noise streams with model checkpoints.""" + states = dict(self._pending_generator_states) + states.update( + {key: generator.get_state().cpu() for key, generator in self._noise_generators.items()} + ) + return {"noise_seed": self._noise_seed, "generator_states": states} + + def set_extra_state(self, state: dict) -> None: + self._noise_seed = int(state["noise_seed"]) + self._pending_generator_states = { + key: value.cpu() for key, value in state.get("generator_states", {}).items() + } + self._noise_generators.clear() + + def _sample_background(self, selection: torch.Tensor) -> torch.Tensor: + device_key = str(selection.device) + generator = self._noise_generators.get(device_key) + if generator is None: + generator = torch.Generator(device=selection.device) + state = self._pending_generator_states.pop(device_key, None) + if state is None: + generator.manual_seed(self._noise_seed) + else: + try: + generator.set_state(state) + except RuntimeError: + # CPU and CUDA generator states are not interchangeable; + # keep inference portable while preserving exact replay + # whenever a checkpoint returns to the same device. + generator.manual_seed(self._noise_seed) + self._noise_generators[device_key] = generator + return torch.randn( + selection.shape, + device=selection.device, + dtype=selection.dtype, + generator=generator, + ) + + def forward( + self, + selection: torch.Tensor, + *, + selection_confidence: torch.Tensor | None = None, + candidate_mask: torch.Tensor | None = None, + noise_scale: float | None = None, + selection_scale: float | None = None, + deterministic: bool = False, + collect_diagnostics: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return hard gates and their pre-commit ignition probabilities. + + ``selection`` is a nonnegative (batch, slot) selection-strength + tensor. The stochastic term is sampled independently per candidate + and remains fixed through the recurrent settling iterations. + """ + if selection.ndim != 2: + raise ValueError("selection must have shape (batch, slots)") + if selection.shape[-1] < 2: + raise ValueError("bistable competition requires at least two slots") + if candidate_mask is not None and tuple(candidate_mask.shape) != (selection.shape[0],): + raise ValueError("candidate_mask must have shape (batch,)") + scale = self.selection_scale if selection_scale is None else selection_scale + noise = self.noise_scale if noise_scale is None else noise_scale + if scale < 0 or noise < 0: + raise ValueError("selection_scale and noise_scale must be >= 0") + + selection = selection.float().clamp_min(0.0) + normalized = selection / selection.sum(dim=-1, keepdim=True).clamp_min(1e-8) + eps = torch.finfo(normalized.dtype).eps + selection_logit = torch.logit(normalized.clamp(eps, 1.0 - eps)) + uniform_logit = torch.logit( + torch.full_like(normalized[:, :1], 1.0 / normalized.shape[-1]).clamp( + eps, 1.0 - eps + ) + ) + selection_logit = selection_logit - uniform_logit + if selection_confidence is None: + confidence_logit = torch.zeros_like(selection_logit[:, :1]) + else: + confidence = selection_confidence.float().reshape(-1, 1).clamp(eps, 1.0 - eps) + confidence_logit = torch.logit(confidence) + sampled_background = self._sample_background(selection) + background = ( + sampled_background * noise + if not deterministic and noise > 0 + else torch.zeros_like(selection) + ) + drive = scale * (selection_logit + confidence_logit) + background + self.bias + self_excitation = F.softplus(self.self_excitation_logit) + inhibition = F.softplus(self.inhibition_logit) + activity = torch.sigmoid(drive / self.temperature) + for _ in range(self.iterations): + population_activity = activity.sum(dim=-1, keepdim=True) + lateral_input = population_activity - activity + activity = torch.sigmoid( + (drive + self_excitation * activity - inhibition * lateral_input) + / self.temperature + ) + + hard = (activity >= 0.5).to(activity.dtype) + committed = hard + activity - activity.detach() + valid = ( + torch.ones(selection.shape[0], device=selection.device, dtype=torch.bool) + if candidate_mask is None + else candidate_mask.to(device=selection.device, dtype=torch.bool) + ) + committed = committed * valid.unsqueeze(-1).to(committed.dtype) + if collect_diagnostics: + with torch.no_grad(): + flat = activity.detach()[valid] + candidate_count = valid.sum().float() * selection.shape[-1] + candidate_denom = candidate_count.clamp_min(1.0) + histogram = ( + torch.histc(flat, bins=10, min=0.0, max=1.0) + if flat.numel() + else activity.new_zeros((10,)) + ) + self._last_diagnostics = { + "gate_candidate_count": candidate_count, + "gate_probability_mean": flat.sum() / candidate_denom, + "gate_histogram": histogram / candidate_denom, + "gate_low_frac": (flat < 0.1).float().sum() / candidate_denom, + "gate_mid_frac": ((flat >= 0.1) & (flat <= 0.9)).float().sum() + / candidate_denom, + "gate_high_frac": (flat > 0.9).float().sum() / candidate_denom, + "write_frequency": hard.detach()[valid].sum() / candidate_denom, + "selection_strength_mean": normalized.detach()[valid].amax(dim=-1).sum() + / valid.sum().clamp_min(1), + } + else: + self._last_diagnostics = {} + return committed.to(selection.dtype), activity + + def last_diagnostics(self) -> dict[str, torch.Tensor]: + return dict(self._last_diagnostics) diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index e7ca473..15e0ce1 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -11,6 +11,8 @@ import torch import torch.nn as nn +from initium.workspace.bistability import BistableWriteGate + class WorkspaceBroadcast(nn.Module): """Competitive MoE-selected writes and all-expert shared-slot readback. @@ -21,24 +23,48 @@ class WorkspaceBroadcast(nn.Module): receives the resulting context before the selected deltas are combined. """ - def __init__(self, d_model: int, num_experts: int, num_slots: int = 5): + def __init__( + self, + d_model: int, + num_experts: int, + num_slots: int = 5, + *, + gate_variant: str = "continuous", + bistable_selection_scale: float = 1.0, + bistable_noise_scale: float = 0.75, + ): super().__init__() if d_model < 1 or num_experts < 1: raise ValueError("d_model and num_experts must be positive") if num_slots < 0: raise ValueError("num_slots must be >= 0") + if gate_variant not in ("continuous", "bistable"): + raise ValueError("gate_variant must be 'continuous' or 'bistable'") self.d_model = d_model self.num_experts = num_experts self.num_slots = num_slots + self.gate_variant = gate_variant + self.bistable_selection_scale = float(bistable_selection_scale) + self.bistable_noise_scale = float(bistable_noise_scale) self.query = nn.Linear(d_model, d_model, bias=False) self.key = nn.Linear(d_model, d_model, bias=False) self.value = nn.Linear(d_model, d_model, bias=False) self.write_gate = nn.Linear(2 * d_model, d_model) self.read_gate = nn.Parameter(torch.zeros(())) self.write_gate_closed = False + self.bistable_gate = ( + BistableWriteGate( + selection_scale=bistable_selection_scale, + noise_scale=bistable_noise_scale, + ) + if gate_variant == "bistable" + else None + ) self.diagnostics_enabled = True self._last_step_diag: dict[str, torch.Tensor] = {} self._last_diag: dict[str, torch.Tensor] = {} + self._bistable_diagnostic_totals: dict[str, torch.Tensor] = {} + self._bistable_diagnostic_calls = 0 nn.init.zeros_(self.write_gate.weight) nn.init.zeros_(self.write_gate.bias) @@ -55,6 +81,9 @@ def forward_step( no_selection: bool = False, reset_each_step: bool = False, bypass_broadcast: bool = False, + bistable_noise_scale: float | None = None, + bistable_selection_scale: float | None = None, + deterministic_gate: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # Keep the zero-slot/closed-gate bypass identical to the selection-only # MoE path. Active workspace attention and recurrent slot updates use @@ -67,6 +96,10 @@ def forward_step( state, no_selection=no_selection, reset_each_step=reset_each_step, + bypass_broadcast=bypass_broadcast, + bistable_noise_scale=bistable_noise_scale, + bistable_selection_scale=bistable_selection_scale, + deterministic_gate=deterministic_gate, ) if specialist_deltas.is_cuda: with torch.autocast(device_type="cuda", enabled=False): @@ -77,6 +110,10 @@ def forward_step( state.float(), no_selection=no_selection, reset_each_step=reset_each_step, + bypass_broadcast=bypass_broadcast, + bistable_noise_scale=bistable_noise_scale, + bistable_selection_scale=bistable_selection_scale, + deterministic_gate=deterministic_gate, ) return self._forward_step( specialist_deltas.float(), @@ -85,6 +122,10 @@ def forward_step( state.float(), no_selection=no_selection, reset_each_step=reset_each_step, + bypass_broadcast=bypass_broadcast, + bistable_noise_scale=bistable_noise_scale, + bistable_selection_scale=bistable_selection_scale, + deterministic_gate=deterministic_gate, ) def _forward_step( @@ -96,6 +137,10 @@ def _forward_step( *, no_selection: bool = False, reset_each_step: bool = False, + bypass_broadcast: bool = False, + bistable_noise_scale: float | None = None, + bistable_selection_scale: float | None = None, + deterministic_gate: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if self.num_slots == 0 or self.write_gate_closed: selected = specialist_deltas.gather( @@ -157,6 +202,24 @@ def _forward_step( update = torch.sigmoid(self.write_gate(torch.cat([old_inputs, slot_inputs], dim=-1))) has_writer = write_weights.sum(dim=-1, keepdim=True).unsqueeze(-1) > 0 update = update * has_writer.to(update.dtype) + bistable_diag = {} + if self.bistable_gate is not None: + committed, _ = self.bistable_gate( + mass, + selection_confidence=write_weights.sum(dim=-1), + candidate_mask=has_writer.reshape(batch), + noise_scale=bistable_noise_scale, + selection_scale=bistable_selection_scale, + deterministic=deterministic_gate, + collect_diagnostics=self.diagnostics_enabled, + ) + update = update * committed.unsqueeze(-1) + bistable_diag = self.bistable_gate.last_diagnostics() + if self.diagnostics_enabled: + writer_weights = torch.zeros( + batch, self.num_experts, device=route_indices.device, dtype=torch.float32 + ).scatter_add_(1, route_indices, route_weights.float()) + writer_frequency = writer_weights.mean(dim=0).detach() next_state = update.mul(slot_inputs).add((1.0 - update).mul(old_inputs)) if bypass_broadcast: @@ -175,6 +238,8 @@ def _forward_step( "slot_rms": next_state.detach().float().square().mean().sqrt(), "broadcast_rms": next_state.new_zeros((), dtype=torch.float32), "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + **bistable_diag, + **({"writer_frequency": writer_frequency} if self.diagnostics_enabled else {}), } return (selected_deltas * route_weights.unsqueeze(-1)).sum(dim=1), next_state @@ -200,6 +265,8 @@ def _forward_step( read_scale.detach().float() * broadcast.detach().float() ).square().mean().sqrt(), "specialist_rms": specialist_deltas.detach().float().square().mean().sqrt(), + **bistable_diag, + **({"writer_frequency": writer_frequency} if self.diagnostics_enabled else {}), } return (enriched * route_weights.unsqueeze(-1)).sum(dim=1), next_state @@ -213,6 +280,9 @@ def forward( no_selection: bool = False, reset_each_step: bool = False, bypass_broadcast: bool = False, + bistable_noise_scale: float | None = None, + bistable_selection_scale: float | None = None, + deterministic_gate: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Process time-major specialists ``(B,T,E,D)`` sequentially.""" outputs = [] @@ -224,6 +294,9 @@ def forward( no_selection=no_selection, reset_each_step=reset_each_step, bypass_broadcast=bypass_broadcast, + bistable_noise_scale=bistable_noise_scale, + bistable_selection_scale=bistable_selection_scale, + deterministic_gate=deterministic_gate, ) outputs.append(out) if step_diags is not None: @@ -238,6 +311,37 @@ def forward( self._last_diag["broadcast_rms"] / self._last_diag["specialist_rms"].clamp_min(1e-8) ) + if "writer_frequency" in step_diags[0]: + self._last_diag["writer_frequency"] = torch.stack( + [diag["writer_frequency"] for diag in step_diags] + ).mean(dim=0).detach() + if "gate_candidate_count" in step_diags[0]: + candidate_counts = torch.stack( + [diag["gate_candidate_count"] for diag in step_diags] + ) + candidate_total = candidate_counts.sum() + candidate_denom = candidate_total.clamp_min(1.0) + self._last_diag["gate_candidate_count"] = candidate_total.detach() + for key in ( + "gate_probability_mean", + "gate_low_frac", + "gate_mid_frac", + "gate_high_frac", + "write_frequency", + "selection_strength_mean", + ): + values = torch.stack([diag[key] for diag in step_diags]) + self._last_diag[key] = ( + (values * candidate_counts).sum() / candidate_denom + ).detach() + histograms = torch.stack( + [diag["gate_histogram"] for diag in step_diags] + ) + self._last_diag["gate_histogram"] = ( + (histograms * candidate_counts.unsqueeze(-1)).sum(dim=0) + / candidate_denom + ).detach() + self._accumulate_bistable_diagnostics(self._last_diag) return torch.stack(outputs, dim=1), state def last_diagnostics(self) -> dict[str, torch.Tensor]: @@ -250,3 +354,71 @@ def set_diagnostics_enabled(self, enabled: bool) -> None: if not self.diagnostics_enabled: self._last_step_diag = {} self._last_diag = {} + self.reset_bistable_diagnostics() + + def _accumulate_bistable_diagnostics(self, diagnostics: dict[str, torch.Tensor]) -> None: + if "gate_histogram" not in diagnostics: + return + candidate_count = diagnostics.get("gate_candidate_count") + if candidate_count is None: + return + keys = ( + "gate_probability_mean", + "gate_histogram", + "gate_low_frac", + "gate_mid_frac", + "gate_high_frac", + "write_frequency", + "selection_strength_mean", + "writer_frequency", + ) + for key in keys: + value = diagnostics.get(key) + if value is not None: + weight = ( + candidate_count / self.num_slots + if key == "writer_frequency" or key == "selection_strength_mean" + else candidate_count + ) + self._bistable_diagnostic_totals[key] = ( + self._bistable_diagnostic_totals[key] + value.detach() * weight + if key in self._bistable_diagnostic_totals + else value.detach().clone() * weight + ) + self._bistable_diagnostic_totals["gate_candidate_count"] = ( + self._bistable_diagnostic_totals.get( + "gate_candidate_count", candidate_count.detach().new_zeros(()) + ) + + candidate_count.detach() + ) + self._bistable_diagnostic_calls += 1 + + def consume_bistable_diagnostics(self) -> dict[str, torch.Tensor]: + """Return candidate diagnostics averaged across forwards since reset.""" + if self._bistable_diagnostic_calls == 0: + return {} + candidate_count = self._bistable_diagnostic_totals["gate_candidate_count"] + denominator = candidate_count.clamp_min(1.0) + result = { + key: value + / ( + denominator / self.num_slots + if key in ("writer_frequency", "selection_strength_mean") + else denominator + ) + for key, value in self._bistable_diagnostic_totals.items() + if key != "gate_candidate_count" + } + result["gate_candidate_count"] = candidate_count + self.reset_bistable_diagnostics() + return result + + def reset_bistable_diagnostics(self) -> None: + self._bistable_diagnostic_totals.clear() + self._bistable_diagnostic_calls = 0 + + def bistable_gate_parameters(self): + """Return experimental gate parameters, or an empty iterator.""" + if self.bistable_gate is None: + return iter(()) + return self.bistable_gate.parameters() diff --git a/tests/unit/test_bistable_workspace_gate.py b/tests/unit/test_bistable_workspace_gate.py new file mode 100644 index 0000000..02ebb17 --- /dev/null +++ b/tests/unit/test_bistable_workspace_gate.py @@ -0,0 +1,105 @@ +import torch + +from initium.workspace.bistability import BistableWriteGate +from initium.workspace.broadcast import WorkspaceBroadcast + + +def test_bistable_gate_commits_hard_writes_with_finite_gradients(): + gate = BistableWriteGate(noise_scale=0.75) + selection = torch.full((1024, 5), 0.2) + confidence = torch.full((1024,), 0.2) + torch.manual_seed(1234) + + committed, probabilities = gate( + selection, + selection_confidence=confidence, + ) + assert torch.all((committed == 0) | (committed == 1)) + assert torch.isfinite(probabilities).all() + + committed.sum().backward() + total_gradient = 0.0 + for parameter in gate.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + total_gradient += float(parameter.grad.abs().sum()) + assert total_gradient > 0.0 + + diagnostics = gate.last_diagnostics() + assert diagnostics["gate_low_frac"] >= 0.1 + assert diagnostics["gate_high_frac"] >= 0.1 + assert diagnostics["gate_mid_frac"] <= 0.2 + + +def test_bistable_gate_noise_is_independent_and_noise_can_be_disabled(): + gate = BistableWriteGate(noise_scale=0.75) + selection = torch.full((128, 5), 0.2) + confidence = torch.full((128,), 0.2) + + torch.manual_seed(10) + _, noisy_a = gate(selection, selection_confidence=confidence) + torch.manual_seed(11) + _, noisy_b = gate(selection, selection_confidence=confidence) + assert not torch.equal(noisy_a, noisy_b) + + torch.manual_seed(10) + _, quiet_a = gate( + selection, + selection_confidence=confidence, + noise_scale=0.0, + ) + torch.manual_seed(11) + _, quiet_b = gate( + selection, + selection_confidence=confidence, + noise_scale=0.0, + ) + assert torch.equal(quiet_a, quiet_b) + + torch.manual_seed(42) + expected_global = torch.rand(4) + torch.manual_seed(42) + gate(selection, selection_confidence=confidence) + actual_global = torch.rand(4) + assert torch.equal(actual_global, expected_global) + + +def test_bistable_gate_noise_stream_roundtrips_in_state_dict(): + gate = BistableWriteGate() + selection = torch.full((16, 5), 0.2) + confidence = torch.full((16,), 0.2) + checkpoint = gate.state_dict() + + _, expected = gate(selection, selection_confidence=confidence) + restored = BistableWriteGate() + restored.load_state_dict(checkpoint) + _, actual = restored(selection, selection_confidence=confidence) + assert torch.equal(actual, expected) + + +def test_workspace_bistable_variant_preserves_state_contract_and_logs_modes(): + torch.manual_seed(3) + workspace = WorkspaceBroadcast( + d_model=8, + num_experts=4, + num_slots=5, + gate_variant="bistable", + ) + specialists = torch.randn(2, 3, 4, 8) + routes = torch.zeros(2, 3, 1, dtype=torch.long) + weights = torch.full((2, 3, 1), 0.2) + state = workspace.init_state(2, device=specialists.device, dtype=specialists.dtype) + + output, next_state = workspace( + specialists, + routes, + weights, + state, + ) + diagnostics = workspace.last_diagnostics() + assert output.shape == (2, 3, 8) + assert next_state.shape == (2, 5, 8) + assert diagnostics["gate_low_frac"] >= 0 + assert diagnostics["gate_mid_frac"] >= 0 + assert diagnostics["gate_high_frac"] >= 0 + assert diagnostics["writer_frequency"].shape == (4,) From e354ef1b8fb3d951f83b7b5379e1415d9c6e2918 Mon Sep 17 00:00:00 2001 From: Samuel Date: Tue, 13 Oct 2026 19:59:37 +0200 Subject: [PATCH 11/14] Optimize DNC memory allocation Add an opt-in exact allocation path that scatters weights through the existing usage permutation, eliminating a redundant sort while preserving state-dict compatibility. Wire the option through training, inference, and memory resizing, and add equivalence tests. --- README.md | 2 + src/initium/core_training.py | 35 ++++++ src/initium/inference/model_loader.py | 2 + .../mamba_controller/mamba_controller.py | 14 ++- .../mamba_controller/split_graph_dnc.py | 5 + .../memory_manipulation/dnc_mechanisms.py | 61 +++++++++-- .../dynamic_memory_resize.py | 11 +- tests/unit/test_optimized_dnc_memory.py | 100 ++++++++++++++++++ 8 files changed, 220 insertions(+), 10 deletions(-) create mode 100644 tests/unit/test_optimized_dnc_memory.py diff --git a/README.md b/README.md index 76cb799..06a89e8 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,8 @@ The split-graph memory exposes independent read and write variants while retaini - The graph curriculum emits ID, OOD, and OOD-offset metrics as before. Graph runs also write `run__relational_probe.csv`, with cosine top-1, selected read-head top-1, and held-out-composition accuracy. `data/graph_traversal/order_probe.py` supplies paired prefix-permutation helpers for the Q99 order-sensitivity evaluation. - Split-graph graph runs also write `run__track_a.csv` and `run__track_b.csv`. Track A records ID/OOD accuracy, relational stress and held-out-composition accuracy, ordinary-input read-score/cosine correlation, and read-content entropy/maximum weight. Track B records the Q99 order-sensitivity gate, ID/OOD posterior variance and write-update diagnostics, and training-window loss, gradient norm, AMP scale, and learning rate. These are experiment logs; no training or evaluation results are checked in. +The DNC implementation defaults to `--base-dnc`. `--optimized-dnc` selects an exact allocation implementation that reuses the stock sorted usage permutation to scatter allocation weights back into memory order, avoiding the second full sort in every write step. It retains dense temporal links, the same controller interface, and the same learned parameters/state-dict keys. Use the mutually exclusive `--base-dnc` flag to make a baseline run explicit; compare matched seeds, episode streams, and step budgets before adopting the optimized mode. This first optimization targets allocation overhead only; it does not claim to remove the dense link matrix's quadratic cost. + Use the selected split-graph baseline with either supported controller combiner: ```bash diff --git a/src/initium/core_training.py b/src/initium/core_training.py index a8cb7e5..73defe7 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -692,6 +692,9 @@ def load_frozen_workspace_warm_start( "workspace_bistable_noise_scale", "workspace_bistable_deterministic", "workspace_stability_max_activation", + # Exact implementation-only optimization; old checkpoints default + # to base and remain weight/state compatible with either backend. + "dnc_implementation", } missing_config = sorted( key @@ -829,6 +832,7 @@ def run( dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, dnc_write_variant: str = DNC_WRITE_VARIANT, dnc_write_observation_variance: float = DNC_WRITE_OBSERVATION_VARIANCE, + dnc_implementation: str = "base", dnc_backward_trace: bool = False, dataset_type: str = DATASET_TYPE, dataset_link: str | None = DATASET_LINK, @@ -1511,6 +1515,7 @@ def run( dnc_read_residual_max_ratio=dnc_read_residual_max_ratio, dnc_write_variant=dnc_write_variant, dnc_write_observation_variance=dnc_write_observation_variance, + dnc_implementation=dnc_implementation, workspace_enabled=workspace_enabled, workspace_num_slots=workspace_num_slots, workspace_no_selection=workspace_no_selection, @@ -1552,9 +1557,19 @@ def run( # equivalence rather than a bit-identical beta=0 substitute. # v7: unaffected by the controller swap -- MambaDNC builds # dnc.memory.Memory identically regardless of rnn_type. + dnc_implementation=dnc_implementation, **mamba_kwargs, ).to(device) + print( + f"[{run_id}] DNC implementation={dnc_implementation}" + + ( + " | exact allocation path avoids one full usage-permutation sort per step" + if dnc_implementation == "optimized" + else "" + ) + ) + if split_graph_enabled: if workspace_enabled: workspace_active = isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None @@ -1623,6 +1638,7 @@ def run( "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, "dnc_write_variant": dnc_write_variant, "dnc_write_observation_variance": dnc_write_observation_variance, + "dnc_implementation": dnc_implementation, "dnc_compute_precision": "fp32", "dnc_backward_trace": dnc_backward_trace, "workspace_enabled": workspace_enabled, @@ -3844,6 +3860,22 @@ def _write_diagnostic_values(diagnostics): help="Trace the first non-finite gradient through DNC internals and report " "windowed memory/read/link magnitudes. Diagnostic only; adds synchronization.", ) + dnc_impl_group = parser.add_mutually_exclusive_group() + dnc_impl_group.add_argument( + "--base-dnc", + dest="dnc_implementation", + action="store_const", + const="base", + help="Use the upstream DNC allocation implementation (default).", + ) + dnc_impl_group.add_argument( + "--optimized-dnc", + dest="dnc_implementation", + action="store_const", + const="optimized", + help="Use exact dense DNC equations with one allocation permutation sort removed.", + ) + parser.set_defaults(dnc_implementation="base") parser.add_argument( "--dataset-type", type=str, @@ -4104,6 +4136,8 @@ def _write_diagnostic_values(diagnostics): run_id += f"_obsvar{args.dnc_write_observation_variance:g}" if args.split_graph_combiner_mode == "controller": run_id = f"{run_id}_combctrl{args.split_graph_combiner_variant.replace('+', '')}" + if args.dnc_implementation == "optimized": + run_id += "_optimizeddnc" if args.dataset_type != "graph": run_id = f"{run_id}_ds{args.dataset_type}" if args.run_id_suffix: @@ -4165,6 +4199,7 @@ def _write_diagnostic_values(diagnostics): dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, dnc_write_variant=args.dnc_write_variant, dnc_write_observation_variance=args.dnc_write_observation_variance, + dnc_implementation=args.dnc_implementation, dnc_backward_trace=args.dnc_backward_trace, dataset_type=args.dataset_type, dataset_link=args.dataset_link, diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index d511f49..927a76b 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -187,6 +187,7 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea dnc_write_observation_variance=c.get( "dnc_write_observation_variance", cc.DNC_WRITE_OBSERVATION_VARIANCE ), + dnc_implementation=c.get("dnc_implementation", "base"), independent_linears=True, moe_enabled=_moe["moe_enabled"], moe_placement=c.get("split_graph_moe_placement", "all"), @@ -252,6 +253,7 @@ def load_model(ckpt: dict, device: torch.device, deterministic_write: bool = Fal batch_first=True, device=device, independent_linears=True, + dnc_implementation=raw.get("dnc_implementation", "base"), **_controller_kwargs(controller, c), **extra, ).to(device) diff --git a/src/initium/mamba_controller/mamba_controller.py b/src/initium/mamba_controller/mamba_controller.py index 7eb0302..d0cd329 100644 --- a/src/initium/mamba_controller/mamba_controller.py +++ b/src/initium/mamba_controller/mamba_controller.py @@ -169,6 +169,7 @@ from dnc.util import cuda from initium.MoE.moe_layer import MoEBlock, MoERNNWrapper +from initium.memory_manipulation.dnc_mechanisms import OptimizedMemory _MAMBA_IMPORT_ERROR: str | None @@ -603,7 +604,11 @@ def __init__( moe_capacity_factor: float = 1.5, moe_load_balance_alpha: float = 0.01, moe_top_k: int = 1, + dnc_implementation: str = "base", ): + if dnc_implementation not in ("base", "optimized"): + raise ValueError("dnc_implementation must be 'base' or 'optimized'") + memory_cls = OptimizedMemory if dnc_implementation == "optimized" else Memory if rnn_type.lower() not in ("mamba", "mamba2", "mamba3", "cfc") and not is_hybrid_rnn_type( rnn_type ): @@ -626,6 +631,10 @@ def __init__( clip=clip, device=device, ) + if dnc_implementation == "optimized": + for memory in self.memories: + memory.__class__ = OptimizedMemory + self.dnc_implementation = dnc_implementation if moe_enabled: # Basic MoE for the plain LSTM/GRU/RNN baseline (per the # roadmap: the shallow baseline controller only needs a @@ -697,6 +706,7 @@ def __init__( self.debug = debug self.clip = clip self.device = device + self.dnc_implementation = dnc_implementation # Mamba-only bookkeeping, stored for __repr__ / checkpoint metadata. self.mamba_d_state = mamba_d_state @@ -887,7 +897,7 @@ def __init__( # memories for each layer -- copied verbatim from dnc.DNC.__init__ if not self.share_memory_between_layers: self.memories.append( - Memory( + memory_cls( input_size=self.output_size, nr_cells=self.nr_cells, cell_size=self.w, @@ -902,7 +912,7 @@ def __init__( # dnc.DNC.__init__ if self.share_memory_between_layers: self.memories.append( - Memory( + memory_cls( input_size=self.output_size, nr_cells=self.nr_cells, cell_size=self.w, diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index ac0dffd..7fad06d 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -204,6 +204,7 @@ def __init__( dnc_read_residual_max_ratio: float = 0.5, dnc_write_variant: str = "learned", dnc_write_observation_variance: float = 1.0, + dnc_implementation: str = "base", workspace_enabled: bool = False, workspace_num_slots: int = 5, workspace_no_selection: bool = False, @@ -240,6 +241,9 @@ def __init__( self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio self.dnc_write_variant = dnc_write_variant self.dnc_write_observation_variance = dnc_write_observation_variance + if dnc_implementation not in ("base", "optimized"): + raise ValueError("dnc_implementation must be 'base' or 'optimized'") + self.dnc_implementation = dnc_implementation self.moe_backbone_residual_scale = moe_backbone_residual_scale if moe_placement not in ("all", "backbone", "combiner"): raise ValueError("moe_placement must be 'all', 'backbone', or 'combiner'") @@ -503,6 +507,7 @@ def __init__( read_residual_scale=dnc_read_residual_scale, read_residual_max_ratio=dnc_read_residual_max_ratio, observation_variance=dnc_write_observation_variance, + dnc_implementation=dnc_implementation, ) ) # See module docstring's API-compatibility section for why this diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py index b8fd456..4b4f229 100644 --- a/src/initium/memory_manipulation/dnc_mechanisms.py +++ b/src/initium/memory_manipulation/dnc_mechanisms.py @@ -1,9 +1,8 @@ """Composable DNC read-addressing and write-value mechanisms. -The stock pytorch-dnc Memory is kept as the baseline. This subclass changes -only the read content score and/or the update applied to the selected rows; -usage, allocation, erase/write addressing, read modes and temporal links keep -their upstream implementations and semantics. +The stock pytorch-dnc Memory is kept as the baseline. MechanismMemory changes +only selected read/write rules; OptimizedMechanismMemory uses the same dense +DNC equations and learned state while avoiding a redundant allocation sort. """ from __future__ import annotations @@ -12,7 +11,7 @@ import torch.nn as nn import torch.nn.functional as F from dnc.memory import Memory -from dnc.util import σ, θ +from dnc.util import δ, σ, θ from initium.memory_manipulation.nvrtc_compat import ( pop_product_numerical_diagnostics, set_product_backward_trace, @@ -20,6 +19,42 @@ READ_VARIANTS = ("cosine", "relational-mlp", "relational-residual") WRITE_VARIANTS = ("learned", "kanerva-closed-form") +DNC_IMPLEMENTATIONS = ("base", "optimized") + + +class _FastAllocationMixin: + """Exact DNC allocation with one full sort instead of two. + + The stock implementation sorts usage with topk(k=N) and then sorts the + returned permutation a second time to invert it. Scattering the sorted + allocation weights back through the first permutation is equivalent and + avoids that second full sort without changing addressing semantics or + module state. + """ + + dnc_implementation = "optimized" + nr_cells: int + + def allocate(self, usage: torch.Tensor, write_gate: torch.Tensor): + usage = δ + (1 - δ) * usage + sorted_usage, order = torch.topk( + usage, self.nr_cells, dim=1, largest=False + ) + leading_one = torch.ones( + (usage.size(0), 1), device=usage.device, dtype=usage.dtype + ) + exclusive_products = torch.cumprod( + torch.cat((leading_one, sorted_usage), dim=1), dim=1 + )[:, :-1] + sorted_weights = (1 - sorted_usage) * exclusive_products + allocation_weights = torch.zeros_like(usage).scatter( + 1, order, sorted_weights + ) + return allocation_weights.unsqueeze(1), usage + + +class OptimizedMemory(_FastAllocationMixin, Memory): + """Stock DNC memory with an exact allocation-sort optimization.""" class MechanismMemory(Memory): @@ -30,6 +65,8 @@ class MechanismMemory(Memory): from low-precision intermediate gradients near the DNC norm epsilon. """ + dnc_implementation = "base" + def __init__( self, *args, @@ -606,6 +643,10 @@ def write( return hidden +class OptimizedMechanismMemory(_FastAllocationMixin, MechanismMemory): + """MechanismMemory with the exact allocation-sort optimization.""" + + def build_memory( *, input_size: int, @@ -619,9 +660,17 @@ def build_memory( read_residual_scale: float = 1.0, read_residual_max_ratio: float = 0.5, observation_variance: float = 1.0, + dnc_implementation: str = "base", ): """Construct a stock-compatible memory for the selected mechanism pair.""" - return MechanismMemory( + if dnc_implementation not in DNC_IMPLEMENTATIONS: + raise ValueError(f"Unknown DNC implementation: {dnc_implementation!r}") + memory_cls = ( + OptimizedMechanismMemory + if dnc_implementation == "optimized" + else MechanismMemory + ) + return memory_cls( input_size=input_size, nr_cells=nr_cells, cell_size=cell_size, diff --git a/src/initium/memory_manipulation/dynamic_memory_resize.py b/src/initium/memory_manipulation/dynamic_memory_resize.py index 98256b3..cbffc88 100644 --- a/src/initium/memory_manipulation/dynamic_memory_resize.py +++ b/src/initium/memory_manipulation/dynamic_memory_resize.py @@ -50,7 +50,11 @@ import torch.nn as nn from dnc.memory import Memory -from initium.memory_manipulation.dnc_mechanisms import MechanismMemory, build_memory +from initium.memory_manipulation.dnc_mechanisms import ( + MechanismMemory, + OptimizedMemory, + build_memory, +) from initium.memory_manipulation.link_matrix_ablation import AblatableSparseLinkMemory @@ -104,6 +108,7 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz independent_linears=True, device=device, ) + dnc_implementation = getattr(old_memory, "dnc_implementation", "base") if isinstance(old_memory, MechanismMemory): new_memory = build_memory( **memory_args, @@ -112,9 +117,11 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz read_residual_scale=old_memory.read_residual_scale, read_residual_max_ratio=old_memory.read_residual_max_ratio, observation_variance=old_memory.observation_variance, + dnc_implementation=dnc_implementation, ).to(device) else: - new_memory = Memory(**memory_args).to(device) + memory_cls = OptimizedMemory if dnc_implementation == "optimized" else Memory + new_memory = memory_cls(**memory_args).to(device) # Transplant every N-independent learned sublayer except # write_vector_transform (that one is a StochasticWriteHead, not a diff --git a/tests/unit/test_optimized_dnc_memory.py b/tests/unit/test_optimized_dnc_memory.py new file mode 100644 index 0000000..76dd683 --- /dev/null +++ b/tests/unit/test_optimized_dnc_memory.py @@ -0,0 +1,100 @@ +import torch +from dnc.memory import Memory + +from initium.memory_manipulation.dnc_mechanisms import ( + MechanismMemory, + OptimizedMechanismMemory, + OptimizedMemory, + build_memory, +) + + +def test_optimized_allocation_matches_base_values_and_gradients(): + base = Memory(input_size=8, nr_cells=19, cell_size=6, read_heads=2) + optimized = OptimizedMemory(input_size=8, nr_cells=19, cell_size=6, read_heads=2) + optimized.load_state_dict(base.state_dict()) + + base_usage = torch.rand(3, 19, requires_grad=True) + optimized_usage = base_usage.detach().clone().requires_grad_(True) + write_gate = torch.ones(3, 1) + + base_weights, base_adjusted_usage = base.allocate(base_usage, write_gate) + optimized_weights, optimized_adjusted_usage = optimized.allocate( + optimized_usage, write_gate + ) + + torch.testing.assert_close(optimized_weights, base_weights) + torch.testing.assert_close(optimized_adjusted_usage, base_adjusted_usage) + + weight = torch.randn_like(base_weights) + (base_weights * weight).sum().backward() + (optimized_weights * weight).sum().backward() + torch.testing.assert_close(optimized_usage.grad, base_usage.grad) + + +def test_optimized_mechanism_memory_preserves_state_dict_and_factory_choice(): + base = build_memory( + input_size=8, + nr_cells=19, + cell_size=6, + read_heads=2, + dnc_implementation="base", + ) + optimized = build_memory( + input_size=8, + nr_cells=19, + cell_size=6, + read_heads=2, + dnc_implementation="optimized", + ) + + assert isinstance(base, MechanismMemory) + assert not isinstance(base, OptimizedMechanismMemory) + assert isinstance(optimized, OptimizedMechanismMemory) + assert base.state_dict().keys() == optimized.state_dict().keys() + optimized.load_state_dict(base.state_dict(), strict=True) + + +def test_optimized_memory_forward_and_parameter_gradients_match_base(): + base = build_memory( + input_size=8, + nr_cells=11, + cell_size=5, + read_heads=2, + dnc_implementation="base", + ) + optimized = build_memory( + input_size=8, + nr_cells=11, + cell_size=5, + read_heads=2, + dnc_implementation="optimized", + ) + optimized.load_state_dict(base.state_dict(), strict=True) + base_input = torch.randn(2, 8, requires_grad=True) + optimized_input = base_input.detach().clone().requires_grad_(True) + + base_read, base_state = base(base_input, base.new(batch_size=2)) + optimized_read, optimized_state = optimized( + optimized_input, optimized.new(batch_size=2) + ) + + torch.testing.assert_close(optimized_read, base_read) + for key in base_state: + torch.testing.assert_close(optimized_state[key], base_state[key]) + + base_loss = base_read.square().sum() + sum( + value.square().sum() for value in base_state.values() + ) + optimized_loss = optimized_read.square().sum() + sum( + value.square().sum() for value in optimized_state.values() + ) + base_loss.backward() + optimized_loss.backward() + + torch.testing.assert_close(optimized_input.grad, base_input.grad) + base_params = dict(base.named_parameters()) + optimized_params = dict(optimized.named_parameters()) + assert base_params.keys() == optimized_params.keys() + for name, parameter in base_params.items(): + torch.testing.assert_close(optimized_params[name].grad, parameter.grad) From 1b0b4390f20dbb2cc26b781047b7d3ca4746cb4f Mon Sep 17 00:00:00 2001 From: Samuel Date: Wed, 14 Oct 2026 10:52:38 +0200 Subject: [PATCH 12/14] Stabilize bistable gate gradients and diagnostics Add configurable truncated or unrolled gradients for bistable settling, defaulting to truncated backpropagation while preserving the full forward trajectory. Extend gradient and activation diagnostics, route CfC steps through module hooks, and validate frozen gate-only training parameters. --- src/initium/LNN_controller/cfc_controller.py | 19 +- src/initium/core_training.py | 169 +++++++++++++++++- src/initium/diagnostics/segment_logging.py | 19 +- src/initium/inference/model_loader.py | 3 + .../mamba_controller/split_graph_dnc.py | 16 +- src/initium/workspace/bistability.py | 18 +- src/initium/workspace/broadcast.py | 2 + 7 files changed, 236 insertions(+), 10 deletions(-) diff --git a/src/initium/LNN_controller/cfc_controller.py b/src/initium/LNN_controller/cfc_controller.py index 3edf67a..6b032ca 100644 --- a/src/initium/LNN_controller/cfc_controller.py +++ b/src/initium/LNN_controller/cfc_controller.py @@ -113,6 +113,15 @@ def step(self, x, state): y = x + out.to(x.dtype) if self.residual else out.to(x.dtype) return y, _state_to_fp32(new_state) + def forward(self, x, state): + """Run one recurrent step through ``Module.__call__``. + + Keeping the step behind ``forward`` preserves the equations while + ensuring module hooks (including segment activation diagnostics) see + the CfC block boundary. + """ + return self.step(x, state) + class CfCControllerWrapper(nn.Module): """Stack of `num_blocks` CfCControllerBlocks with the nn.LSTM-compatible @@ -247,6 +256,12 @@ def init_state(self, batch_size, device=None, dtype=None): return [blk.init_state(batch_size, device=device) for blk in blocks] def forward(self, input, hx): + if isinstance(input, list): + if self.moe_source_dims is None: + raise RuntimeError( + "list input is supported only when moe_source_dims is configured" + ) + return self.forward_multi_source(input, hx) assert input.dim() == 3 and input.size(1) == 1, ( "CfCControllerWrapper only supports single-timestep calls " f"(got shape {tuple(input.shape)})." @@ -262,7 +277,7 @@ def forward(self, input, hx): new_hx = [] blocks = cast(Iterable[CfCControllerBlock], self.blocks) for i, (block, state) in enumerate(zip(blocks, hx)): - x, new_state = block.step(x, state) + x, new_state = block(x, state) if self.moe_enabled: assert self.moe_blocks is not None x = self.moe_blocks[i](x) @@ -303,7 +318,7 @@ def forward_multi_source(self, sources: list[torch.Tensor], hx): new_hx = [] blocks = cast(Iterable[CfCControllerBlock], self.blocks) for block, state in zip(blocks, hx): - x, new_state = block.step(x, state) + x, new_state = block(x, state) new_hx.append(new_state) return x.unsqueeze(1), new_hx diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 73defe7..9f16266 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -690,6 +690,7 @@ def load_frozen_workspace_warm_start( "workspace_gate_variant", "workspace_bistable_selection_scale", "workspace_bistable_noise_scale", + "workspace_bistable_gradient_mode", "workspace_bistable_deterministic", "workspace_stability_max_activation", # Exact implementation-only optimization; old checkpoints default @@ -822,6 +823,7 @@ def run( workspace_gate_variant: str = "continuous", workspace_bistable_selection_scale: float = 1.0, workspace_bistable_noise_scale: float = 0.75, + workspace_bistable_gradient_mode: str = "truncated", workspace_bistable_deterministic: bool = False, workspace_stability_max_activation: float = 100_000.0, eval_only: bool = False, @@ -905,6 +907,13 @@ def run( or checkpoint_config.get("workspace_train_mode") != workspace_train_mode or checkpoint_config.get("workspace_gate_variant", "continuous") != workspace_gate_variant + or ( + workspace_gate_variant == "bistable" + and checkpoint_config.get( + "workspace_bistable_gradient_mode", "unrolled" + ) + != workspace_bistable_gradient_mode + ) ) and not eval_only ) @@ -1524,6 +1533,7 @@ def run( workspace_gate_variant=workspace_gate_variant, workspace_bistable_selection_scale=workspace_bistable_selection_scale, workspace_bistable_noise_scale=workspace_bistable_noise_scale, + workspace_bistable_gradient_mode=workspace_bistable_gradient_mode, device=device, ).to(device) if dnc_backward_trace: @@ -1585,6 +1595,7 @@ def run( f"slots={workspace_num_slots} | write_selection={write_mode} | " f"persistence={persistence} | train_mode={workspace_train_mode} | " f"gate={workspace_gate_variant} | " + f"bistable_gradient={workspace_bistable_gradient_mode} | " f"compute_precision={'fp32' if workspace_active else 'bypass'} | " "scope=final-backbone-MoE-bank; each specialist reads shared slots" ) @@ -1650,6 +1661,7 @@ def run( "workspace_train_scope": workspace_train_scope, "workspace_bistable_selection_scale": workspace_bistable_selection_scale, "workspace_bistable_noise_scale": workspace_bistable_noise_scale, + "workspace_bistable_gradient_mode": workspace_bistable_gradient_mode, "workspace_bistable_deterministic": workspace_bistable_deterministic, "workspace_stability_max_activation": workspace_stability_max_activation, } @@ -1752,7 +1764,27 @@ def run( else None ), ) + if workspace_train_scope == "gate" and isinstance(rnn, SplitGraphDNC): + rnn._initium_segment_observer = segment_diagnostics.observe_tensor segment_diagnostics.add_module(output_proj, "task_output") + if segment_diagnostics_enabled and isinstance(rnn, SplitGraphDNC): + combiner = rnn.combiner_wrapper + if combiner is not None and hasattr(combiner, "in_adapter"): + segment_diagnostics.add_module( + combiner.in_adapter, "controller.cfc.in_adapter" + ) + for index, adapter in enumerate( + getattr(combiner, "source_in_adapters", None) or () + ): + segment_diagnostics.add_module( + adapter, f"controller.cfc.source_adapter.{index}" + ) + for index, block in enumerate(getattr(combiner, "blocks", ())): + cfc_cell = getattr(block, "cfc", None) + if isinstance(cfc_cell, nn.Module): + segment_diagnostics.add_module( + cfc_cell, f"controller.cfc.cell.{index}" + ) moe_layer_ids = {id(layer) for layer in getattr(rnn, "moe_layers", [])} for module_name, module in rnn.named_modules(): if id(module) in moe_layer_ids: @@ -1841,13 +1873,37 @@ def report_bistable_evaluation(mode: str) -> None: parameter.requires_grad_(True) for parameter in output_proj.parameters(): parameter.requires_grad_(False) - trainable_parameters = [ - parameter - for parameter in list(rnn.parameters()) + list(output_proj.parameters()) + trainable_named_parameters = [ + (f"rnn.{name}", parameter) + for name, parameter in rnn.named_parameters() + if parameter.requires_grad + ] + [ + (f"output_proj.{name}", parameter) + for name, parameter in output_proj.named_parameters() if parameter.requires_grad ] + trainable_parameters = [parameter for _, parameter in trainable_named_parameters] if not trainable_parameters: raise ValueError("training configuration has no trainable parameters") + if ( + workspace_enabled + and workspace_train_mode == "frozen" + and workspace_train_scope == "gate" + and workspace_gate_variant in ("continuous", "bistable") + ): + gate_prefix = ( + "rnn.workspace.bistable_gate." + if workspace_gate_variant == "bistable" + else "rnn.workspace.write_gate." + ) + unexpected_trainable = [ + name for name, _ in trainable_named_parameters if not name.startswith(gate_prefix) + ] + if unexpected_trainable: + raise RuntimeError( + "Frozen gate-only training has parameters outside the selected write gate: " + + ", ".join(unexpected_trainable) + ) optimizer = torch.optim.Adam(trainable_parameters, lr=LR) # Option 3 (Concept 24/LB-17): multi-axis efficiency accounting. # param_count is fixed for the life of this run -- static Option 2's @@ -2200,6 +2256,29 @@ def set_lr(step_now): running_task_loss, running_kl_loss, running_div, running_grad_norm = 0.0, 0.0, 0, 0.0 running_finite_grad_steps = 0 running_nonfinite_grad_steps = 0 + bistable_grad_probe = None + if workspace_train_scope == "gate" and workspace_gate_variant == "bistable": + # SegmentDiagnostics reports parameter gradients after clipping. Keep + # a separate pre-clip trace so an unexpectedly large global norm can + # be attributed to the exact trainable tensors that produced it. + bistable_grad_probe = { + name: { + "sum_sq": torch.zeros((), device=device, dtype=torch.float64), + "max": torch.zeros((), device=device, dtype=torch.float64), + "max_step": torch.full((), -1, device=device, dtype=torch.int64), + } + for name, _ in trainable_named_parameters + } + bistable_grad_probe_steps = 0 + bistable_grad_probe_global_max = torch.zeros( + (), device=device, dtype=torch.float64 + ) + bistable_grad_probe_global_max_step = torch.full( + (), -1, device=device, dtype=torch.int64 + ) + bistable_grad_probe_relative_error = torch.zeros( + (), device=device, dtype=torch.float64 + ) # Dynamic-beta controller state (no-op / unused when beta_mode == "static"). # Initialized post-resume so a resumed dynamic run starts its health check @@ -2338,6 +2417,41 @@ def set_lr(step_now): ) scaler.scale(loss).backward() scaler.unscale_(optimizer) + independent_grad_norm = None + if bistable_grad_probe is not None: + step_parameter_norms = [] + for name, parameter in trainable_named_parameters: + if parameter.grad is None: + continue + parameter_norm = torch.linalg.vector_norm( + parameter.grad.detach().to(dtype=torch.float64) + ) + probe = bistable_grad_probe[name] + probe["sum_sq"].add_(parameter_norm.square()) + probe["max_step"].copy_( + torch.where( + parameter_norm > probe["max"], + torch.full_like(probe["max_step"], step), + probe["max_step"], + ) + ) + probe["max"].copy_(torch.maximum(probe["max"], parameter_norm)) + step_parameter_norms.append(parameter_norm) + if step_parameter_norms: + independent_grad_norm = torch.linalg.vector_norm( + torch.stack(step_parameter_norms) + ) + bistable_grad_probe_global_max_step.copy_( + torch.where( + independent_grad_norm > bistable_grad_probe_global_max, + torch.full_like(bistable_grad_probe_global_max_step, step), + bistable_grad_probe_global_max_step, + ) + ) + bistable_grad_probe_global_max.copy_( + torch.maximum(bistable_grad_probe_global_max, independent_grad_norm) + ) + bistable_grad_probe_steps += 1 bistable_gate_grad_norm = 0.0 if workspace_gate_variant == "bistable": assert isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None @@ -2381,6 +2495,13 @@ def set_lr(step_now): # error_if_nonfinite raises before applying the clip coefficient, # preserving the original bad gradients for the diagnostic below. grad_norm = torch.tensor(float("nan"), device=device) + if independent_grad_norm is not None and math.isfinite(float(grad_norm)): + relative_error = ( + independent_grad_norm - grad_norm.detach().to(dtype=torch.float64) + ).abs() / independent_grad_norm.clamp_min(1e-12) + bistable_grad_probe_relative_error.copy_( + torch.maximum(bistable_grad_probe_relative_error, relative_error) + ) if not math.isfinite(float(grad_norm)): grad_was_finite_since_eval = False consecutive_nonfinite_grad_steps += 1 @@ -2587,6 +2708,36 @@ def set_lr(step_now): print(f"[{run_id}] [workspace] writer_frequency={writer_frequency}") if dyn_n_str: parts.append(dyn_n_str.strip(" |")) + if bistable_grad_probe is not None: + probe_steps = max(bistable_grad_probe_steps, 1) + per_parameter = [] + for name, probe in bistable_grad_probe.items(): + rms = (probe["sum_sq"] / probe_steps).sqrt() + maximum = probe["max"] + maximum_step = int(probe["max_step"].cpu()) + per_parameter.append( + (float(maximum.cpu()), name, float(rms.cpu()), maximum_step) + ) + per_parameter.sort(reverse=True) + parameter_summary = ", ".join( + f"{name}:rms={rms:.3g},max={maximum:.3g}@step{maximum_step}" + for maximum, name, rms, maximum_step in per_parameter[:8] + ) or "none" + parts.append( + "bistable_grad_probe[" + f"trainable={len(trainable_named_parameters)} steps={bistable_grad_probe_steps} " + f"independent_global_max={float(bistable_grad_probe_global_max.cpu()):.3g}" + f"@step{int(bistable_grad_probe_global_max_step.cpu())} " + f"max_norm_relative_error={float(bistable_grad_probe_relative_error.cpu()):.3g} " + f"parameters={parameter_summary}]" + ) + for probe in bistable_grad_probe.values(): + probe["sum_sq"].zero_() + probe["max"].zero_() + bistable_grad_probe_steps = 0 + bistable_grad_probe_global_max.zero_() + bistable_grad_probe_global_max_step.fill_(-1) + bistable_grad_probe_relative_error.zero_() print(" | ".join(parts)) segment_diagnostics.emit("train", include_gradients=True) if dnc_backward_trace and split_graph_enabled: @@ -3698,6 +3849,16 @@ def _write_diagnostic_values(diagnostics): default=0.75, help="Independent background-noise scale for bistable ignition; set to 0 for the noise ablation.", ) + parser.add_argument( + "--workspace-bistable-gradient-mode", + choices=("truncated", "unrolled"), + default="truncated", + help=( + "Backward mode for bistable settling. truncated preserves the full " + "forward iteration but detaches intermediate settling states; " + "unrolled backpropagates through every iteration." + ), + ) parser.add_argument( "--workspace-bistable-deterministic", action="store_true", @@ -4116,6 +4277,7 @@ def _write_diagnostic_values(diagnostics): run_id += ( f"_sel{args.workspace_bistable_selection_scale:g}" f"_noise{args.workspace_bistable_noise_scale:g}" + f"_grad{args.workspace_bistable_gradient_mode}" ) if args.workspace_bistable_deterministic: run_id += "_detgate" @@ -4189,6 +4351,7 @@ def _write_diagnostic_values(diagnostics): workspace_gate_variant=args.workspace_gate_variant, workspace_bistable_selection_scale=args.workspace_bistable_selection_scale, workspace_bistable_noise_scale=args.workspace_bistable_noise_scale, + workspace_bistable_gradient_mode=args.workspace_bistable_gradient_mode, workspace_bistable_deterministic=args.workspace_bistable_deterministic, workspace_stability_max_activation=args.workspace_stability_max_activation, eval_only=args.eval_only, diff --git a/src/initium/diagnostics/segment_logging.py b/src/initium/diagnostics/segment_logging.py index 97cf528..5cfaff9 100644 --- a/src/initium/diagnostics/segment_logging.py +++ b/src/initium/diagnostics/segment_logging.py @@ -147,7 +147,10 @@ def _accumulate_tensor(record: dict[str, Any], prefix: str, tensor: torch.Tensor f"{prefix}_bad": (~finite).sum().float(), } for key, item in values.items(): - record[key] = record[key] + item if key in record else item + if key in record and key.endswith("_max"): + record[key] = torch.maximum(record[key], item) + else: + record[key] = record[key] + item if key in record else item record[f"{prefix}_numel"] = record.get(f"{prefix}_numel", 0) + value.numel() def add_module(self, module: nn.Module, label: str) -> None: @@ -160,6 +163,20 @@ def add_module(self, module: nn.Module, label: str) -> None: def set_channel(self, channel: str | None) -> None: self.channel = channel + def observe_tensor(self, label: str, tensor: torch.Tensor) -> None: + """Record a named tensor at a call site that has no module boundary. + + This keeps diagnostics on-device and lets recurrent models expose + exact per-step values without adding synchronization to their loop. + """ + if not self.enabled or self.channel is None: + return + record = self._records.setdefault(self.channel, OrderedDict()).setdefault( + label, {"path": label, "calls": 0} + ) + record["calls"] += 1 + self._accumulate_tensor(record, "out", tensor) + def _gradient_records(self) -> OrderedDict[str, tuple[torch.Tensor, int]]: records: OrderedDict[str, list[tuple[torch.Tensor, int]]] = OrderedDict() unassigned: list[tuple[torch.Tensor, int]] = [] diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 927a76b..b671371 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -220,6 +220,9 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea "workspace_bistable_selection_scale", 1.0 ), workspace_bistable_noise_scale=c.get("workspace_bistable_noise_scale", 0.75), + workspace_bistable_gradient_mode=c.get( + "workspace_bistable_gradient_mode", "truncated" + ), device=device, ).to(device) model.workspace_bistable_deterministic = c.get( diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 7fad06d..e1f98e4 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -213,6 +213,7 @@ def __init__( workspace_gate_variant: str = "continuous", workspace_bistable_selection_scale: float = 1.0, workspace_bistable_noise_scale: float = 0.75, + workspace_bistable_gradient_mode: str = "truncated", ): super().__init__() if not independent_linears: @@ -268,6 +269,7 @@ def __init__( self.workspace_reset_each_step = workspace_reset_each_step self.workspace_train_mode = workspace_train_mode self.workspace_bistable_noise_scale = float(workspace_bistable_noise_scale) + self.workspace_bistable_gradient_mode = workspace_bistable_gradient_mode self.workspace_bistable_selection_scale = float(workspace_bistable_selection_scale) self.workspace_bistable_deterministic = False # Runtime-only counterfactual for evaluation. It is deliberately not @@ -534,6 +536,7 @@ def __init__( gate_variant=workspace_gate_variant, bistable_selection_scale=workspace_bistable_selection_scale, bistable_noise_scale=workspace_bistable_noise_scale, + bistable_gradient_mode=workspace_bistable_gradient_mode, ) if workspace_enabled and workspace_num_slots > 0 else None @@ -621,6 +624,11 @@ def forward( raise ValueError("start_step>0 requires a resumed hx=(chx, mhx, last_read)") device = input.device + def observe_tensor(label: str, value: torch.Tensor) -> None: + observer = getattr(self, "_initium_segment_observer", None) + if observer is not None: + observer(label, value) + # ---- (a) PARALLEL BACKBONE ------------------------------------ # Single call, whole sequence, zero dependency on memory state. # This is the entire "removed from the sequential critical path" @@ -664,6 +672,7 @@ def forward( H[:, start_step:] = base[:, start_step:] + ( self.moe_backbone_residual_scale * broadcast_delta ) + observe_tensor("split_graph.h_after_workspace", H) else: backbone_stages = ( list(self.backbone) @@ -717,6 +726,8 @@ def forward( outputs = [] for t in range(start_step, T): h_t = H[:, t, :] # (B, hidden_size) + observe_tensor("split_graph.backbone_source", h_t) + observe_tensor("split_graph.read_source", read_vec) if self.combiner_mode == "controller": assert self.combiner_wrapper is not None @@ -725,7 +736,9 @@ def forward( and getattr(self.combiner_wrapper, "moe_source_dims", None) is not None ) if _cfc_multi_source: - xi_out, combiner_hx = self.combiner_wrapper.forward_multi_source( + # Call the nn.Module entry point so registered segment + # diagnostics observe the CfC combiner boundary. + xi_out, combiner_hx = self.combiner_wrapper( [h_t, read_vec], combiner_hx ) else: @@ -748,6 +761,7 @@ def forward( xi_t = h_t if pass_through_memory: + observe_tensor("split_graph.dnc_input", xi_t) read_vec, mem_state = self.memories[0](xi_t, mem_state) read_vec = read_vec.reshape( B, -1 diff --git a/src/initium/workspace/bistability.py b/src/initium/workspace/bistability.py index e463322..f0fad3b 100644 --- a/src/initium/workspace/bistability.py +++ b/src/initium/workspace/bistability.py @@ -24,6 +24,7 @@ def __init__( noise_scale: float = 0.75, iterations: int = 6, temperature: float = 0.4, + gradient_mode: str = "truncated", ) -> None: super().__init__() if selection_scale < 0 or noise_scale < 0: @@ -32,10 +33,13 @@ def __init__( raise ValueError("iterations must be >= 1") if temperature <= 0: raise ValueError("temperature must be > 0") + if gradient_mode not in ("truncated", "unrolled"): + raise ValueError("gradient_mode must be 'truncated' or 'unrolled'") self.selection_scale = float(selection_scale) self.noise_scale = float(noise_scale) self.iterations = int(iterations) self.temperature = float(temperature) + self.gradient_mode = gradient_mode self.bias = nn.Parameter(torch.tensor(1.1)) self.self_excitation_logit = nn.Parameter(torch.tensor(1.1)) self.inhibition_logit = nn.Parameter(torch.tensor(0.5)) @@ -137,10 +141,18 @@ def forward( inhibition = F.softplus(self.inhibition_logit) activity = torch.sigmoid(drive / self.temperature) for _ in range(self.iterations): - population_activity = activity.sum(dim=-1, keepdim=True) - lateral_input = population_activity - activity + # Keep the competitive settling trajectory intact in the forward + # pass. Its recurrent Jacobian can be expansive by design, so the + # default backward pass truncates gradients at each settling step + # instead of multiplying them through the full positive-feedback + # loop. ``unrolled`` remains available for controlled comparisons. + previous_activity = ( + activity.detach() if self.gradient_mode == "truncated" else activity + ) + population_activity = previous_activity.sum(dim=-1, keepdim=True) + lateral_input = population_activity - previous_activity activity = torch.sigmoid( - (drive + self_excitation * activity - inhibition * lateral_input) + (drive + self_excitation * previous_activity - inhibition * lateral_input) / self.temperature ) diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index 15e0ce1..838c9f1 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -32,6 +32,7 @@ def __init__( gate_variant: str = "continuous", bistable_selection_scale: float = 1.0, bistable_noise_scale: float = 0.75, + bistable_gradient_mode: str = "truncated", ): super().__init__() if d_model < 1 or num_experts < 1: @@ -56,6 +57,7 @@ def __init__( BistableWriteGate( selection_scale=bistable_selection_scale, noise_scale=bistable_noise_scale, + gradient_mode=bistable_gradient_mode, ) if gate_variant == "bistable" else None From c2c236751525681568b759eb88dee40fbec6d8f0 Mon Sep 17 00:00:00 2001 From: Samuel Date: Wed, 14 Oct 2026 21:52:46 +0200 Subject: [PATCH 13/14] Add configurable workspace BPTT truncation Allow gradients through persistent workspace slots to be truncated at a configurable horizon without changing the forward trajectory. Wire the setting through training, checkpoints, and inference; reset optimizer state when the horizon changes; and extend recurrent diagnostics with state-path metrics and gradient clipping counts. --- src/initium/core_training.py | 62 +++- src/initium/diagnostics/segment_logging.py | 267 ++++++++++++++++-- src/initium/inference/model_loader.py | 1 + .../mamba_controller/split_graph_dnc.py | 3 + src/initium/workspace/broadcast.py | 12 + 5 files changed, 308 insertions(+), 37 deletions(-) diff --git a/src/initium/core_training.py b/src/initium/core_training.py index 9f16266..dd2a97a 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -691,6 +691,7 @@ def load_frozen_workspace_warm_start( "workspace_bistable_selection_scale", "workspace_bistable_noise_scale", "workspace_bistable_gradient_mode", + "workspace_bptt_horizon", "workspace_bistable_deterministic", "workspace_stability_max_activation", # Exact implementation-only optimization; old checkpoints default @@ -824,6 +825,7 @@ def run( workspace_bistable_selection_scale: float = 1.0, workspace_bistable_noise_scale: float = 0.75, workspace_bistable_gradient_mode: str = "truncated", + workspace_bptt_horizon: int = 0, workspace_bistable_deterministic: bool = False, workspace_stability_max_activation: float = 100_000.0, eval_only: bool = False, @@ -924,6 +926,14 @@ def run( and checkpoint_config.get("workspace_train_mode") != workspace_train_mode and not eval_only ) + workspace_bptt_warm_start = bool( + resume_requested + and workspace_enabled + and checkpoint_workspace_enabled + and int(checkpoint_config.get("workspace_bptt_horizon", 0)) + != workspace_bptt_horizon + and not eval_only + ) gate_variant_warm_start = bool( resume_requested and workspace_enabled @@ -932,9 +942,10 @@ def run( and checkpoint_config.get("workspace_gate_variant", "continuous") == "continuous" and not eval_only ) - phase4_gate_warm_start = bool( + workspace_condition_warm_start = bool( phase4_gate_control_warm_start or workspace_mode_warm_start + or workspace_bptt_warm_start or gate_variant_warm_start ) frozen_workspace_warm_start = bool( @@ -956,7 +967,7 @@ def run( resuming = ( resume_requested and not frozen_workspace_warm_start - and not phase4_gate_warm_start + and not workspace_condition_warm_start and not eval_only ) log_file = open(log_path, "a" if resuming else "w", newline="") @@ -1534,6 +1545,7 @@ def run( workspace_bistable_selection_scale=workspace_bistable_selection_scale, workspace_bistable_noise_scale=workspace_bistable_noise_scale, workspace_bistable_gradient_mode=workspace_bistable_gradient_mode, + workspace_bptt_horizon=workspace_bptt_horizon, device=device, ).to(device) if dnc_backward_trace: @@ -1596,6 +1608,7 @@ def run( f"persistence={persistence} | train_mode={workspace_train_mode} | " f"gate={workspace_gate_variant} | " f"bistable_gradient={workspace_bistable_gradient_mode} | " + f"workspace_bptt_horizon={workspace_bptt_horizon or 'full'} | " f"compute_precision={'fp32' if workspace_active else 'bypass'} | " "scope=final-backbone-MoE-bank; each specialist reads shared slots" ) @@ -1662,6 +1675,7 @@ def run( "workspace_bistable_selection_scale": workspace_bistable_selection_scale, "workspace_bistable_noise_scale": workspace_bistable_noise_scale, "workspace_bistable_gradient_mode": workspace_bistable_gradient_mode, + "workspace_bptt_horizon": workspace_bptt_horizon, "workspace_bistable_deterministic": workspace_bistable_deterministic, "workspace_stability_max_activation": workspace_stability_max_activation, } @@ -1795,6 +1809,10 @@ def run( if isinstance(rnn, SplitGraphDNC) and rnn.workspace is not None: rnn.workspace.set_diagnostics_enabled(moe_diagnostics_enabled) rnn.workspace_bistable_deterministic = workspace_bistable_deterministic + if workspace_train_scope == "gate": + rnn.workspace._initium_recurrent_state_observer = ( + segment_diagnostics.observe_recurrent_state + ) def emit_failure_diagnostics(channel: str, *, include_gradients: bool = False) -> None: segment_diagnostics.emit(channel, include_gradients=include_gradients) @@ -2102,7 +2120,7 @@ def set_lr(step_now): "workspace tensors loaded, gate parameters initialized fresh, " "optimizer reset." ) - elif phase4_gate_warm_start: + elif workspace_condition_warm_start: rnn.load_state_dict(ckpt["rnn_state_dict"]) output_proj.load_state_dict(ckpt["output_proj_state_dict"]) print( @@ -2118,9 +2136,9 @@ def set_lr(step_now): # of leaving `scaler` at its freshly-constructed init_scale -- see # header note. `.get(...)` instead of ["..."] so a pre-v3 checkpoint # (saved before this key existed) still loads, just without this fix. - if not frozen_workspace_warm_start and not phase4_gate_warm_start and "scaler_state_dict" in ckpt: + if not frozen_workspace_warm_start and not workspace_condition_warm_start and "scaler_state_dict" in ckpt: scaler.load_state_dict(ckpt["scaler_state_dict"]) - elif not frozen_workspace_warm_start and not phase4_gate_warm_start: + elif not frozen_workspace_warm_start and not workspace_condition_warm_start: print( f"[{run_id}] WARNING: checkpoint predates GradScaler-state " f"saving (v3 fix) -- scaler starting fresh at " @@ -2256,6 +2274,8 @@ def set_lr(step_now): running_task_loss, running_kl_loss, running_div, running_grad_norm = 0.0, 0.0, 0, 0.0 running_finite_grad_steps = 0 running_nonfinite_grad_steps = 0 + running_clipped_grad_steps = 0 + max_grad_norm = 10.0 bistable_grad_probe = None if workspace_train_scope == "gate" and workspace_gate_variant == "bistable": # SegmentDiagnostics reports parameter gradients after clipping. Keep @@ -2316,6 +2336,7 @@ def set_lr(step_now): optimizer.zero_grad(set_to_none=True) segment_diagnostics.set_channel("train") + segment_diagnostics.set_step(step) if dnc_backward_trace and split_graph_enabled: rnn.memories[0].set_backward_trace(True, step=step) with torch.amp.autocast("cuda", enabled=amp_enabled): @@ -2487,7 +2508,7 @@ def set_lr(step_now): # unbounded while its gradients were still being applied by Adam. try: grad_norm = torch.nn.utils.clip_grad_norm_( - trainable_parameters, max_norm=10.0, error_if_nonfinite=True + trainable_parameters, max_norm=max_grad_norm, error_if_nonfinite=True ) except RuntimeError as error: if "non-finite" not in str(error): @@ -2495,14 +2516,15 @@ def set_lr(step_now): # error_if_nonfinite raises before applying the clip coefficient, # preserving the original bad gradients for the diagnostic below. grad_norm = torch.tensor(float("nan"), device=device) - if independent_grad_norm is not None and math.isfinite(float(grad_norm)): + grad_norm_value = float(grad_norm) + if independent_grad_norm is not None and math.isfinite(grad_norm_value): relative_error = ( independent_grad_norm - grad_norm.detach().to(dtype=torch.float64) ).abs() / independent_grad_norm.clamp_min(1e-12) bistable_grad_probe_relative_error.copy_( torch.maximum(bistable_grad_probe_relative_error, relative_error) ) - if not math.isfinite(float(grad_norm)): + if not math.isfinite(grad_norm_value): grad_was_finite_since_eval = False consecutive_nonfinite_grad_steps += 1 running_nonfinite_grad_steps += 1 @@ -2542,9 +2564,11 @@ def set_lr(step_now): running_task_loss += task_loss.detach() running_kl_loss += kl_loss.detach() running_div += dataset.diversity(output.detach(), target_digits, answer_mask) - if math.isfinite(float(grad_norm)): - running_grad_norm += float(grad_norm) + if math.isfinite(grad_norm_value): + running_grad_norm += grad_norm_value running_finite_grad_steps += 1 + if grad_norm_value > max_grad_norm: + running_clipped_grad_steps += 1 step += 1 current_lr = set_lr(step) # Phase 1 FIX: clamped manual LR, replaces scheduler.step() @@ -2654,6 +2678,7 @@ def set_lr(step_now): f"LR {current_lr:.6f}", f"grad_norm {avg_grad_norm:.4f}", f"nonfinite_grad_steps {running_nonfinite_grad_steps}", + f"clipped_grad_steps {running_clipped_grad_steps}/{running_finite_grad_steps}", ] if amp_enabled: parts.append(f"amp_scale {amp_scale:.1f}") @@ -2817,6 +2842,7 @@ def set_lr(step_now): running_task_loss, running_kl_loss, running_div, running_grad_norm = 0.0, 0.0, 0, 0.0 running_finite_grad_steps = 0 running_nonfinite_grad_steps = 0 + running_clipped_grad_steps = 0 if kl_on and step % PRIOR_SNAPSHOT_EVERY == 0: # v5 (Phase 2): refit (mu_g, Sigma_g) from the writes @@ -3859,6 +3885,17 @@ def _write_diagnostic_values(diagnostics): "unrolled backpropagates through every iteration." ), ) + parser.add_argument( + "--workspace-bptt-horizon", + type=int, + default=0, + metavar="STEPS", + help=( + "Truncate gradients carried by persistent workspace slots every STEPS " + "sequence positions while preserving the forward trajectory; 0 keeps " + "full-sequence BPTT." + ), + ) parser.add_argument( "--workspace-bistable-deterministic", action="store_true", @@ -4197,6 +4234,8 @@ def _write_diagnostic_values(diagnostics): parser.error("--workspace requires backbone MoE; use --moe-placement backbone or all") if args.workspace_num_slots < 0: parser.error("--workspace-num-slots must be >= 0") + if args.workspace_bptt_horizon < 0: + parser.error("--workspace-bptt-horizon must be >= 0") if args.workspace_train_scope == "gate" and not args.workspace: parser.error("--workspace-train-scope=gate requires --workspace") if args.workspace_gate_variant == "bistable" and ( @@ -4281,6 +4320,8 @@ def _write_diagnostic_values(diagnostics): ) if args.workspace_bistable_deterministic: run_id += "_detgate" + if args.workspace_bptt_horizon: + run_id += f"_wbptt{args.workspace_bptt_horizon}" if args.workspace_no_selection: run_id += "_randomwrite" if args.workspace_reset_each_step: @@ -4352,6 +4393,7 @@ def _write_diagnostic_values(diagnostics): workspace_bistable_selection_scale=args.workspace_bistable_selection_scale, workspace_bistable_noise_scale=args.workspace_bistable_noise_scale, workspace_bistable_gradient_mode=args.workspace_bistable_gradient_mode, + workspace_bptt_horizon=args.workspace_bptt_horizon, workspace_bistable_deterministic=args.workspace_bistable_deterministic, workspace_stability_max_activation=args.workspace_stability_max_activation, eval_only=args.eval_only, diff --git a/src/initium/diagnostics/segment_logging.py b/src/initium/diagnostics/segment_logging.py index 5cfaff9..3740b1e 100644 --- a/src/initium/diagnostics/segment_logging.py +++ b/src/initium/diagnostics/segment_logging.py @@ -30,6 +30,28 @@ def _first_tensor(value: Any) -> torch.Tensor | None: return None +def _all_tensors(value: Any) -> list[torch.Tensor]: + """Return distinct tensor leaves in output order for structured modules.""" + tensors: list[torch.Tensor] = [] + seen: set[int] = set() + + def visit(item: Any) -> None: + if torch.is_tensor(item): + identity = id(item) + if identity not in seen: + seen.add(identity) + tensors.append(item) + elif isinstance(item, (tuple, list)): + for child in item: + visit(child) + elif isinstance(item, dict): + for child in item.values(): + visit(child) + + visit(value) + return tensors + + def _segment_label(path: str, module: nn.Module) -> str | None: class_name = type(module).__name__ lower_name = class_name.lower() @@ -93,6 +115,8 @@ def __init__( self.enabled = bool(enabled) self.activation_absmax_limit = activation_absmax_limit self.channel: str | None = "train" + self._current_step: int | None = None + self._step_call_counts: dict[str, int] = {} self._records: dict[str, OrderedDict[str, dict[str, Any]]] = {} self._segments: list[tuple[str, str]] = [] self._external_modules: list[tuple[str, nn.Module]] = [] @@ -112,32 +136,115 @@ def hook(_module, inputs, output): return channel = self.channel input_tensor = _first_tensor(inputs) - output_tensor = _first_tensor(output) + tracks_secondary_outputs = label.startswith( + ("workspace:", "workspace_bistable_gate:") + ) + output_tensors = ( + _all_tensors(output) + if tracks_secondary_outputs + else [tensor for tensor in (_first_tensor(output),) if tensor is not None] + ) + output_tensor = output_tensors[0] if output_tensors else None record = self._records.setdefault(channel, OrderedDict()).setdefault( label, {"path": path, "calls": 0} ) record["calls"] += 1 + call_index = self._step_call_counts.get(label, 0) + self._step_call_counts[label] = call_index + 1 + location = ( + (self._current_step, call_index) + if self._current_step is not None and channel == "train" + else None + ) + input_context = None + output_context = None + if ( + location is not None + and label.startswith("workspace_bistable_gate:") + and input_tensor is not None + and input_tensor.requires_grad + ): + selection = input_tensor.detach().float().clamp_min(0.0) + normalized = selection / selection.sum(dim=-1, keepdim=True).clamp_min(1e-8) + selection_min, selection_max = torch.aminmax(normalized) + input_context = { + "selection_min": selection_min, + "selection_max": selection_max, + } + if ( + location is not None + and label.startswith("workspace_bistable_gate:") + and isinstance(output, (tuple, list)) + and len(output) > 1 + and torch.is_tensor(output[1]) + and output_tensor is not None + and output_tensor.requires_grad + ): + probability = output[1].detach().float() + probability_min, probability_max = torch.aminmax(probability) + output_context = { + "prob_min": probability_min, + "prob_max": probability_max, + } if input_tensor is not None: - self._accumulate_tensor(record, "in", input_tensor) + self._accumulate_tensor(record, "in", input_tensor, location=location) if input_tensor.requires_grad: input_tensor.register_hook( - lambda grad, target=record: self._accumulate_tensor( - target, "grad_in", grad + lambda grad, target=record, call_location=location, context=input_context: self._accumulate_tensor( + target, + "grad_in", + grad, + location=call_location, + context=context, ) ) if output_tensor is not None: - self._accumulate_tensor(record, "out", output_tensor) + self._accumulate_tensor(record, "out", output_tensor, location=location) if output_tensor.requires_grad: output_tensor.register_hook( - lambda grad, target=record: self._accumulate_tensor( - target, "grad_out", grad + lambda grad, target=record, call_location=location, context=output_context: self._accumulate_tensor( + target, + "grad_out", + grad, + location=call_location, + context=context, + ) + ) + + # Structured recurrent modules can return a large persistent state + # alongside their visible output. Keep these gradients separate so + # a state-path spike cannot be hidden by the first returned tensor. + secondary_outputs = output_tensors[1:] if tracks_secondary_outputs else () + for output_index, tensor in enumerate(secondary_outputs, start=1): + output_prefix = f"out_{output_index}" + gradient_prefix = f"grad_out_{output_index}" + self._accumulate_tensor(record, output_prefix, tensor, location=location) + if tensor.requires_grad: + tensor.register_hook( + lambda grad, + target=record, + call_location=location, + prefix=gradient_prefix, + context=output_context if output_index == 1 else None: self._accumulate_tensor( + target, + prefix, + grad, + location=call_location, + context=context, ) ) self._handles.append(module.register_forward_hook(hook)) @staticmethod - def _accumulate_tensor(record: dict[str, Any], prefix: str, tensor: torch.Tensor) -> None: + def _accumulate_tensor( + record: dict[str, Any], + prefix: str, + tensor: torch.Tensor, + *, + location: tuple[int, int] | None = None, + context: dict[str, torch.Tensor] | None = None, + ) -> None: value = tensor.detach().float().reshape(-1) finite = torch.isfinite(value) safe = torch.where(finite, value, torch.zeros_like(value)) @@ -148,9 +255,35 @@ def _accumulate_tensor(record: dict[str, Any], prefix: str, tensor: torch.Tensor } for key, item in values.items(): if key in record and key.endswith("_max"): + previous = record[key] + is_new_max = item > previous record[key] = torch.maximum(record[key], item) + if location is not None: + step_key = f"{key}_step" + call_key = f"{key}_call" + old_step = record.get(step_key, torch.full_like(item, -1, dtype=torch.long)) + old_call = record.get(call_key, torch.full_like(item, -1, dtype=torch.long)) + record[step_key] = torch.where( + is_new_max, torch.full_like(old_step, location[0]), old_step + ) + record[call_key] = torch.where( + is_new_max, torch.full_like(old_call, location[1]), old_call + ) + for context_name, context_value in (context or {}).items(): + context_key = f"{key}_context_{context_name}" + previous_context = record.get( + context_key, torch.zeros_like(context_value) + ) + record[context_key] = torch.where( + is_new_max, context_value.to(previous_context), previous_context + ) else: record[key] = record[key] + item if key in record else item + if key.endswith("_max") and location is not None: + record[f"{key}_step"] = torch.full_like(item, location[0], dtype=torch.long) + record[f"{key}_call"] = torch.full_like(item, location[1], dtype=torch.long) + for context_name, context_value in (context or {}).items(): + record[f"{key}_context_{context_name}"] = context_value.detach().clone() record[f"{prefix}_numel"] = record.get(f"{prefix}_numel", 0) + value.numel() def add_module(self, module: nn.Module, label: str) -> None: @@ -163,6 +296,11 @@ def add_module(self, module: nn.Module, label: str) -> None: def set_channel(self, channel: str | None) -> None: self.channel = channel + def set_step(self, step: int | None) -> None: + """Set the outer training step used to locate recurrent call maxima.""" + self._current_step = step + self._step_call_counts.clear() + def observe_tensor(self, label: str, tensor: torch.Tensor) -> None: """Record a named tensor at a call site that has no module boundary. @@ -177,6 +315,28 @@ def observe_tensor(self, label: str, tensor: torch.Tensor) -> None: record["calls"] += 1 self._accumulate_tensor(record, "out", tensor) + def observe_recurrent_state(self, tensor: torch.Tensor, *, call_index: int) -> None: + """Track forward and backward values for an internal recurrent state.""" + if not self.enabled or self.channel != "train": + return + label = "workspace.slot_state_step" + record = self._records.setdefault(self.channel, OrderedDict()).setdefault( + label, {"path": label, "calls": 0} + ) + record["calls"] += 1 + location = ( + (self._current_step, call_index) + if self._current_step is not None + else None + ) + self._accumulate_tensor(record, "out", tensor, location=location) + if tensor.requires_grad: + tensor.register_hook( + lambda grad, target=record, call_location=location: self._accumulate_tensor( + target, "grad_out", grad, location=call_location + ) + ) + def _gradient_records(self) -> OrderedDict[str, tuple[torch.Tensor, int]]: records: OrderedDict[str, list[tuple[torch.Tensor, int]]] = OrderedDict() unassigned: list[tuple[torch.Tensor, int]] = [] @@ -219,19 +379,64 @@ def emit(self, channel: str, *, include_gradients: bool = False) -> None: return scalar_tensors: list[torch.Tensor] = [] - record_shapes: list[tuple[str, int, int, int, int, int]] = [] + record_shapes: list[ + tuple[ + str, + int, + list[tuple[str, int, str]], + dict[str, tuple[int, int, dict[str, float]]], + ] + ] = [] for label, record in (records or {}).items(): - record_shapes.append( - ( - label, - record["calls"], - record.get("in_numel", 0), - record.get("out_numel", 0), - record.get("grad_in_numel", 0), - record.get("grad_out_numel", 0), + locations = {} + gradient_prefixes = ["grad_in", "grad_out"] + secondary_output_indices = sorted( + int(key[4:-6]) + for key in record + if key.startswith("out_") and key.endswith("_numel") + and key[4:-6].isdigit() + ) + gradient_prefixes.extend( + f"grad_out_{output_index}" for output_index in secondary_output_indices + ) + location_prefixes = ["in", "out", *gradient_prefixes] + location_prefixes.extend( + f"out_{output_index}" for output_index in secondary_output_indices + ) + for prefix in location_prefixes: + maximum_key = f"{prefix}_max" + step_value = record.get(f"{maximum_key}_step") + call_value = record.get(f"{maximum_key}_call") + if step_value is not None and call_value is not None: + context = { + name.removeprefix(f"{maximum_key}_context_"): float(value.cpu()) + for name, value in record.items() + if name.startswith(f"{maximum_key}_context_") + } + locations[prefix] = ( + int(step_value.cpu()), int(call_value.cpu()), context + ) + metrics = [ + ("in", record.get("in_numel", 0), "in"), + ("out", record.get("out_numel", 0), "out"), + ("grad_in", record.get("grad_in_numel", 0), "g_in"), + ("grad_out", record.get("grad_out_numel", 0), "g_out"), + ] + for output_index in secondary_output_indices: + metrics.extend( + ( + (f"out_{output_index}", record.get(f"out_{output_index}_numel", 0), f"out{output_index}"), + ( + f"grad_out_{output_index}", + record.get(f"grad_out_{output_index}_numel", 0), + f"g_out{output_index}", + ), + ) ) + record_shapes.append( + (label, record["calls"], metrics, locations) ) - for prefix in ("in", "out", "grad_in", "grad_out"): + for prefix, _, _ in metrics: for suffix in ("sum_sq", "max", "bad"): item = record.get(f"{prefix}_{suffix}") if item is None: @@ -246,20 +451,25 @@ def emit(self, channel: str, *, include_gradients: bool = False) -> None: parts = [] index = 0 - for label, calls, in_numel, out_numel, grad_in_numel, grad_out_numel in record_shapes: + for label, calls, metrics, locations in record_shapes: summaries = [] - for prefix, numel, display in ( - ("in", in_numel, "in"), - ("out", out_numel, "out"), - ("grad_in", grad_in_numel, "g_in"), - ("grad_out", grad_out_numel, "g_out"), - ): + for prefix, numel, display in metrics: sum_sq, maximum, bad = values[index : index + 3] index += 3 if numel: + location_text = "" + if prefix in locations: + location_step, location_call, context = locations[prefix] + context_text = ",".join( + f"{name}={value:.3g}" for name, value in context.items() + ) + context_suffix = f",{context_text}" if context_text else "" + location_text = ( + f"@step{location_step}/call{location_call}{context_suffix}" + ) summaries.append( f"{display}:rms={(sum_sq / numel) ** 0.5:.3g}," - f"max={maximum:.3g},nonfinite={bad / numel:.2%}" + f"max={maximum:.3g}{location_text},nonfinite={bad / numel:.2%}" ) parts.append(f"{label}[calls={calls} {' '.join(summaries)}]") if gradient_shapes: @@ -278,7 +488,10 @@ def activation_absmax_exceeded(self, channel: str) -> tuple[bool, float, str | N records = self._records.get(channel, {}) candidates: list[tuple[str, torch.Tensor]] = [] for label, record in records.items(): - for key in ("in_max", "out_max"): + secondary_output_maxima = ( + key for key in record if key.startswith("out_") and key.endswith("_max") + ) + for key in ("in_max", "out_max", *secondary_output_maxima): value = record.get(key) if torch.is_tensor(value): candidates.append((label, value)) diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index b671371..027269a 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -223,6 +223,7 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea workspace_bistable_gradient_mode=c.get( "workspace_bistable_gradient_mode", "truncated" ), + workspace_bptt_horizon=c.get("workspace_bptt_horizon", 0), device=device, ).to(device) model.workspace_bistable_deterministic = c.get( diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index e1f98e4..9876b93 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -214,6 +214,7 @@ def __init__( workspace_bistable_selection_scale: float = 1.0, workspace_bistable_noise_scale: float = 0.75, workspace_bistable_gradient_mode: str = "truncated", + workspace_bptt_horizon: int = 0, ): super().__init__() if not independent_linears: @@ -271,6 +272,7 @@ def __init__( self.workspace_bistable_noise_scale = float(workspace_bistable_noise_scale) self.workspace_bistable_gradient_mode = workspace_bistable_gradient_mode self.workspace_bistable_selection_scale = float(workspace_bistable_selection_scale) + self.workspace_bptt_horizon = int(workspace_bptt_horizon) self.workspace_bistable_deterministic = False # Runtime-only counterfactual for evaluation. It is deliberately not # part of checkpoint config/state so a saved model can be scored with @@ -537,6 +539,7 @@ def __init__( bistable_selection_scale=workspace_bistable_selection_scale, bistable_noise_scale=workspace_bistable_noise_scale, bistable_gradient_mode=workspace_bistable_gradient_mode, + bptt_horizon=workspace_bptt_horizon, ) if workspace_enabled and workspace_num_slots > 0 else None diff --git a/src/initium/workspace/broadcast.py b/src/initium/workspace/broadcast.py index 838c9f1..7fdf6fb 100644 --- a/src/initium/workspace/broadcast.py +++ b/src/initium/workspace/broadcast.py @@ -33,6 +33,7 @@ def __init__( bistable_selection_scale: float = 1.0, bistable_noise_scale: float = 0.75, bistable_gradient_mode: str = "truncated", + bptt_horizon: int = 0, ): super().__init__() if d_model < 1 or num_experts < 1: @@ -41,12 +42,15 @@ def __init__( raise ValueError("num_slots must be >= 0") if gate_variant not in ("continuous", "bistable"): raise ValueError("gate_variant must be 'continuous' or 'bistable'") + if bptt_horizon < 0: + raise ValueError("bptt_horizon must be >= 0 (0 disables truncation)") self.d_model = d_model self.num_experts = num_experts self.num_slots = num_slots self.gate_variant = gate_variant self.bistable_selection_scale = float(bistable_selection_scale) self.bistable_noise_scale = float(bistable_noise_scale) + self.bptt_horizon = int(bptt_horizon) self.query = nn.Linear(d_model, d_model, bias=False) self.key = nn.Linear(d_model, d_model, bias=False) self.value = nn.Linear(d_model, d_model, bias=False) @@ -289,7 +293,13 @@ def forward( """Process time-major specialists ``(B,T,E,D)`` sequentially.""" outputs = [] step_diags = [] if self.diagnostics_enabled else None + state_observer = getattr(self, "_initium_recurrent_state_observer", None) for step in range(specialist_deltas.shape[1]): + # Truncate only the gradient carried by persistent slots. Detach + # preserves the exact forward state and therefore the full write / + # read trajectory; zero keeps the historical full-sequence BPTT. + if self.bptt_horizon and step and step % self.bptt_horizon == 0: + state = state.detach() out, state = self.forward_step( specialist_deltas[:, step], route_indices[:, step], route_weights[:, step], state, @@ -300,6 +310,8 @@ def forward( bistable_selection_scale=bistable_selection_scale, deterministic_gate=deterministic_gate, ) + if state_observer is not None: + state_observer(state, call_index=step) outputs.append(out) if step_diags is not None: step_diags.append(self._last_step_diag) From 02b0ff6ea6d30a2b06d5fcd40496b59611293d69 Mon Sep 17 00:00:00 2001 From: Samuel Date: Thu, 15 Oct 2026 09:08:34 +0200 Subject: [PATCH 14/14] Add DNC execution and addressing options Add compiled memory forwarding, fast and top-k allocation, temporal-link sharpening, and configurable batch sizes. Skip dense traversal products when links are ablated, preserve these settings when resizing memory, and document the tradeoffs. --- README.md | 8 +- src/initium/core_training.py | 114 ++++++++++++++++- src/initium/inference/model_loader.py | 2 + .../mamba_controller/split_graph_dnc.py | 6 + .../memory_manipulation/dnc_mechanisms.py | 116 +++++++++++++++++- .../dynamic_memory_resize.py | 8 ++ .../link_matrix_ablation.py | 28 +++-- 7 files changed, 262 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 06a89e8..f35e31f 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,13 @@ The split-graph memory exposes independent read and write variants while retaini - The graph curriculum emits ID, OOD, and OOD-offset metrics as before. Graph runs also write `run__relational_probe.csv`, with cosine top-1, selected read-head top-1, and held-out-composition accuracy. `data/graph_traversal/order_probe.py` supplies paired prefix-permutation helpers for the Q99 order-sensitivity evaluation. - Split-graph graph runs also write `run__track_a.csv` and `run__track_b.csv`. Track A records ID/OOD accuracy, relational stress and held-out-composition accuracy, ordinary-input read-score/cosine correlation, and read-content entropy/maximum weight. Track B records the Q99 order-sensitivity gate, ID/OOD posterior variance and write-update diagnostics, and training-window loss, gradient norm, AMP scale, and learning rate. These are experiment logs; no training or evaluation results are checked in. -The DNC implementation defaults to `--base-dnc`. `--optimized-dnc` selects an exact allocation implementation that reuses the stock sorted usage permutation to scatter allocation weights back into memory order, avoiding the second full sort in every write step. It retains dense temporal links, the same controller interface, and the same learned parameters/state-dict keys. Use the mutually exclusive `--base-dnc` flag to make a baseline run explicit; compare matched seeds, episode streams, and step budgets before adopting the optimized mode. This first optimization targets allocation overhead only; it does not claim to remove the dense link matrix's quadratic cost. +The DNC implementation defaults to `--base-dnc`. `--dnc-fast-allocation` uses the exact allocation implementation that reuses the stock sorted usage permutation and removes its second full sort; `--dnc-compile-memory` compiles the per-timestep memory forward with `torch.compile`. These preserve DNC equations and state-dict keys. `--dnc-allocation-topk K` is an independent approximation that allocates among only the K least-used cells. It can reduce selection work, but omits allocation mass outside the selected cells and may affect learning. `--dnc-link-sharpness [EXPONENT]` is a parameter-free learning variant inspired by arXiv:1904.10278 that sharpens the shape of forward/backward temporal-link reads while preserving their original total mass; it defaults to exponent 2. Both functional variants can be started from a lesson-14 checkpoint. `--optimized-dnc` enables fast allocation and compiled memory forward; for split-graph runs it also defaults to top-64 allocation and exponent-2 link sharpening. Compilation is lazy, so compare sustained steps/s after warm-up. The exact execution options do not remove dense temporal-link cost. + +`--link-matrix-mode ablated` bypasses temporal-link updates and traversal products, which saves dense link work but removes ordered recall; `sparse_topk` currently masks a dense update and therefore is not a speed optimization. The paper evidence is task-dependent: rsDNC reports that temporal links contribute little on its QA workload, while the original DNC reports sparse links working on copy but link removal failing there. Split-graph traversal chains multiple hops, so link removal stays an isolated comparison, never part of `--optimized-dnc`. Use `--batch-size` to sweep larger episode batches; this can improve utilization but increases memory use and changes gradient noise. Compare equal seeds and optimizer-step budgets, and report both steps/s and steps-to-target accuracy. + +Gradient checkpointing is omitted from the speed bundle because it saves activation memory by recomputing work in backward. Truncated BPTT is omitted because graph facts needed for a later query can precede it by many steps; detaching DNC state can remove that learning signal. A hybrid scan over guessed addressing weights is plausible mathematically, since memory updates become affine once weights are fixed, but the addressing depends on the evolving memory. It would need an iterative solver and convergence/accuracy evidence; at the current short lesson-1 sequences, its repeated passes may cost more than the sequential loop. These are follow-up research directions rather than robust drop-in optimizations. + +The paper review supports this staged approach. Sparse Access Memory reports much faster sparse reads/writes, including a temporal-association variant, but it changes the addressing/indexing mechanism and is not a drop-in exact DNC optimization. BPPSA and parallel-scan work apply to linear recurrences; the DNC loop here has a nonlinear read-to-next-write dependency, so parallelizing timesteps would require a changed model. The rsDNC work's slimmer memory and link changes were developed for question answering, and the local link ablation already gives us a direct task-specific control. For now, compilation and exact allocation are the faithful speed tests; sparse addressing remains a distinct research branch. Use the selected split-graph baseline with either supported controller combiner: diff --git a/src/initium/core_training.py b/src/initium/core_training.py index dd2a97a..5605e40 100644 --- a/src/initium/core_training.py +++ b/src/initium/core_training.py @@ -164,6 +164,7 @@ from initium.mamba_controller.split_graph_dnc import SplitGraphDNC from initium.memory_manipulation.dynamic_memory_resize import resize_memory from initium.memory_manipulation.dynamic_n_controller import DynamicNController +from initium.memory_manipulation.dnc_mechanisms import compile_memory_forward from initium.memory_manipulation.link_matrix_ablation import patch_link_matrix from initium.memory_manipulation.stochastic_write_head_v2 import ( get_prior_state, @@ -697,6 +698,10 @@ def load_frozen_workspace_warm_start( # Exact implementation-only optimization; old checkpoints default # to base and remain weight/state compatible with either backend. "dnc_implementation", + "dnc_compile_memory", # exact execution optimization; state-dict compatible + "dnc_link_sharpness", # parameter-free read-addressing experiment + "dnc_allocation_topk", # parameter-free allocation experiment + "batch_size", # training hyperparameter, not model architecture } missing_config = sorted( key @@ -836,13 +841,17 @@ def run( dnc_read_residual_max_ratio: float = DNC_READ_RESIDUAL_MAX_RATIO, dnc_write_variant: str = DNC_WRITE_VARIANT, dnc_write_observation_variance: float = DNC_WRITE_OBSERVATION_VARIANCE, + dnc_link_sharpness: float = 1.0, + dnc_allocation_topk: int | None = None, dnc_implementation: str = "base", + dnc_compile_memory: bool = False, dnc_backward_trace: bool = False, dataset_type: str = DATASET_TYPE, dataset_link: str | None = DATASET_LINK, test_dataset_link: str | None = None, prepared_dataset=None, ogs_weights: tuple[float, float, float] = OGS_WEIGHTS, + batch_size: int = BATCH_SIZE, ): # Deliberately does NOT touch LR_DECAY_STEPS - that's a separate # module-level constant, fixed at import time from the *original* @@ -1535,6 +1544,8 @@ def run( dnc_read_residual_max_ratio=dnc_read_residual_max_ratio, dnc_write_variant=dnc_write_variant, dnc_write_observation_variance=dnc_write_observation_variance, + dnc_link_sharpness=dnc_link_sharpness, + dnc_allocation_topk=dnc_allocation_topk, dnc_implementation=dnc_implementation, workspace_enabled=workspace_enabled, workspace_num_slots=workspace_num_slots, @@ -1548,6 +1559,7 @@ def run( workspace_bptt_horizon=workspace_bptt_horizon, device=device, ).to(device) + if dnc_backward_trace: rnn.memories[0].set_backward_trace(True) print( @@ -1583,6 +1595,14 @@ def run( **mamba_kwargs, ).to(device) + if dnc_compile_memory: + for memory in rnn.memories: + compile_memory_forward(memory) + print( + f"[{run_id}] DNC memory forward compiled with torch.compile | " + "first-call compilation overhead is expected" + ) + print( f"[{run_id}] DNC implementation={dnc_implementation}" + ( @@ -1590,6 +1610,12 @@ def run( if dnc_implementation == "optimized" else "" ) + + (" | compiled memory forward" if dnc_compile_memory else "") + + ( + f" | temporal-link sharpness exponent={dnc_link_sharpness:g}" + if dnc_link_sharpness != 1.0 + else "" + ) ) if split_graph_enabled: @@ -1662,9 +1688,13 @@ def run( "dnc_read_residual_max_ratio": dnc_read_residual_max_ratio, "dnc_write_variant": dnc_write_variant, "dnc_write_observation_variance": dnc_write_observation_variance, + "dnc_link_sharpness": dnc_link_sharpness, + "dnc_allocation_topk": dnc_allocation_topk, "dnc_implementation": dnc_implementation, + "dnc_compile_memory": dnc_compile_memory, "dnc_compute_precision": "fp32", "dnc_backward_trace": dnc_backward_trace, + "batch_size": batch_size, "workspace_enabled": workspace_enabled, "workspace_num_slots": workspace_num_slots, "workspace_no_selection": workspace_no_selection, @@ -2320,7 +2350,7 @@ def set_lr(step_now): t_run_start = time.time() while step < total_steps and not eval_only: - input_seq, target_digits, answer_mask = dataset.sample_batch(curriculum, BATCH_SIZE) + input_seq, target_digits, answer_mask = dataset.sample_batch(curriculum, batch_size) input_seq = input_seq.to(device, non_blocking=True) target_digits = target_digits.to(device, non_blocking=True) @@ -3609,6 +3639,12 @@ def _write_diagnostic_values(diagnostics): import argparse import sys + class _OptimizedDNCAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + namespace.dnc_implementation = "optimized" + namespace.dnc_compile_memory = True + namespace._dnc_optimized_bundle = True + parser = argparse.ArgumentParser( description="Phase 2 KL sweep (learned, periodically-snapshotted " "prior). Positional beta for single-beta mode; " @@ -4052,6 +4088,28 @@ def _write_diagnostic_values(diagnostics): default=DNC_WRITE_OBSERVATION_VARIANCE, help="Fixed positive observation variance for the Kanerva-style write.", ) + parser.add_argument( + "--dnc-link-sharpness", + type=float, + nargs="?", + const=2.0, + default=1.0, + metavar="EXPONENT", + help="Sharpen the shape of forward/backward temporal-link reads using a " + "fixed exponent inspired by arXiv:1904.10278 while preserving each " + "vector's original mass (default exponent 2 when specified; 1 leaves " + "the DNC unchanged). Adds no parameters, so it can resume from a " + "baseline checkpoint.", + ) + parser.add_argument( + "--dnc-allocation-topk", + type=int, + default=None, + metavar="K", + help="Approximate DNC allocation using only the K least-used cells. " + "This can reduce allocation selection work, but omits allocation mass " + "outside the selected cells and may change learning behavior.", + ) parser.add_argument( "--dnc-backward-trace", action="store_true", @@ -4068,12 +4126,25 @@ def _write_diagnostic_values(diagnostics): ) dnc_impl_group.add_argument( "--optimized-dnc", + action=_OptimizedDNCAction, + nargs=0, + help="Enable exact fast allocation and compiled memory forward. For " + "--split-graph, also enable top-64 allocation and exponent-2 temporal " + "link sharpening as separate functional experiments.", + ) + parser.set_defaults(dnc_implementation="base", _dnc_optimized_bundle=False) + dnc_impl_group.add_argument( + "--dnc-fast-allocation", dest="dnc_implementation", action="store_const", const="optimized", - help="Use exact dense DNC equations with one allocation permutation sort removed.", + help="Use the exact allocation path with one usage-permutation sort removed (without compilation).", + ) + parser.add_argument( + "--dnc-compile-memory", + action="store_true", + help="Compile the DNC memory forward with torch.compile; keeps DNC equations unchanged.", ) - parser.set_defaults(dnc_implementation="base") parser.add_argument( "--dataset-type", type=str, @@ -4087,6 +4158,14 @@ def _write_diagnostic_values(diagnostics): ], help="Dataset plugged into the core loop (data/dataset_registry.py).", ) + parser.add_argument( + "--batch-size", + type=int, + default=BATCH_SIZE, + help=f"Training episodes per optimizer step (default: {BATCH_SIZE}); " + "larger values may improve device utilization at higher memory use and " + "change the effective gradient batch.", + ) parser.add_argument( "--dataset-link", type=str, @@ -4141,6 +4220,21 @@ def _write_diagnostic_values(diagnostics): parser.error("--prepare-only cannot be combined with --load-prepared") if args.dnc_write_observation_variance <= 0: parser.error("--dnc-write-observation-variance must be positive") + if args._dnc_optimized_bundle and args.split_graph: + if args.dnc_link_sharpness == 1.0: + args.dnc_link_sharpness = 2.0 + if args.dnc_allocation_topk is None: + args.dnc_allocation_topk = 64 + if args.dnc_link_sharpness <= 0: + parser.error("--dnc-link-sharpness must be positive") + if args.dnc_allocation_topk is not None and args.dnc_allocation_topk <= 0: + parser.error("--dnc-allocation-topk must be positive") + if args.dnc_allocation_topk is not None and not args.split_graph: + parser.error("--dnc-allocation-topk requires --split-graph") + if args.batch_size <= 0: + parser.error("--batch-size must be positive") + if args.dnc_link_sharpness != 1.0 and not args.split_graph: + parser.error("--dnc-link-sharpness requires --split-graph") if args.dnc_backward_trace and not args.split_graph: parser.error("--dnc-backward-trace requires --split-graph") if args.dnc_read_residual_scale < 0: @@ -4337,12 +4431,20 @@ def _write_diagnostic_values(diagnostics): run_id += f"_rmax{args.dnc_read_residual_max_ratio:g}" if args.dnc_write_variant == "kanerva-closed-form": run_id += f"_obsvar{args.dnc_write_observation_variance:g}" + if args.dnc_link_sharpness != 1.0: + run_id += f"_linksharp{args.dnc_link_sharpness:g}" + if args.dnc_allocation_topk is not None: + run_id += f"_allocK{args.dnc_allocation_topk}" if args.split_graph_combiner_mode == "controller": run_id = f"{run_id}_combctrl{args.split_graph_combiner_variant.replace('+', '')}" if args.dnc_implementation == "optimized": - run_id += "_optimizeddnc" + run_id += "_optimizeddnc" if args.dnc_compile_memory else "_dncfastalloc" + elif args.dnc_compile_memory: + run_id += "_dnccompiled" if args.dataset_type != "graph": run_id = f"{run_id}_ds{args.dataset_type}" + if args.batch_size != BATCH_SIZE: + run_id += f"_bs{args.batch_size}" if args.run_id_suffix: run_id = f"{run_id}_{args.run_id_suffix}" summary = run( @@ -4404,13 +4506,17 @@ def _write_diagnostic_values(diagnostics): dnc_read_residual_max_ratio=args.dnc_read_residual_max_ratio, dnc_write_variant=args.dnc_write_variant, dnc_write_observation_variance=args.dnc_write_observation_variance, + dnc_link_sharpness=args.dnc_link_sharpness, + dnc_allocation_topk=args.dnc_allocation_topk, dnc_implementation=args.dnc_implementation, + dnc_compile_memory=args.dnc_compile_memory, dnc_backward_trace=args.dnc_backward_trace, dataset_type=args.dataset_type, dataset_link=args.dataset_link, test_dataset_link=args.test_dataset_link, prepared_dataset=prepared_dataset, ogs_weights=tuple(args.ogs_weights), + batch_size=args.batch_size, ) all_summaries.append(summary) diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 027269a..81861d9 100644 --- a/src/initium/inference/model_loader.py +++ b/src/initium/inference/model_loader.py @@ -187,6 +187,8 @@ def _build_split_graph(c: _Cfg, input_dim, hidden, nr_cells, cell_size, read_hea dnc_write_observation_variance=c.get( "dnc_write_observation_variance", cc.DNC_WRITE_OBSERVATION_VARIANCE ), + dnc_link_sharpness=c.get("dnc_link_sharpness", 1.0), + dnc_allocation_topk=c.get("dnc_allocation_topk", None), dnc_implementation=c.get("dnc_implementation", "base"), independent_linears=True, moe_enabled=_moe["moe_enabled"], diff --git a/src/initium/mamba_controller/split_graph_dnc.py b/src/initium/mamba_controller/split_graph_dnc.py index 9876b93..ec52433 100644 --- a/src/initium/mamba_controller/split_graph_dnc.py +++ b/src/initium/mamba_controller/split_graph_dnc.py @@ -204,6 +204,8 @@ def __init__( dnc_read_residual_max_ratio: float = 0.5, dnc_write_variant: str = "learned", dnc_write_observation_variance: float = 1.0, + dnc_link_sharpness: float = 1.0, + dnc_allocation_topk: int | None = None, dnc_implementation: str = "base", workspace_enabled: bool = False, workspace_num_slots: int = 5, @@ -243,6 +245,8 @@ def __init__( self.dnc_read_residual_max_ratio = dnc_read_residual_max_ratio self.dnc_write_variant = dnc_write_variant self.dnc_write_observation_variance = dnc_write_observation_variance + self.dnc_link_sharpness = dnc_link_sharpness + self.dnc_allocation_topk = dnc_allocation_topk if dnc_implementation not in ("base", "optimized"): raise ValueError("dnc_implementation must be 'base' or 'optimized'") self.dnc_implementation = dnc_implementation @@ -511,6 +515,8 @@ def __init__( read_residual_scale=dnc_read_residual_scale, read_residual_max_ratio=dnc_read_residual_max_ratio, observation_variance=dnc_write_observation_variance, + link_sharpness=dnc_link_sharpness, + allocation_topk=dnc_allocation_topk, dnc_implementation=dnc_implementation, ) ) diff --git a/src/initium/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py index 4b4f229..50b628f 100644 --- a/src/initium/memory_manipulation/dnc_mechanisms.py +++ b/src/initium/memory_manipulation/dnc_mechanisms.py @@ -22,6 +22,24 @@ DNC_IMPLEMENTATIONS = ("base", "optimized") +def compile_memory_forward(memory: Memory) -> Memory: + """Compile the per-timestep Memory forward for kernel fusion. + + This is an execution-only optimization: it does not change the module's + parameters, state dict, or addressing equations. Compilation is lazy, so + the first call pays compile cost; callers should use it for sustained + training rather than short profiling snippets. + """ + if getattr(memory, "_initium_compiled_forward", False): + return memory + compiler = getattr(torch, "compile", None) + if compiler is None: + raise RuntimeError("--dnc-compile-memory requires a PyTorch build with torch.compile") + memory.forward = compiler(memory.forward, dynamic=False) + memory._initium_compiled_forward = True + return memory + + class _FastAllocationMixin: """Exact DNC allocation with one full sort instead of two. @@ -53,6 +71,35 @@ def allocate(self, usage: torch.Tensor, write_gate: torch.Tensor): return allocation_weights.unsqueeze(1), usage +class _TopKAllocationMixin: + """Approximate allocation over only the K least-used memory cells. + + The selected cells retain the ordinary DNC allocation formula in sorted + order. Omitted cells receive zero allocation mass, so this is an explicit + addressing approximation rather than a state-dict or implementation-only + optimization. + """ + + nr_cells: int + allocation_topk: int + + def allocate(self, usage: torch.Tensor, write_gate: torch.Tensor): + usage = δ + (1 - δ) * usage + k = min(self.allocation_topk, self.nr_cells) + sorted_usage, order = torch.topk(usage, k, dim=1, largest=False) + leading_one = torch.ones( + (usage.size(0), 1), device=usage.device, dtype=usage.dtype + ) + exclusive_products = torch.cumprod( + torch.cat((leading_one, sorted_usage), dim=1), dim=1 + )[:, :-1] + sorted_weights = (1 - sorted_usage) * exclusive_products + allocation_weights = torch.zeros_like(usage).scatter( + 1, order, sorted_weights + ) + return allocation_weights.unsqueeze(1), usage + + class OptimizedMemory(_FastAllocationMixin, Memory): """Stock DNC memory with an exact allocation-sort optimization.""" @@ -76,6 +123,8 @@ def __init__( read_residual_max_ratio: float = 0.5, relational_hidden_size: int | None = None, observation_variance: float = 1.0, + link_sharpness: float = 1.0, + allocation_topk: int | None = None, **kwargs, ): if read_variant not in READ_VARIANTS: @@ -88,12 +137,18 @@ def __init__( raise ValueError("read_residual_max_ratio must be non-negative") if observation_variance <= 0: raise ValueError("observation_variance must be positive") + if link_sharpness <= 0: + raise ValueError("link_sharpness must be positive") + if allocation_topk is not None and allocation_topk <= 0: + raise ValueError("allocation_topk must be positive") super().__init__(*args, **kwargs) self.read_variant = read_variant self.write_variant = write_variant self.read_residual_scale = float(read_residual_scale) self.read_residual_max_ratio = float(read_residual_max_ratio) self.observation_variance = float(observation_variance) + self.link_sharpness = float(link_sharpness) + self.allocation_topk = allocation_topk self._collect_stage_diagnostics = False self._stage_diagnostic_sums: dict[str, torch.Tensor] = {} self._stage_diagnostic_counts: dict[str, int] = {} @@ -132,6 +187,38 @@ def __init__( if self.device is not None: self.relational_score.to(self.device) + def directional_weightings(self, link_matrix, read_weights): + """Optionally sharpen DNC temporal-link traversal distributions. + + This uses a fixed exponent inspired by the link-distribution + sharpness method in arXiv:1904.10278. It preserves each directional + vector's original mass for stable DNC sub-probability semantics. + Exponent 1 is an exact no-op, preserving the upstream DNC path. The + paper's controller-predicted exponent is deliberately not used here, + so the experiment adds no parameters and can resume from a baseline + checkpoint. + """ + forward, backward = super().directional_weightings(link_matrix, read_weights) + if self.link_sharpness == 1.0: + return forward, backward + + def sharpen(distribution): + eps = torch.finfo(distribution.dtype).tiny + nonnegative = distribution.clamp_min(0) + mass = nonnegative.sum(dim=-1, keepdim=True) + + # Sharpen only the normalized shape and then restore the original + # mass. DNC directional weights are sub-probability vectors: their + # total mass reflects how much temporal-link signal is available. + # Normalizing the sharpened vector to 1 would amplify arbitrarily + # weak, nearly empty link traversals and can destabilize gradients. + probabilities = nonnegative / mass.clamp_min(eps) + powered = probabilities.pow(self.link_sharpness) + shape = powered / powered.sum(dim=-1, keepdim=True).clamp_min(eps) + return shape * mass + + return sharpen(forward), sharpen(backward) + def forward(self, xi: torch.Tensor, hidden: dict[str, torch.Tensor]): """Run all DNC interface, addressing, and state math in fp32. @@ -647,6 +734,14 @@ class OptimizedMechanismMemory(_FastAllocationMixin, MechanismMemory): """MechanismMemory with the exact allocation-sort optimization.""" +class TopKMechanismMemory(_TopKAllocationMixin, MechanismMemory): + """MechanismMemory with approximate allocation over its least-used K cells.""" + + +class TopKOptimizedMechanismMemory(_TopKAllocationMixin, OptimizedMechanismMemory): + """Top-K allocation variant retaining the optimized implementation tag.""" + + def build_memory( *, input_size: int, @@ -660,16 +755,25 @@ def build_memory( read_residual_scale: float = 1.0, read_residual_max_ratio: float = 0.5, observation_variance: float = 1.0, + link_sharpness: float = 1.0, + allocation_topk: int | None = None, dnc_implementation: str = "base", ): """Construct a stock-compatible memory for the selected mechanism pair.""" if dnc_implementation not in DNC_IMPLEMENTATIONS: raise ValueError(f"Unknown DNC implementation: {dnc_implementation!r}") - memory_cls = ( - OptimizedMechanismMemory - if dnc_implementation == "optimized" - else MechanismMemory - ) + if allocation_topk is not None: + memory_cls = ( + TopKOptimizedMechanismMemory + if dnc_implementation == "optimized" + else TopKMechanismMemory + ) + else: + memory_cls = ( + OptimizedMechanismMemory + if dnc_implementation == "optimized" + else MechanismMemory + ) return memory_cls( input_size=input_size, nr_cells=nr_cells, @@ -682,4 +786,6 @@ def build_memory( read_residual_scale=read_residual_scale, read_residual_max_ratio=read_residual_max_ratio, observation_variance=observation_variance, + link_sharpness=link_sharpness, + allocation_topk=allocation_topk, ) diff --git a/src/initium/memory_manipulation/dynamic_memory_resize.py b/src/initium/memory_manipulation/dynamic_memory_resize.py index cbffc88..20efabb 100644 --- a/src/initium/memory_manipulation/dynamic_memory_resize.py +++ b/src/initium/memory_manipulation/dynamic_memory_resize.py @@ -54,6 +54,7 @@ MechanismMemory, OptimizedMemory, build_memory, + compile_memory_forward, ) from initium.memory_manipulation.link_matrix_ablation import AblatableSparseLinkMemory @@ -117,6 +118,8 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz read_residual_scale=old_memory.read_residual_scale, read_residual_max_ratio=old_memory.read_residual_max_ratio, observation_variance=old_memory.observation_variance, + link_sharpness=old_memory.link_sharpness, + allocation_topk=old_memory.allocation_topk, dnc_implementation=dnc_implementation, ).to(device) else: @@ -155,6 +158,11 @@ def resize_memory(model, new_nr_cells: int, device=None, layer: int = 0, optimiz if hasattr(old_memory, "link_matrix_topk"): new_memory.link_matrix_topk = old_memory.link_matrix_topk + # Compilation is execution-only state and is not represented in a state + # dict. Preserve it when the curriculum replaces the size-dependent memory. + if getattr(old_memory, "_initium_compiled_forward", False): + compile_memory_forward(new_memory) + 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 diff --git a/src/initium/memory_manipulation/link_matrix_ablation.py b/src/initium/memory_manipulation/link_matrix_ablation.py index 50f5d10..85e9df0 100644 --- a/src/initium/memory_manipulation/link_matrix_ablation.py +++ b/src/initium/memory_manipulation/link_matrix_ablation.py @@ -60,16 +60,12 @@ class AblatableSparseLinkMemory(Memory): 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. + "sparse_topk" -- link matrix is updated densely as normal, then each + row keeps only its link_matrix_topk largest-magnitude + entries. This is a functional sparsity experiment only: + it still allocates and updates a dense N x N tensor, + performs dense traversal matmuls, and adds a top-k pass. + It does not currently provide a speed or memory saving. """ def get_link_matrix( @@ -98,6 +94,18 @@ def get_link_matrix( raise ValueError(f"Unknown link_matrix_mode: {mode!r}") + def directional_weightings(self, link_matrix: torch.Tensor, read_weights: torch.Tensor): + # In ablated mode the link matrix is invariantly zero for the whole + # episode. Return the mathematically identical zero traversals without + # launching the two dense O(B * heads * N^2) matrix products. + if getattr(self, "link_matrix_mode", "dense") == "ablated": + # Match Memory.directional_weightings' (B, heads, 1, cells) + # layout. read_weightings combines read modes over that singleton + # write-head axis before producing its (B, heads, cells) result. + zeros = torch.zeros_like(read_weights).unsqueeze(2) + return zeros, zeros + return super().directional_weightings(link_matrix, read_weights) + 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,