Fix EK-FAC gradient store overflow above ~537M tracked params

[davidoj]

Fix EK-FAC gradient store overflow above ~537M tracked params

Problem

bergson ekfac (the EK-FAC influence pipeline) cannot attribute any model with
more than ~537M tracked-linear parameters. The apply-Hessian stage loads a structured memmap which contains unprojected grads so is beyond the structured memmap size limit and aborts with ValueError: integer won't fit into a C int.

So sum(grad_sizes.values()) equals
the full per-example gradient dimension, and at fp32 the record crosses 2**31
bytes once the model has more than 2**31 / 4 = 536,870,912 tracked params.

The exact failing calls are in the apply stage:

• bergson/hessians/apply_hessian.py:79 — load_gradients(self.gradient_path)
  (structured) maps the unprojected query gradient → overflow.
• bergson/hessians/apply_hessian.py:83 — create_index(...) (structured,
  default) allocates the transformed-gradient output → overflow.

Fix

You can avoid the crash by loading with the flag structured=False.

For oversized stores only, create_index / load_gradients now serve the data
as a flat 2D memmap wrapped in a small field-access view that preserves the
structured interface the rest of the codebase uses. Small models are completely
unaffected — they still get a genuine numpy structured memmap.

• bergson/data.py — added FlatGradientView: exposes view[module_name]
  (returns the module's (num_grads, size) column block as a live, writable view
  into the memmap), view.dtype.names, len(view), and view.flush().
• bergson/data.py — added _structured_dtype_fits(); create_index
  (with_structure=True) and load_gradients (structured=True) fall back to
  FlatGradientView when the record would overflow the C-int cap, and are
  otherwise unchanged. The persisted info.json is unchanged (still records the
  structured dtype, grad_sizes, and base_dtype), so existing stores and
  metadata readers keep working.
• bergson/query/faiss_index.py — gradients_loader now routes through
  load_gradients (so the approximate-query path is overflow-safe too), and the
  loader loop skips the structured_to_unstructured copy for an already-flat
  view.

No public API or signature changes. apply_hessian's field-access reads/writes
(mmap[k], grad_buffer[k][:] = ..., len, .flush) work unchanged against
FlatGradientView.

Test

tests/test_data.py (CPU-only, no CUDA, no model):

• test_large_gradient_store_roundtrips — builds a store whose combined fp32
  record exceeds 2**31 bytes (> 2**31 / 4 float32 elements across two
  modules), writes sentinels at each module's start/end via structured field
  access, and asserts they round-trip through create_index +
  load_gradients(structured=True). It also cross-checks the flat
  (structured=False) view of the same bytes to confirm the byte-identical
  layout. The store is a sparse file (truncate only sets the size), so the
  test touches a handful of pages instead of materializing ~2 GB.
• test_small_gradient_store_stays_structured — asserts sub-cap stores still
  return a genuine np.memmap structured array with unchanged behavior.

Before this fix the oversized case raises ValueError: integer won't fit into a C int in create_index; after, it round-trips. The previous test_data.py
(which asserted the failure as expected behavior) is replaced accordingly.

Verified in the pinned venv (numpy 2.5.1, torch 2.9.1+cu128):

pytest tests/test_data.py -v        # 2 passed
pytest tests/test_score.py -v       # 8 passed (incl. test_large_gradients_query)

Scope and limitations

• This raises the storage ceiling for the EK-FAC gradient store past
  ~537M params; it does not change results for models below the old ceiling.
• EK-FAC influence has a separate, larger ceiling that this fix does not
  address: the query gradient is materialized unprojected in full model
  space on the GPU during apply_hessian (and again when reading it back for
  scoring). That is an architectural memory ceiling (order ~7B params on an 80 GB
  GPU, and it also implies a large host/disk footprint), independent of the
  numpy itemsize limit fixed here. Projecting or streaming the unprojected query
  gradient to lift that ceiling is out of scope for this bug fix.
• FlatGradientView implements only the subset of the structured-array
  interface Bergson actually uses (field access, dtype.names, len, flush,
  and integer/slice row indexing on the underlying flat array). It is not a
  general drop-in for every numpy structured-array operation.

---

🤖 Generated with Claude Code

https://claude.ai/code/session_016rrSNbhfYPhfuxuBrdJeFG

[luciaquirke]

For oversized stores only, create_index / load_gradients now serve the data
as a flat 2D memmap wrapped in a small field-access view that preserves the
structured interface the rest of the codebase uses

Please merge whenever makes sense to you to maximize development speed. Long term it seems like we could maybe just pass structured=False here rather than writing a bunch of new code unless we need the structure for something.

[luciaquirke]

@claude are you still wired up here haha

[claude]

Claude finished @luciaquirke's task in 15s —— View job

---

👋 Yep, still here!

Just confirming I'm wired up and watching this PR. Nothing to act on from your "haha" ping specifically — but happy to help with anything on this PR. A couple of things worth flagging while I'm here:

• Lucia's earlier suggestion (just passing structured=False instead of adding FlatGradientView) is a reasonable simplification to consider long-term, but I won't implement it unless you explicitly ask — it wasn't a direct instruction to me.
• Let me know if you'd like me to: review the current diff, run the test suite to re-verify, or implement any specific change (e.g., the structured=False simplification).

Just say the word (@claude ...) and I'll get to work.

[luciaquirke]

@claude I'm happy with this work, can you fix the typecheck so we can merge

[claude]

Claude encountered an error after 2m 11s —— View job

---

Working on it

• Gather context (PR diff, CLAUDE.md)
• Run pyright to reproduce typecheck failures
• Fix typecheck errors
• Re-run pyright to confirm clean
• Run test suite to confirm no regressions
• Commit and push fix