Repository navigation
Conversation
…reaches The end-of-backward reconciliation (drop forward captures whose backward never came, re-evaluate completion) was queued only from a hooked layer's own backward hook, and only when a hooked layer had a trainable parameter. With frozen layers (include_frozen=True) it therefore never ran: under non-reentrant gradient checkpointing the forward fires twice and the backward once, so the step never completed and the buffers grew step over step; a step whose hooked layers the backward never reached stalled the same way, silently. - A forward hook on the model root arms a tensor hook on the output, which queues the end-of-backward callback from the outermost backward task whatever the hooked layers' parameters. - At the end of the backward a step in which no hooked layer received a gradient raises, naming the layers and the reason (the loss must depend on their output through a parameter that requires grad; frozen layers need a trainable parameter upstream, typically the input embedding). Layers left out of an otherwise complete step are reported once per layer set. The streamer's own "no gradient captured" message is reduced to a guard for a loss_fn that never runs the model. - Registration warns when include_frozen=True hooks layers without a trainable parameter (and what that requires of the model), and when include_frozen=False left nothing hooked because every candidate was frozen. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Hooks store a layer's factors raw, and every product kernel preprocesses them on each call (the bias column of a linear layer, the normalized activation of a norm layer, the unfolded patches of a convolution). A query layer that stayed factorized was therefore re-preprocessed for every train block, and a train block once per test block it met. finalize_factors (attribution/utils.py) replaces each raw factorized layer by its final factors (Factorized with module_kwargs=None), releasing the raw payload as it goes. _route_test_rep applies it to the query layers that stay factorized, under the same cache budget as the dense copies, and the base score_sources applies it to every transformed train block, so every inner-product attributor inherits it. Embedding layers stay raw: their materialization needs module_kwargs. ops.effective_dims reads final factors directly instead of misreading a finalized convolution. Scores are unchanged (exact); at 32-token sequences the scoring loop is about 5% faster. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
cross_gram's factorized route issued three einsums per layer (two 4-D token grams and their contraction) and the materialized route two. For short sequences the kernels take microseconds and an einsum's equation parsing and view building dominated: scoring a block of 96 layers cost 22 ms against 8 ms for the same math as two matrix products, a multiply and a sum. The 4-D einsums were also slow on the device: at 2048 tokens the factorized route drops from 433 ms to 167 ms per block. Both routes are now plain matrix products. Results are unchanged, the two sides may still differ in token count, and the measured crossover between the routes stays where the FLOP rule puts it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
With a live source, asking for the next block runs the next forward and backward pass, which is where the memory peak sits. A finished block was still resident then, up to three times over: * rebatch_blocks ran torch.cat over the pending tensors of every dense block even when there was one (the usual case, block size == batch size) -- a full copy; * it kept the whole previous block as its metadata template, and the last layer through a loop variable; * the scoring loop still held the finished batch when it requested the next one. A single-block batch now reuses the block's tensors, only the metadata is remembered, and the references are dropped before the next request. At 2048 tokens and batch 4 the peak falls from 29.2 GB to 20.2 GB (the forward/backward floor), at 512 tokens from 10.8 GB to 6.8 GB; scores are unchanged and the step is a few percent faster where blocks are dense. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
At full dimension the fitted K-FAC / EK-FAC factors are as large as the squared layer widths, and only one layer's factors are needed at a time. KroneckerAttributor(factor_cache_residency=...) holds them in a TensorCache the attributor owns: "memory" (host memory), "disk" (files), or "tiered" (host memory that spills to files once its budget, a share of the free host memory, is used). None, the default, keeps them on the device. The fit, the damping and a persisted fit loaded from fisher_dir place each layer's factors in the cache; damping and scoring bring one layer to the device at a time. A tensor the damping leaves unchanged (EK-FAC's eigenbases) is not copied again, EK-FAC's intermediate eigenbases use the same cache, and a cache is closed once a newer fit has replaced it, which deletes any spilled files. The Kronecker accumulators hand their buffers over on result(consume=True) instead of returning a normalized copy, so the fit never holds the covariances twice. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
K-FAC and EK-FAC attributors take covariances_at_capture (default True). When the attributor collects the train gradients into a store -- cache(), or attribute() with a gradient_cache_residency -- it attaches a KroneckerCovarianceCallback to that pass. The covariances are reduced across ranks, kept for the store and, for a disk store, saved next to the train gradients (kronecker_covariances.pt); prepare_scoring and fit use them instead of sweeping the store, as long as the source reads every stored step. A materialized "logra" store can therefore be fitted and scored through the one-call workflows. With the option off the fit is swept from the stored factors as before. The callback's module docstring now describes this use and states that the covariances follow the capture's projection. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
attribute() with a gradient_cache_residency and cache() sharded the test set across the ranks, so under DDP/FSDP each rank's store held only its own queries and the assembled score matrix was missing the columns of the other ranks. The store path now captures the whole test set on every rank, as the live path does. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The per-run options of the K-FAC family are arguments of attribute, attribute_from_cache, fit, cache and cache_preconditioned_test, handled by _set_options, and no longer of the constructor: covariances_at_capture, factor_cache_residency and the new relative_damping, which reads damping as a multiple of the mean eigenvalue of whatever it is added to (each covariance factor, the corrected EK-FAC spectrum, each dense Fisher block). fisher_dir is accepted by the live attribute() too. A directory that holds a fit is loaded; an empty one receives the fit of the call, so later calls over the same training set skip the pre-pass. The file records what defines the fit (algorithm, mode, non_kfac_strategy, direct_fim_max_params, layer_name, selected_training_steps) and a mismatch raises; damping and relative_damping are recorded but applied after loading, so one fit serves any damping. The attributor keeps the fit it last saved to or loaded from a directory, and a call naming the same unchanged file does not read it again. The file is written under a per-rank name and renamed into place. The warning of the dense-Fisher fallback states how many of the K-FAC-eligible layers it covers. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
OptimizerSnapshot reads a layer's optimizer state under FullyShardedDataParallel (use_orig_params=True) and fully_shard. Under the wrapper each rank lays its slice of a state tensor out on zeros and the slices are summed across ranks, one all_reduce per layer read; under fully_shard a DTensor state is gathered with full_tensor(). A layer is bound to the parameters the optimizer holds from inside its backward, where the wrapper exposes the flat parameter and fully_shard the unsharded view; the HookManager's preconditioner and the OptimizerStateCallback bind there. Only use_orig_params=False is refused. OptimizerStateCallback now copies the moments it records: a dense read of a bias-free layer on CPU was a view of the optimizer's tensor, so the recorded pre-step moments followed the next update. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
fully_shard without a mesh builds its device mesh on the default accelerator, so with a GPU visible the test's collectives ran on CUDA against a gloo process group and the worker died with SIGSEGV. The test passes an explicit CPU mesh. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ut and routing benchmarks Fidelity: the setting runs on GPT-2, Qwen2.5 (0.5B, 1.5B, 3B) and, under OLMo-core's trainer, OLMo-2-1B (FSDP over four GPUs) and OLMo-3-7B (one GPU); masked AdamW-influence at 512, 2048 and 8192 coordinates per layer; the scripts split into utils/truth (the ground truth) and utils/attribution (one script per method), over the shared setting, protocol and helpers in utils/. Benchmark: the throughput ladder on four GPUs at each library's largest batch, the routing sweep across sequence lengths, timing of Bergson's worker start-up, attribution-only timing for every table. The ordinary-versus-invasive capture benchmark moves to experiments/capture, with its own adapter, over the runner, data and logger of experiments/benchmark. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
HookManager(sample_hash_fields=[...]) restricts a sample's content hash to the named forward arguments (a kwarg by name, a positional argument by index), an alternative to sample_id_key. A trainer that passes per-sample inputs which are not part of a sample's identity, such as a label mask or the sample's position in the dataset, can then hash the same fields as a query loop that passes only the labels, so one identity serves both stores. A named field missing from a step raises; the scheme stays the content hash, so records are stamped sample_id_key=None as before. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
A batch is padded to its longest sequence, so the padded tokens, the attention mask and the -100 label fill carried the batch's length into a sample's identity: the same sample hashed differently in two batches. With an attention_mask in the inputs, hash_sample keeps only the unmasked positions of every field of the mask's length (left or right padding alike) and leaves the mask, all ones after that, out of the digest. A sample's hash is then that of its unpadded tokens and equals the hash of the raw dataset item. Without a mask nothing is stripped; a field of another length, such as a shifted labels row, is kept whole. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
DataSelectionCallback takes scoring_kwargs and selection_kwargs; the example passed threshold, threshold_mode and score_mode flat and raised TypeError. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ttings make check runs ruff's lint and format checks over the library, the tests and the examples; make test runs the suite on one GPU; make darglint checks docstrings against signatures. pytest.ini declares the gpu marker the tests use; .darglint holds darglint's google-style settings. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
A non-blocking copy from the device to pageable host memory returns before the data has arrived, so a test block on the CPU read stale factors. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…transformers, not dattri
…each method's capture requirements
… memory store budget crashes
… altering training
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.