jeli04/hubble-ekfac-preconditioners

EK-FAC preconditioners for the Hubble perturbed models, fit on DCLM

EK-FAC approximations of the Hessian (eigenvalue-corrected K-FAC) of the Hubble 500B-token perturbed models, fit on text from DCLM-baseline, Hubble's pretraining distribution. They are stored undamped, in bergson 2.4.0's on-disk format.

folder	model	modules	DCLM tokens	labels	size
8b-500b.mlp_down_proj.ekfac-dclm10M	allegrolab/hubble-8b-500b_toks-perturbed-hf	mlp.down_proj of each block (36)	9,994,240	sampled from the model (true Fisher)	76 GB
1b-500b.all.ekfac-dclm1M	allegrolab/hubble-1b-500b_toks-perturbed-hf	every linear module of the blocks: q, k, v, o, gate, up, down (112)	999,424	sampled from the model (true Fisher)	36 GB

The 1b preconditioner is fit on 1M DCLM tokens, not 10M.

How they were fit

• Data: the first documents of mlfoundations/dclm-baseline-1.0-parquet, packed as bergson packs them (EOS separators) into 2048-token sequences, Hubble's training length: 4,880 sequences for the 8b fit and 488 for the 1b fit, each cut to a multiple of the GPU count.
• bergson.approximate_hessians with method="kfac", ev_correction=True: the K-FAC covariances, their eigendecomposition, then a second pass over the same sequences for the eigenvalue correction. Model in bf16, token batch of 2048 (one sequence), 4 GPUs.
• Never covered: embeddings, norms and lm_head. A preconditioner can only be applied to gradients of modules it covers.

Layout

Each folder holds an ekfac/ directory as bergson writes it: activation_sharded/, gradient_sharded/ (K-FAC covariances), eigen_activation_sharded/, eigen_gradient_sharded/ (their eigenvectors), eigenvalue_sharded/, eigenvalue_correction_sharded/, factor_eig_a/, factor_eig_g/, and total_processed.pt. Each *_sharded/ directory holds shard_<i>.safetensors files whose rows concatenate (in shard order) to each module's full tensor, keyed by module name (e.g. layers.0.mlp.down_proj).

Loading

Load a folder's ekfac/ with bergson's load_preconditioner. Damping is applied at load time, not stored; we damp each module by a factor times its own mean eigenvalue.

The 8b factors are about 38 GB for down_proj alone, which is more than a 48 GB GPU holds once bergson concatenates the shards. To apply them module by module, read one module's rows from each shard:

from pathlib import Path
import torch
from safetensors import safe_open

def module_rows(ekfac_dir, subdir, name):
    shards = sorted(Path(ekfac_dir, subdir).glob("shard_*.safetensors"),
                    key=lambda shard: int(shard.stem.split("_")[1]))
    parts = []
    for shard in shards:
        with safe_open(shard, framework="pt") as handle:
            parts.append(handle.get_tensor(name))
    return torch.cat(parts).float()

eigvecs_a = module_rows("ekfac", "eigen_activation_sharded", "layers.0.mlp.down_proj")

Download one folder only with, e.g.:

hf download jeli04/hubble-ekfac-preconditioners --include "8b-500b.mlp_down_proj.ekfac-dclm10M/*" --local-dir .