diff --git a/bergson/hessians/apply_hessian.py b/bergson/hessians/apply_hessian.py index 18e8a363..0a844b9d 100644 --- a/bergson/hessians/apply_hessian.py +++ b/bergson/hessians/apply_hessian.py @@ -3,12 +3,14 @@ import warnings from dataclasses import dataclass from pathlib import Path +from typing import Literal import numpy as np import torch import torch.distributed as dist from simple_parsing import ArgumentParser +from bergson.collector.collector import create_projection_matrix from bergson.config import InversionConfig from bergson.data import column_offsets, create_index, load_gradients from bergson.distributed import init_dist @@ -28,6 +30,10 @@ class EkfacConfig: `HessianConfig.ev_correction=True`.""" apply_batch_size: int = 32 """Number of query gradients moved on-device and preconditioned at a time.""" + projection_dim: int = 0 + """When set, compress each module's IVHP output to a ``[p, p]`` Kronecker + random projection (``P_S @ (H^-1 G) @ P_A^T``).""" + projection_type: Literal["normal", "rademacher"] = "rademacher" debug: bool = False @@ -52,6 +58,13 @@ def __init__( if inversion_cfg is not None and apply_fn is not None: raise ValueError("Pass either inversion_cfg or apply_fn, not both.") + if cfg.projection_dim > 0 and cfg.ev_correction: + raise ValueError( + "projection_dim compression is not supported with EK-FAC " + "(ev_correction=True); set ev_correction=False or " + "projection_dim=0." + ) + self.cfg = cfg self.path = cfg.hessian_method_path self.gradient_path = cfg.gradient_path @@ -77,9 +90,18 @@ def compute_ivhp_sharded(self): ev_correction=self.cfg.ev_correction, ) - grad_sizes = { + o_dims = { name: preconditioner.eigen_g[name].shape[1] - * preconditioner.eigen_a[name].shape[1] + for name in preconditioner.eigen_a + } + i_dims = { + name: preconditioner.eigen_a[name].shape[1] + for name in preconditioner.eigen_a + } + + p = self.cfg.projection_dim + grad_sizes = { + name: p * p if p > 0 else o_dims[name] * i_dims[name] for name in preconditioner.eigen_a } @@ -126,6 +148,30 @@ def compute_ivhp_sharded(self): self.logger.debug("Finished H^{-1} G = Q_S @ (G' / lambda) @ Q_A^T batch") + if p > 0: + for name, flat in transformed.items(): + if name not in o_dims: + continue + g = flat.view(-1, o_dims[name], i_dims[name]) + P_l = create_projection_matrix( + f"{name}/left", + p, + o_dims[name], + g.dtype, + g.device, + self.cfg.projection_type, + ) + P_r = create_projection_matrix( + f"{name}/right", + p, + i_dims[name], + g.dtype, + g.device, + self.cfg.projection_type, + ) + transformed[name] = torch.einsum("ps,nsa,ra->npr", P_l, g, P_r) + self.logger.debug("Compressed IVHP output to [p, p] per module") + # Stage the async D2H copies, synchronize once, then read them # into the numpy buffer. `.numpy()` is a host read so the # sync must sit between it and the copies. diff --git a/bergson/hessians/pipeline.py b/bergson/hessians/pipeline.py index cba8ea38..c7addf46 100644 --- a/bergson/hessians/pipeline.py +++ b/bergson/hessians/pipeline.py @@ -56,6 +56,14 @@ def hessian_pipeline( 3. Apply the inverse Hessian to the mean query gradient. 4. Score each training example against the transformed query gradient. """ + if preprocess_cfg.unit_normalize: + raise ValueError( + "preprocess_cfg.unit_normalize (cosine similarity) is not " + "supported with Kronecker-factored Hessians; hessian_pipeline " + "only ever fits and applies a factored (kfac/tkfac/shampoo) " + "Hessian." + ) + run_path = index_cfg.run_path method = hessian_cfg.method query_path = f"{run_path}/query" @@ -110,6 +118,8 @@ def _validate(cfg: IndexConfig): gradient_path=query_path, run_path=transformed_query_path, ev_correction=hessian_cfg.ev_correction, + projection_dim=index_cfg.projection_dim, + projection_type=index_cfg.projection_type, ) launch_distributed_run( "apply_hessian", @@ -123,7 +133,6 @@ 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_cfg.query_path = transformed_query_path score_cfg.higher_is_better = True _validate(score_index_cfg) diff --git a/bergson/score/score.py b/bergson/score/score.py index ccbbbdf3..3b2dffb3 100644 --- a/bergson/score/score.py +++ b/bergson/score/score.py @@ -378,6 +378,16 @@ def score_dataset( and index_cfg.projection_dim != 0 and is_factored_hessian(preprocess_cfg.hessian_path) ): + if preprocess_cfg.unit_normalize: + raise ValueError( + f"Scoring with a factored (EKFAC) hessian at " + f"{preprocess_cfg.hessian_path} and unit_normalize=True (cosine " + f"similarity) requires projection_dim=0: split (two-sided " + f"H^-1/2) scoring needs full, unprojected index gradients, but " + f"index_cfg.projection_dim={index_cfg.projection_dim}. Rebuild " + f"the index and query with projection_dim=0, or use a dense " + f"(autocorrelation) hessian." + ) raise ValueError( f"Scoring with a factored (EKFAC) hessian at " f"{preprocess_cfg.hessian_path} requires projection_dim=0, but " diff --git a/tests/ekfac_tests/test_factored_preconditioner.py b/tests/ekfac_tests/test_factored_preconditioner.py index 244b6e58..66f7049e 100644 --- a/tests/ekfac_tests/test_factored_preconditioner.py +++ b/tests/ekfac_tests/test_factored_preconditioner.py @@ -16,6 +16,7 @@ import torch from safetensors.torch import load_file, save_file +from bergson.collector.collector import create_projection_matrix from bergson.config import InversionConfig from bergson.data import create_index, load_module_gradients from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig @@ -239,3 +240,73 @@ def test_factored_apply_invariant_to_shard_count(tmp_path, inversion: str): assert torch.allclose( out1[name], out2[name], atol=1e-5 ), f"{inversion}: apply on {name} depends on shard count" + + +def test_apply_hessian_compresses_per_module(tmp_path): + """``EkfacConfig.projection_dim > 0`` compresses the IVHP output to a + ``[p, p]`` per-module Kronecker random projection, applied *after* H^-1 + on the full gradient. It must use the exact same + ``create_projection_matrix`` convention (identifiers ``f"{name}/left"`` + / ``f"{name}/right"``) that ``bergson build`` uses on training + gradients, so a compressed query is directly comparable to a compressed + training-side gradient store. + """ + modules = {"a": (4, 6), "b": (5, 3)} # (O, I) + hessian_path = tmp_path / "hessian" + _write_factored_hessian(hessian_path, modules, num_shards=1, seed=0) + + grad_sizes = {name: o * i for name, (o, i) in modules.items()} + num_grads = 3 + query_path = str(tmp_path / "query") + _make_query_gradients(query_path, grad_sizes, num_grads) + + inversion_cfg = InversionConfig(inversion="damped_inverse", damping_factor=0.1) + + # Reference: full, uncompressed IVHP output. + ref = _apply( + str(hessian_path), + query_path, + str(tmp_path / "out_full"), + "damped_inverse", + ev_correction=False, + ) + + p = 2 + compressed_cfg = EkfacConfig( + hessian_method_path=str(hessian_path), + gradient_path=query_path, + run_path=str(tmp_path / "out_compressed"), + ev_correction=False, + projection_dim=p, + ) + EkfacApplicator(compressed_cfg, inversion_cfg=inversion_cfg).compute_ivhp_sharded() + got = load_module_gradients(str(tmp_path / "out_compressed")) + + for name, (o, i) in modules.items(): + full = torch.from_numpy(np.asarray(ref[name][:])).view(num_grads, o, i) + P_l = create_projection_matrix( + f"{name}/left", p, o, full.dtype, full.device, "rademacher" + ) + P_r = create_projection_matrix( + f"{name}/right", p, i, full.dtype, full.device, "rademacher" + ) + expected = torch.einsum("ps,nsa,ra->npr", P_l, full, P_r) + + compressed = torch.from_numpy(np.asarray(got[name][:])).view(num_grads, p, p) + assert compressed.shape == (num_grads, p, p) + assert torch.allclose( + compressed, expected, atol=1e-4, rtol=1e-4 + ), f"{name}: compressed IVHP output doesn't match manual post-projection" + + +def test_apply_hessian_rejects_compression_with_ev_correction(): + """projection_dim compression is not supported with EK-FAC (ev_correction).""" + cfg = EkfacConfig( + hessian_method_path="unused", + gradient_path="unused", + run_path="unused", + ev_correction=True, + projection_dim=8, + ) + with pytest.raises(ValueError, match="EK-FAC"): + EkfacApplicator(cfg, inversion_cfg=InversionConfig()) diff --git a/tests/ekfac_tests/test_hessian_pipeline.py b/tests/ekfac_tests/test_hessian_pipeline.py new file mode 100644 index 00000000..72b2e4c1 --- /dev/null +++ b/tests/ekfac_tests/test_hessian_pipeline.py @@ -0,0 +1,23 @@ +import pytest + +from bergson.config import ( + HessianConfig, + HessianPipelineConfig, + IndexConfig, + PreprocessConfig, + ScoreConfig, +) +from bergson.hessians.pipeline import hessian_pipeline + + +def test_hessian_pipeline_rejects_unit_normalize(tmp_path): + """Cosine similarity (unit_normalize) is not supported with the + Kronecker-factored Hessians hessian_pipeline fits and applies.""" + with pytest.raises(ValueError, match="unit_normalize"): + hessian_pipeline( + IndexConfig(run_path=str(tmp_path / "run")), + HessianConfig(method="kfac"), + ScoreConfig(), + PreprocessConfig(unit_normalize=True), + HessianPipelineConfig(), + ) diff --git a/tests/test_score.py b/tests/test_score.py index a85ba3fd..aded4094 100644 --- a/tests/test_score.py +++ b/tests/test_score.py @@ -690,3 +690,21 @@ def test_score_factored_hessian_rejects_projection(tmp_path: Path): ScoreConfig(query_path=str(tmp_path / "q")), PreprocessConfig(hessian_path=hessian_path), ) + + +def test_score_factored_hessian_rejects_projection_with_unit_normalize( + tmp_path: Path, +): + """Kronecker-factored hessian + projection + unit_normalize (cosine + similarity) together fail fast, with a message naming the actual cause.""" + modules = {"mod_a": (4, 6)} + hessian_path = str(tmp_path / "hessian") + _write_factored_hessian(Path(hessian_path), modules) + + index_cfg = IndexConfig(run_path=str(tmp_path / "out"), projection_dim=16) + with pytest.raises(ValueError, match="unit_normalize=True"): + score_dataset( + index_cfg, + ScoreConfig(query_path=str(tmp_path / "q")), + PreprocessConfig(hessian_path=hessian_path, unit_normalize=True), + )