From 892a9f71c0d4e78f74cef3e81b16bb0d109f7055 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Tue, 5 May 2026 11:05:05 -0700 Subject: [PATCH 01/10] Add first-class K-FAC support to build Bergson supports random-projection compression of gradients in `bergson build` and `bergson score`. This wires the same per-layer projection through the K-FAC / EK-FAC Hessian path so compressed gradient stores can be scored against a compressed Hessian. - New `HessianConfig.projection_dim` / `projection_type` fields. Default `projection_dim=0` keeps existing behavior. - `compute_eigendecomposition` now takes optional `projection_dim` / `projection_type` / `side` and compresses each gathered covariance to `P @ M @ P.T` before `eigh`, using the same per-layer projection identifier convention as `collector.py`. - `LambdaCollector` projects activations (right) and output gradients (left) before applying the now-`[p, p]` eigenvectors so the eigenvalue corrections are in the compressed space. - Hessian worker validates that compression is only enabled with `method='kfac'` and `projection_target='per_module'`. - One-sided application only: query gets projected and `apply_hessian` runs unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/config.py | 11 +++ bergson/hessians/eigenvectors.py | 87 ++++++++++++++++++++-- bergson/hessians/hessian_approximations.py | 26 ++++++- 3 files changed, 118 insertions(+), 6 deletions(-) diff --git a/bergson/config.py b/bergson/config.py index ffca3a2d..600134bc 100644 --- a/bergson/config.py +++ b/bergson/config.py @@ -548,6 +548,17 @@ class HessianConfig(Serializable): """Whether to use dataset labels for Hessian (empirical Fisher) approximation. If false, the model predictions will be used.""" + projection_dim: int = 0 + """If > 0, compress the K-FAC factors (A and S) and the EK-FAC eigenvalue + correction via the same per-layer random projection used by ``bergson build``. + Set this to the same value as the gradient store these factors will be applied + to, with matching ``projection_type``. Only supported with ``method='kfac'`` and + ``projection_target='per_module'``.""" + + projection_type: Literal["normal", "rademacher"] = "rademacher" + """Type of random projection used when ``projection_dim > 0``. Must match the + value used in ``bergson build``.""" + @dataclass class FaissConfig: diff --git a/bergson/hessians/eigenvectors.py b/bergson/hessians/eigenvectors.py index 9c7f2209..d72e6aae 100644 --- a/bergson/hessians/eigenvectors.py +++ b/bergson/hessians/eigenvectors.py @@ -1,6 +1,7 @@ import gc import os from dataclasses import dataclass +from typing import Literal import torch import torch.distributed as dist @@ -10,7 +11,7 @@ from torch import Tensor from tqdm import tqdm -from bergson.collector.collector import HookCollectorBase +from bergson.collector.collector import HookCollectorBase, create_projection_matrix from bergson.hessians.sharded_computation import ShardedMul from bergson.utils.logger import get_logger from bergson.utils.utils import ( @@ -76,10 +77,22 @@ class LambdaCollector(HookCollectorBase): Transforms activations and gradients using precomputed eigenvectors, then computes outer products for diagonal correction terms. + + When ``projection_dim > 0`` the precomputed eigenvectors live in the + compressed [p, p] space, so we project the raw activations and output + gradients with the same per-layer random matrices used by ``bergson build`` + before applying the eigenvectors. """ path: str + projection_dim: int = 0 + """If > 0, project activations (right) and output gradients (left) before + transforming them with the (compressed) eigenvectors.""" + + projection_type: Literal["normal", "rademacher"] = "rademacher" + """Random projection family used to compress activations and gradients.""" + def setup(self) -> None: """Load eigenvectors and initialize storage.""" self.shard_computer = ShardedMul() @@ -108,10 +121,21 @@ def forward_hook(self, module: nn.Module, a: Tensor) -> None: name = assert_type(str, module._name) # a shape: [N, S, I] + if self.projection_dim > 0: + P = create_projection_matrix( + identifier=f"{name}/right", + m=self.projection_dim, + n=a.shape[-1], + dtype=a.dtype, + device=a.device, + projection_type=self.projection_type, + ) + a = a @ P.T # [N, S, p] + # Transform: a @ eigen_a transformed = self.shard_computer._matmul( vector_nsa=a, matrix_cb=self.eigen_a[name] - ) # shape [N, S, I] + ) # shape [N, S, I] (or [N, S, p] when compressed) # Cache for use in backward pass self.transformed_a_cache[name] = transformed @@ -121,10 +145,21 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None: name = assert_type(str, module._name) # g shape: [N, S, O] + if self.projection_dim > 0: + P = create_projection_matrix( + identifier=f"{name}/left", + m=self.projection_dim, + n=g.shape[-1], + dtype=g.dtype, + device=g.device, + projection_type=self.projection_type, + ) + g = g @ P.T # [N, S, p] + # Transform: g @ eigen_g transformed_g = self.shard_computer._matmul( vector_nsa=g, matrix_cb=self.eigen_g[name] - ) # shape [N, S, O] + ) # shape [N, S, O] (or [N, S, p] when compressed) # Compute outer product: sum_n (transformed_a_n^T @ transformed_g_n) # Einstein notation: [N, S, I] x [N, S, O] -> [N, O, I] @@ -215,6 +250,10 @@ def _compute_full_matrix( def compute_eigendecomposition( covariance_path: str, total_processed: int | Tensor, + *, + projection_dim: int = 0, + projection_type: Literal["normal", "rademacher"] = "rademacher", + side: Literal["left", "right"] | None = None, ) -> None: """ Compute eigendecomposition from covariance matrices (Eq. 18 from paper). @@ -224,14 +263,32 @@ def compute_eigendecomposition( via _compute_full_matrix(). Output sharding is inferred from the eigenvector shapes. + When ``projection_dim > 0`` and ``side`` is set, each gathered covariance + matrix M of shape [d, d] is compressed to ``P @ M @ P.T`` of shape + ``[projection_dim, projection_dim]`` before eigendecomposition. ``P`` is the + same per-layer random projection matrix (keyed by ``f"{name}/{side}"``) + that ``bergson build`` uses to project gradients, so the resulting + eigenvectors operate in the same compressed space as projected gradients. + Args: covariance_path: Full path to the covariance sharded directory. total_processed: Number of samples used to compute covariance. + projection_dim: If > 0, compress each covariance matrix using a + per-layer random projection before eigendecomposing. + projection_type: Random projection family to use when compressing. + side: Which side projection to apply ("right" for activation + covariance, "left" for gradient covariance). Required when + ``projection_dim > 0``. """ rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist.is_initialized() else 1 device = get_device(rank) + if projection_dim > 0 and side is None: + raise ValueError( + "side must be specified ('left' or 'right') when projection_dim > 0." + ) + # Handle total_processed as tensor if needed if isinstance(total_processed, int): total_processed = torch.tensor(total_processed, device=device) @@ -246,8 +303,16 @@ def compute_eigendecomposition( # Get dimensions for fair distribution (columns not sharded, shape[-1]=d) key_dimensions = {key: f.get_tensor(key).shape[-1] for key in all_keys} + if projection_dim > 0: + # After compression each matrix is [projection_dim, projection_dim], so + # eigendecomposition cost is identical across keys. We still use the + # same fair_distribute helper below for consistent ordering. + eigen_dimensions = {key: projection_dim for key in all_keys} + else: + eigen_dimensions = key_dimensions + # Distribute keys fairly based on O(d³) eigendecomposition cost - all_assignments = fair_distribute_by_cost(key_dimensions, world_size) + all_assignments = fair_distribute_by_cost(eigen_dimensions, world_size) keys_for_this_rank = all_assignments[rank] covariance_eigenvectors = {} @@ -266,6 +331,18 @@ def compute_eigendecomposition( world_size=world_size, ) + if projection_dim > 0: + assert side is not None + P = create_projection_matrix( + identifier=f"{key}/{side}", + m=projection_dim, + n=matrix.shape[0], + dtype=matrix.dtype, + device=matrix.device, + projection_type=projection_type, + ) + matrix = P @ matrix @ P.T + # original_dtype = matrix.dtype matrix_normalized = matrix.to(torch.float64) / total_processed matrix_normalized = (matrix_normalized + matrix_normalized.T).div(2) @@ -289,7 +366,7 @@ def compute_eigendecomposition( covariance_eigenvectors = _merge_and_shard_eigenvectors( input_dict=covariance_eigenvectors, all_keys=all_keys, - key_dimensions=key_dimensions, + key_dimensions=eigen_dimensions, dtype=original_dtype, # type: ignore rank=rank, world_size=world_size, diff --git a/bergson/hessians/hessian_approximations.py b/bergson/hessians/hessian_approximations.py index 41fda238..18cf3d26 100644 --- a/bergson/hessians/hessian_approximations.py +++ b/bergson/hessians/hessian_approximations.py @@ -145,6 +145,20 @@ def hessian_worker( model, target_modules = setup_model_and_peft(index_cfg) + if hessian_cfg.projection_dim > 0: + if hessian_cfg.method != "kfac": + raise ValueError( + "K-FAC factor compression (projection_dim > 0) is only supported with " + f"method='kfac', got method='{hessian_cfg.method}'." + ) + if index_cfg.projection_target != "per_module": + raise ValueError( + "K-FAC factor compression requires projection_target='per_module' " + "(global projection cannot compose with the per-layer Kronecker " + "factorization). Got projection_target=" + f"'{index_cfg.projection_target}'." + ) + attention_cfgs = { module: index_cfg.attention for module in index_cfg.split_attention_modules } @@ -173,10 +187,16 @@ def hessian_worker( compute_eigendecomposition( os.path.join(index_cfg.partial_run_path, "activation_sharded"), total_processed=total_processed, + projection_dim=hessian_cfg.projection_dim, + projection_type=hessian_cfg.projection_type, + side="right", ) compute_eigendecomposition( os.path.join(index_cfg.partial_run_path, "gradient_sharded"), total_processed=total_processed, + projection_dim=hessian_cfg.projection_dim, + projection_type=hessian_cfg.projection_type, + side="left", ) dist.barrier() if dist.is_initialized() else None @@ -212,7 +232,11 @@ def collect_hessians( } desc = f"Approximating Hessians with {hessian_cfg.method}" if ev_correction: - collector = LambdaCollector(**collector_args) + collector = LambdaCollector( + **collector_args, + projection_dim=hessian_cfg.projection_dim, + projection_type=hessian_cfg.projection_type, + ) desc += " (eigenvalue correction)" else: collector_args["dtype"] = hessian_dtype From 4da719279b568469d9c77422726c20abc22dc071 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Wed, 6 May 2026 14:58:46 -0700 Subject: [PATCH 02/10] Stage 1: revert projection-before-eigh, broaden method support Per Louis's review (#275): the random projection should be the *last* step in the K-FAC pipeline, folded into the damped inverse rather than applied to covariances before eigh. This commit tears out the project-before-eigh path so the next stages can introduce the correct post-inverse compression. - compute_eigendecomposition: drop projection_dim/projection_type/side kwargs and the per-key P @ M @ P.T compression. Returns the same eigenvalue dict (from #273) but always in the full-rank space. - LambdaCollector: drop projection_dim/projection_type fields and the pre-multiplication branches in forward_hook/backward_hook. - hessian_approximations.py: drop method != 'kfac' validation (Louis: any Kronecker-factored method works). Replace it with a check that ev_correction=True is incompatible with projection_dim > 0, since the EK-FAC eigenvalue correction does not decompose into per-side row/col ops (Lucia + Louis confirmed). - HessianConfig.projection_dim doc: drop the EK-FAC mention, broaden to any Kronecker-factored method, note ev_correction incompatibility. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/config.py | 10 ++- bergson/hessians/eigenvectors.py | 89 ++-------------------- bergson/hessians/hessian_approximations.py | 26 +++---- 3 files changed, 21 insertions(+), 104 deletions(-) diff --git a/bergson/config.py b/bergson/config.py index 600134bc..e4b5994b 100644 --- a/bergson/config.py +++ b/bergson/config.py @@ -549,11 +549,13 @@ class HessianConfig(Serializable): If false, the model predictions will be used.""" projection_dim: int = 0 - """If > 0, compress the K-FAC factors (A and S) and the EK-FAC eigenvalue - correction via the same per-layer random projection used by ``bergson build``. + """If > 0, fold the same per-layer random projection used by ``bergson build`` + into the damped inverse of the Kronecker factors (A and S) so each per-side + factor is compressed to a ``[d, projection_dim]`` whitening-projection matrix. Set this to the same value as the gradient store these factors will be applied - to, with matching ``projection_type``. Only supported with ``method='kfac'`` and - ``projection_target='per_module'``.""" + to, with matching ``projection_type``. Works for any Kronecker-factored method + (``kfac``, ``shampoo``, ...). Requires ``projection_target='per_module'`` and + is incompatible with ``ev_correction=True``.""" projection_type: Literal["normal", "rademacher"] = "rademacher" """Type of random projection used when ``projection_dim > 0``. Must match the diff --git a/bergson/hessians/eigenvectors.py b/bergson/hessians/eigenvectors.py index d1c8afe6..5a8530e6 100644 --- a/bergson/hessians/eigenvectors.py +++ b/bergson/hessians/eigenvectors.py @@ -1,7 +1,6 @@ import gc import os from dataclasses import dataclass -from typing import Literal import torch import torch.distributed as dist @@ -11,7 +10,7 @@ from torch import Tensor from tqdm import tqdm -from bergson.collector.collector import HookCollectorBase, create_projection_matrix +from bergson.collector.collector import HookCollectorBase from bergson.hessians.sharded_computation import ShardedMul from bergson.utils.logger import get_logger from bergson.utils.utils import ( @@ -77,22 +76,10 @@ class LambdaCollector(HookCollectorBase): Transforms activations and gradients using precomputed eigenvectors, then computes outer products for diagonal correction terms. - - When ``projection_dim > 0`` the precomputed eigenvectors live in the - compressed [p, p] space, so we project the raw activations and output - gradients with the same per-layer random matrices used by ``bergson build`` - before applying the eigenvectors. """ path: str - projection_dim: int = 0 - """If > 0, project activations (right) and output gradients (left) before - transforming them with the (compressed) eigenvectors.""" - - projection_type: Literal["normal", "rademacher"] = "rademacher" - """Random projection family used to compress activations and gradients.""" - def setup(self) -> None: """Load eigenvectors and initialize storage.""" self.shard_computer = ShardedMul() @@ -121,21 +108,10 @@ def forward_hook(self, module: nn.Module, a: Tensor) -> None: name = assert_type(str, module._name) # a shape: [N, S, I] - if self.projection_dim > 0: - P = create_projection_matrix( - identifier=f"{name}/right", - m=self.projection_dim, - n=a.shape[-1], - dtype=a.dtype, - device=a.device, - projection_type=self.projection_type, - ) - a = a @ P.T # [N, S, p] - # Transform: a @ eigen_a transformed = self.shard_computer._matmul( vector_nsa=a, matrix_cb=self.eigen_a[name] - ) # shape [N, S, I] (or [N, S, p] when compressed) + ) # shape [N, S, I] # Cache for use in backward pass self.transformed_a_cache[name] = transformed @@ -145,21 +121,10 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None: name = assert_type(str, module._name) # g shape: [N, S, O] - if self.projection_dim > 0: - P = create_projection_matrix( - identifier=f"{name}/left", - m=self.projection_dim, - n=g.shape[-1], - dtype=g.dtype, - device=g.device, - projection_type=self.projection_type, - ) - g = g @ P.T # [N, S, p] - # Transform: g @ eigen_g transformed_g = self.shard_computer._matmul( vector_nsa=g, matrix_cb=self.eigen_g[name] - ) # shape [N, S, O] (or [N, S, p] when compressed) + ) # shape [N, S, O] # Compute outer product: sum_n (transformed_a_n^T @ transformed_g_n) # Einstein notation: [N, S, I] x [N, S, O] -> [N, O, I] @@ -250,10 +215,6 @@ def _compute_full_matrix( def compute_eigendecomposition( covariance_path: str, total_processed: int | Tensor, - *, - projection_dim: int = 0, - projection_type: Literal["normal", "rademacher"] = "rademacher", - side: Literal["left", "right"] | None = None, ) -> dict[str, Tensor]: """ Compute eigendecomposition from covariance matrices (Eq. 18 from paper). @@ -263,22 +224,9 @@ def compute_eigendecomposition( via _compute_full_matrix(). Output sharding is inferred from the eigenvector shapes. - When ``projection_dim > 0`` and ``side`` is set, each gathered covariance - matrix M of shape [d, d] is compressed to ``P @ M @ P.T`` of shape - ``[projection_dim, projection_dim]`` before eigendecomposition. ``P`` is the - same per-layer random projection matrix (keyed by ``f"{name}/{side}"``) - that ``bergson build`` uses to project gradients, so the resulting - eigenvectors operate in the same compressed space as projected gradients. - Args: covariance_path: Full path to the covariance sharded directory. total_processed: Number of samples used to compute covariance. - projection_dim: If > 0, compress each covariance matrix using a - per-layer random projection before eigendecomposing. - projection_type: Random projection family to use when compressing. - side: Which side projection to apply ("right" for activation - covariance, "left" for gradient covariance). Required when - ``projection_dim > 0``. Returns: Per-key eigenvalue shards (each `[m/world_size]`) on CPU. The @@ -290,11 +238,6 @@ def compute_eigendecomposition( world_size = dist.get_world_size() if dist.is_initialized() else 1 device = get_device(rank) - if projection_dim > 0 and side is None: - raise ValueError( - "side must be specified ('left' or 'right') when projection_dim > 0." - ) - # Handle total_processed as tensor if needed if isinstance(total_processed, int): total_processed = torch.tensor(total_processed, device=device) @@ -309,16 +252,8 @@ def compute_eigendecomposition( # Get dimensions for fair distribution (columns not sharded, shape[-1]=d) key_dimensions = {key: f.get_tensor(key).shape[-1] for key in all_keys} - if projection_dim > 0: - # After compression each matrix is [projection_dim, projection_dim], so - # eigendecomposition cost is identical across keys. We still use the - # same fair_distribute helper below for consistent ordering. - eigen_dimensions = {key: projection_dim for key in all_keys} - else: - eigen_dimensions = key_dimensions - # Distribute keys fairly based on O(d³) eigendecomposition cost - all_assignments = fair_distribute_by_cost(eigen_dimensions, world_size) + all_assignments = fair_distribute_by_cost(key_dimensions, world_size) keys_for_this_rank = all_assignments[rank] covariance_eigenvectors = {} @@ -338,18 +273,6 @@ def compute_eigendecomposition( world_size=world_size, ) - if projection_dim > 0: - assert side is not None - P = create_projection_matrix( - identifier=f"{key}/{side}", - m=projection_dim, - n=matrix.shape[0], - dtype=matrix.dtype, - device=matrix.device, - projection_type=projection_type, - ) - matrix = P @ matrix @ P.T - # original_dtype = matrix.dtype matrix_normalized = matrix.to(torch.float64) / total_processed matrix_normalized = (matrix_normalized + matrix_normalized.T).div(2) @@ -374,7 +297,7 @@ def compute_eigendecomposition( covariance_eigenvectors = _gather_and_shard_along_dim0( input_dict=covariance_eigenvectors, - full_shape_per_key={k: (m, m) for k, m in eigen_dimensions.items()}, + full_shape_per_key={k: (m, m) for k, m in key_dimensions.items()}, dtype=original_dtype, # type: ignore rank=rank, world_size=world_size, @@ -382,7 +305,7 @@ def compute_eigendecomposition( ) covariance_eigenvalues = _gather_and_shard_along_dim0( input_dict=covariance_eigenvalues, - full_shape_per_key={k: (m,) for k, m in eigen_dimensions.items()}, + full_shape_per_key={k: (m,) for k, m in key_dimensions.items()}, dtype=original_dtype, # type: ignore rank=rank, world_size=world_size, diff --git a/bergson/hessians/hessian_approximations.py b/bergson/hessians/hessian_approximations.py index 4adbf164..0de18cc0 100644 --- a/bergson/hessians/hessian_approximations.py +++ b/bergson/hessians/hessian_approximations.py @@ -150,18 +150,20 @@ def hessian_worker( model, target_modules = setup_model_and_peft(index_cfg) if hessian_cfg.projection_dim > 0: - if hessian_cfg.method != "kfac": - raise ValueError( - "K-FAC factor compression (projection_dim > 0) is only supported with " - f"method='kfac', got method='{hessian_cfg.method}'." - ) if index_cfg.projection_target != "per_module": raise ValueError( - "K-FAC factor compression requires projection_target='per_module' " + "Hessian factor compression requires projection_target='per_module' " "(global projection cannot compose with the per-layer Kronecker " "factorization). Got projection_target=" f"'{index_cfg.projection_target}'." ) + if hessian_cfg.ev_correction: + raise ValueError( + "Hessian factor compression (projection_dim > 0) is incompatible " + "with ev_correction=True: the EK-FAC eigenvalue correction does not " + "decompose into per-side row/column operations and so cannot be " + "folded into the Kronecker compression scheme." + ) attention_cfgs = { module: index_cfg.attention for module in index_cfg.split_attention_modules @@ -191,16 +193,10 @@ def hessian_worker( eva_a = compute_eigendecomposition( os.path.join(index_cfg.partial_run_path, "activation_sharded"), total_processed=total_processed, - projection_dim=hessian_cfg.projection_dim, - projection_type=hessian_cfg.projection_type, - side="right", ) eva_g = compute_eigendecomposition( os.path.join(index_cfg.partial_run_path, "gradient_sharded"), total_processed=total_processed, - projection_dim=hessian_cfg.projection_dim, - projection_type=hessian_cfg.projection_type, - side="left", ) dist.barrier() if dist.is_initialized() else None @@ -245,11 +241,7 @@ def collect_hessians( } desc = f"Approximating Hessians with {hessian_cfg.method}" if ev_correction: - collector = LambdaCollector( - **collector_args, - projection_dim=hessian_cfg.projection_dim, - projection_type=hessian_cfg.projection_type, - ) + collector = LambdaCollector(**collector_args) desc += " (eigenvalue correction)" else: collector_args["dtype"] = hessian_dtype From 11a4d8e75eac3daff6c612cec604cda22d835029 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Wed, 6 May 2026 15:00:21 -0700 Subject: [PATCH 03/10] Stage 2: persist per-side eigenvalues to disk Stage 3 (precompute the damped-inverse-projection M matrices) needs E_A and E_S separately. The #273 merge gave us their outer product on disk via `save_uncorrected_eigenvalues` (-> `eigenvalue_sharded/`), but not the per-side spectra. `compute_eigendecomposition` now writes the per-key eigenvalue shards alongside the eigenvectors, mirroring the existing naming: eigen_activation_sharded/ (Q_A) eigenvalue_activation_sharded/ (E_A) eigen_gradient_sharded/ (Q_S) eigenvalue_gradient_sharded/ (E_S) The function still returns the eigenvalue dict so the in-process `save_uncorrected_eigenvalues` call can use them without reloading. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/hessians/eigenvectors.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/bergson/hessians/eigenvectors.py b/bergson/hessians/eigenvectors.py index 5a8530e6..efe26f7c 100644 --- a/bergson/hessians/eigenvectors.py +++ b/bergson/hessians/eigenvectors.py @@ -229,10 +229,11 @@ def compute_eigendecomposition( total_processed: Number of samples used to compute covariance. Returns: - Per-key eigenvalue shards (each `[m/world_size]`) on CPU. The - eigenvectors are written to disk; the eigenvalues are returned so - callers (e.g. `save_uncorrected_eigenvalues`) can use them without - reloading. + Per-key eigenvalue shards (each `[m/world_size]`) on CPU. Both the + eigenvectors and eigenvalues are persisted to disk (under + ``eigen_/`` and ``eigenvalue_/`` respectively); + the eigenvalues are also returned so in-process callers (e.g. + `save_uncorrected_eigenvalues`) can use them without reloading. """ rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -312,20 +313,28 @@ def compute_eigendecomposition( device=device, ) - # Generic output path by adding eigen_prefix to the path + # Generic output paths by prefixing the covariance basename dirname = os.path.dirname(covariance_path) basename = os.path.basename(covariance_path) - output_path = os.path.join(dirname, "eigen_" + basename) + eigvec_path = os.path.join(dirname, "eigen_" + basename) + eigval_path = os.path.join(dirname, "eigenvalue_" + basename) - os.makedirs(output_path, exist_ok=True) + os.makedirs(eigvec_path, exist_ok=True) save_file( covariance_eigenvectors, - os.path.join(output_path, f"shard_{rank}.safetensors"), + os.path.join(eigvec_path, f"shard_{rank}.safetensors"), + ) + + os.makedirs(eigval_path, exist_ok=True) + save_file( + covariance_eigenvalues, + os.path.join(eigval_path, f"shard_{rank}.safetensors"), ) gc.collect() - get_logger().info(f"Saved eigenvectors to {output_path}") + get_logger().info(f"Saved eigenvectors to {eigvec_path}") + get_logger().info(f"Saved eigenvalues to {eigval_path}") return covariance_eigenvalues From 3eab445c4562d96a3ade0206fdc432c7750f68ec Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Wed, 6 May 2026 15:07:30 -0700 Subject: [PATCH 04/10] Stage 3: precompute damped-inverse-projection matrices M_A, M_S MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For each layer, fold the per-side damped curvature inverse and the per-layer random projection into a single [d, projection_dim] matrix: M = Q · diag((E + α)^{-1}) · Q^T · P^T where Q, E come from the saved eigendecomposition, P is the same Rademacher / normal projection used by `bergson build` (identifier `f"{name}/{side}"`), and α = lambda_damp_factor * mean(E) is the adaptive damping (mirrors `ShardedMul._sharded_hadamard`). `compute_whitening_projection_matrices` reads Q and E shards from disk, distributes layers across ranks via `fair_distribute_by_cost`, computes M in float64 on the assigned rank, then re-shards along dim 0 (rows of d) so each rank ends up with `M[rank*d/W : (rank+1)*d/W, :]` — matching the existing Q sharding so `ShardedMul._matmul(activations, M_shard)` is the natural way to use these in stage 4. Outputs land in: whitening_projection_activation_sharded/ (M_A, side='right') whitening_projection_gradient_sharded/ (M_S, side='left') `HessianConfig.lambda_damp_factor` (new field, default 0.1) controls the damping baked into M. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/config.py | 6 + bergson/hessians/eigenvectors.py | 128 ++++++++++++++++++++- bergson/hessians/hessian_approximations.py | 34 ++++++ 3 files changed, 167 insertions(+), 1 deletion(-) diff --git a/bergson/config.py b/bergson/config.py index e4b5994b..6df72566 100644 --- a/bergson/config.py +++ b/bergson/config.py @@ -561,6 +561,12 @@ class HessianConfig(Serializable): """Type of random projection used when ``projection_dim > 0``. Must match the value used in ``bergson build``.""" + lambda_damp_factor: float = 0.1 + """Damping factor baked into the precomputed whitening-projection matrices + when ``projection_dim > 0``. The damped inverse uses + ``α = lambda_damp_factor * mean(E)`` per side. Has no effect when + ``projection_dim == 0``.""" + @dataclass class FaissConfig: diff --git a/bergson/hessians/eigenvectors.py b/bergson/hessians/eigenvectors.py index efe26f7c..fafa7508 100644 --- a/bergson/hessians/eigenvectors.py +++ b/bergson/hessians/eigenvectors.py @@ -1,6 +1,7 @@ import gc import os from dataclasses import dataclass +from typing import Literal import torch import torch.distributed as dist @@ -10,7 +11,7 @@ from torch import Tensor from tqdm import tqdm -from bergson.collector.collector import HookCollectorBase +from bergson.collector.collector import HookCollectorBase, create_projection_matrix from bergson.hessians.sharded_computation import ShardedMul from bergson.utils.logger import get_logger from bergson.utils.utils import ( @@ -429,3 +430,128 @@ def _gather_and_shard_along_dim0( gc.collect() return result_dict + + +def compute_whitening_projection_matrices( + eigvec_path: str, + eigval_path: str, + output_path: str, + *, + projection_dim: int, + projection_type: Literal["normal", "rademacher"], + side: Literal["left", "right"], + lambda_damp_factor: float = 0.1, +) -> None: + """ + For each layer, compute the damped-inverse-projection matrix + + M = (C + α I)^{-1} P^T = Q · diag((E + α)^{-1}) · Q^T · P^T + + of shape ``[d, projection_dim]``, where: + + - ``C`` is the per-layer covariance (A for ``side='right'``, S for + ``side='left'``) with eigendecomposition ``C = Q E Q^T``. + - ``P`` is the same per-layer Rademacher / normal projection matrix used + by ``bergson build``, identified by ``f"{name}/{side}"`` and shape + ``[projection_dim, d]``. + - ``α = lambda_damp_factor * mean(E)`` is the adaptive damping (mirrors + the convention used by ``ShardedMul._sharded_hadamard``). + + The output is row-sharded along ``d`` so each rank stores + ``M[rank*d/W : (rank+1)*d/W, :]``, matching the ``ShardedMul`` convention + used elsewhere in the K-FAC pipeline. With this sharding, + ``ShardedMul._matmul(activations, M_shard)`` produces whitened-and-projected + activations in one matmul. + + Args: + eigvec_path: Directory containing eigenvector shards (e.g. + ``.../eigen_activation_sharded``). + eigval_path: Directory containing eigenvalue shards (e.g. + ``.../eigenvalue_activation_sharded``). + output_path: Directory to write ``shard_{rank}.safetensors`` files + containing the per-rank row-shard of ``M`` for every layer. + projection_dim: Width ``p`` of the random projection (must match the + value used by ``bergson build``). + projection_type: Random projection family (must match build). + side: Which side of the per-layer factorization this is. ``'right'`` + for the activation covariance ``A``, ``'left'`` for the gradient + covariance ``S``. + lambda_damp_factor: Multiplier on ``mean(E)`` used to damp the inverse. + """ + rank = dist.get_rank() if dist.is_initialized() else 0 + world_size = dist.get_world_size() if dist.is_initialized() else 1 + device = get_device(rank) + + # Discover keys / dimensions from the first eigenvector shard. Each rank's + # eigenvector shard has shape [d/W, d], so column count is the full layer + # dim. + first_eigvec = os.path.join(eigvec_path, "shard_0.safetensors") + with safe_open(first_eigvec, framework="pt") as f: + all_keys = list(f.keys()) + original_dtype = f.get_tensor(all_keys[0]).dtype + key_dimensions = {key: f.get_tensor(key).shape[1] for key in all_keys} + + # Distribute keys across ranks. Cost is dominated by the full-rank matmul + # Q · diag · Q^T · P^T which is O(d^3) just like the eigendecomposition. + all_assignments = fair_distribute_by_cost(key_dimensions, world_size) + keys_for_this_rank = all_assignments[rank] + + M_dict: dict[str, Tensor] = {} + for key in tqdm( + keys_for_this_rank, + disable=False, + desc=f"Rank {rank}: Computing M_{side}", + position=rank, + leave=False, + ): + d = key_dimensions[key] + + # Gather full Q [d, d] and E [d] for this key from disk. + Q = _compute_full_matrix( + name=key, shard_path=eigvec_path, rank=rank, world_size=world_size + ).to(device=device, dtype=torch.float64) + E = _compute_full_matrix( + name=key, shard_path=eigval_path, rank=rank, world_size=world_size + ).to(device=device, dtype=torch.float64) + + # Adaptive damping: α = lambda_damp_factor * mean(E). + alpha = lambda_damp_factor * E.mean() + E_inv = (E + alpha).reciprocal() # [d] + + # Per-layer random projection, identified the same way bergson build + # does it (see HookCollectorBase.projection). + P = create_projection_matrix( + identifier=f"{key}/{side}", + m=projection_dim, + n=d, + dtype=Q.dtype, + device=Q.device, + projection_type=projection_type, + ) # [p, d] + + # M = Q · diag(E_inv) · Q^T · P^T, computed as Q @ ((E_inv ⊙ (Q^T P^T))). + QtPt = Q.T @ P.T # [d, p] + QtPt.mul_(E_inv.unsqueeze(-1)) # [d, p] + M = Q @ QtPt # [d, p] + + M_dict[key] = M.to(dtype=original_dtype, device="cpu").contiguous() + + # Reshard along dim 0 so every rank ends up with M[rank*d/W : (rank+1)*d/W] + # for every layer. This matches the row-sharding used for Q. + M_dict = _gather_and_shard_along_dim0( + input_dict=M_dict, + full_shape_per_key={k: (d, projection_dim) for k, d in key_dimensions.items()}, + dtype=original_dtype, # type: ignore + rank=rank, + world_size=world_size, + device=device, + ) + + os.makedirs(output_path, exist_ok=True) + save_file( + M_dict, + os.path.join(output_path, f"shard_{rank}.safetensors"), + ) + + gc.collect() + get_logger().info(f"Saved whitening-projection matrices to {output_path}") diff --git a/bergson/hessians/hessian_approximations.py b/bergson/hessians/hessian_approximations.py index 0de18cc0..11d4b500 100644 --- a/bergson/hessians/hessian_approximations.py +++ b/bergson/hessians/hessian_approximations.py @@ -17,6 +17,7 @@ from bergson.hessians.eigenvectors import ( LambdaCollector, compute_eigendecomposition, + compute_whitening_projection_matrices, save_uncorrected_eigenvalues, ) from bergson.hessians.kfac import CovarianceCollector @@ -210,6 +211,39 @@ def hessian_worker( world_size=world_size, ) + if hessian_cfg.projection_dim > 0: + compute_whitening_projection_matrices( + eigvec_path=os.path.join( + index_cfg.partial_run_path, "eigen_activation_sharded" + ), + eigval_path=os.path.join( + index_cfg.partial_run_path, "eigenvalue_activation_sharded" + ), + output_path=os.path.join( + index_cfg.partial_run_path, "whitening_projection_activation_sharded" + ), + projection_dim=hessian_cfg.projection_dim, + projection_type=hessian_cfg.projection_type, + side="right", + lambda_damp_factor=hessian_cfg.lambda_damp_factor, + ) + compute_whitening_projection_matrices( + eigvec_path=os.path.join( + index_cfg.partial_run_path, "eigen_gradient_sharded" + ), + eigval_path=os.path.join( + index_cfg.partial_run_path, "eigenvalue_gradient_sharded" + ), + output_path=os.path.join( + index_cfg.partial_run_path, "whitening_projection_gradient_sharded" + ), + projection_dim=hessian_cfg.projection_dim, + projection_type=hessian_cfg.projection_type, + side="left", + lambda_damp_factor=hessian_cfg.lambda_damp_factor, + ) + dist.barrier() if dist.is_initialized() else None + if hessian_cfg.ev_correction: collect_hessians(**kwargs, ev_correction=True) From 71172a5ae5b0725a6363df8e153d57757bb90721 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Wed, 6 May 2026 15:12:15 -0700 Subject: [PATCH 05/10] Stage 4: use precomputed M in apply_hessian when present MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When hessian build saves whitening-projection matrices (stage 3), the IVHP computation collapses to a single sharded matmul per side per layer: G̃_q = M_Sᵀ · G_q · M_A shape: [N, p_S, p_A] No eigenbasis rotation, no eigenvalue divide, no rotation back — the damped curvature inverse and the random projection are already folded into M. `compute_ivhp_sharded` now dispatches: if ``whitening_projection_activation_sharded/`` exists under the hessian path, take the new path; otherwise fall through to the legacy rotate-divide-rotate flow. The new path also rejects ``ev_correction=True`` at runtime, mirroring the build-time guard added in stage 1. Output shape is ``[p_S, p_A]`` flattened per layer (down from ``[d_S, d_A]`` in the legacy path), which is what makes this a compression of the IVHP store, not just a reformulation. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/hessians/apply_hessian.py | 95 +++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 2841c287..644235cd 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -48,6 +48,18 @@ def __init__(self, cfg: EkfacConfig): self.sharded_computer = ShardedMul() def compute_ivhp_sharded(self): + if os.path.isdir( + os.path.join(self.path, "whitening_projection_activation_sharded") + ): + if self.cfg.ev_correction: + raise ValueError( + "ev_correction=True is incompatible with precomputed " + "whitening-projection matrices (the EK-FAC eigenvalue " + "correction does not decompose into per-side row/column ops " + "and so was disallowed at hessian-build time)." + ) + return self._compute_ivhp_with_whitening_projection() + eigen_a = load_file( self.path + f"/eigen_activation_sharded/shard_{self.rank}.safetensors", device=self.device, @@ -151,6 +163,89 @@ def compute_ivhp_sharded(self): self.logger.info(f"Saved IVHP gradients to {self.cfg.run_path}") + def _compute_ivhp_with_whitening_projection(self): + """IVHP via precomputed `M = (cov + αI)^{-1} P^T` matrices. + + For each layer, compute ``M_S^T G M_A`` of shape ``[N, p, p]`` in two + sharded matmuls — the per-side damped curvature inverse and the + random projection are both already folded into the ``M`` matrices + (built by ``compute_whitening_projection_matrices`` at hessian time), + so no eigenbasis rotation or eigenvalue divide is needed here. + """ + M_a = load_file( + self.path + + f"/whitening_projection_activation_sharded/shard_{self.rank}.safetensors", + device=self.device, + ) + M_g = load_file( + self.path + + f"/whitening_projection_gradient_sharded/shard_{self.rank}.safetensors", + device=self.device, + ) + + for k in M_a: + M_a[k] = M_a[k].to(dtype=torch.float32) + M_g[k] = M_g[k].to(dtype=torch.float32) + + # Per-layer output is [p_S, p_A] flattened; M shards are [d/W, p]. + grad_sizes = {name: M_g[name].shape[1] * M_a[name].shape[1] for name in M_a} + + mmap = load_gradients(self.gradient_path) + with open(os.path.join(self.gradient_path, "info.json")) as f: + info = json.load(f) + + grad_buffer = create_index( + Path(self.cfg.run_path), + num_grads=info["num_grads"], + grad_sizes=grad_sizes, + dtype=np.float32, + ) + + self.logger.info( + f"Loaded gradients for {len(mmap)} queries and computing IVHP " + "via whitening-projection matrices..." + ) + + transformed_gradients: dict[str, Tensor] = {} + for k, M_a_shard in M_a.items(): + M_g_shard = M_g[k] + + # Recover full per-layer dims from shard shapes. + d_S = M_g_shard.shape[0] * self.world_size + d_A = M_a_shard.shape[0] * self.world_size + + gradients_noi = torch.from_numpy(mmap[k][:]).to( + device=self.device, dtype=torch.float32 + ) + gradients_noi = gradients_noi.view(-1, d_S, d_A) + + # Step 1: G @ M_A : [N, d_S, d_A] @ [d_A, p] -> [N, d_S, p] + intermediate = self.sharded_computer._matmul( + vector_nsa=gradients_noi, matrix_cb=M_a_shard + ) + + # Step 2: M_S^T @ (G @ M_A) -> [N, p_S, p_A]. + # Computed as ((G @ M_A)^T @ M_S)^T to fit the row-sharded M_g + # shape via _matmul. + intermediate = self.sharded_computer._matmul( + vector_nsa=intermediate.transpose(-2, -1), matrix_cb=M_g_shard + ).transpose(-2, -1) + + transformed_gradients[k] = intermediate + + del M_a, M_g + gc.collect() + + self.logger.debug("Finished M_S^T @ G @ M_A") + + torch.cuda.synchronize() + for k, v in transformed_gradients.items(): + grad_buffer[k][:] = v.to(device="cpu", non_blocking=True).flatten(1).numpy() + + grad_buffer.flush() + + self.logger.info(f"Saved IVHP gradients to {self.cfg.run_path}") + def apply_worker( rank: int, From 10943ac1f81df774f1f242b0c2b983422d650b3b Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Wed, 6 May 2026 15:40:04 -0700 Subject: [PATCH 06/10] Stage 6: propagate projection settings into the score step Step 4 of `hessian_pipeline` always reset `score_index_cfg.projection_dim = 0`, which was correct for the legacy IVHP path (apply_hessian writes the transformed query at full `[d_S, d_A]`, and training gradients must match that for a meaningful inner product at score time). With Hessian-factor compression on (`hessian_cfg.projection_dim > 0`), apply_hessian now writes the query at `[p, p]` per layer, so training gradients also need to be projected to `[p, p]` per layer. Otherwise the score step compares `[d_S, d_A]` against `[p, p]` and either errors out or computes a meaningless inner product. Inherit the matching `projection_dim`, `projection_type`, and `projection_target='per_module'` from `hessian_cfg` when compression is on; keep the `projection_dim=0` override for the legacy path. This is the last piece needed to run `bergson ekfac` end-to-end with compression enabled. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/hessians/pipeline.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/bergson/hessians/pipeline.py b/bergson/hessians/pipeline.py index 1c52652c..36aedded 100644 --- a/bergson/hessians/pipeline.py +++ b/bergson/hessians/pipeline.py @@ -98,8 +98,18 @@ def _validate(cfg: IndexConfig): if not _step_complete(scores_path, resume): score_index_cfg = deepcopy(index_cfg) score_index_cfg.run_path = scores_path - score_index_cfg.projection_dim = 0 score_index_cfg.skip_hessians = True + if hessian_cfg.projection_dim > 0: + # Hessian-factor compression is on, so apply_hessian writes the + # transformed query at [p, p] per layer. Training gradients must + # match that shape, which means projecting them with the same + # per-module random matrix used to build M (same projection_dim, + # projection_type, and projection_target). + score_index_cfg.projection_dim = hessian_cfg.projection_dim + score_index_cfg.projection_type = hessian_cfg.projection_type + score_index_cfg.projection_target = "per_module" + else: + score_index_cfg.projection_dim = 0 score_cfg.query_path = transformed_query_path _validate(score_index_cfg) From 9a7325f93e22aef10bf4dbe57c78ab69233d59b7 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Thu, 7 May 2026 05:35:57 -0700 Subject: [PATCH 07/10] =?UTF-8?q?Stage=207:=20simplify=20per=20Lucia=20?= =?UTF-8?q?=E2=80=94=20post-IVHP=20projection,=20no=20precompute?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per Lucia's review (#275), the build-time `M` precompute was overkill. The brief always asked for: K-FAC computed on uncompressed gradients (unchanged), and compression added at apply time. This commit restores that minimal shape. Replaces the `M = (cov + αI)⁻¹ Pᵀ` precompute (stages 2-4) with a trivial post-projection of the legacy IVHP output: G̃_q = P_S · (H⁻¹ G_q) · P_Aᵀ ∈ [N, p, p] The legacy rotate-divide-rotate is unchanged; the new compression block runs after it, using the same `f"{name}/{side}"` per-layer projection identifiers as `bergson build`. Train-side stored gradients are already `P_S · G_t · P_Aᵀ` (collector.py's `g_proj.T @ a_proj` collapses to this), so the Frobenius inner product approximates `` up to a known `(p/d)²` constant. Deletions vs current branch: - `eigenvectors.py`: revert stage 2's per-side eigenvalue saving; delete `compute_whitening_projection_matrices` entirely. - `hessian_approximations.py`: delete the M precompute calls and the `projection_target == 'per_module'` + `ev_correction` validation block (the post-projection design composes with EV correction). - `config.py`: delete `HessianConfig.projection_dim`, `projection_type`, `lambda_damp_factor` — compression is now controlled entirely by `IndexConfig.projection_dim` (already used by build/score) and the new `EkfacConfig.projection_dim`. - `apply_hessian.py`: delete `_compute_ivhp_with_whitening_projection` and the dispatch. Additions: - `EkfacConfig`: new `projection_dim`, `projection_type` fields. - `apply_hessian.py:compute_ivhp_sharded`: ~15 lines projecting the rotated-and-divided IVHP output per layer when `projection_dim > 0`, plus a `grad_sizes` adjustment. - `pipeline.py` step 3: pass `index_cfg.projection_dim` and `projection_type` into `EkfacConfig`. Step 4 stops overriding `score_index_cfg.projection_dim = 0` — the deepcopy of `index_cfg` now correctly carries the projection settings to the score step, so training-side gradient shapes match the apply output automatically. Net: -258 lines vs current branch, +~25 lines vs main. Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/config.py | 19 --- bergson/hessians/apply_hessian.py | 144 +++++++------------ bergson/hessians/eigenvectors.py | 155 ++------------------- bergson/hessians/hessian_approximations.py | 50 ------- bergson/hessians/pipeline.py | 17 +-- 5 files changed, 63 insertions(+), 322 deletions(-) diff --git a/bergson/config.py b/bergson/config.py index 6df72566..ffca3a2d 100644 --- a/bergson/config.py +++ b/bergson/config.py @@ -548,25 +548,6 @@ class HessianConfig(Serializable): """Whether to use dataset labels for Hessian (empirical Fisher) approximation. If false, the model predictions will be used.""" - projection_dim: int = 0 - """If > 0, fold the same per-layer random projection used by ``bergson build`` - into the damped inverse of the Kronecker factors (A and S) so each per-side - factor is compressed to a ``[d, projection_dim]`` whitening-projection matrix. - Set this to the same value as the gradient store these factors will be applied - to, with matching ``projection_type``. Works for any Kronecker-factored method - (``kfac``, ``shampoo``, ...). Requires ``projection_target='per_module'`` and - is incompatible with ``ev_correction=True``.""" - - projection_type: Literal["normal", "rademacher"] = "rademacher" - """Type of random projection used when ``projection_dim > 0``. Must match the - value used in ``bergson build``.""" - - lambda_damp_factor: float = 0.1 - """Damping factor baked into the precomputed whitening-projection matrices - when ``projection_dim > 0``. The damped inverse uses - ``α = lambda_damp_factor * mean(E)`` per side. Has no effect when - ``projection_dim == 0``.""" - @dataclass class FaissConfig: diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 644235cd..877cbe90 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from datetime import timedelta from pathlib import Path +from typing import Literal import numpy as np import torch @@ -12,6 +13,7 @@ from simple_parsing import ArgumentParser from torch import Tensor +from bergson.collector.collector import create_projection_matrix from bergson.data import create_index, load_gradients from bergson.hessians.sharded_computation import ShardedMul from bergson.utils.logger import get_logger @@ -29,6 +31,15 @@ class EkfacConfig: `HessianConfig.ev_correction=True`.""" debug: bool = False lambda_damp_factor: float = 0.1 + projection_dim: int = 0 + """If > 0, compress the IVHP output per layer by ``P_S · (H⁻¹G) · P_Aᵀ`` + using the same per-layer ``f"{name}/{side}"`` random projections that + ``bergson build`` uses on training gradients. The result is ``[p, p]`` + per layer, suitable for inner-product scoring against a compressed + gradient store.""" + projection_type: Literal["normal", "rademacher"] = "rademacher" + """Random projection family for compression. Must match the value used + in the matching ``bergson build``.""" class EkfacApplicator: @@ -48,18 +59,6 @@ def __init__(self, cfg: EkfacConfig): self.sharded_computer = ShardedMul() def compute_ivhp_sharded(self): - if os.path.isdir( - os.path.join(self.path, "whitening_projection_activation_sharded") - ): - if self.cfg.ev_correction: - raise ValueError( - "ev_correction=True is incompatible with precomputed " - "whitening-projection matrices (the EK-FAC eigenvalue " - "correction does not decompose into per-side row/column ops " - "and so was disallowed at hessian-build time)." - ) - return self._compute_ivhp_with_whitening_projection() - eigen_a = load_file( self.path + f"/eigen_activation_sharded/shard_{self.rank}.safetensors", device=self.device, @@ -83,9 +82,14 @@ def compute_ivhp_sharded(self): eigen_g[k] = eigen_g[k].to(dtype=torch.float32) lambda_factor[k] = v.to(dtype=torch.float32) - grad_sizes = { - name: eigen_g[name].shape[1] * eigen_a[name].shape[1] for name in eigen_a - } + p = self.cfg.projection_dim + if p > 0: + grad_sizes = {name: p * p for name in eigen_a} + else: + grad_sizes = { + name: eigen_g[name].shape[1] * eigen_a[name].shape[1] + for name in eigen_a + } mmap = load_gradients(self.gradient_path) with open(os.path.join(self.gradient_path, "info.json")) as f: @@ -155,88 +159,34 @@ def compute_ivhp_sharded(self): del eigen_a gc.collect() - torch.cuda.synchronize() - for k, v in transformed_gradients.items(): - grad_buffer[k][:] = v.to(device="cpu", non_blocking=True).flatten(1).numpy() - - grad_buffer.flush() - - self.logger.info(f"Saved IVHP gradients to {self.cfg.run_path}") - - def _compute_ivhp_with_whitening_projection(self): - """IVHP via precomputed `M = (cov + αI)^{-1} P^T` matrices. - - For each layer, compute ``M_S^T G M_A`` of shape ``[N, p, p]`` in two - sharded matmuls — the per-side damped curvature inverse and the - random projection are both already folded into the ``M`` matrices - (built by ``compute_whitening_projection_matrices`` at hessian time), - so no eigenbasis rotation or eigenvalue divide is needed here. - """ - M_a = load_file( - self.path - + f"/whitening_projection_activation_sharded/shard_{self.rank}.safetensors", - device=self.device, - ) - M_g = load_file( - self.path - + f"/whitening_projection_gradient_sharded/shard_{self.rank}.safetensors", - device=self.device, - ) - - for k in M_a: - M_a[k] = M_a[k].to(dtype=torch.float32) - M_g[k] = M_g[k].to(dtype=torch.float32) - - # Per-layer output is [p_S, p_A] flattened; M shards are [d/W, p]. - grad_sizes = {name: M_g[name].shape[1] * M_a[name].shape[1] for name in M_a} - - mmap = load_gradients(self.gradient_path) - with open(os.path.join(self.gradient_path, "info.json")) as f: - info = json.load(f) - - grad_buffer = create_index( - Path(self.cfg.run_path), - num_grads=info["num_grads"], - grad_sizes=grad_sizes, - dtype=np.float32, - ) - - self.logger.info( - f"Loaded gradients for {len(mmap)} queries and computing IVHP " - "via whitening-projection matrices..." - ) - - transformed_gradients: dict[str, Tensor] = {} - for k, M_a_shard in M_a.items(): - M_g_shard = M_g[k] - - # Recover full per-layer dims from shard shapes. - d_S = M_g_shard.shape[0] * self.world_size - d_A = M_a_shard.shape[0] * self.world_size - - gradients_noi = torch.from_numpy(mmap[k][:]).to( - device=self.device, dtype=torch.float32 - ) - gradients_noi = gradients_noi.view(-1, d_S, d_A) - - # Step 1: G @ M_A : [N, d_S, d_A] @ [d_A, p] -> [N, d_S, p] - intermediate = self.sharded_computer._matmul( - vector_nsa=gradients_noi, matrix_cb=M_a_shard - ) - - # Step 2: M_S^T @ (G @ M_A) -> [N, p_S, p_A]. - # Computed as ((G @ M_A)^T @ M_S)^T to fit the row-sharded M_g - # shape via _matmul. - intermediate = self.sharded_computer._matmul( - vector_nsa=intermediate.transpose(-2, -1), matrix_cb=M_g_shard - ).transpose(-2, -1) - - transformed_gradients[k] = intermediate - - del M_a, M_g - gc.collect() - - self.logger.debug("Finished M_S^T @ G @ M_A") + # Compress the IVHP output per layer by ``P_S · (H⁻¹G) · P_Aᵀ`` so + # the saved gradients are ``[N, p, p]`` and can be directly inner- + # producted against a `bergson build` store that used the same + # per-layer projections. + if p > 0: + for k, v in transformed_gradients.items(): + d_S = v.shape[-2] + d_A = v.shape[-1] + P_left = create_projection_matrix( + identifier=f"{k}/left", + m=p, + n=d_S, + dtype=v.dtype, + device=v.device, + projection_type=self.cfg.projection_type, + ) + P_right = create_projection_matrix( + identifier=f"{k}/right", + m=p, + n=d_A, + dtype=v.dtype, + device=v.device, + projection_type=self.cfg.projection_type, + ) + transformed_gradients[k] = torch.einsum( + "ps,nsa,ra->npr", P_left, v, P_right + ) + self.logger.debug(f"Finished P_S · (H^{{-1}} G) · P_Aᵀ at p={p}") torch.cuda.synchronize() for k, v in transformed_gradients.items(): diff --git a/bergson/hessians/eigenvectors.py b/bergson/hessians/eigenvectors.py index fafa7508..5a8530e6 100644 --- a/bergson/hessians/eigenvectors.py +++ b/bergson/hessians/eigenvectors.py @@ -1,7 +1,6 @@ import gc import os from dataclasses import dataclass -from typing import Literal import torch import torch.distributed as dist @@ -11,7 +10,7 @@ from torch import Tensor from tqdm import tqdm -from bergson.collector.collector import HookCollectorBase, create_projection_matrix +from bergson.collector.collector import HookCollectorBase from bergson.hessians.sharded_computation import ShardedMul from bergson.utils.logger import get_logger from bergson.utils.utils import ( @@ -230,11 +229,10 @@ def compute_eigendecomposition( total_processed: Number of samples used to compute covariance. Returns: - Per-key eigenvalue shards (each `[m/world_size]`) on CPU. Both the - eigenvectors and eigenvalues are persisted to disk (under - ``eigen_/`` and ``eigenvalue_/`` respectively); - the eigenvalues are also returned so in-process callers (e.g. - `save_uncorrected_eigenvalues`) can use them without reloading. + Per-key eigenvalue shards (each `[m/world_size]`) on CPU. The + eigenvectors are written to disk; the eigenvalues are returned so + callers (e.g. `save_uncorrected_eigenvalues`) can use them without + reloading. """ rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist.is_initialized() else 1 @@ -314,28 +312,20 @@ def compute_eigendecomposition( device=device, ) - # Generic output paths by prefixing the covariance basename + # Generic output path by adding eigen_prefix to the path dirname = os.path.dirname(covariance_path) basename = os.path.basename(covariance_path) - eigvec_path = os.path.join(dirname, "eigen_" + basename) - eigval_path = os.path.join(dirname, "eigenvalue_" + basename) + output_path = os.path.join(dirname, "eigen_" + basename) - os.makedirs(eigvec_path, exist_ok=True) + os.makedirs(output_path, exist_ok=True) save_file( covariance_eigenvectors, - os.path.join(eigvec_path, f"shard_{rank}.safetensors"), - ) - - os.makedirs(eigval_path, exist_ok=True) - save_file( - covariance_eigenvalues, - os.path.join(eigval_path, f"shard_{rank}.safetensors"), + os.path.join(output_path, f"shard_{rank}.safetensors"), ) gc.collect() - get_logger().info(f"Saved eigenvectors to {eigvec_path}") - get_logger().info(f"Saved eigenvalues to {eigval_path}") + get_logger().info(f"Saved eigenvectors to {output_path}") return covariance_eigenvalues @@ -430,128 +420,3 @@ def _gather_and_shard_along_dim0( gc.collect() return result_dict - - -def compute_whitening_projection_matrices( - eigvec_path: str, - eigval_path: str, - output_path: str, - *, - projection_dim: int, - projection_type: Literal["normal", "rademacher"], - side: Literal["left", "right"], - lambda_damp_factor: float = 0.1, -) -> None: - """ - For each layer, compute the damped-inverse-projection matrix - - M = (C + α I)^{-1} P^T = Q · diag((E + α)^{-1}) · Q^T · P^T - - of shape ``[d, projection_dim]``, where: - - - ``C`` is the per-layer covariance (A for ``side='right'``, S for - ``side='left'``) with eigendecomposition ``C = Q E Q^T``. - - ``P`` is the same per-layer Rademacher / normal projection matrix used - by ``bergson build``, identified by ``f"{name}/{side}"`` and shape - ``[projection_dim, d]``. - - ``α = lambda_damp_factor * mean(E)`` is the adaptive damping (mirrors - the convention used by ``ShardedMul._sharded_hadamard``). - - The output is row-sharded along ``d`` so each rank stores - ``M[rank*d/W : (rank+1)*d/W, :]``, matching the ``ShardedMul`` convention - used elsewhere in the K-FAC pipeline. With this sharding, - ``ShardedMul._matmul(activations, M_shard)`` produces whitened-and-projected - activations in one matmul. - - Args: - eigvec_path: Directory containing eigenvector shards (e.g. - ``.../eigen_activation_sharded``). - eigval_path: Directory containing eigenvalue shards (e.g. - ``.../eigenvalue_activation_sharded``). - output_path: Directory to write ``shard_{rank}.safetensors`` files - containing the per-rank row-shard of ``M`` for every layer. - projection_dim: Width ``p`` of the random projection (must match the - value used by ``bergson build``). - projection_type: Random projection family (must match build). - side: Which side of the per-layer factorization this is. ``'right'`` - for the activation covariance ``A``, ``'left'`` for the gradient - covariance ``S``. - lambda_damp_factor: Multiplier on ``mean(E)`` used to damp the inverse. - """ - rank = dist.get_rank() if dist.is_initialized() else 0 - world_size = dist.get_world_size() if dist.is_initialized() else 1 - device = get_device(rank) - - # Discover keys / dimensions from the first eigenvector shard. Each rank's - # eigenvector shard has shape [d/W, d], so column count is the full layer - # dim. - first_eigvec = os.path.join(eigvec_path, "shard_0.safetensors") - with safe_open(first_eigvec, framework="pt") as f: - all_keys = list(f.keys()) - original_dtype = f.get_tensor(all_keys[0]).dtype - key_dimensions = {key: f.get_tensor(key).shape[1] for key in all_keys} - - # Distribute keys across ranks. Cost is dominated by the full-rank matmul - # Q · diag · Q^T · P^T which is O(d^3) just like the eigendecomposition. - all_assignments = fair_distribute_by_cost(key_dimensions, world_size) - keys_for_this_rank = all_assignments[rank] - - M_dict: dict[str, Tensor] = {} - for key in tqdm( - keys_for_this_rank, - disable=False, - desc=f"Rank {rank}: Computing M_{side}", - position=rank, - leave=False, - ): - d = key_dimensions[key] - - # Gather full Q [d, d] and E [d] for this key from disk. - Q = _compute_full_matrix( - name=key, shard_path=eigvec_path, rank=rank, world_size=world_size - ).to(device=device, dtype=torch.float64) - E = _compute_full_matrix( - name=key, shard_path=eigval_path, rank=rank, world_size=world_size - ).to(device=device, dtype=torch.float64) - - # Adaptive damping: α = lambda_damp_factor * mean(E). - alpha = lambda_damp_factor * E.mean() - E_inv = (E + alpha).reciprocal() # [d] - - # Per-layer random projection, identified the same way bergson build - # does it (see HookCollectorBase.projection). - P = create_projection_matrix( - identifier=f"{key}/{side}", - m=projection_dim, - n=d, - dtype=Q.dtype, - device=Q.device, - projection_type=projection_type, - ) # [p, d] - - # M = Q · diag(E_inv) · Q^T · P^T, computed as Q @ ((E_inv ⊙ (Q^T P^T))). - QtPt = Q.T @ P.T # [d, p] - QtPt.mul_(E_inv.unsqueeze(-1)) # [d, p] - M = Q @ QtPt # [d, p] - - M_dict[key] = M.to(dtype=original_dtype, device="cpu").contiguous() - - # Reshard along dim 0 so every rank ends up with M[rank*d/W : (rank+1)*d/W] - # for every layer. This matches the row-sharding used for Q. - M_dict = _gather_and_shard_along_dim0( - input_dict=M_dict, - full_shape_per_key={k: (d, projection_dim) for k, d in key_dimensions.items()}, - dtype=original_dtype, # type: ignore - rank=rank, - world_size=world_size, - device=device, - ) - - os.makedirs(output_path, exist_ok=True) - save_file( - M_dict, - os.path.join(output_path, f"shard_{rank}.safetensors"), - ) - - gc.collect() - get_logger().info(f"Saved whitening-projection matrices to {output_path}") diff --git a/bergson/hessians/hessian_approximations.py b/bergson/hessians/hessian_approximations.py index 11d4b500..71678f9f 100644 --- a/bergson/hessians/hessian_approximations.py +++ b/bergson/hessians/hessian_approximations.py @@ -17,7 +17,6 @@ from bergson.hessians.eigenvectors import ( LambdaCollector, compute_eigendecomposition, - compute_whitening_projection_matrices, save_uncorrected_eigenvalues, ) from bergson.hessians.kfac import CovarianceCollector @@ -150,22 +149,6 @@ def hessian_worker( model, target_modules = setup_model_and_peft(index_cfg) - if hessian_cfg.projection_dim > 0: - if index_cfg.projection_target != "per_module": - raise ValueError( - "Hessian factor compression requires projection_target='per_module' " - "(global projection cannot compose with the per-layer Kronecker " - "factorization). Got projection_target=" - f"'{index_cfg.projection_target}'." - ) - if hessian_cfg.ev_correction: - raise ValueError( - "Hessian factor compression (projection_dim > 0) is incompatible " - "with ev_correction=True: the EK-FAC eigenvalue correction does not " - "decompose into per-side row/column operations and so cannot be " - "folded into the Kronecker compression scheme." - ) - attention_cfgs = { module: index_cfg.attention for module in index_cfg.split_attention_modules } @@ -211,39 +194,6 @@ def hessian_worker( world_size=world_size, ) - if hessian_cfg.projection_dim > 0: - compute_whitening_projection_matrices( - eigvec_path=os.path.join( - index_cfg.partial_run_path, "eigen_activation_sharded" - ), - eigval_path=os.path.join( - index_cfg.partial_run_path, "eigenvalue_activation_sharded" - ), - output_path=os.path.join( - index_cfg.partial_run_path, "whitening_projection_activation_sharded" - ), - projection_dim=hessian_cfg.projection_dim, - projection_type=hessian_cfg.projection_type, - side="right", - lambda_damp_factor=hessian_cfg.lambda_damp_factor, - ) - compute_whitening_projection_matrices( - eigvec_path=os.path.join( - index_cfg.partial_run_path, "eigen_gradient_sharded" - ), - eigval_path=os.path.join( - index_cfg.partial_run_path, "eigenvalue_gradient_sharded" - ), - output_path=os.path.join( - index_cfg.partial_run_path, "whitening_projection_gradient_sharded" - ), - projection_dim=hessian_cfg.projection_dim, - projection_type=hessian_cfg.projection_type, - side="left", - lambda_damp_factor=hessian_cfg.lambda_damp_factor, - ) - dist.barrier() if dist.is_initialized() else None - if hessian_cfg.ev_correction: collect_hessians(**kwargs, ev_correction=True) diff --git a/bergson/hessians/pipeline.py b/bergson/hessians/pipeline.py index 36aedded..d869c6f5 100644 --- a/bergson/hessians/pipeline.py +++ b/bergson/hessians/pipeline.py @@ -85,6 +85,8 @@ def _validate(cfg: IndexConfig): run_path=transformed_query_path, ev_correction=hessian_cfg.ev_correction, lambda_damp_factor=hessian_pipeline_cfg.lambda_damp_factor, + projection_dim=index_cfg.projection_dim, + projection_type=index_cfg.projection_type, ) launch_distributed_run( "apply_hessian", @@ -94,22 +96,15 @@ def _validate(cfg: IndexConfig): ) # ── Step 4: Score training examples ─────────────────────────────────── + # The deepcopy carries projection_dim / projection_type / projection_target + # from index_cfg, so training-side gradients are stored at the same shape + # as the IVHP output (apply_hessian compresses iff index_cfg.projection_dim + # > 0). No override needed on either branch. print("Step 4/4: Scoring training data against transformed query...") if not _step_complete(scores_path, resume): score_index_cfg = deepcopy(index_cfg) score_index_cfg.run_path = scores_path score_index_cfg.skip_hessians = True - if hessian_cfg.projection_dim > 0: - # Hessian-factor compression is on, so apply_hessian writes the - # transformed query at [p, p] per layer. Training gradients must - # match that shape, which means projecting them with the same - # per-module random matrix used to build M (same projection_dim, - # projection_type, and projection_target). - score_index_cfg.projection_dim = hessian_cfg.projection_dim - score_index_cfg.projection_type = hessian_cfg.projection_type - score_index_cfg.projection_target = "per_module" - else: - score_index_cfg.projection_dim = 0 score_cfg.query_path = transformed_query_path _validate(score_index_cfg) From b761933d09ae3fd4d15811743b54d34b8126d3f2 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Thu, 7 May 2026 14:18:09 -0700 Subject: [PATCH 08/10] Removes unnecessary comments --- bergson/collector/gradient_collectors.py | 11 ----------- bergson/hessians/apply_hessian.py | 11 ----------- bergson/hessians/pipeline.py | 4 ---- 3 files changed, 26 deletions(-) diff --git a/bergson/collector/gradient_collectors.py b/bergson/collector/gradient_collectors.py index 19982264..a1c817e9 100644 --- a/bergson/collector/gradient_collectors.py +++ b/bergson/collector/gradient_collectors.py @@ -126,13 +126,6 @@ def backward_hook(self, module: nn.Module, g: Float[Tensor, "N S O"]): def global_project(self) -> None: """Concatenate per-module per-example gradients and project. Sets ``self.mod_grads`` to ``{"gradients": projected}``. - - Projects in row-chunks. A naive ``flat = torch.cat(..., dim=1)`` of - all module gradients can need tens of GiB contiguous on rank 0 when - the bin-packer assigns many short examples to a single batch (e.g. - flan_v2 with token_batch_size=2048 packs ~80 rows of ~525 MB each). - The projector is per-row, so chunking is exact; chunk size is sized - to a fixed GPU-byte budget. """ # backward_hook fires in reverse forward order, so insertion order in # mod_grads is deterministic for a given model. @@ -150,9 +143,6 @@ def global_project(self) -> None: projection_type=self.processor.projection_type, ) - # Cap chunk_flat at ~4 GiB per chunk to leave headroom for fast_jl's - # internal scratch buffers alongside the per-module tensors that stay - # alive until the loop exits. bytes_per_row = total_grad_dim * parts[0].element_size() chunk_rows = max(1, (4 * 1024**3) // max(bytes_per_row, 1)) chunk_rows = min(chunk_rows, n_rows) @@ -171,7 +161,6 @@ def global_project(self) -> None: ) del chunk_flat, chunk_projected - # Free per-module GPU tensors now that all chunks are projected. self.mod_grads = {} self.mod_grads = {"gradients": torch.cat(chunks_cpu, dim=0)} diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 877cbe90..28b4891c 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -32,14 +32,7 @@ class EkfacConfig: debug: bool = False lambda_damp_factor: float = 0.1 projection_dim: int = 0 - """If > 0, compress the IVHP output per layer by ``P_S · (H⁻¹G) · P_Aᵀ`` - using the same per-layer ``f"{name}/{side}"`` random projections that - ``bergson build`` uses on training gradients. The result is ``[p, p]`` - per layer, suitable for inner-product scoring against a compressed - gradient store.""" projection_type: Literal["normal", "rademacher"] = "rademacher" - """Random projection family for compression. Must match the value used - in the matching ``bergson build``.""" class EkfacApplicator: @@ -159,10 +152,6 @@ def compute_ivhp_sharded(self): del eigen_a gc.collect() - # Compress the IVHP output per layer by ``P_S · (H⁻¹G) · P_Aᵀ`` so - # the saved gradients are ``[N, p, p]`` and can be directly inner- - # producted against a `bergson build` store that used the same - # per-layer projections. if p > 0: for k, v in transformed_gradients.items(): d_S = v.shape[-2] diff --git a/bergson/hessians/pipeline.py b/bergson/hessians/pipeline.py index d869c6f5..625c6888 100644 --- a/bergson/hessians/pipeline.py +++ b/bergson/hessians/pipeline.py @@ -96,10 +96,6 @@ def _validate(cfg: IndexConfig): ) # ── Step 4: Score training examples ─────────────────────────────────── - # The deepcopy carries projection_dim / projection_type / projection_target - # from index_cfg, so training-side gradients are stored at the same shape - # as the IVHP output (apply_hessian compresses iff index_cfg.projection_dim - # > 0). No override needed on either branch. print("Step 4/4: Scoring training data against transformed query...") if not _step_complete(scores_path, resume): score_index_cfg = deepcopy(index_cfg) From 56e736bcc9f1a378b233f4f97ff30f8246421e84 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Thu, 7 May 2026 14:24:44 -0700 Subject: [PATCH 09/10] Tighten the post-projection block Drop one-arg-per-line `create_projection_matrix(...)` calls in favor of positional args, fold the `grad_sizes` if/else into a dict-comprehension ternary, and remove the trailing debug log. Same behavior, ~17 fewer lines. Brings this PR to net +14 vs `kfac-full-support` (#273). Co-Authored-By: Claude Opus 4.7 (1M context) --- bergson/hessians/apply_hessian.py | 38 +++++++++---------------------- 1 file changed, 11 insertions(+), 27 deletions(-) diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 28b4891c..a808e8c9 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -76,13 +76,10 @@ def compute_ivhp_sharded(self): lambda_factor[k] = v.to(dtype=torch.float32) p = self.cfg.projection_dim - if p > 0: - grad_sizes = {name: p * p for name in eigen_a} - else: - grad_sizes = { - name: eigen_g[name].shape[1] * eigen_a[name].shape[1] - for name in eigen_a - } + grad_sizes = { + name: p * p if p > 0 else eigen_g[name].shape[1] * eigen_a[name].shape[1] + for name in eigen_a + } mmap = load_gradients(self.gradient_path) with open(os.path.join(self.gradient_path, "info.json")) as f: @@ -153,29 +150,16 @@ def compute_ivhp_sharded(self): gc.collect() if p > 0: + pt = self.cfg.projection_type for k, v in transformed_gradients.items(): - d_S = v.shape[-2] - d_A = v.shape[-1] - P_left = create_projection_matrix( - identifier=f"{k}/left", - m=p, - n=d_S, - dtype=v.dtype, - device=v.device, - projection_type=self.cfg.projection_type, - ) - P_right = create_projection_matrix( - identifier=f"{k}/right", - m=p, - n=d_A, - dtype=v.dtype, - device=v.device, - projection_type=self.cfg.projection_type, + d_S, d_A = v.shape[-2:] + P_l = create_projection_matrix( + f"{k}/left", p, d_S, v.dtype, v.device, pt ) - transformed_gradients[k] = torch.einsum( - "ps,nsa,ra->npr", P_left, v, P_right + P_r = create_projection_matrix( + f"{k}/right", p, d_A, v.dtype, v.device, pt ) - self.logger.debug(f"Finished P_S · (H^{{-1}} G) · P_Aᵀ at p={p}") + transformed_gradients[k] = torch.einsum("ps,nsa,ra->npr", P_l, v, P_r) torch.cuda.synchronize() for k, v in transformed_gradients.items(): From c3201ea71f6b64a9e04bcd8dfc952cdad3ba4928 Mon Sep 17 00:00:00 2001 From: Girish Gupta Date: Thu, 21 May 2026 11:50:29 -0700 Subject: [PATCH 10/10] Renames variable --- bergson/hessians/apply_hessian.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index f382b9ba..ea226851 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -158,14 +158,14 @@ def compute_ivhp_sharded(self): gc.collect() if p > 0: - pt = self.cfg.projection_type + projection_type = self.cfg.projection_type for k, v in transformed_gradients.items(): d_S, d_A = v.shape[-2:] P_l = create_projection_matrix( - f"{k}/left", p, d_S, v.dtype, v.device, pt + f"{k}/left", p, d_S, v.dtype, v.device, projection_type ) P_r = create_projection_matrix( - f"{k}/right", p, d_A, v.dtype, v.device, pt + f"{k}/right", p, d_A, v.dtype, v.device, projection_type ) transformed_gradients[k] = torch.einsum("ps,nsa,ra->npr", P_l, v, P_r)