Repository navigation
Add MoE workspace and diagnostics - #38
Open
samueleternity wants to merge 14 commits into
Open
samueleternity wants to merge 14 commits into
samueleternity wants to merge 14 commits into
Conversation
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.
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.
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.
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.
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.
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.
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.
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 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.
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.
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.
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.
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 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.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.