F.1 — finite-difference gradchecks for tanh, sigmoid, and the exact gated
composite (1-z)*m + z*c (z=sigmoid, c=tanh) added to tape_cpu_gradcheck.rs.
All pass (rel-err < 2e-2) — the activation VJPs that already existed are now
proven correct on the real CpuBackend (E0 discipline), so the gated cell can
rely on them.
F.2 — GatedMemoryUpdater: a GRU-style gated recurrent memory cell.
z = sigmoid([M‖x]·W_z); c = tanh([M‖x]·W_c); M_t = (1-z)⊙M + z⊙c
The convex update is a non-expansion (|M_t| ≤ max(|M_0|, 1)), so free rollout
stays bounded with NO clamp — and the learned gate can both jump at a regime
switch (z≈1) and hold+denoise in steady state (z≈0), which the residual+leaky
ClonedMemoryUpdater cannot. Same method surface (new/step/predict/train_step/
train_step_memory) so the omni-think facade is cell-generic.
Tests: gated cell trains (loss drops); 500-step free rollout stays in [-1,1]
without a clamp. clippy(-D)/fmt clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
SetEncoderConfig gains an opt-in `recency` flag (default false = unchanged
mean-pool stationary behavior). When set, the encoder adds sinusoidal timestamp
embeddings to the token embeddings (so attention can reason about position) and
pools with an exponential-decay (recent-weighted) reduction instead of a uniform
mean — a recent-window sufficient statistic that tracks non-stationary signals.
Both use only existing-VJP tape ops (add + matmul/softmax); no new params.
Test: recency-mode teacher trains end-to-end (loss >5x drop). Existing
stationary tests unchanged. clippy(-D)/fmt clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Two robustness additions for on-policy distillation, which trains on the
policy's own (clamped, sometimes extreme) visited states:
- train_step now skips non-finite-loss steps and clips gradients, so an extreme
rollout state can't poison every weight via Adam.
- train_step_memory: a memory-only correction at a caller-chosen lr — trains
just the recurrence (W_in, W_mem) to map (prev,x)->target_mem, leaving the
BC-trained readout (W_read) intact. This is the key to DAgger working: it
pulls the free-running memory back toward the oracle trajectory without the
readout retraining that otherwise blows the predictions up.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Make the cloned updater's recurrence M_t = ρ·M_{t-1} + Δ a contraction (ρ=0.9)
and clamp the free-running state (|M| ≤ 4). One-step BC has no signal against
autoregressive drift (that is E4/DAgger's job); these keep a free rollout
numerically bounded — no 1e25 blow-up — so E4 has a stable base to refine.
Training is teacher-forced on the bounded oracle memory, so neither ever binds
during training.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
A trainable recurrent cell g(M_{t-1}, x_t) -> M_t (gelu MLP, residual memory
update) plus a predict-the-future readout, trained on Autodiff<CpuBackend> with
the same host-Adam pattern as SetEncoderTeacher.
train_step is one-step behavioral cloning, TEACHER-FORCED on the oracle memory:
each example feeds the oracle's previous memory and regresses (new_mem ->
target_mem) + (readout -> next latent) in a single tape graph. No rollout, no
backprop-through-time. Off-tape step/predict use a gelu byte-identical to the
CPU tape's (tanh approx, sqrt(2/pi), 0.044715) so rollout matches training.
The residual update M_t = M_{t-1} + delta lets the recurrence rule stay
length-invariant — the basis for extrapolating past the teacher's horizon (E3b
generalization test, omni-cortex side).
Tests: updater learns a one-step transition (>10x loss drop); off-tape
step/predict shapes + finiteness. clippy(-D)/fmt clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
MambaRecurrence reconstructs MambaBlock's selective-scan as an explicit
step(state, x_t) -> (state', y_t) recurrence over host f32 weights, with the
recurrent state (h, conv-ring) carried between steps. This is the substrate the
SMT memory updater is behaviorally cloned on (E3b) — no BPTT, one step at a time.
Dimensions are derived from the persisted tensor shapes, so mamba.rs is left
untouched (it is near the 1250-LOC cap). silu/softplus are byte-identical copies
of the forward's.
Equivalence pin (the plan's highest-risk item): stepping a window one token at a
time from a zero state reproduces the full-sequence forward EXACTLY —
max_abs_diff = 0.0 (bit-identical), on an active_block with wide Delta so the
scan genuinely drives the output. Plus a fresh-state determinism test.
clippy(-D warnings) + fmt clean on the new module and test.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Promotes the proven attention set-encoder smoke test (PR #4,
tape_train_smoke.rs) into a reusable rtx-transformers layer for Supervised
Memory Training. The teacher maps a window of past tokens to a fixed-size
memory M via embed → self-attention → residual → mean-pool → memory
projection, with a decoder head supervised by predict-the-future MSE so that
M becomes a sufficient statistic of the past.
- Trains end-to-end on Autodiff<CpuBackend> (the gradient-correct real backend
from PR #3), with a self-contained deterministic host-side Adam.
- Time-parallel by construction (one window → one memory, no recurrence to
unroll) — this is the oracle whose trajectory the recurrent Mamba updater is
later behaviorally cloned against, so the recurrent net never needs BPTT.
- Exposes named_params/set_named_params so the caller (omni-think's
PredictiveStateTeacher facade) owns safetensors persistence + BLAKE3 sealing.
- Adds rtx-backend + rtx-backend-cpu deps (the tape needs a concrete backend).
Tests: teacher trains (loss >5x drop), encode is deterministic + fixed-size,
params round-trip. fmt + clippy(-D warnings) clean on the new module.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Persist accumulated WIP across rtx-tensor / rtx-transformers / rtx-runtime.
Two related feature groups:
GPU enablement (unblocks the Phase-3 spec §8 CUDA Mamba path):
- rtx-tensor: `Tensor::cuda_device_ptr()` + storage GPU-buffer accessors
(`raw_ptr.rs`) expose the raw CUdeviceptr that kernel launches need —
the "rtx-tensor GPU memory access API" the Mamba CUDA kernels were
blocked on.
- rtx-transformers: `MambaBlock::forward_cuda` runs the four linear
projections through cuBLAS on-device (in/out/x/dt_proj), keeping the
selective scan + conv1d + activations on CPU; dispatched automatically
from `forward` when on a CUDA device under the `cuda` feature. Updated
`mamba_cuda_kernels.rs` accordingly.
- supporting plumbing in rtx-runtime stream/bridge and rtx-tensor
storage/conversion/concatenation/creation + rtx-flash-attention.
Linear algebra (rtx-tensor):
- `linalg/matrix_exp.rs`: real matrix exponential via scaling-and-squaring
with a degree-13 Padé approximant (Higham 2005), f64 internally.
- `complex/linalg.rs`: complex matmul/adjoint, Hermitian eigendecomposition
(`ComplexEigenResult`), and the complex matrix exponential, nalgebra-backed.
- tests for both.
Builds verified on the CPU path (`cargo check -p rtx-tensor -p rtx-transformers
-p rtx-runtime -p rtx-flash-attention` clean). The `cuda` feature and the
rtx-backend-cuda NVCC build remain unbuildable on this host (CUDA/glibc header
mismatch) — pre-existing and unrelated to these changes.
End-to-end demonstration that training the real selective-scan Mamba buys
a genuine temporal-modeling win — the payoff of M1–M3.
Task: next-token prediction on a multi-regime sequence (x[t] = μ_regime +
noise). Predicting x[t+1] inside a regime requires integrating recent
history to average out the noise — a memoryless model can't.
Result (same sequence, all three):
persistence (memoryless) MSE = 0.0895
untrained Mamba MSE = 0.2167
trained Mamba (400 Adam) MSE = 0.0002
The trained backbone integrates history to de-noise the regime mean,
beating the memoryless persistence baseline by ~450× and improving
~1000× over its untrained self. (Single-sequence fit: demonstrates the
SSM's temporal-modeling capacity, not held-out generalization.)
This replaces the old non-result ("random Mamba 6.5% vs linear 48%") with
a real "trained SSM exploits temporal structure" demonstration.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Plain randn(0,1) init made A=-exp(A_log) and Δ wildly large, overflowing
the real exp(Δ·A) scan to NaN (the stub never cared — it discarded these
weights). new()/new_seeded() now share a build() with canonical S6 init:
- A_log = ln(1..=d_state) ⇒ A = -(1..=d_state), bounded
- dt_bias so softplus(dt_bias) ≈ 0.01 (small, stable Δ; near-identity
scan at init — intentional for gradient flow)
- D = 1, zero conv bias, projections scaled by 1/√fan_in (capped 0.5)
This fixes the NaN that broke omni-cortex's d231 action-conditioned
predictor training (now green). Seeded determinism preserved.
Tests: active_block helper (Δ overridden to ≈0.69) exercises the
scan-active regime so the liveness check can observe each weight; the
training test asserts a seed-varying backbone weight (conv1d_weight)
moves. All 6 selective-scan tests green incl. the finite-diff grad check.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Demonstrates end-to-end trainability of the real selective scan. A
self-contained Adam loop (the rtx-transformers AdamOptimizer has no
public gradient setter — the spec permits a bespoke loop) fits a teacher
block's output on a fixed input: forward → MSE → analytic backward →
Adam step → rebuild. Over 200 steps the loss drops >50% and the backbone
weight A_log moves, confirming gradients actually train the model (not
just the head). All 6 selective-scan tests green.
The production AdamOptimizer can be wired once it exposes a gradient
setter; the M2 backward already returns grads in its HashMap shape.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Hand-written VJP of the selective-scan forward (Phase-3 spec §5/§7):
MambaBlock::backward(x, d_out) -> per-parameter gradients keyed by
persistence name, summed over the batch. Differentiates the scan
analytically (reverse-time recurrence over the cached h trajectory)
rather than through the immature rtx-tensor autograd tape.
Covers every parameter: in_proj, conv1d_weight, conv1d_bias, A_log
(via A=-exp(A_log) ⇒ dA_log = dA·A), x_proj, dt_proj, dt_bias, D,
out_proj. Adds stable sigmoid_f32 / silu_grad_f32 helpers.
New test analytic_gradients_match_finite_differences: on a small
well-conditioned instance, ≥30 sampled grad elements across all 9
params match central finite differences within (5e-3 + 5e-2·|fd|).
All 5 selective-scan tests green.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
The MambaBlock forward was a stub — SelectiveScan::forward passed input
through, discretize returned zeros, conv1d was a no-op, and B/C were
randn per call, leaving conv1d_weight/dt_proj/A_log as dead weights with
zero temporal mixing. This implements a genuine S6 selective-scan:
- New params: x_proj [d_inner, dt_rank+2*d_state] (data-dependent dt,B,C),
dt_bias [d_inner], D [d_inner] (skip). Added to new()/new_seeded() and
the persistence contract (persistence_tensors/from_persistence_tensors).
- Real forward (CPU f32, looped — backbone is small): in_proj -> causal
depthwise conv1d -> SiLU -> x_proj->(dt,B,C) -> delta=softplus(dt.dt_proj
+dt_bias) -> A=-exp(A_log) -> sequential scan h=dA.h+dBu, y=sum C.h + D.u
-> gate by SiLU(z) -> out_proj. Residual moved OUT (canonical).
Numerically-stable silu_f32/softplus_f32 helpers.
- The scan runs inline (not via the immature rtx-tensor autograd tape);
the analytic backward lands in M2 per docs/phase3_real_ssm_spec.md.
New tests/real_selective_scan.rs (4 cases, all green): liveness (each
formerly-dead weight now moves the output), causality (no future
leakage), seeded determinism, and finite/non-constant output.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
Two coordinated additions for the omni-cortex D249/D250 work:
rtx-tensor: Tensor::randn_seeded(shape, device, seed) — like
randn() but routes through StdRng::seed_from_u64(seed) so two
calls with (shape, device, seed) produce bit-exact identical
tensors. CPU is the canonical generator; GPU calls go via to_device
transfer. Required for reproducible model init.
rtx-transformers: MambaBlock gains:
- new_seeded(config, device, seed) — every internal weight tensor
initialised via randn_seeded() with per-tensor SplitMix64-derived
seeds. Two calls with the same (config, device, seed) → bit-
exact identical block.
- persistence_tensors() -> Vec<(&'static str, &Tensor)> — read-only
view of the six (or seven, with conv_bias) internal weight
tensors with canonical names (in_proj, conv1d_weight,
conv1d_bias?, A_log, dt_proj, out_proj). Stable across versions
so safetensors round-trip works.
- from_persistence_tensors(config, device, HashMap<String, Tensor>)
— rebuild a MambaBlock from a name → tensor map. Validates each
tensor's shape against the config and surfaces clean errors on
mismatch (so wrong-DIM safetensors loads fail explicitly).
These three primitives together give omni-cortex's D249 (operator-
seeded determinism) and D250 (safetensors round-trip + BLAKE3 hash
pin) clean library hooks without exposing MambaBlock's private
fields.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
ring_allreduce() contained its own simulation that multiplied each gradient
value by world_size (to fake an AllReduce sum), bypassing the ProcessGroup
dispatch entirely. This meant the overlapped synchronization path never used
NCCL or RNCCL, even when those features were compiled in.
Replace the hand-rolled simulation with a call to
self.process_group.allreduce(tensor, ReduceOp::Sum) so the overlapped path
uses the same backend as synchronize_gradients_sequential. The communication
latency sleep is kept for benchmarking purposes.
Add test_ring_allreduce_matches_sequential_path to verify both paths produce
identical gradient values under CPU simulation.
Closes#10
The workspace root was upgraded to thiserror = "2" in an earlier commit,
but 56 per-crate Cargo.toml files still independently declared "1.0".
These crates do not use workspace.dependencies inheritance for thiserror.
All updated to thiserror = "2" for complete fleet alignment.
Includes: rtx-backend, rtx-tensor, rtx-losses, rtx-backend-cuda/rocm/metal,
all training crates (rtx-auto, rtx-rl, rtx-distributed, rtx-federated, etc.),
specialized crates (rtx-science, rtx-platform, rtx-nmf, rtx-neuro-*),
production crates (rtx-streaming, rtx-serving-api), and all demo crates.
cargo check --workspace: PASSES.
osobh