Command Line Interface

Bergson’s post-hoc attribution exposes three building block commands — build, reduce, and score — plus supporting commands for querying, preconditioner computation, and end-to-end pipelines.

build and query are designed for working with compressed gradients stored on disk and queried multiple times.

reduce and score are designed for working with uncompressed gradients primarily on GPUs, with a single predetermined query set. Use reduce to accumulate a dataset into a single query gradient (mean or sum), and score to map over an arbitrarily large dataset, computing the gradient of each item and scoring it against precomputed query gradients.

hessian and preconditioners compute preconditioner statistics independently of gradient collection. trackstar runs preconditioners, build, and score as a single pipeline (see Trackstar).

bergson {build,query,preconditioners,reduce,score,hessian,trackstar} [OPTIONS]

class bergson.__main__.Build(index_cfg, preprocess_cfg)

Bases: Serializable

Build a gradient index.

decode_into_subclasses: ClassVar[bool] = False

execute()

Build the gradient index.

index_cfg: IndexConfig

preprocess_cfg: PreprocessConfig

Example:

bergson build runs/my-index \
    --model EleutherAI/pythia-14m \
    --dataset NeelNanda/pile-10k \
    --truncation

class bergson.__main__.Query(index='', model='', text_field='text', unit_norm=True, device_map_auto=False, faiss=False, top_k=5, record='')

Bases: QueryConfig

Query an existing gradient index.

decode_into_subclasses: ClassVar[bool] = False

execute()

Query an existing gradient index.

Example:

bergson query \
    --index runs/my-index

class bergson.__main__.Reduce(index_cfg, preprocess_cfg)

Bases: Serializable

Reduce a gradient index.

decode_into_subclasses: ClassVar[bool] = False

execute()

Reduce a gradient index.

index_cfg: IndexConfig

preprocess_cfg: PreprocessConfig

Example:

bergson reduce runs/my-index \
    --model EleutherAI/pythia-14m \
    --dataset NeelNanda/pile-10k \
    --truncation \
    --aggregation mean \
    --unit_normalize \
    --projection_dim 0 \
    --skip_preconditioners

class bergson.__main__.Score(score_cfg, index_cfg, preprocess_cfg)

Bases: Serializable

Score a dataset against an existing gradient index.

decode_into_subclasses: ClassVar[bool] = False

execute()

Score a dataset against an existing gradient index.

index_cfg: IndexConfig

preprocess_cfg: PreprocessConfig

score_cfg: ScoreConfig

Example:

bergson score runs/my-scores \
    --model EleutherAI/pythia-14m \
    --dataset NeelNanda/pile-10k \
    --truncation \
    --query_path runs/my-index \
    --projection_dim 16

class bergson.__main__.Hessian(hessian_cfg, index_cfg)

Bases: Serializable

Approximate Hessian matrices using KFAC or EKFAC.

decode_into_subclasses: ClassVar[bool] = False

execute()

Compute Hessian approximation.

hessian_cfg: HessianConfig

index_cfg: IndexConfig

Example:

bergson hessian runs/my-hessian \
    --model EleutherAI/pythia-14m \
    --dataset NeelNanda/pile-10k \
    --truncation \
    --method kfac

class bergson.__main__.Preconditioners(run_path, model='EleutherAI/pythia-160m', precision='fp32', revision=None, distributed=<factory>, fsdp=False, peft_init_kwargs='', model_kwargs='', data=<factory>, tokenizer='', drop_columns=True, max_tokens=None, overwrite=False, use_tf32_matmuls=False, debug=False, projection_dim=16, include_bias=False, reshape_to_square=False, projection_type='rademacher', token_batch_size=2048, auto_batch_size=False, processor_path='', optimizer_state_path='', skip_preconditioners=False, skip_index=False, stats_sample_size=10000, loss_fn='ce', loss_reduction='sum', label_smoothing=0.0, stream_shard_size=400000, split_attention_modules=<factory>, attention=<factory>, profile=False, filter_modules=None, force_math_sdp=False, attribute_tokens=False, modules=<factory>)

Bases: IndexConfig

Compute normalizers and preconditioners without gradient collection.

decode_into_subclasses: ClassVar[bool] = False

execute()

Compute normalizers and preconditioners.

Example:

bergson preconditioners runs/my-processor \
    --model EleutherAI/pythia-14m \
    --dataset NeelNanda/pile-10k \
    --truncation \
    --normalizer adam

class bergson.__main__.Trackstar(index_cfg, trackstar_cfg)

Bases: Serializable

Run preconditioners, build, and score as a single pipeline.

decode_into_subclasses: ClassVar[bool] = False

execute()

index_cfg: IndexConfig

trackstar_cfg: TrackstarConfig

Example:

bergson trackstar runs/my-trackstar \
    --model EleutherAI/pythia-14m \
    --data.dataset NeelNanda/pile-10k \
    --data.truncation \
    --query.dataset NeelNanda/pile-10k \
    --query.truncation \
    --projection_dim 16