Releases: EleutherAI/bergson
Release list
v0.13.3
v0.13.3 (2026-07-27)
This release is published under the MIT License License.
Performance Improvements
- fix: guard torch.compile so bergson imports on Python 3.14
torch.compile raises RuntimeError at decoration time on interpreters Dynamo doesn't support (Python 3.14+ with torch <= 2.9). The normalizers in bergson/gradients.py apply it at import time, so importing bergson crashed outright on 3.14.
- Add compile_if_supported: falls back to the eager function with a warning when torch.compile is unavailable - Add tests covering the fallback path, including a subprocess test that simulates the 3.14 condition on any interpreter - Add an import-py314 CI job that installs bergson on Python 3.14 and imports it, catching future import-time regressions
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
-
ci: run compile-fallback tests on the Python 3.14 job
-
perf: remove torch.compile from normalizers
A100 benchmarks (torch 2.7.1, py3.11) show the compiled normalizers are net slower in the common case: ~50us of dynamo guard overhead per call dominates the sub-100us kernels at 768-dim shapes and on every bias path, and each run pays 0.3-1.5s compile latency per graph (17 graphs observed from bin-packing's varying batch sizes). Compiled kernels only win on very large tensors (up to ~2x on [4, 50304, 768] bf16), a regime worth a few percent of end-to-end build time at most, and only when --optimizer_state is set.
Removing the decorators also makes bergson importable on Python 3.14, where torch.compile raises at decoration time (torch <= 2.9). The import-py314 CI job guards against reintroducing import-time compiles.
Co-authored-by: Claude Fable 5 noreply@anthropic.com
Detailed Changes: v0.13.2...v0.13.3
v0.13.2
v0.13.2 (2026-07-24)
This release is published under the MIT License License.
Bug Fixes
hessian_dtype only reached the covariance collectors — the eigenvalue-correction pass ignored it and accumulated in the activations' dtype. Pass it to LambdaCollector in both the standard and per-checkpoint (SOURCE) paths, and add an fp64 option for kronfluence-parity precision studies.
Detailed Changes: v0.13.1...v0.13.2
v0.13.1
v0.13.1 (2026-07-22)
This release is published under the MIT License License.
Bug Fixes
- Scale random projections by 1/sqrt(p) (Johnson-Lindenstrauss)
create_projection_matrix normalized each row over its n entries, giving entry variance 1/n. JL requires 1/p (p = projected dim) so that E[A^T A] = I and projected inner products are unbiased. The consequence is not the overall scale — a constant cancels from any ranking — but that the factor p/n depends on the module's SHAPE, so modules are silently reweighted against each other in the summed score. Attention (d x d) and MLP (d x 4d) blocks differ by 4x.
Measured E[proj]/true, 3000 draws per shape:
shape p before predicted p^2/(o*i) after 32x32 16 0.24972 0.25000 0.99889 64x64 32 0.25042 0.25000 1.00167 128x128 32 0.06247 0.06250 0.99949 256x64 32 0.06253 0.06250 1.00045
This matches the TrackStar paper (arXiv 2410.17413 §A.1.2), which specifies the same two-sided per-block projection and states "Projection matrix entries are sampled i.i.d. from N(0, 1/d)" — variance 1/(number of rows), exactly A /= sqrt(m). It also matches docs/preprocessing.rst's claim that projections preserve inner products, which the old scaling did not.
Existing tests could not catch this: they compare the collector against L @ G @ R^T using the SAME matrices, so they are invariant to how those are scaled. test_global_projection_linearity uses one R sliced into blocks — the paper's construction, not the code's. tests/test_projection_inner_products.py now pins projected against unprojected (12 tests, all 12 fail without this fix), including that the scale does not depend on module shape.
BREAKING: changes all stored index scales. Existing indices built with projection_dim > 0 must be regenerated, and any published numbers re-run. runs/test_build_cache.npy is regenerated here; the measured ratio is exactly 0.5000, matching theory (sqrt(o*i)/p = 8/16 for tiny-Phi3).
THREE TESTS FAIL ON THIS BRANCH AND ARE LEFT FAILING, PENDING A DECISION:
tests/test_batch_size_invariance.py::test_gradient_scale_invariance[100-100] tests/test_batch_size_invariance.py::test_gradient_scale_invariance[50-150] Relative error is UNCHANGED by this fix (6.83e-06 on main vs 7.58e-06 here) and already exceeds the assert_close default rtol=1.3e-06 on main. It passes on main only because the values are ~150x smaller, so the absolute difference stays under atol=1e-05. Restoring the correct scale trips atol. The test is magnitude-sensitive rather than wrong about the invariant.
tests/test_gradients.py::test_gradient_collector_proj_norm 1 of 256 elements, greatest relative difference 2.46e-04 against the test's explicit 1e-4 tolerance. Probably a near-cancelling element landing differently at the new scale — the test already carries a comment about accumulating numerical error — but this is NOT proven.
These are regression guards; loosening them is a call for a human, not something to do to make a branch green.
Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com
- Make two projection tests independent of gradient magnitude
test_gradient_scale_invariance compared std separately-vs-combined with assert_close's defaults. The invariant is scale-free, but the default atol=1e-5 let a relative disagreement pass whenever the operands were small enough: on main the relative error is already 6.83e-06 against rtol=1.3e-06, and the test passes only because the values are ~150x smaller than the correctly-scaled ones. Use rtol=1e-4 with atol=0 so the check tracks the invariant rather than the magnitude.
test_gradient_collector_proj_norm compares A @ normalize(g) @ B.T against the collector, which projects activations and output grads separately. The two orderings are mathematically identical and differ only in fp32 rounding; at the corrected scale that noise reached 2.46e-04 relative on an O(1) element, over the existing 1e-4. Widen to 1e-3 and say what the tolerance is measuring.
- Add projection_scale to read indexes built with the old scaling
Indexes store their projected gradients, so changing the projection scaling makes existing ones unreadable at the correct weighting. projection_scale selects the convention: "jl" (default, variance 1/projection_dim) or "row_norm".
GradientProcessor.load resolves a missing projection_scale key to "row_norm", so existing indexes keep working with no user action, and new ones record "jl" in processor_config.yaml. Threaded through IndexConfig, create_processor and EkfacConfig, which builds matching matrices to compress the IVHP output.
- [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Co-authored-by: Claude Opus 4.8 noreply@anthropic.com
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Detailed Changes: v0.13.0...v0.13.1
v0.13.0
v0.13.0 (2026-07-19)
This release is published under the MIT License License.
Features
- Compress K-FAC IVHP output to match compressed gradient stores (
157fa1b)
EKFAC's H^-1 must be applied to the full, unprojected gradient (the eigenbasis rotation only makes sense at the true parameter dimension), so apply_hessian.py's IVHP path and the query build step have always forced projection_dim=0. But the training-side gradient store built by bergson build/bergson score uses index_cfg.projection_dim's per-module Kronecker random projection by default, so the hessian_pipeline's scoring step had to force projection_dim=0 there too -- silently ignoring whatever compression the user asked for and scoring against full dense gradients.
EkfacConfig gains projection_dim/projection_type; when set, EkfacApplicator.compute_ivhp_sharded compresses each module's H^-1 G output to [p, p] as a post-processing step, using the same create_projection_matrix identifiers (f"{name}/left" / f"{name}/right") that bergson build already uses on training gradients. The hessian_pipeline now threads index_cfg.projection_dim/projection_type into the apply step and no longer overrides the score step's projection_dim to 0, so a compressed query and a compressed training-side store are directly comparable end to end.
Supersedes #275, which predates the FactoredPreconditioner refactor and no longer applies; ports the same idea (compress the IVHP output rather than the pre-Hessian gradient) onto the current apply path.
Co-Authored-By: Girish Gupta girish@girishgupta.com
Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Detailed Changes: v0.12.1...v0.13.0
v0.12.1
v0.12.1 (2026-07-18)
This release is published under the MIT License License.
Bug Fixes
- Correctly differentiate cross-rank gradient sync in MAGIC's backward-through-training (
bfc2938)
Trainer.step() synchronizes per-rank parameter gradients with a differentiable all-reduce so the resulting graph can be differentiated again during backward-through-training. Rebased onto main, which had independently landed a different fix for the same underlying bug (51a822d, "Update DDP MAGIC all reduce") using torch.distributed.nn.functional.all_reduce plus a single /world_size correction applied once at the end of Trainer.backward().
This commit replaces that mechanism with a custom _ReplicatedAllReduceSum autograd Function, differentiable to arbitrary order via a recursive backward. Unlike torch.distributed.nn's all_reduce, this one already applies its /world_size correction inline (dividing before the reduce, once per occurrence), so the leftover unconditional bwd_state.weight_grads /= dist.get_world_size() at the end of backward() — carried over from main's fix, which does need it — was double-correcting and silently deflating scores by another 1/world_size. Confirmed via tests/test_ddp.py::test_ddp_matches_single_process, which should be exact with no cross-rank correction needed for a single step: it failed at exactly a 1/world_size ratio with the line still in place, and passes with it removed. Both CPU (test_ddp.py) and real multi-GPU (test_distributed_magic.py, FSDP-vs-DDP with and without grad clipping) suites pass after the fix.
Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Detailed Changes: v0.12.0...v0.12.1
v0.12.0
v0.12.0 (2026-07-18)
This release is published under the MIT License License.
Features
- Add mix_hessians toggle to TrackStar pipeline (
08e9e23)
Lets TrackStar skip computing and mixing the query autocorrelation hessian (steps 2-3) and precondition with the value/train hessian alone. Useful when query and train data are IID and the query set is too small to estimate its own hessian reliably.
Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Refactoring
- Use invert_psd_matrix in semantic scoring examples (
d005091)
Updates example scripts to call bergson.hessians.inversion.invert_psd_matrix instead of the old bergson.utils.math.damped_psd_power, matching the already-landed inversion API. Where a custom regularizer was passed in, fold it into H + damping_factor*regularizer before inverting rather than passing it as a separate argument.
Co-Authored-By: Claude Sonnet 5 noreply@anthropic.com
Detailed Changes: v0.11.1...v0.12.0
v0.11.1
v0.11.1 (2026-07-18)
This release is published under the MIT License License.
Bug Fixes
- Optimizer-state loading — orientation, HF param groups, PEFT, FSDP (
342a216)
Four fixes to the optimizer.pt loading used for attribution normalizers, plus matching export hardening:
- Orient second moments by module class (LayerAdapter.weight_transposed: HF Conv1D stores [in, out]) instead of shape matching, which cannot detect transposed storage for SQUARE weights — GPT-2's attn.c_proj normalizers were silently transposed. - Map state indices group-aware: HF Trainer writes two decay/no-decay param groups whose index order differs from named_parameters(); the old positional mapping assigned moments to the wrong modules. Reconstructed with transformers' own utilities and shape-verified. - Apply the group-aware mapping to PEFT checkpoints too (HF Trainer splits LoRA params across the same groups). - Make save_second_moments_as_optimizer_pt collective-safe (FSDP DTensor moments are gathered by every rank; rank 0 writes — previously a rank-0 -only call would hang) and record each entry's canonical param_name, since FSDP wrapping renames and reorders parameters; readers prefer the recorded name. Optional standard step/betas/eps fields make snapshot exports
self-describing. Model type widened to nn.Module.
Tests: square-Conv1D orientation, decay-split identity vs a real two-group AdamW, PEFT + two-group checkpoint, and the real-checkpoint round-trips now look entries up group-aware with shape assertions the old convention fails.
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
Detailed Changes: v0.11.0...v0.11.1
v0.11.0
v0.11.0 (2026-07-17)
This release is published under the MIT License License.
Features
- Custom gradient store class (
3ae519c)
Trim comments referencing the removed structured-dtype layout.
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
Detailed Changes: v0.10.2...v0.11.0
v0.10.2
v0.10.2 (2026-07-11)
This release is published under the MIT License License.
Bug Fixes
- Update bergson (
ab4b933)
Documentation
- Add Known limitations section (MoE fused experts, FSDP host-RAM load) (
290140a)
Document two issues we are not fixing now, each with the responsible source file:line so maintainers can locate them:
- MoE fused-parameter experts and router are bare nn.Parameters on custom
modules, outside LayerAdapter.supported_modules (bergson/gradients.py:238;
module walk in bergson/collector/collector.py:154-155), so they are silently
skipped. Only attention projections and lm_head (~1-2% of params) are tracked. - Under FSDP the load path uses device_map="cpu" per rank
(bergson/utils/worker_utils.py:166), so every rank replicates a full
dequantized model in host RAM before sharding, which can OOM host memory
independent of GPU VRAM.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012NMtUnc12k72cAXd4uRMzN
- Compact MoE limitation, drop FSDP host-RAM limitation (
1f573ba)
Trim the MoE fused-experts note to what users need (affected model families, consequence, legacy-layout caveat) and remove the FSDP host-RAM limitation, which is a straightforward fix pending coordination.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_0149bxnr7vz4KcR8HpT4ipgn
Detailed Changes: v0.10.1...v0.10.2
v0.10.1
v0.10.1 (2026-07-09)
This release is published under the MIT License License.
Bug Fixes
- Load mmap'd FAISS ANN indices on CPU without crashing (
12e54ac)
index_to_device(index, "cpu") unconditionally called faiss.index_gpu_to_cpu, which clones the index and raises RuntimeError: clone not supported ... OnDiskInvertedLists for any IVF/ANN index mmap'd from disk (the shipped mmap_index=False default). Only exact Flat survived. The CPU branch is now a no-op; the GPU->CPU conversion after a GPU build is done explicitly in create_index only when device != "cpu".
Also expose the FAISS index config on the query CLI (nested QueryConfig.faiss_cfg) so ANN / mmap_index=True is actually reachable, and add a regression test for the on-disk ANN CPU load path.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_012NMtUnc12k72cAXd4uRMzN
Documentation
- Drop redundant inline comments at index_to_device call sites (
2ebc293)
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
- Trim index_to_device docstring (
5ecc0b0)
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
Refactoring
- Keep symmetric index_to_device, fix the bug at the init call site (
50345cf)
Revert the index_to_gpu rename. index_to_device stays a symmetric CPU<->GPU move; the actual bug was only that FaissIndex.__init__ routed an already-CPU mmap'd shard through it with "cpu", hitting the index_gpu_to_cpu clone that crashes on OnDiskInvertedLists. The loader now guards device != "cpu" so an already-CPU shard is never moved, and create_index uses the same helper for its genuine GPU->CPU conversion (also guarded). No behavioural change on any tested path; the crash stays fixed.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
- Make index_to_device self-guarding (no-op when already on target) (
fc43de6)
Rather than requiring call sites to know "don't ask for cpu if it's already cpu", index_to_device now detects GPU residency (_is_gpu_resident: a direct GpuIndex, or the IndexShards/IndexReplicas container a multi-GPU move returns) and treats a CPU->CPU request as a no-op, returning the index unchanged instead of cloning via index_gpu_to_cpu (which crashes on mmap'd OnDiskInvertedLists). Both call sites (FaissIndex.__init__, create_index) drop their device != "cpu" guards and just call the helper.
Adds unit tests for the guard: it no-ops on an in-memory CPU index and on the mmap'd OnDisk index that used to crash (asserting the raw index_gpu_to_cpu still raises, so the regression case is real).
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
- Rename index_to_device -> index_to_gpu, drop dead CPU branch (
080bbf0)
The CPU destination in index_to_device was only ever a no-op after the prior fix: FaissIndex.__init__ reads shards from disk (already CPU) and only ever needs to push up to a GPU, and the one genuine GPU->CPU move lives in create_index. Rename to index_to_gpu (no-op on "cpu") and guard the loader call with device != "cpu" so the direction is explicit and the OnDiskInvertedLists clone trap cannot be reintroduced. Also trims a verbose config docstring.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
Testing
- Cover the index_to_device op path, not just the no-op (
64484d8)
Add a test that a GPU-resident index is actually converted (not returned as-is): an IndexShards container is what a multi-GPU move returns and what _is_gpu_resident flags for conversion, so bringing it to CPU must yield a new, non-container, still-searchable index. Complements the two no-op tests.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com
Claude-Session: https://claude.ai/code/session_01EeaHeis3YkmXwc5emPJnXo
Detailed Changes: v0.10.0...v0.10.1