You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
This commit was created on GitHub.com and signed with GitHub’s verified signature.
Add MoE workspace and diagnostics (#38)
* Add MoE workspace and diagnostics
Adds an optional shared workspace for split-graph MoE models, including competitive slot writes and broadcast readback to specialists. This also adds per-bank MoE diagnostics, segment-level activation/gradient logging, paired memory-on/off eval replay, and stricter finite-gradient/NaN safeguards for recurrent training. The change keeps existing behavior unchanged unless the new --workspace and related flags are enabled.
* Improve split-graph MoE routing and diagnostics
This patch keeps MoE additive around the original controller signal instead of replacing it at initialization, zero-initializes expert output projections, and fixes top-1 routing so the router still receives useful task gradients. Capacity is now assigned by route confidence rather than flattened token order, which avoids first-source bias during multi-source CfC routing and adds accepted/dropped source diagnostics for training logs and inference reports.
It also adds split-graph MoE placement controls (`all`, `backbone`, `combiner`) so backbone and combiner MoE effects can be isolated cleanly, and updates the README with the placement behavior and a short diagnostic smoke-test matrix for comparing MoE effects on memory-learning transitions. The change keeps the baseline CfC fusion path intact while letting source-aware experts provide learned residual corrections.
* Add backbone MoE isolation diagnostics
Adds optional split-graph backbone MoE isolation per zero-based block and wires it through the training/run, loader, and model configs. This enables backbone-only MoE experiments, validates incompatible options (workspace/bad ranges), and emits valid-vs-padded routing diagnostics for graph padding without masking the actual execution path.
* Add backbone MoE residual scaling control
This change adds a configurable backbone MoE residual scale for split-graph runs, defaulting to 1.0 to preserve existing behavior. The scale is validated as finite and non-negative, only allowed with backbone/all placement, rejected with the workspace broadcast path, and included in run IDs and checkpoint metadata.
The residual is applied before the backbone residual addition as x + scale * moe_delta, while leaving routing and auxiliary load-balancing behavior active. The new option is wired through the MoE layer, Mamba and CfC parallel backbone builders, and inference config loading. The README also documents the experimental usage and follow-up commands for conservative backbone-only scale testing.
* Enable fp32 DNC tracing and diagnostics
This commit forces DNC interface, addressing, and recurrent memory math onto fp32 even under CUDA autocast, and records the precision choice in startup logs/checkpoint metadata. It adds --dnc-backward-trace for split-graph runs to watch DNC read/write addressing, temporal links, usage/allocation, prod/cumprod numerics, and the first non-finite backward boundary with batch-row context; collected diagnostics are emitted during training and failure windows and include product-input stats, read-weight mass checks, and memory/read/link magnitudes.
* Enable residual-scaled workspace MoE broadcasts
This change allows the workspace broadcast path to run with the split-graph backbone MoE residual scale instead of rejecting it. The final-bank broadcast delta is scaled with the matched backbone setting, and active workspace attention/slot updates run in fp32 outside CUDA autocast while the zero-slot/closed-gate bypass remains unchanged. Startup reporting and the README were updated to describe the new phase-1 configuration, fp32 precision behavior, and the relevant MoE/DNC validation commands.
* Add frozen workspace warm-start and MoE diagnostics
This commit adds a warm-start path for frozen workspace experiments: a compatible workspace-disabled checkpoint can be resumed into a new frozen-workspace run without reinitializing backbone, MoE, combiner, DNC, stochastic-write, or task-head weights. It validates model metadata and shared tensor shapes, keeps newly introduced workspace parameters freshly initialized, and resets optimizer, scaler, RNG, curriculum, and step counters to start the workspace from lesson 1. Existing checkpoints with workspace state still follow the normal exact-resume behavior.
It also adds optional MoE routing diagnostics control. Detailed routing/activation diagnostics can be disabled to reduce overhead while keeping the auxiliary load-balancing loss and routing logic active. The code resets cached diagnostics cleanly when toggled, and the training loop now conditionally emits the more expensive MoE diagnostics. The README is updated to document the new warm-start workflow and expected behavior.
* Add optional segment/workspace diagnostics
Adds a segment-diagnostics flag to the training CLI and passes it into SegmentDiagnostics, which now skips hook registration when disabled. WorkspaceBroadcast also gains a diagnostics toggle to avoid expensive per-timestep reductions, and the training entrypoint reports whether MoE and segment diagnostics are enabled. This keeps long recurrent runs lighter while preserving the optional trace output when needed.
* Add eval-only checkpoint evaluation mode
Add evaluation-only resume flows that load checkpoint weights without restoring optimizer state or training updates. The new flags support deterministic eval RNGs and workspace counterfactual modes (checkpoint/reset-each-step/bypass), while preventing lesson advancement during evaluation-only passes so checkpoint scoring stays comparable.
* Add experimental bistable workspace gate
Introduce an opt-in recurrent stochastic write gate with checkpointed noise, gate-only warm starts, stability checks, and training/evaluation diagnostics. Document Phase 4 usage and add unit tests; preserve the model's prior training mode during traversal evaluation.
* Optimize DNC memory allocation
Add an opt-in exact allocation path that scatters weights through the existing usage permutation, eliminating a redundant sort while preserving state-dict compatibility. Wire the option through training, inference, and memory resizing, and add equivalence tests.
* Stabilize bistable gate gradients and diagnostics
Add configurable truncated or unrolled gradients for bistable settling, defaulting to truncated backpropagation while preserving the full forward trajectory. Extend gradient and activation diagnostics, route CfC steps through module hooks, and validate frozen gate-only training parameters.
* Add configurable workspace BPTT truncation
Allow gradients through persistent workspace slots to be truncated at a configurable horizon without changing the forward trajectory. Wire the setting through training, checkpoints, and inference; reset optimizer state when the horizon changes; and extend recurrent diagnostics with state-path metrics and gradient clipping counts.
* Add DNC execution and addressing options
Add compiled memory forwarding, fast and top-k allocation, temporal-link sharpening, and configurable batch sizes. Skip dense traversal products when links are ablated, preserve these settings when resizing memory, and document the tradeoffs.
* Add workspace reconsideration and diagnostics
Allow high-innovation proposals to reopen writes to occupied workspace slots, with configurable thresholds and a rate-matched random trigger. Add bistable, CFC, and DNC gradient tracing and training diagnostics to help investigate numerical issues.
* formatted and linted
* Preserve exact bistable gate outputs
Avoid floating-point round-off in straight-through hard writes and tighten type and nullability handling across MoE, diagnostics, and controller code.
* versioning
* format
* typecheck fix
* linter fix
* Fix typing for segment diagnostic hooks
Mark runtime callback assignments on PyTorch modules as intentional typing exceptions and reformat tensor-gradient hook lambdas without changing their behavior.
* typecheck fix