From 7a76dc3cf98c8c1d380d0d86275b094c22b081c8 Mon Sep 17 00:00:00 2001 From: TorgemanTarak <91268951+TorgemanTarak@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:54:38 +0200 Subject: [PATCH 1/2] feat: add optional sparse aggregation for graded synapses --- benchmarks/sparse_current.py | 131 +++++++++++++++++++ docs/docs/reference/network.md | 41 ++++++ flyvis/network/dynamics.py | 18 +-- flyvis/network/network.py | 43 +++++- flyvis/network/sparse.py | 154 ++++++++++++++++++++++ tests/test_sparse_current.py | 231 +++++++++++++++++++++++++++++++++ 6 files changed, 606 insertions(+), 12 deletions(-) create mode 100644 benchmarks/sparse_current.py create mode 100644 flyvis/network/sparse.py create mode 100644 tests/test_sparse_current.py diff --git a/benchmarks/sparse_current.py b/benchmarks/sparse_current.py new file mode 100644 index 0000000..98cc2b0 --- /dev/null +++ b/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() diff --git a/docs/docs/reference/network.md b/docs/docs/reference/network.md index 97f08e1..822f3d6 100644 --- a/docs/docs/reference/network.md +++ b/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 diff --git a/flyvis/network/dynamics.py b/flyvis/network/dynamics.py index 7a41fbe..d1c4989 100644 --- a/flyvis/network/dynamics.py +++ b/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( diff --git a/flyvis/network/network.py b/flyvis/network/network.py index db55efa..fb71f0d 100644 --- a/flyvis/network/network.py +++ b/flyvis/network/network.py @@ -8,6 +8,7 @@ import logging import warnings from contextlib import contextmanager +from functools import partial from typing import Any, Callable, Dict, Iterable, Literal, Optional, Union import numpy as np @@ -29,8 +30,9 @@ from flyvis.utils.nn_utils import n_params, simulation from flyvis.utils.tensor_utils import AutoDeref, RefTensor -from .dynamics import NetworkDynamics +from .dynamics import NetworkDynamics, PPNeuronIGRSynapses, activation_fns from .initialization import Parameter +from .sparse import SparseCurrent from .stimulus import init_stimulus logger = logging.getLogger(__name__) @@ -46,6 +48,9 @@ class Network(nn.Module): dynamics: Network dynamics configuration. node_config: Node parameter configuration. edge_config: Edge parameter configuration. + backend: ``"scatter"`` (default) or ``"sparse_trainable"``. The latter + avoids edge-by-batch activations for the default PPNeuronIGRSynapses + dynamics, with first-order float32/float64 gradients on CPU/CUDA. Attributes: connectome (Connectome): Connectome directory. @@ -120,6 +125,7 @@ def __init__( ), ), stimulus_config: Dict[str, Any] = Namespace(type="Stimulus", init_buffer=False), + backend: str = "scatter", ): super().__init__() @@ -133,6 +139,21 @@ def __init__( # Store the connectome, dynamics, and parameters. self.connectome = init_connectome(**connectome) self.dynamics = forward_subclass(NetworkDynamics, dynamics) + if backend not in ("scatter", "sparse_trainable"): + raise ValueError(f"Unknown propagation backend: {backend!r}") + if backend == "sparse_trainable" and ( + type(self.dynamics) is not PPNeuronIGRSynapses + or type(self.dynamics.activation) not in tuple(activation_fns.values()) + or getattr(self.dynamics.activation, "inplace", False) + ): + raise ValueError( + "sparse_trainable requires PPNeuronIGRSynapses with a built-in, " + "out-of-place pointwise activation" + ) + self.backend = backend + # Leave old configurations unchanged unless the new backend is selected. + if backend != "scatter": + self.config.backend = backend # Load constant indices into memory. # Store source/target indices. @@ -141,6 +162,11 @@ def __init__( self.n_nodes = len(self.connectome.nodes.index) self.n_edges = len(self.connectome.edges.source_index) + self._sparse_current = ( + SparseCurrent(self._source_indices, self._target_indices, self.n_nodes) + if backend == "sparse_trainable" + else None + ) # Optional way of parameter sharing is averaging at every call across # precomputed masks. This can be useful for e.g. symmetric electrical @@ -293,6 +319,8 @@ def _param_api(self) -> AutoDeref[str, AutoDeref[str, RefTensor]]: params[route][param_name] = RefTensor(values, indices) # Add derived parameters. self.dynamics.write_derived_params(params) + if self._sparse_current is not None: + params.edges.csr_weight = self._sparse_current.prepare(params.edges.weight) for k, v in params.nodes.items(): if k not in params.sources: params.sources[k] = self._source_gather(v) @@ -397,8 +425,13 @@ def _next_state( """ vel = AutoDeref(nodes=AutoDeref(), edges=AutoDeref()) + kwargs = {} + if self._sparse_current is not None: + kwargs["current_sum"] = partial( + self._sparse_current, values=params.edges.csr_weight + ) self.dynamics.write_state_velocity( - vel, state, params, self.target_sum, x_t, dt=dt + vel, state, params, self.target_sum, x_t, dt=dt, **kwargs ) next_state = AutoDeref( @@ -863,13 +896,15 @@ def handle_stim(stim, fade_in_state): states = self(stimulus(), dt, state=fade_in_state, as_states=True) return ( stim.cpu().numpy().squeeze(), - torch.stack( + torch + .stack( [s.nodes.activity.cpu() for s in states], dim=1, ) .numpy() .squeeze(), - torch.stack( + torch + .stack( [self.dynamics.currents(s, params).cpu() for s in states], dim=1, ) diff --git a/flyvis/network/sparse.py b/flyvis/network/sparse.py new file mode 100644 index 0000000..10ccf0f --- /dev/null +++ b/flyvis/network/sparse.py @@ -0,0 +1,154 @@ +"""Opt-in CSR aggregation for instantaneous, graded-release synapses. + +Only first-order differentiation is supported. The edge gradient is sampled at +existing coordinates, avoiding both dense adjacency gradients and edge-by-batch +messages. Layouts are derived from the connectome and are not checkpoint data. +""" + +import torch +from torch import Tensor, nn +from torch.autograd.function import once_differentiable + + +class _CSRMM(torch.autograd.Function): + @staticmethod + def forward(ctx, values, activity, crow, col, transpose_crow, transpose_col, order): + n = activity.shape[-1] + adjacency = torch.sparse_csr_tensor( + crow, + col, + values, + size=(n, n), + device=values.device, + dtype=values.dtype, + check_invariants=False, + ) + ctx.save_for_backward( + values, activity, crow, col, transpose_crow, transpose_col, order + ) + return torch.sparse.mm(adjacency, activity.T).T + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + values, activity, crow, col, transpose_crow, transpose_col, order = ( + ctx.saved_tensors + ) + n = activity.shape[-1] + adjacency = torch.sparse_csr_tensor( + crow, + col, + values, + size=(n, n), + device=values.device, + dtype=values.dtype, + check_invariants=False, + ) + grad_values = grad_activity = None + if ctx.needs_input_grad[0]: + grad_values = torch.sparse.sampled_addmm( + adjacency, grad_output.T, activity, beta=0 + ).values() + if ctx.needs_input_grad[1]: + # CSC @ dense would sort/convert to CSR at every backward step. + # Cache the transposed topology and only gather the current values. + transposed = torch.sparse_csr_tensor( + transpose_crow, + transpose_col, + values.index_select(0, order), + size=(n, n), + device=values.device, + dtype=values.dtype, + check_invariants=False, + ) + grad_activity = torch.sparse.mm(transposed, grad_output.T).T + return grad_values, grad_activity, None, None, None, None, None + + +class SparseCurrent(nn.Module): + """Fixed connectivity with dynamic weights and batch-major neuron activity. + + Unsorted and parallel edges are supported. Parallel weights are summed once per + network forward, with gradients routed back to every original edge. Preparing + weights once also avoids retaining a reordered edge vector per recurrent step. + """ + + def __init__(self, source: Tensor, target: Tensor, num_nodes: int): + super().__init__() + if num_nodes < 1 or source.ndim != 1 or source.shape != target.shape: + raise ValueError( + "sparse_trainable requires 1D source/target arrays and nodes" + ) + if source.dtype not in (torch.int32, torch.int64) or target.dtype not in ( + torch.int32, + torch.int64, + ): + raise ValueError("sparse_trainable requires integer edge indices") + if source.device != target.device: + raise ValueError("sparse_trainable source/target indices must share a device") + if source.numel() and bool( + ( + (source < 0) + | (source >= num_nodes) + | (target < 0) + | (target >= num_nodes) + ).any() + ): + raise ValueError("sparse_trainable edge index is out of bounds") + self.num_nodes = num_nodes + self.num_edges = source.numel() + coordinates, inverse = torch.unique( + target.to(torch.int64) * num_nodes + source.to(torch.int64), + sorted=True, + return_inverse=True, + ) + rows = coordinates // num_nodes + counts = torch.bincount(rows, minlength=num_nodes) + crow = torch.cat((counts.new_zeros(1), counts.cumsum(0))) + self.register_buffer("crow", crow, persistent=False) + self.register_buffer("col", coordinates % num_nodes, persistent=False) + self.register_buffer("inverse", inverse, persistent=False) + order = torch.argsort((coordinates % num_nodes) * num_nodes + rows) + transpose_counts = torch.bincount(self.col, minlength=num_nodes) + self.register_buffer( + "transpose_crow", + torch.cat((transpose_counts.new_zeros(1), transpose_counts.cumsum(0))), + persistent=False, + ) + self.register_buffer("transpose_col", rows[order], persistent=False) + self.register_buffer("transpose_order", order, persistent=False) + + def prepare(self, edge_weight: Tensor) -> Tensor: + """Map current edge weights to unique CSR coordinates, preserving autograd.""" + if edge_weight.shape != (self.num_edges,): + raise ValueError("sparse_trainable requires one weight per original edge") + if edge_weight.device != self.inverse.device: + raise ValueError("sparse_trainable weights and topology must share a device") + return edge_weight.new_zeros(self.col.numel()).index_add( + 0, self.inverse, edge_weight + ) + + def forward(self, activity: Tensor, values: Tensor) -> Tensor: + """Sum weighted activated inputs, returning ``[batch, nodes]``.""" + if activity.ndim != 2 or activity.shape[1] != self.num_nodes: + raise ValueError("sparse_trainable expects activity shaped [batch, nodes]") + if activity.device.type not in ("cpu", "cuda") or activity.dtype not in ( + torch.float32, + torch.float64, + ): + raise ValueError("sparse_trainable supports float32/float64 on CPU and CUDA") + if activity.device != values.device or activity.dtype != values.dtype: + raise ValueError( + "sparse_trainable activity and weights must share device/dtype" + ) + if values.shape != self.col.shape or values.device != self.col.device: + raise ValueError("sparse_trainable values must match the prepared CSR layout") + return _CSRMM.apply( + values, + activity, + self.crow, + self.col, + self.transpose_crow, + self.transpose_col, + self.transpose_order, + ) diff --git a/tests/test_sparse_current.py b/tests/test_sparse_current.py new file mode 100644 index 0000000..f43ebe4 --- /dev/null +++ b/tests/test_sparse_current.py @@ -0,0 +1,231 @@ +"""Sparse aggregation parity, including the real FlyVis recurrent runtime.""" + +import pytest +import torch + +from flyvis import Network +from flyvis.network.dynamics import activation_fns +from flyvis.network.sparse import SparseCurrent + +DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +@pytest.fixture(autouse=True) +def defaults(): + device, dtype = torch.get_default_device(), torch.get_default_dtype() + torch.set_default_dtype(torch.float64) + yield + torch.set_default_device(device) + torch.set_default_dtype(dtype) + + +def make_network(backend, activation="relu"): + return Network( + connectome=dict( + type="ConnectomeFromAvgFilters", + file="fib25-fib19_v2.2.json", + extent=1, + n_syn_fill=1, + ), + dynamics=dict(type="PPNeuronIGRSynapses", activation=dict(type=activation)), + node_config=dict( + bias=dict( + type="RestingPotential", + groupby=["type"], + initial_dist="Normal", + mode="mean", + requires_grad=True, + mean=0.5, + std=0.05, + ), + time_const=dict( + type="TimeConstant", + groupby=["type"], + initial_dist="Value", + value=0.05, + requires_grad=True, + ), + ), + backend=backend, + ).to(dtype=torch.get_default_dtype()) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize( + "train_values,train_state", [(True, True), (True, False), (False, True)] +) +def test_unsorted_parallel_edges_and_gradients(device, train_values, train_state): + torch.set_default_device(device) + source = torch.tensor([2, 0, 2, 4, 1, 0]) + target = torch.tensor([1, 3, 1, 0, 4, 0]) + op = SparseCurrent(source, target, 6) # Also includes an isolated node. + w0 = torch.tensor([0.0, 1.0, -0.5, 0.3, -0.2, 0.1]) + h0 = torch.randn(3, 12)[:, ::2] # Noncontiguous input. + results = [] + for sparse in (False, True): + w = w0.clone().requires_grad_(train_values) + h = h0.detach().requires_grad_(train_state) + if sparse: + y = op(h, op.prepare(w)) + else: + y = torch.zeros(3, 6).scatter_add( + -1, target.expand(3, -1), w * h.index_select(-1, source) + ) + y.square().sum().backward() + results.append((y.detach(), w.grad, h.grad)) + for reference, actual in zip(*results): + if reference is not None: + torch.testing.assert_close(actual, reference, atol=1e-10, rtol=1e-9) + + +def test_finite_difference_gradcheck(): + torch.set_default_device("cpu") + op = SparseCurrent(torch.tensor([2, 0, 2]), torch.tensor([1, 2, 1]), 4) + w = torch.randn(3, requires_grad=True) + h = torch.randn(2, 4, requires_grad=True) + assert torch.autograd.gradcheck(lambda w, h: op(h, op.prepare(w)), (w, h)) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("activation", list(activation_fns)) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_real_network_outputs_and_all_gradients(device, activation, dtype): + torch.set_default_device(device) + torch.set_default_dtype(dtype) + reference = make_network("scatter", activation) + actual = make_network("sparse_trainable", activation) + actual.load_state_dict(reference.state_dict(), strict=True) + x0 = torch.randn(2, 5, reference.n_nodes) * 0.2 + outputs, inputs, gradients = [], [], [] + for model in (reference, actual): + x = x0.clone().requires_grad_(True) + y = model(x, dt=0.01) + y.square().mean().backward() + outputs.append(y.detach()) + inputs.append(x.grad) + gradients.append({ + n: p.grad for n, p in model.named_parameters() if p.requires_grad + }) + tolerances = ( + dict(atol=1e-6, rtol=1e-4) + if dtype == torch.float32 + else dict(atol=1e-9, rtol=1e-8) + ) + torch.testing.assert_close(outputs[1], outputs[0], **tolerances) + torch.testing.assert_close(inputs[1], inputs[0], **tolerances) + assert gradients[0].keys() == gradients[1].keys() + for name in gradients[0]: + assert gradients[0][name] is not None, name + torch.testing.assert_close(gradients[1][name], gradients[0][name], **tolerances) + + +@pytest.mark.parametrize("device", DEVICES) +def test_hooks_initial_state_currents_and_optimizer_updates(device): + torch.set_default_device(device) + reference, actual = make_network("scatter"), make_network("sparse_trainable") + actual.load_state_dict(reference.state_dict()) + optimizers = [torch.optim.Adam(m.parameters(), lr=1e-4) for m in (reference, actual)] + + def hook(state): + state.nodes.activity = state.nodes.activity * 0.9 + 0.02 + return state + + for model in (reference, actual): + model.register_state_hook(hook) + x = torch.randn(2, 4, reference.n_nodes) * 0.2 + for _ in range(3): + results = [] + for model, optimizer in zip((reference, actual), optimizers): + optimizer.zero_grad(set_to_none=True) + params = model._param_api() + state = model._initial_state(params, 2) + h = torch.full((2, model.n_nodes), 0.3, requires_grad=True) + state.nodes.activity = h + states = model(x, dt=0.01, state=model._state_api(state), as_states=True) + y = torch.stack([s.nodes.activity for s in states], dim=1) + # Per-edge inspection remains in original connectome order. + currents = model.dynamics.currents(states[-1], params) + loss = y.square().mean() + currents.square().mean() + loss.backward() + results.append((y.detach(), h.grad.clone(), currents.detach())) + optimizer.step() + for got, expected in zip(results[1], results[0]): + torch.testing.assert_close(got, expected, atol=1e-9, rtol=1e-7) + rebuilt = Network(**actual.config) + rebuilt.load_state_dict(actual.state_dict(), strict=True) + assert rebuilt.backend == "sparse_trainable" + assert set(actual.state_dict()) == set(reference.state_dict()) + + +def test_no_edge_batch_activations_saved(): + torch.set_default_device("cpu") + model = make_network("sparse_trainable") + saved_shapes = [] + + def pack(tensor): + saved_shapes.append(tuple(tensor.shape)) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack, lambda t: t): + model(torch.randn(3, 5, model.n_nodes), dt=0.01).square().mean().backward() + assert (3, model.n_edges) not in saved_shapes + assert (model.n_edges, 3) not in saved_shapes + assert (model.n_nodes, model.n_nodes) not in saved_shapes + + +@pytest.mark.parametrize("backend", ["scatter", "sparse_trainable"]) +def test_forward_inference_then_training(backend): + torch.set_default_device("cpu") + model = make_network(backend) + with torch.inference_mode(): + model(torch.randn(2, 3, model.n_nodes), dt=0.01) + model(torch.randn(2, 3, model.n_nodes), dt=0.01).sum().backward() + assert model.edges_syn_strength.grad is not None + + +def test_unsupported_dynamics_and_inplace_activation_rejected(): + torch.set_default_device("cpu") + for dynamics in ( + dict(type="NetworkDynamics", activation=dict(type="relu")), + dict(type="PPNeuronIGRSynapses", activation=dict(type="relu", inplace=True)), + ): + with pytest.raises(ValueError, match="requires PPNeuronIGRSynapses"): + Network( + connectome=dict(type="ConnectomeFromAvgFilters", extent=1), + dynamics=dynamics, + backend="sparse_trainable", + ) + + +@pytest.mark.parametrize("device", DEVICES) +def test_empty_edges(device): + torch.set_default_device(device) + indices = torch.empty(0, dtype=torch.int64) + op = SparseCurrent(indices, indices, 4) + values = torch.empty(0, requires_grad=True) + h = torch.randn(2, 4, requires_grad=True) + y = op(h, op.prepare(values)) + y.sum().backward() + torch.testing.assert_close(y, torch.zeros_like(y)) + torch.testing.assert_close(h.grad, torch.zeros_like(h)) + assert values.grad is not None and values.grad.numel() == 0 + + +@pytest.mark.parametrize("source,target", [([4], [0]), ([-1], [0]), ([0], [4])]) +def test_invalid_edge_coordinates_rejected(source, target): + torch.set_default_device("cpu") + with pytest.raises(ValueError, match="out of bounds"): + SparseCurrent(torch.tensor(source), torch.tensor(target), 4) + + +def test_sparse_layout_moves_with_module_and_stays_out_of_checkpoint(): + torch.set_default_device("cpu") + op = SparseCurrent(torch.tensor([2, 0]), torch.tensor([1, 2]), 3) + assert not op.state_dict() + for device in DEVICES: + op.to(device) + h = torch.randn(2, 3, device=device, requires_grad=True) + w = torch.randn(2, device=device, requires_grad=True) + op(h, op.prepare(w)).sum().backward() + assert all(t.device.type == device for t in op.buffers()) + assert h.grad is not None and w.grad is not None From f8c48cf35946b3205dcd3428c929e8de01336605 Mon Sep 17 00:00:00 2001 From: TorgemanTarak <91268951+TorgemanTarak@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:56:54 +0200 Subject: [PATCH 2/2] docs: record RTX 5090 recurrent memory and timing measurements --- benchmarks/rtx5090-sparse-training.jsonl | 8 ++++ benchmarks/rtx5090-sparse-training.md | 54 ++++++++++++++++++++++++ 2 files changed, 62 insertions(+) create mode 100644 benchmarks/rtx5090-sparse-training.jsonl create mode 100644 benchmarks/rtx5090-sparse-training.md diff --git a/benchmarks/rtx5090-sparse-training.jsonl b/benchmarks/rtx5090-sparse-training.jsonl new file mode 100644 index 0000000..ae96ebe --- /dev/null +++ b/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}} diff --git a/benchmarks/rtx5090-sparse-training.md b/benchmarks/rtx5090-sparse-training.md new file mode 100644 index 0000000..45de025 --- /dev/null +++ b/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.