diff --git a/README.md b/README.md index f009511..f35e31f 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,14 @@ 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`. `--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: ```bash @@ -104,6 +112,98 @@ 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. 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 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 --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. + +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. + +### 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. + +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. + +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 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 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. + +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. + +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 +``` + +`--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 +``` + +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. ### Controllers diff --git a/src/initium/LNN_controller/cfc_backbone_parallel.py b/src/initium/LNN_controller/cfc_backbone_parallel.py index a6f7035..24b9297 100644 --- a/src/initium/LNN_controller/cfc_backbone_parallel.py +++ b/src/initium/LNN_controller/cfc_backbone_parallel.py @@ -100,6 +100,8 @@ 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, ): @@ -134,32 +136,68 @@ 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, + residual_scale=moe_residual_scale, + 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, *, 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, token_mask=token_mask + ) + return base, experts, indices, weights + 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") + def build_parallel_backbone( in_dim, @@ -177,6 +215,8 @@ 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, ): @@ -204,6 +244,8 @@ 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 {}), ) ) @@ -226,6 +268,8 @@ 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, ) ) cur = d_model @@ -246,7 +290,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/LNN_controller/cfc_controller.py b/src/initium/LNN_controller/cfc_controller.py index 012f84b..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) @@ -283,15 +298,27 @@ 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 = [] 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/MoE/moe_layer.py b/src/initium/MoE/moe_layer.py index fd06de4..d790ea4 100644 --- a/src/initium/MoE/moe_layer.py +++ b/src/initium/MoE/moe_layer.py @@ -81,7 +81,7 @@ from __future__ import annotations -from typing import cast +import math import torch import torch.nn as nn @@ -102,13 +102,21 @@ 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))) 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 @@ -149,12 +157,26 @@ 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 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) -> torch.Tensor: + 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, ( @@ -162,10 +184,27 @@ 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) + 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 + # 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 @@ -178,78 +217,242 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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. - # 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. + # 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 + # 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) - - # 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) - + 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) + + 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_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 - # 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) + 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 + # 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) + if collect_diagnostics + else None + ) - 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) + 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, + ) - return output.reshape(orig_shape) + # 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) + + 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: + 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,10 +462,27 @@ 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 + 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. @@ -273,6 +493,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 @@ -282,6 +511,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 @@ -303,10 +537,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, @@ -323,15 +563,33 @@ 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, + token_mask: torch.Tensor | None = None, + ): + if not return_details: + 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 + self.residual_scale * routed, experts, indices, weights 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() + 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 +616,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 @@ -376,6 +634,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, @@ -406,6 +666,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) @@ -428,29 +690,39 @@ 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 - if topk_idx is not None: + 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: - 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 + ] + 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).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()), - } + 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 - self._last_source_diag = source_diag outs, offset = [], 0 for i, b in enumerate(batch_sizes): @@ -461,9 +733,18 @@ 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() + 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 +753,60 @@ 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 = [] + 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]), + "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): @@ -526,27 +860,88 @@ 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"]) + 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"])) + 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", + "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", + "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 {}), + } + ) 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..5605e40 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,10 +159,12 @@ 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 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, @@ -166,7 +173,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}") @@ -204,6 +211,107 @@ 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", []) + ) + 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}] " + 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}" + f"{split_summary}" + f"{(' ' + source_summary) if source_summary else ''}]" + ) + 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)) @@ -545,6 +653,133 @@ 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", + "workspace_train_scope", + "workspace_gate_variant", + "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 + # 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 + 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 # ========================================== @@ -566,6 +801,11 @@ 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, + 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, moe_num_experts: int = MOE_NUM_EXPERTS, moe_expert_dim: int | None = MOE_EXPERT_DIM, moe_capacity_factor: float = MOE_CAPACITY_FACTOR, @@ -580,16 +820,38 @@ 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, + 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_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, + 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, 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* @@ -617,7 +879,106 @@ 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 + 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 + ) + 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 + ) + 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 + or ( + workspace_gate_variant == "bistable" + and checkpoint_config.get( + "workspace_bistable_gradient_mode", "unrolled" + ) + != workspace_bistable_gradient_mode + ) + ) + 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 + ) + 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 + and workspace_gate_variant == "bistable" + and checkpoint_workspace_enabled + and checkpoint_config.get("workspace_gate_variant", "continuous") == "continuous" + and not eval_only + ) + 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( + resume_requested + and workspace_enabled + 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. + # 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 workspace_condition_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: @@ -628,12 +989,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", @@ -645,6 +1008,7 @@ def run( "moe_cv_importance", "moe_cv_load", "moe_max_load_frac", + "moe_capacity_drop_frac", ] ) @@ -972,6 +1336,16 @@ 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=" + 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) @@ -1089,14 +1463,34 @@ 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)." ) + 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}" + ) + 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} | " + "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 # reusing the Mamba-1 constants here -- harmless previously (any @@ -1136,19 +1530,47 @@ def run( combiner_num_blocks=split_graph_combiner_num_blocks, 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, 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, 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, + 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, + workspace_bistable_gradient_mode=workspace_bistable_gradient_mode, + workspace_bptt_horizon=workspace_bptt_horizon, 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, @@ -1169,9 +1591,59 @@ 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) + 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}" + + ( + " | exact allocation path avoids one full usage-permutation sort per step" + 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: + 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" + 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"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" + ) + 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 @@ -1208,11 +1680,111 @@ 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, + "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, "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, + "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_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, + } + 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": @@ -1225,6 +1797,95 @@ 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, + activation_absmax_limit=( + workspace_stability_max_activation + if workspace_train_scope == "gate" + 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: + 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) + 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) + 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 and moe_diagnostics_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)) + + 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 @@ -1243,7 +1904,55 @@ 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 + 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_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 # nr_cells resize never changes it (every Memory sublayer is sized by @@ -1344,8 +2053,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 @@ -1371,16 +2082,93 @@ 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 + 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) + 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." + ) + 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 workspace_condition_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"]) + 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 # (saved before this key existed) still loads, just without this fix. - if "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"]) - else: + 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 " @@ -1394,7 +2182,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 @@ -1404,7 +2193,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: @@ -1425,9 +2214,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) " @@ -1440,18 +2229,29 @@ 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 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( f"[{run_id}] beta_target changed on resume " @@ -1459,28 +2259,76 @@ 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 step >= total_steps: + 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 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"{'[resumed]' if resuming else ''} ===" + 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). 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 + # 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 @@ -1488,6 +2336,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() @@ -1500,66 +2349,240 @@ 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: - input_seq, target_digits, answer_mask = dataset.sample_batch(curriculum, BATCH_SIZE) + 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) target_digits = target_digits.to(device, non_blocking=True) 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") + 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): - output, hidden = rnn(input_seq, hidden, reset_experience=True) - if check_finite_this_step: - _first_nonfinite_report(output, "rnn_output", step) + 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 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 + ): + 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=( + moe_diagnostics_enabled and (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) + 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 + 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 # 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) - if not math.isfinite(float(grad_norm)): + # 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=max_grad_norm, 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) + 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(grad_norm_value): 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 +2594,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) - 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() @@ -1647,7 +2674,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,19 +2707,103 @@ 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}", + f"clipped_grad_steps {running_clipped_grad_steps}/{running_finite_grad_steps}", ] 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%}", ] + 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 "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(" |")) + 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: + 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 @@ -1752,12 +2863,16 @@ 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 + 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 @@ -1807,9 +2922,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 +3014,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 +3022,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 +3032,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 +3190,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 +3233,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, @@ -2184,13 +3321,28 @@ 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" + 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") _, 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") + report_bistable_evaluation("ID") final_id_stage_diagnostics = ( final_id_stage_memory.finish_stage_diagnostics() if hasattr(final_id_stage_memory, "finish_stage_diagnostics") @@ -2220,8 +3372,11 @@ 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") ood_triple_acc, ood_perfect_frac, ood_hop_breakdown = dataset.evaluate_ood( rnn, device, @@ -2230,6 +3385,8 @@ def set_lr(step_now): verbose_n=10, 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") @@ -2269,6 +3426,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}/" @@ -2279,8 +3444,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: @@ -2339,10 +3503,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 +3628,7 @@ def _write_diagnostic_values(diagnostics): if dynamic_n_ctrl is not None: dynamic_n_log_file.close() + segment_diagnostics.close() return summary @@ -2472,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; " @@ -2488,9 +3661,27 @@ 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( + "--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", @@ -2638,6 +3829,120 @@ 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, 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", + 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( + "--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( + "--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). " + "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.", + ) + 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( + "--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-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-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", + 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, @@ -2783,6 +4088,63 @@ 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", + 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", + 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 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.add_argument( "--dataset-type", type=str, @@ -2796,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, @@ -2850,6 +4220,23 @@ 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: parser.error("--dnc-read-residual-scale must be non-negative") if args.dnc_read_residual_max_ratio < 0: @@ -2888,9 +4275,86 @@ 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") + 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.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.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 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 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") + 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 ( + 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] @@ -2930,8 +4394,36 @@ 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: + 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_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}" + f"_grad{args.workspace_bistable_gradient_mode}" + ) + 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: + 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": @@ -2939,10 +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" 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( @@ -2965,6 +4467,11 @@ 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, + 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, moe_num_experts=args.moe_num_experts, moe_expert_dim=args.moe_expert_dim, moe_capacity_factor=args.moe_capacity_factor, @@ -2978,16 +4485,38 @@ 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, + 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_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, + 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, 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/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..796e136 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]: @@ -439,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 @@ -525,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/__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..3740b1e --- /dev/null +++ b/src/initium/diagnostics/segment_logging.py @@ -0,0 +1,509 @@ +"""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 _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() + 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 == "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): + 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, + *, + 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._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]] = [] + self._handles = [] + if not self.enabled: + return + 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) + 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, location=location) + if input_tensor.requires_grad: + input_tensor.register_hook( + 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, location=location) + if output_tensor.requires_grad: + output_tensor.register_hook( + 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, + *, + 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)) + 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(): + 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: + """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)) + + 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. + + 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 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]] = [] + 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: + 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: + return + + scalar_tensors: list[torch.Tensor] = [] + 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(): + 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 metrics: + 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, metrics, locations in record_shapes: + summaries = [] + 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}{location_text},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 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(): + 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)) + 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() + self._handles.clear() diff --git a/src/initium/inference/model_loader.py b/src/initium/inference/model_loader.py index 0b68989..81861d9 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, @@ -187,8 +187,18 @@ 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"], + moe_placement=c.get("split_graph_moe_placement", "all"), + 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"], @@ -198,8 +208,30 @@ 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), + 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), + 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( + "workspace_bistable_deterministic", False + ) + return model def load_model(ckpt: dict, device: torch.device, deterministic_write: bool = False) -> LoadedModel: @@ -227,6 +259,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/inference/run_inference.py b/src/initium/inference/run_inference.py index 0eaa6b6..7ab96e0 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() @@ -574,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/mamba_backbone_parallel.py b/src/initium/mamba_controller/mamba_backbone_parallel.py index d109ffd..2f1c35e 100644 --- a/src/initium/mamba_controller/mamba_backbone_parallel.py +++ b/src/initium/mamba_controller/mamba_backbone_parallel.py @@ -224,6 +224,8 @@ 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, ): @@ -233,6 +235,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 +270,28 @@ 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, + residual_scale=moe_residual_scale, + 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,7 +299,38 @@ 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, *, 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 + 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") + 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, token_mask=token_mask + ) + return base, experts, indices, weights + 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/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 11e3af8..ec52433 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): @@ -180,11 +181,14 @@ def __init__( independent_linears: bool = True, 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, 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 -- @@ -200,6 +204,19 @@ 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, + 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, + workspace_bistable_gradient_mode: str = "truncated", + workspace_bptt_horizon: int = 0, ): super().__init__() if not independent_linears: @@ -228,6 +245,43 @@ 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 + 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") + 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" + ) + 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 + 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 + # 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). @@ -246,12 +300,14 @@ 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, 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, ) @@ -318,7 +374,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, @@ -334,7 +390,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, @@ -357,7 +413,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, @@ -370,7 +426,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, @@ -381,7 +437,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, @@ -397,7 +453,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, @@ -411,7 +467,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, @@ -424,7 +480,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, @@ -459,6 +515,9 @@ 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, ) ) # See module docstring's API-compatibility section for why this @@ -467,16 +526,37 @@ 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, + 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, + bptt_horizon=workspace_bptt_horizon, + ) + 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. - self.moe_enabled = moe_enabled + # Keep a flat list of every active MoE block so the training loop can + # collect auxiliary losses and routing diagnostics uniformly. + 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 @@ -501,6 +581,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 @@ -518,6 +599,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, @@ -525,6 +609,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. @@ -535,11 +633,64 @@ 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" # 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, 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, token_mask=sequence_valid_mask + ) + ) + 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, + 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 + # 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 + ) + observe_tensor("split_graph.h_after_workspace", H) + else: + 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) --- @@ -576,14 +727,16 @@ 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) ) 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 @@ -592,7 +745,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: @@ -615,6 +770,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 @@ -632,6 +788,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/memory_manipulation/dnc_mechanisms.py b/src/initium/memory_manipulation/dnc_mechanisms.py index 5a7571a..50b628f 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,13 +11,108 @@ 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") +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. + + 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 _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.""" 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. + """ + + dnc_implementation = "base" def __init__( self, @@ -29,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: @@ -41,17 +137,35 @@ 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] = {} 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 +187,56 @@ 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. + + 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 +276,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 +554,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 +645,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 +670,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,9 +722,26 @@ 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 +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, @@ -338,9 +755,26 @@ 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.""" - return MechanismMemory( + if dnc_implementation not in DNC_IMPLEMENTATIONS: + raise ValueError(f"Unknown DNC implementation: {dnc_implementation!r}") + 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, cell_size=cell_size, @@ -352,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 98256b3..20efabb 100644 --- a/src/initium/memory_manipulation/dynamic_memory_resize.py +++ b/src/initium/memory_manipulation/dynamic_memory_resize.py @@ -50,7 +50,12 @@ 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, + compile_memory_forward, +) from initium.memory_manipulation.link_matrix_ablation import AblatableSparseLinkMemory @@ -104,6 +109,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 +118,13 @@ 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: - 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 @@ -148,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, 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 {})) 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/bistability.py b/src/initium/workspace/bistability.py new file mode 100644 index 0000000..f0fad3b --- /dev/null +++ b/src/initium/workspace/bistability.py @@ -0,0 +1,194 @@ +"""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, + gradient_mode: str = "truncated", + ) -> 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") + 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)) + 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): + # 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 * previous_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 new file mode 100644 index 0000000..7fdf6fb --- /dev/null +++ b/src/initium/workspace/broadcast.py @@ -0,0 +1,438 @@ +"""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 + +from initium.workspace.bistability import BistableWriteGate + + +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, + *, + gate_variant: str = "continuous", + 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: + 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'") + 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) + 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, + gradient_mode=bistable_gradient_mode, + ) + 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) + + 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, + 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 + # 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, + 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): + return self._forward_step( + specialist_deltas.float(), + route_indices, + route_weights.float(), + 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(), + route_indices, + route_weights.float(), + 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( + 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, + 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( + 1, route_indices.unsqueeze(-1).expand(-1, -1, self.d_model) + ) + 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 + 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) + 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: + # 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(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(write_candidates) + 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) + 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: + # 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(), + **bistable_diag, + **({"writer_frequency": writer_frequency} if self.diagnostics_enabled else {}), + } + 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) + 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 = 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(), + **bistable_diag, + **({"writer_frequency": writer_frequency} if self.diagnostics_enabled else {}), + } + return (enriched * 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, + 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 = [] + 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, + 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 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) + 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) + ) + 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]: + """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 = {} + 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,) 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)