Repository navigation
Add distributed alignment search for interchange interventions - #1893
Conversation
a625238 to
15dddc6
Compare
|
Hi @dipo101, please rebase this onto |
Add fit_das and LearnedSubspace to tools.analysis.interchange: DAS trains an orthonormal basis so that interchanging it makes the model produce the outputs a high-level causal model predicts for each base/source pair. The model is frozen during training and restored afterwards, and batches can be split to save memory. Includes unit tests on a planted direction, a GPT-2 test on held-out nouns, and a docs section. Part of TransformerLensOrg#1879.
15dddc6 to
73ffb4b
Compare
|
@jlarson4 Rebase done, ptal |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for the rebase! Just a handful of comments below, let me know if you have any questions
… device fit_das, collect_activations and intervened_logits now reject a model with any submodule in training mode, via require_eval_mode. fit_das seeds its initialisation inside torch.random.fork_rng, so the caller's random stream is unchanged, and creates the subspace on the model's device unless told otherwise. The splitting test now checks that no forward pass exceeds the micro-batch size, and the docs examples fit on some nouns and report interchange accuracy on held-out ones.
| model, hook_name, source_inputs, task_logits, positions=source_positions, batch_size=micro | ||
| ) | ||
| # Seed the initialisation without moving the caller's random stream | ||
| with torch.random.fork_rng(devices=[]): |
There was a problem hiding this comment.
torch.manual_seed also reseeds the CUDA and MPS generators, but fork_rng(devices=[]) only restores the CPU one. On a GPU, the caller's device random stream still restarts from seed after fit_das, and the new test only checks the CPU stream. _preserve_model_rng in backward_lens.py handles the same problem by also restoring the CUDA and MPS state for the model's devices. Lets make sure to follow that pattern
There was a problem hiding this comment.
Good point. torch.manual_seed reseeds every CUDA device as well as MPS and XPU, so restoring only the model’s devices would still reset streams on other GPUs. Instead, fit_das no longer calls torch.manual_seed. the subspace is built on the CPU, so it seeds only torch.default_generator, inside fork_rng(devices=[]), which restores it. No device generator is touched, and the same seed still gives the same basis.
I also have a new test fail if torch.manual_seed or any CUDA or MPS seeding function is called during fit_das, so this holds on GPUs and MPS even though CI is CPU-only. Happy to switch to _preserve_model_rng if you’d rather keep the patterns identical thoughj.
There was a problem hiding this comment.
Looks great as is! I am going to approve and merge this one. We will need a rebase on the final PR before I review, whenever you have a moment. Just ping me over there when it is ready. I appreciate the responsiveness!
|
Just one additional comment! |
torch.manual_seed reseeds every CUDA, MPS and XPU generator as well as the CPU one, and fork_rng(devices=[]) restores only the CPU one. The subspace is built on the CPU, so fit_das now seeds only the CPU generator, inside the fork that restores it. A test fails if torch.manual_seed or a CUDA or MPS seeding function is called during fit_das.
2 participants
Description
Part 2 of 3 for #1879. Stacked on #1892; only the last commit is new.
Adds
fit_dasandLearnedSubspace: distributed alignment search trains an orthonormal basis so that interchanging it makes the model produce the outputs a high-level causal model predicts for each base/source pair. Tested on a planted direction and on GPT-2 subject number with held-out nouns; docs section added.Type of change
Checklist: