Skip to content

Repository files navigation

memory-native

📝 Preprint: Training Without Master Weights: a 6-bit Finite-State Synapse with Optimizer-in-State (v0.1 draft). A 6-bit code holds the entire per-parameter training state — no FP master weight, no Adam moments. Composed with reversible activations + int8/bf16 compute, a 1.21B-parameter model trains on a single T4 in 2.25 GiB peak, where the dense+Adam equivalent cannot allocate its 18 GiB of state.

Finite-state counter synapses + reversible activations for memory-efficient training, in pure PyTorch — no custom engine, runs on stock CPU/CUDA. Pure-Python package: pip install -e . and train.

GPU-validated (Tesla T4 / T4×2): results/KERNEL.md (fused update ×45.9 / step ×1.26) · results/SCALE_1B.md (1.21B params on one T4 at 2.25 GiB; dense+Adam OOMs at 18 GiB) · results/POOLS.md (all four memory pools) · results/SHOOTOUT.md (vs AdamW / 8-bit Adam / GaLore / LoMo).

Validated value proposition (real runs on a Tesla T4 — results/SHOOTOUT.md): the lowest training-memory of every contestant (AdamW, 8-bit Adam, GaLore, LoMo) at competitive-to-better quality, for ~2× slower steps — and the advantage grows with scale. At d=768 counter+int4 beats AdamW and 8-bit Adam on memory and quality at once (peak 1.32× below AdamW / 1.62× below 8-bit Adam; val −4.7% vs AdamW; speed gap narrowing). The method wins because it cuts both the optimizer pool (zero state) and activations (int4), while the memory-efficient optimizers only shrink optimizer state — a small slice of the activation-bound peak.

🍎 MLX / Apple-silicon port: the method now runs on MLX (Metal on macOS; CPU backend anywhere) — same 6-bit state, same packed 0.75 B/weight layout, same deterministic hash-SR update, validated bit-for-bit against this torch implementation. Counter models cross torch↔MLX losslessly, so you can PTQ/train on CUDA and fine-tune on a MacBook's unified memory. See docs/MLX_PORT.md, package src/memory_native_mlx/, demo scripts/mlx_demo.py.

Project map

Start here Purpose
PROJECT_STATUS.md What is verified, experimental, or still open
REPRODUCIBILITY.md Reproduce CPU, CUDA, scale, and recovery witnesses
results/README.md Evidence index and raw run artifacts
paper/MEMORY_NATIVE_PREPRINT.md Method, assumptions, and open scientific questions
CONTRIBUTING.md Contribution and review workflow
SECURITY.md Private vulnerability reporting and support scope

This repository is the reference research implementation and primary public project. MotifCL is a separate, supporting native runtime used to explore Vulkan deployment and compact-state execution on legacy GPUs. Results from one repository are not treated as proof for the other unless the corresponding witness is linked explicitly.

The method attacks all four memory pools of training at once:

Pool Lever What it is
Parameters + optimizer + gradients CompactCounterLinear / RMSCounterLinear a ternary weight whose optimizer state lives inside a per-synapse finite-state automaton; the update is fused into backward — no FP master weight, no Adam moments, no full gradient buffer
Activations ReversibleCouplingBlock activations are recomputed in backward from the output instead of stored — depth-independent activation memory

Install

pip install -e .            # from this directory
# or, once published: pip install memory-native

Dependencies are torch>=2.1 and numpy>=1.21.

Quickstart

import torch
from memory_native import RMSCounterLinear

layer = RMSCounterLinear(256, 256, C=11, lr=3e-3, lr_scale=2e-4)
x = torch.randn(64, 256)                 # note: no requires_grad needed
for _ in range(200):
    loss = (layer(x) - target).pow(2).mean()
    loss.backward()                       # the layer self-updates here; no optimizer.step()

The counter layer exposes no nn.Parameter — its weight is a uint8 state buffer and it updates itself during backward(). Mix it with normal modules; an AdamW over the rest of the model (embeddings, norms, head) trains those as usual.

Command-line gates

# char-LM parity: counter+RMS vs ternary-QAT vs dense AdamW, same arch/data/seed.
# Falls back to a synthetic corpus offline; use --data-path for real tinyshakespeare.
memory-native-charlm --config tiny --steps 600 --device cuda

# compare the dense baseline under a memory-efficient optimizer instead of FP32 AdamW:
memory-native-charlm --config tiny --kinds dense --optimizer galore --device cuda

# training-peak memory gate: counter_rms vs dense across optimizers (real peak on CUDA)
memory-native-memgate --config s512 --optimizers adamw,galore,lomo,bnb8 --device cuda

Baselines (the honest comparison set)

--optimizer / --optimizers selects what the dense baseline is trained with, so the memory claim is measured against real memory-efficient training, not only FP32 AdamW:

name what optimizer-state memory deps
adamw FP32 AdamW 2×params (fp32 m,v) torch
bnb8 8-bit AdamW (bitsandbytes) ~0.5×params bitsandbytes + CUDA
galore low-rank projected AdamW ~2×rank×dim torch (built in)
lomo fused-backward SGD zero moments torch (built in)

GaLore and LoMo are implemented in plain PyTorch here (run on CPU, no extra deps); bnb8 is optional and used where bitsandbytes+CUDA are present (skips cleanly otherwise).

Scale validation

The biggest open question is whether parity holds beyond micro/tiny. One command runs the full sweep (parity across kinds + dense-vs-optimizers + memory gate) at d=512 and saves logs:

scripts/run_scale_validation.sh s512 2000 cuda      # config steps device
DATA_PATH=/path/to/tinyshakespeare.txt scripts/run_scale_validation.sh small 4000 cuda

Logs land in results/ (not committed pre-filled — capture them on your GPU).

A 60-step micro run already reproduces the method's story: counter_rms lands within ~0.1% of dense, beats vanilla counter, and isolates to ~+2.5% over ternary-QAT (the counter- optimizer cost) — matching the larger-scale numbers in results/SUMMARY.md.

What is measured vs what is roadmap

Honest framing — read before quoting memory numbers.

  • Realized now (pure PyTorch, verified on CPU): the learning dynamics, the optimizer-state saving (no FP master weight, no Adam moments), and — with PackedRMSCounterLinear (kind="counter_packed") — genuinely packed 0.75 byte/weight persistent state (4 codes / 3 bytes, bit-identical to the engine's packing; round-trip and identical-dynamics tested). memory_report / memory-native-memgate quantify it (real torch.cuda.max_memory_allocated on CUDA, byte accounting on CPU).
  • Verified on GPU (Tesla T4): the Triton forward kernel (triton_counter.py, TritonCounterLinear) decodes the packed state inside the GEMM (no dense weight) and matches the dense reference within f32 tol (err ≤ 3e-6). And at d=512 counter+RMS keeps parity with — slightly beats — dense AdamW (val −1.7%). See results/SUMMARY.md.
  • Activation pool — collapsed (T4): ReversibleSequence is a single whole-chain RevNet Function (stores only the final output, reconstructs the rest), so activation memory is O(1) in depth — 13.5× below plain at depth 256 (5.29 GiB → 392 MiB), gradients identical to the per-block version. With act_save_bits (unbiased int8/int4 saved activations) the saved-X cost drops too. See results/POOLS.md.
  • Fused update kernel — done (T4): memory_native.fused_update collapses the per-element RMS+stochastic-rounding update into one launch (deterministic hash-SR, bit-quantified against a CPU reference), ×45.9 on the update / ×1.26 on the step, wired into PackedRMSCounterLinear. See results/KERNEL.md.
  • Scale — demonstrated (T4): the full method (counter + reversible) trains a 1.21B-param model on a single 14.6 GiB T4 at 2.25 GiB peak, where dense+Adam needs 18 GiB of state and OOMs before step 0. See results/SCALE_1B.md.
  • Strict update-from-IO — implemented and T4-validated: both the row-scale update_from_io.triton_counter_update_from_io path and the group-scale solver-v3 path can form the update directly from (state, scale, v, x, grad_out) without materializing a dense [out,in] grad_w. The result is correct and establishes the strict-memory bound, but the hand-written correlation is dramatically slower than cuBLAS at real token counts: the row-scale witness was ~860× slower, while the group-scale Qwen-shape path took 30–46 seconds per layer at M=4096. The practical default therefore remains cuBLAS correlation plus the fused update, optionally row-tiled to bound transient memory. See results/ACCELERATION.md and results/group_kernel_opt_stage01.md.

Roadmap

  1. Practical low-memory correlation — retain the verified strict from-IO kernel as the zero-grad_w bound, while improving the useful frontier between full cuBLAS correlation and row-tiled correlation. Promote a path only when end-to-end training throughput and peak memory both beat the current default on a named model and GPU.
  2. Multi-session scale — 2×T4 data-parallel + checkpoint/resume (scripts/fineweb_1b_2xt4.py); accumulate tokens across Kaggle sessions on real web text (FineWeb-Edu, BPE).
  3. Reversible at depth — anchors (store every k blocks) to trade a little memory for less recompute, and a depth-sweep for float-reconstruction error beyond the tested range.

Constraints (same eager-only contract as the engine)

  • The in-backward update is eager-only: incompatible with gradient accumulation, weight sharing of a counter module, activation checkpointing, and torch.compile graph capture without an explicit scheduler. One forward → one backward per step; measure under torch.no_grad(). Data-parallel is supported a different way than DDP: the counter gradient is all-reduced inside backward (the optimizer is an in-place state update, so there is no Parameter .grad), keeping every replica's packed state bit-identical — validated on 2×T4 (0 bytes differ across ranks). See scripts/fineweb_1b_2xt4.py.
  • ReversibleCouplingBlock needs deterministic F/G (no dropout/RNG). Float reconstruction is exact enough at tested depth (≤12 blocks); deeper stacks need a depth-sweep and possibly anchors.

Layout

src/memory_native/
  counter.py        CompactCounterLinear, RMSCounterLinear, encode/decode, stochastic rounding
  packed.py         PackedRMSCounterLinear — real 0.75 byte/weight storage (4 codes / 3 bytes)
  triton_counter.py TritonCounterLinear + in-GEMM packed decode (CUDA; T4-verified, net-negative)
  fused_update.py   one-launch RMS+SR counter update kernel (CUDA; T4-verified, ×45.9)
  reversible.py     ReversibleCouplingBlock, ReversibleSequential (recompute backward)
  actquant.py       unbiased low-bit saved activations (the counter layer's act_save_bits)
  budget.py         training_budget — symbolic 4-pool memory model (deep-v2, reproduced exact)
  baselines.py      TernaryQATLinear, make_linear factory (kinds incl. counter_packed)
  optimizers.py     build_optimizer: adamw / bnb8 / galore / lomo
  models.py         swappable-linear GPT harness + configs (micro/tiny/s512/small)
  memory.py         memory_report, peak_training_memory, compare_training_peak
  data.py           char corpus loader with offline synthetic fallback
  cli.py            memory-native-charlm / memory-native-memgate entry points
src/memory_native_mlx/
  counter.py        RMSCounterLinear on MLX (custom-VJP self-update, deterministic hash-SR)
  packed.py         PackedRMSCounterLinear — same 0.75 B/weight packed layout, bit-identical
  reversible.py     whole-chain reversible sequence (O(1)-in-depth) + anchored mode
  metal_update.py   fused RMS+SR update as a custom Metal kernel (mirror of the Triton one)
  interop.py        lossless torch <-> MLX counter-state transfer (the CUDA->MacBook handoff)
scripts/            run_scale_validation.sh (one-command GPU sweep -> results/)
                    mlx_demo.py (counter+reversible char-LM on MLX; Metal or CPU)
tests/              pytest: encode/decode, learning, reversible grad-check, packed round-trip,
                    optimizers, memory gate, triton (CUDA-skipped)

How this was built

Built entirely by AI (Claude) under sustained human direction, over months, by someone with no formal CS/math background. Working protocol: every claim needs an executable witness — tests, frozen pre-run forecasts, raw logs committed next to results, negative results reported first-class. The repo, not the author, answers technical questions. Status: frozen (July 2026) — out of money and hardware, not out of ideas. Everything reproduces from a cold clone.

License

MIT.

About

Finite-state counter synapses + reversible activations for memory-efficient training. 6-bit code holds the entire per-parameter training state — no FP master weight, no Adam moments. 1.21B params trains on a single T4 in 2.25 GiB peak.

Topics

Resources

Contributing

Security policy

Stars

Watchers

Forks

Releases

Packages

Used by

Contributors

Languages