Conversation
The backward pass builds `grad_values[:, :, None] * input[:, None, :]`, an `[N, k, D]` tensor holding every top-k contribution at once. It is the largest single allocation in SAE training: 288 MiB at batch 8, ctx 512, d_model 576, k 32, and 1 GiB in the shape `tests/test_encode.py` already exercises. It scales with k, so the `--k 192` configuration the README recommends pays six times that. Consume it in row-blocks sized to a fixed byte budget instead. The arithmetic is unchanged; only the accumulation order inside `index_add_` differs, and that was already unspecified on CUDA. Peak allocated per rank on SmolLM2-135M with 30 hookpoints at 2 GPUs: 9.19 -> 9.04 GiB under DDP, 4.24 -> 3.86 GiB under --distribute_modules. Step time is unchanged within run-to-run variance.
|
|
Author
|
recheck |
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.
What
FusedEncoder.backwardbuildsgrad_values[:, :, None] * input[:, None, :], an[N, k, D]tensor holding every top-k contribution at once, and then folds it intograd_weightwith a singleindex_add_. This consumes it in row-blocks sized to a fixed byte budget instead.Why
It is the largest single allocation in SAE training. Some sizes, all measured on an A100:
d_model576, k 32 (repo defaults)tests/test_encode.pyalready runs)It is linear in
k, so the--k 192configuration the README recommends for Llama 3 8B pays six times the first row.Effect
Peak allocated per rank, SmolLM2-135M, 30 hookpoints, 2x A100-40GB, batch 8:
--distribute_modulesMedian step time is unchanged within run-to-run variance, which I measured at about 25% on this setup, so I would not read a throughput difference either way from these runs.
Correctness
The arithmetic is unchanged. Only the accumulation order inside
index_add_differs, and that was already unspecified on CUDA.tests/test_encode.pycompares the encoder gradient against a naive autograd reference at N=8192, D=1024, k=32 and passes.Full suite on an A100, this branch against
main: same results, with the one pre-existingtest_auxk_lossfailure on both (that one is #142).