add gradient accumulation support

.gitignore

@@ -215,4 +215,8 @@ __marimo__/
 .vscode/settings.json
 
 # Ignore logs
-logs/
+logs/
+*/logs/
+*ckpt
+.nfs**
+examples/cifar-10-batches-py/*

CLAUDE.md

@@ -0,0 +1,74 @@
+# perspic
+
+A tool to study neural network training dynamics. Package `perspic`, v0.0.1, authors Konstantin Nikolaou and Jonas Scheunemann, requires Python >=3.11. Built on PyTorch and PyTorch Lightning.
+
+## Overview
+
+perspic wraps a user's `pytorch_lightning.LightningModule` with an `analyzer()` factory, transparently instrumenting training to compute per-sample gradient-norm and linear-response metrics — quantities used to study training dynamics (gradient coupling, sensitivity, collective variables) — without requiring the user to modify their model code. The wrapped module trains normally under a standard `pytorch_lightning.Trainer`; the analysis runs and logs metrics (`chi_net`, `chi_loss`, `chi_coup`, `grad_norm_squared`, `loss`, `batch_size`, ...) alongside it.
+
+## Architecture
+
+`perspic/analyzer.py` — `analyzer(lightning_module, sample_wise_engine="opacus", disable_analyzer=False, log_metrics=True, opacus_strict=False, opacus_approximate_with_n=None, analyze_every=None, analysis_schedule=None, cross_response=False, micro_batch_size=None, effective_batch_size=None, measure_dataloader=None, measure_batch_size=None, measure_subset_seed=None, **model_kwargs)`. A factory function that dynamically subclasses the given `LightningModule` into an `Analyzer` class. `training_step` is overridden to: run pre-step analysis (`_before_training_step`) → delegate to the wrapped module's own `training_step` → manual backward and optimizer step → step any `interval='step'` LR schedulers → log results. `on_train_epoch_end` steps `interval='epoch'` LR schedulers. The wrapped module must define a `criterion` attribute. Supports `cross_response=True`, which expects a dict batch `{"train": ..., "measure": ...}` (via `CombinedLoader`) to additionally measure the model's linear response against a held-out batch, and sparse analysis scheduling via `analyze_every` or `analysis_schedule`.
+
+`micro_batch_size`/`effective_batch_size` enable train-side gradient accumulation: the DataLoader yields `micro_batch_size` micro-batches, and the optimizer steps only every `effective_batch_size // micro_batch_size` of them, with per-sample gradient-norm metrics accumulated across the cycle and combined (via the correct extensive/intensive scaling — see `examples/batch_accumulation_notes.md`) before logging.
+
+`measure_dataloader`/`measure_batch_size`/`measure_subset_seed` give the cross-response measurement side an **independent** batch size, decoupled from the train/`CombinedLoader` path above: the analyzer holds a persistent iterator over `measure_dataloader` and, on each analyzed step, gathers a pool of `max(measure_batch_size)` samples. `measure_batch_size` accepts an int or a `list[int]` to sweep multiple sizes per step (largest processed first; smaller sizes are seed-fixable random subsets of the same pool via `measure_subset_seed`). `measure_dataloader.batch_size` is the maximum single-pass batch size; each swept size `S <= measure_dataloader.batch_size` is measured with a single direct pass (no accumulation), while each `S >` it must be an exact multiple and is measured via gradient accumulation (`S // measure_dataloader.batch_size` passes combined into one measurement) — logging `cross_*` metrics with a `@bs{S}` suffix per swept size. See `examples/measurement_batch_sweep.md` for the full design and usage.
+
+`perspic/calculator/` — the analysis engines used by `analyzer()`:
+- `coupling.py` — `CouplingCalculator`: computes the coupling value `chi_coup = ||grad_L||^2 / (chi_loss * chi_net)`.
+- `linearizer.py` — `Linearizer`: computes the exact first-order linear response of the loss via gradient dot products. `compute(model, criterion, x1, y1, x2=None, y2=None)` returns `(loss, perturbed_loss, delta_loss)`, with an optional cross term against a second batch.
+- `samplewise.py` — `SamplewiseCalculator`: abstract base class defining the interface for per-sample gradient-norm calculators (`compute`, network- and loss-level norm helpers), plus shared helpers.
+- `samplewise_functorch.py` — `SamplewiseCalculatorFunctorch`: per-sample gradient norms via `torch.func` (`vmap` + `jacrev`).
+- `samplewise_opacus.py` — `SamplewiseCalculatorOpacus`: per-sample gradient norms via Opacus ghost clipping (`GradSampleModuleFastGradientClipping`), with custom BatchNorm grad/norm samplers for eval-mode frozen stats and an optional Hutchinson trace-estimator approximation (`approximate_with_n`) for the network-level norm.
+
+`perspic/logger.py` — `LogarithmicWindowSchedule` (dataclass) and `logarithmic_windows(max_steps, points_per_decade=10, base_window=5, adaptive_scale=0.0)`, for building logarithmically-spaced step schedules so analysis can run sparsely over long training runs (pass as `analyzer(..., analysis_schedule=...)`).
+
+`perspic/utils.py` — `BatchStatSnapshot`, a context manager that freezes BatchNorm running stats (with Bessel's correction) so per-sample `vmap` gradients match a train-mode forward pass; `MultiEpochsDataLoader` and `RepeatSampler`, a `DataLoader` subclass that reuses its workers/iterator across epochs.
+
+Class design relies on: ABCs for the calculator hierarchy, a dataclass for the logging schedule, and a factory-returns-dynamic-subclass pattern for `analyzer()` (it builds `class Analyzer(lightning_module): ...` at call time rather than requiring users to subclass anything themselves).
+
+## Public API
+
+`from perspic import ...`:
+- `analyzer` — the primary entry point; wraps and instantiates an `Analyzer` from a `LightningModule`.
+- `Linearizer` — standalone linear-response calculator.
+- `SamplewiseCalculatorFunctorch`, `SamplewiseCalculatorOpacus` — standalone per-sample gradient-norm calculators, also selectable inside `analyzer()` via `sample_wise_engine`.
+- `LogarithmicWindowSchedule`, `logarithmic_windows` — sparse analysis scheduling.
+- `MultiEpochsDataLoader` — performance-oriented DataLoader.
+
+Also available via `perspic.calculator`: `CouplingCalculator`, `SamplewiseCalculator` (base class).
+
+## Design notes
+
+`analyzer()` sets `automatic_optimization = False` because its analysis metrics require extra forward/backward passes separate from the training pass: the `Linearizer` does its own `zero_grad -> forward -> backward` to get `||grad_L||^2`, and `SamplewiseCalculatorOpacus` runs `output_dim` extra forward+backward passes through the Opacus ghost-clipping hooks for `chi_net` (plus one extra forward for `chi_loss`). Manual optimization keeps these extra passes from corrupting Lightning's own gradient/optimizer state. If the wrapped module already uses manual optimization itself, `analyzer()` delegates to it instead of double-handling the optimizer step.
+
+## Development
+
+Install: `pip install -r requirements.txt -r dev-requirements.txt`
+
+Test: `pytest`
+- `tests/unit/` mirrors each `perspic` module 1:1 (`test_analyzer.py`, `test_linearizer.py`, `test_logger.py`, `test_samplewise.py`, `test_samplewise_functorch.py`, `test_samplewise_opacus.py`, `test_utils.py`), using `unittest.mock` and small synthetic Lightning modules.
+- `tests/integration/` covers end-to-end/deployment scenarios: `test_analyzer_deployment.py`, `test_hutchinson_approximation.py`, `test_linearizer_deployment.py`, `test_lna_ntk_ground_truth.py`, `test_normalization_scaling.py`, `test_training_deployment.py`.
+
+Lint/format: `pre-commit run --all-files` (runs `black`, `isort`, `flake8` in that order, `fail_fast: true`), or individually `black .`, `isort .`, `flake8`.
+
+## Conventions
+
+- Line length 88 (`black`, preview mode; `flake8` matches).
+- `isort` uses the `black` profile.
+- `flake8` ignores `W503, E203, E402, C901`; `examples/` is excluded from flake8 checks.
+- Docstrings are predominantly Google-style (`Args:`, `Returns:`, `Raises:`).
+- Type hints are used throughout function/method signatures.
+
+## Examples
+
+`examples/`:
+- `cifar10.ipynb` — basic `analyzer()` usage: wrap a `LightningModule`, train on CIFAR-10, plot logged metrics.
+- `cifar10_CrossReseponse.ipynb` — `cross_response=True` demo, measuring linear response against a held-out batch via a `CombinedLoader`.
+- `batch_accumulation.ipynb` / `batch_accumulation_notes.md` — `micro_batch_size`/`effective_batch_size` gradient-accumulation demo and the design notes explaining the χ_loss aggregation fix and `effective_step`-based logging.
+- `batch_size_scaling_analysis.ipynb` — synthetic sweep over batch size (repeated identical samples) verifying `chi_net`/`chi_loss` batch-size invariance directly against `SamplewiseCalculatorFunctorch`, independent of `analyzer()`.
+- `measurement_batch_sweep.md` — design notes for `measure_dataloader`/`measure_batch_size`/`measure_subset_seed`, the independent (and sweepable) cross-response measurement batch size.
+- `logging_scheduler.ipynb` — `logarithmic_windows` / `LogarithmicWindowSchedule` demo combined with LR scheduling.
+- `mup_integration.ipynb` — integrating Maximal Update Parametrization (mup) with perspic.
+- `core/hutchinson_convergence.py` — convergence of the Opacus Hutchinson trace-estimator approximation toward the exact per-sample gradient norm.
+- `models/` — shared model zoo used by the notebooks above (`cnns.py`, `mlps.py`, `lightning_modules.py`, `utils.py`).

PHASE1_batch_accumulation_review.md

@@ -0,0 +1,44 @@
+# PHASE 1 — Batch-Accumulation Feature Audit
+
+## Data-flow trace
+
+Per micro-batch, [analyzer.py:262-365](perspic/analyzer.py#L262-L365) executes: `opt.zero_grad()` only when `_accumulation_count == 0` → analysis before-hook → wrapped `training_step` → `manual_backward(output / K)` → increment count → when `count >= K`: `opt.step()`, increment `_optimizer_step_count`, reset count, step per-step schedulers.
+
+The analysis path ([_analyze_accumulated_step](perspic/analyzer.py#L503-L584)) decides `_analysis_active` on micro-batch 0 using `effective_step`, saves/restores `p.grad` clones around the grad-clobbering analysis passes, accumulates chi metrics and unscaled linearizer gradients per micro-batch, and finalizes on micro-batch K−1 — **before** the optimizer step, consistent with the single-step path, which also measures pre-step state. Exactly K entries accumulate per cycle; **no off-by-one in the accumulation or flush ordering itself**.
+
+Math verified end-to-end:
+- Gradient: `Σₖ ∇(Lₖ/K) = ∇(mean loss over N=K·B)` ✓ (assuming mean-reduced criterion, equal micro-batch sizes)
+- `grad_norm_squared = ‖Σgₖ‖²/K² = ‖∇L_full‖²` ✓ — confirmed numerically by [test G](tests/unit/test_analyzer.py#L1248-L1307)
+- `chi_net_eff = chi_loss_eff = mean over K` ✓ per the derivation in [batch_accumulation_notes.md](examples/batch_accumulation_notes.md). Note: **the old `K·sum` chi_loss bug (a K² error, 16× at K=4) is fixed only in the uncommitted working tree** — the committed branch tip still contains it. That fix needs to be committed.
+
+## Verdict: ⚠️ Minor issues — core math correct; one label misalignment and several edge cases
+
+## Correctness issues
+
+1. **`effective_step` is off by one against everything it's compared with** — [analyzer.py:329-334](perspic/analyzer.py#L329-L334). Cycle N (0-indexed) tags its micro-batches `effective_step = N+1`, but the same cycle's analysis metrics log `analysis_step = N` ([analyzer.py:765](perspic/analyzer.py#L765)), and the full-batch comparison run's `train_loss` lands at Lightning `step = N`. The notebook plots loss by `effective_step` and chi by `analysis_step` on shared axes, so the accumulation loss curve is shifted +1 — visible at early steps on log axes.
+   **Fix**: `self.log("effective_step", float(self._optimizer_step_count), ...)` — drop the `+1`.
+
+2. **`_accum_step_losses` is dead code** — initialized at [analyzer.py:251](perspic/analyzer.py#L251), appended at [323](perspic/analyzer.py#L323), cleared at [345-346](perspic/analyzer.py#L345-L346) and [735](perspic/analyzer.py#L735), never read. Retains detached GPU scalars for nothing. **Fix**: delete all four sites.
+
+3. **No guard against ragged micro-batches** — `output / K` assumes every micro-batch has exactly `micro_batch_size` samples. With `drop_last=False`, a smaller final batch makes the accumulated gradient a weighted (wrong) mean, and [analyzer.py:628](perspic/analyzer.py#L628) uses the *last* micro-batch's `x.shape[0]` for the logged `batch_size`/`effective_batch_size`. **Fix**: warn once in `training_step` when `batch[0].shape[0] != micro_batch_size`; document `drop_last=True`.
+
+4. **Interrupted / partial cycles**: `_accumulation_count` isn't reset at epoch end, so cycles span epoch boundaries (self-consistent — gradients and analysis both span — but undocumented); a trailing partial cycle at end of training silently drops its gradients; `on_train_epoch_end` steps epoch-interval schedulers even mid-cycle. Additionally `_optimizer_step_count`/`_accumulation_count` are **not checkpointed**, so a resumed run resets `effective_step` to 0 and misaligns `analysis_schedule`. **Fix**: `on_save_checkpoint`/`on_load_checkpoint` for both counters; document epoch-spanning behavior.
+
+5. **Question — mean-reduction assumption**: loss/K scaling *and* the mean-aggregation of `chi_loss` are only valid for a mean-reduced criterion (your notes acknowledge this). Nothing checks `criterion.reduction`; a sum-reduced criterion silently yields wrong gradients and a K²-wrong `chi_loss_eff`. Suggest a best-effort warning via `getattr(criterion, "reduction", "mean")`.
+
+6. **Question**: `effective_step` is logged even when `log_metrics=False` — intentional (training metadata, not an analysis metric)? It's the only unconditional `self.log` in the class.
+
+No race conditions: single-process Lightning, no threading/async anywhere in the accumulation flow.
+
+## Efficiency issues
+
+1. **Train-side linearizer passes are redundant** — [_accumulate_linearizer_grads](perspic/analyzer.py#L586-L623) adds one extra forward+backward per micro-batch, but at cycle end `p.grad` already holds `Σ∇(Lₖ)/K`, whose squared norm equals exactly the `‖Σgₖ‖²/K²` that [finalize](perspic/analyzer.py#L644-L647) computes. Reading `p.grad` just before `opt.step()` eliminates **K forward+backward passes per analysis cycle plus one full-model gradient buffer**. (Matches your existing hook-based-refactor design notes; measure side genuinely needs its own passes.) ~15% of analysis cost at 10 output dims; much more for low-output-dim models.
+2. **`BatchStatSnapshot` pays a full dummy forward even with zero BatchNorm layers** — [utils.py:145-146](perspic/utils.py#L145-L146). For BN-free models: one wasted full-batch forward per analysis micro-batch. Early-out (but keep the `model.eval()` switch so dropout behavior is unchanged) — cheapest win in the feature.
+3. **One GPU sync per parameter** — `.item()` inside per-parameter generator sums at [analyzer.py:644-647](perspic/analyzer.py#L644-L647), [675-682](perspic/analyzer.py#L675-L682), and [linearizer.py:78-79](perspic/calculator/linearizer.py#L78-L79), [103-107](perspic/calculator/linearizer.py#L103-L107). Sum on-device, call `.item()` once.
+4. The loop itself is **O(K)** with in-place `acc.add_()` — no quadratic behavior, no unneeded copies beyond the required `saved_grads` clone (necessary because `sample_calc.compute` calls `model.zero_grad()`). Peak: ~3 concurrent full-gradient buffers (`saved_grads`, `_accum_grad_train`, `_accum_grad_measure`) — acceptable, worth a docstring note. `itertools`/`deque`/vectorization don't apply here; the buffers are parameter-shaped tensors handled correctly.
+
+## Best-practice suggestions
+
+- The train/measure branches of `_accumulate_linearizer_grads` are copy-paste duplicates; unify.
+- Several tests drive `_before_training_step` by poking `_accumulation_count` directly — brittle coupling to private state; the `training_step`-driven style of `test_analysis_does_not_corrupt_training_gradients` is more robust.
+- Typo "the the" in the manual-optimization warning ([analyzer.py:185](perspic/analyzer.py#L185)).