Add opt-in memory-efficient CSR aggregation for graded synapses

benchmarks/rtx5090-sparse-training.jsonl

@@ -0,0 +1,8 @@
+{"backend": "scatter", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 4, "steps": 8, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 6.658950005657971, "repetition_ms": [12.325900024734437, 6.494899978861213, 7.438400061801076, 6.689100060611963, 6.978000048547983, 7.02719995751977, 6.510999985039234, 6.6287999507039785, 6.534100044518709, 6.385100074112415], "baseline_bytes": 119839232, "peak_bytes": 440835584, "incremental_peak_bytes": 320996352, "loss": 0.2890213429927826, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "sparse_trainable", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 4, "steps": 8, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 9.959999995771796, "repetition_ms": [27.376099955290556, 9.837899939157069, 10.078800027258694, 10.149299982003868, 9.966399986296892, 10.204300051555037, 9.953600005246699, 9.730400051921606, 9.147300035692751, 9.25999996252358], "baseline_bytes": 170425344, "peak_bytes": 260221440, "incremental_peak_bytes": 89796096, "loss": 0.2890213429927826, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "scatter", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 4, "steps": 32, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 20.696200022939593, "repetition_ms": [54.0492000291124, 22.04800001345575, 21.971500013023615, 19.167799968272448, 20.561299985274673, 19.632500014267862, 18.653900013305247, 20.831100060604513, 18.73799995519221, 28.39949994813651], "baseline_bytes": 137375744, "peak_bytes": 1118399488, "incremental_peak_bytes": 981023744, "loss": 0.2865479290485382, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "sparse_trainable", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 4, "steps": 32, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 29.006949975155294, "repetition_ms": [34.638999961316586, 35.13699991162866, 26.536599965766072, 27.772999950684607, 29.86629994120449, 25.52969998214394, 28.1476000091061, 30.274000018835068, 25.780600029975176, 31.48280002642423], "baseline_bytes": 188438528, "peak_bytes": 382202880, "incremental_peak_bytes": 193764352, "loss": 0.2865479290485382, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "scatter", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 16, "steps": 8, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 13.377799943555146, "repetition_ms": [29.85140006057918, 13.411999912932515, 13.18919996265322, 12.390799936838448, 13.343599974177778, 15.434899949468672, 14.01330006774515, 13.886100030504167, 12.91369996033609, 12.86340004298836], "baseline_bytes": 137375744, "peak_bytes": 1294208512, "incremental_peak_bytes": 1156832768, "loss": 0.288987934589386, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "sparse_trainable", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 16, "steps": 8, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 8.459850039798766, "repetition_ms": [12.31670007109642, 8.641700027510524, 8.04660003632307, 8.791400003246963, 8.373800083063543, 8.452400099486113, 9.07190004363656, 8.081800071522593, 7.650499930605292, 8.46729998011142], "baseline_bytes": 188438528, "peak_bytes": 373998592, "incremental_peak_bytes": 185560064, "loss": 0.288987934589386, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "scatter", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 16, "steps": 32, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 48.0207500513643, "repetition_ms": [96.65060008410364, 48.018300090916455, 47.19750001095235, 47.827200032770634, 48.02320001181215, 49.035799922421575, 48.76949999015778, 49.54509995877743, 46.662600012496114, 46.610499965026975], "baseline_bytes": 208365056, "peak_bytes": 3909596672, "incremental_peak_bytes": 3701231616, "loss": 0.2865470051765442, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}
+{"backend": "sparse_trainable", "extent": 15, "nodes": 45669, "edges": 1513231, "trainable_parameters": 734, "batch": 16, "steps": 32, "dtype": "torch.float32", "dt": 0.02, "repeats": 10, "warmup": 3, "forward_backward_ms": 30.952000059187412, "repetition_ms": [38.178599905222654, 33.12720009125769, 27.798900031484663, 27.643300010822713, 30.70800006389618, 27.18470001127571, 31.554799992591143, 28.717399924062192, 31.196000054478645, 32.88720001000911], "baseline_bytes": 257521152, "peak_bytes": 951452672, "incremental_peak_bytes": 693931520, "loss": 0.2865470051765442, "gradients_finite": true, "environment": {"torch": "2.14.0+cu130", "cuda": "13.0", "device": "NVIDIA GeForce RTX 5090", "platform": "Windows-11-10.0.26200-SP0", "python": "3.12.11", "threads": 4, "source_commit": "7a76dc3cf98c8c1d380d0d86275b094c22b081c8", "source_dirty": false}}

benchmarks/rtx5090-sparse-training.md

@@ -0,0 +1,54 @@
+# Sparse training on the FlyVis visual-system graph
+
+Measured 2026-09-26 on an RTX 5090, Windows 11, Python 3.12.11, PyTorch
+2.14.0+cu130 (CUDA 13.0), float32, four CPU threads. Source commit:
+`7a76dc3cf98c8c1d380d0d86275b094c22b081c8` (clean at benchmark start).
+
+The supplied `fib25-fib19_v2.2.json` graph at extent 15 has **45,669 nodes,
+1,513,231 edges and 734 trainable shared parameters**. This measures the
+actual FlyVis `Network` forward/backward with its default dynamics, synthetic
+input, `dt=0.02`, and a squared-activity loss. It excludes graph compilation,
+data loading, a task decoder, and optimizer updates. It does not measure
+optical-flow accuracy or learning speed to a target score.
+
+Times are medians of 10 synchronized wall-clock repetitions after 3 warmups.
+Peak is `torch.cuda.max_memory_allocated`, including live model tensors,
+input, activations and gradients. It is not allocator-reserved or total
+process VRAM. The JSONL includes every timing, the baseline allocation,
+incremental peak, loss, and environment. Initialization uses seed 20260926;
+input uses seed 20260927 for both backends.
+
+| Batch | Steps | Backend | Forward + backward (ms) | Peak allocated (GiB) |
+|---:|---:|---|---:|---:|
+| 4 | 8 | `scatter` | 6.66 | 0.411 |
+| 4 | 8 | `sparse_trainable` | 9.96 | 0.242 |
+| 4 | 32 | `scatter` | 20.70 | 1.042 |
+| 4 | 32 | `sparse_trainable` | 29.01 | 0.356 |
+| 16 | 8 | `scatter` | 13.38 | 1.205 |
+| 16 | 8 | `sparse_trainable` | 8.46 | 0.348 |
+| 16 | 32 | `scatter` | 48.02 | 3.641 |
+| 16 | 32 | `sparse_trainable` | 30.95 | 0.886 |
+
+At batch 16 and 32 steps, the sparse backend used **4.11x less peak allocated memory**
+(75.7% reduction) and was **1.55x faster** in this run.
+At batch 4, it saved memory but was slower. This supports an opt-in
+backend, not a universal speedup claim. Timing varies across runs;
+the individual repetitions are retained for inspection. All gradients
+were finite and paired final losses matched at the recorded precision.
+
+Reproduce:
+
+```bash
+python benchmarks/sparse_current.py --extent 15 --batches 4 16 --steps 8 32 --repeats 10 --warmup 3 --out results.jsonl
+```
+
+Local Windows setup note: datamate 1.0.0 needed its HDF5 handle closed
+before fallback file deletion when building the connectome cache.
+That environment-only workaround does not affect timed tensor operations
+and is not included in this FlyVis source patch. Additional environment
+versions: NumPy 2.5.3, pandas 3.0.6, SciPy 1.18.1, torchvision 0.29.0+cu130.
+
+Validation: 198 passed, 12 skipped across the new sparse tests and
+existing network, dynamics and initialization suites. Ruff lint and
+format checks passed for changed Python files. This is targeted
+validation, not the full suite requiring pretrained assets.

benchmarks/sparse_current.py

@@ -0,0 +1,131 @@
+"""Benchmark actual FlyVis Network forward/backward on its visual connectome.
+
+Run after installing FlyVis: python benchmarks/sparse_current.py --out results.jsonl
+Model compilation and data loading are excluded. No training dataset is needed.
+"""
+
+import argparse
+import gc
+import json
+import platform
+import statistics
+import subprocess
+import time
+from pathlib import Path
+
+import torch
+
+from flyvis import Network
+
+
+def measure(backend, extent, batch, steps, repeats, warmup):
+    torch.manual_seed(20260926)
+    model = Network(
+        connectome=dict(
+            type="ConnectomeFromAvgFilters",
+            file="fib25-fib19_v2.2.json",
+            extent=extent,
+            n_syn_fill=1,
+        ),
+        backend=backend,
+    )
+    torch.manual_seed(20260927)
+    stimulus = torch.randn(batch, steps, model.n_nodes) * 0.1
+
+    def step():
+        model.zero_grad(set_to_none=True)
+        output = model(stimulus, dt=0.02)
+        loss = output.square().mean()
+        loss.backward()
+        return loss
+
+    for _ in range(warmup):
+        step()
+    model.zero_grad(set_to_none=True)
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+    baseline = torch.cuda.memory_allocated()
+    torch.cuda.reset_peak_memory_stats()
+    measurements = []
+    for _ in range(repeats):
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+        loss = step()
+        torch.cuda.synchronize()
+        measurements.append((time.perf_counter() - start) * 1000)
+    peak = torch.cuda.max_memory_allocated()
+    result = dict(
+        backend=backend,
+        extent=extent,
+        nodes=model.n_nodes,
+        edges=model.n_edges,
+        trainable_parameters=sum(
+            p.numel() for p in model.parameters() if p.requires_grad
+        ),
+        batch=batch,
+        steps=steps,
+        dtype=str(stimulus.dtype),
+        dt=0.02,
+        repeats=repeats,
+        warmup=warmup,
+        forward_backward_ms=statistics.median(measurements),
+        repetition_ms=measurements,
+        baseline_bytes=baseline,
+        peak_bytes=peak,
+        incremental_peak_bytes=peak - baseline,
+        loss=float(loss.detach()),
+        gradients_finite=all(
+            bool(torch.isfinite(p.grad).all())
+            for p in model.parameters()
+            if p.requires_grad and p.grad is not None
+        ),
+    )
+    model.zero_grad(set_to_none=True)
+    gc.collect()
+    torch.cuda.empty_cache()
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--extent", type=int, default=15)
+    parser.add_argument("--batches", type=int, nargs="+", default=[4, 16])
+    parser.add_argument("--steps", type=int, nargs="+", default=[8, 32])
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--warmup", type=int, default=2)
+    parser.add_argument("--out", type=Path, required=True)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        parser.error("This peak-memory benchmark requires CUDA")
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(torch.float32)
+    torch.set_num_threads(4)
+    environment = dict(
+        torch=torch.__version__,
+        cuda=torch.version.cuda,
+        device=torch.cuda.get_device_name(),
+        platform=platform.platform(),
+        python=platform.python_version(),
+        threads=torch.get_num_threads(),
+        source_commit=subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], text=True
+        ).strip(),
+        source_dirty=bool(subprocess.check_output(["git", "status", "--porcelain"])),
+    )
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    with args.out.open("w", encoding="utf-8") as stream:
+        for batch in args.batches:
+            for steps in args.steps:
+                for backend in ("scatter", "sparse_trainable"):
+                    result = measure(
+                        backend, args.extent, batch, steps, args.repeats, args.warmup
+                    )
+                    result["environment"] = environment
+                    stream.write(json.dumps(result) + "\n")
+                    stream.flush()
+                    print(json.dumps(result), flush=True)
+
+
+if __name__ == "__main__":
+    main()

docs/docs/reference/network.md

@@ -1,5 +1,46 @@
 # Network
 
+## Optional sparse training
+
+`Network(backend="sparse_trainable")` selects CSR aggregation for the default
+`PPNeuronIGRSynapses` dynamics. The default remains `backend="scatter"`.
+For Hydra-based training, add `+network.backend=sparse_trainable` to the overrides.
+
+The default scatter path expands activity to `[batch, edges]` before activation
+and weighted summation. The sparse backend instead activates `[batch, nodes]`
+and computes the same summed current with sparse matrix multiplication. Its
+custom backward samples the weight gradient only at existing connections; it
+does not construct a dense `[nodes, nodes]` gradient. Forward and transposed CSR
+layouts are cached, so backward does not sort the graph at each time step.
+
+Shared synaptic strengths, signs, counts, biases, time constants, Euler updates,
+and node state hooks retain their existing meaning. Unsorted and parallel edges
+are supported: parallel weights are summed once per forward and gradients flow
+to every original edge. Per-edge current inspection still uses the original
+connectome ordering and will allocate edge-sized tensors when requested.
+
+Backend activation storage scales with `batch * nodes * steps` instead of
+`batch * edges * steps`. Graph layouts, parameters, output trajectories and
+temporary workspace still consume memory. This is not a guarantee of faster
+execution: benchmark your graph and batch size before selecting a backend.
+
+Requirements and limitations:
+
+- CPU or CUDA, float32 or float64, with PyTorch CSR `sampled_addmm` support.
+- First-order gradients only; higher-order derivatives are unsupported.
+- The built-in `PPNeuronIGRSynapses` class and its out-of-place pointwise
+  activations. Other dynamics continue to use scatter.
+- Fixed topology. Construct a new network for a different connectome.
+- No AMP, `torch.compile`, MPS or arbitrary custom-activation support is claimed.
+
+CSR layouts are nonpersistent buffers. A sparse network and a scatter network
+constructed from the same configuration have compatible parameter state dicts.
+The selected backend is also retained in `network.config` for reconstruction.
+
+Run `python benchmarks/sparse_current.py --out results.jsonl` for a bounded
+forward/backward benchmark on the supplied visual-system connectome. It does
+not require downloading the training dataset or running an optimization job.
+
 ::: flyvis.network.network.Network
 
 ## Stimulus

flyvis/network/dynamics.py

@@ -190,6 +190,7 @@ def write_state_velocity(
         target_sum: Callable,
         x_t: torch.Tensor,
         dt: float,
+        current_sum: Callable = None,
         **kwargs,
     ) -> None:
         """
@@ -202,19 +203,20 @@ def write_state_velocity(
             target_sum: Function to sum edge values for each target node.
             x_t: External input at time t.
             dt: Time step.
+            current_sum: Optional node-level weighted aggregation. The sparse
+                backend activates each neuron before propagation instead of
+                materializing an activated copy for every outgoing edge.
             **kwargs: Additional keyword arguments.
         """
+        internal_current = (
+            target_sum(params.edges.weight * self.activation(state.sources.activity))
+            if current_sum is None
+            else current_sum(self.activation(state.nodes.activity))
+        )
         vel.nodes.activity = (
             1
             / torch.max(params.nodes.time_const, torch.tensor(dt).float())
-            * (
-                -state.nodes.activity
-                + params.nodes.bias
-                + target_sum(
-                    params.edges.weight * self.activation(state.sources.activity)
-                )  # internal chemical current
-                + x_t
-            )
+            * (-state.nodes.activity + params.nodes.bias + internal_current + x_t)
         )
 
     def currents(