Skip to content

Support simple-stories lms #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
danbraunai-apollo merged 76 commits into from
Apr 22, 2025
Merged

Support simple-stories lms #3

danbraunai-apollo merged 76 commits into from
Apr 22, 2025

Conversation

@danbraunai-apollo
Copy link
Contributor

@danbraunai-apollo danbraunai-apollo commented Apr 22, 2025

Description

  • Allows for running APD on simple-stories models.
  • Only includes the loss functions necessary for our initial exploration, and not everything that is supported for tms and resid_mlp.
  • Adds a streamlit dashboard at spd/experiments/lm/app.py to explore the components for a trained model.

How Has This Been Tested?

None! Notably, no tests yet made for whether the tokens accurately match the text.

Does this PR introduce a breaking change?

No

@danbraunai-apollo
Rename some topk_mask vars to mask
b53df86
@danbraunai-apollo
Implement gating (untested)
0f4e7f8
@danbraunai-apollo
Fix grad attributions and calc_recon_mse
c784489
@danbraunai-apollo
Init gate with bias=1 and weights normal dist mean=0 std=0.2
e3c3eb0
@danbraunai-apollo
Fix lp sparsity loss
15b310c
@danbraunai-apollo
Add random mask loss
3aff69b
@danbraunai-apollo
Use relud masks for lp sparsity loss
13b8097
@danbraunai-apollo
Use masked_target_component_acts in calc_act_recon_mse
0923c0f
@danbraunai-apollo
Comment out grad attribution calculation so people don't use now
3aceb8a
@danbraunai-apollo
Store gates in model class
61247dc
@danbraunai-apollo
Remove buggy tms deprecated params replacement
64c3a23
@danbraunai-apollo
Tie the gates for TMS
ed32237
@danbraunai-apollo
Plot masks
60cc056
@danbraunai-apollo
Fix resid_mlp test (sensitive to float precision)
bc9505c
@danbraunai-apollo
Add init_from_target for tms
01a03bc
@danbraunai-apollo
Support init_from_target for resid_mlp
6d6d99f
@danbraunai-apollo
Normalise lp sparsity by batch size
c303c14
@danbraunai-apollo
Don't copy biases in init_spd_model_from_target_model
41bd85b
@danbraunai-apollo
Fix resid_mlp init_from_target test
befac1d
@danbraunai-apollo
Add randrecon to run label
e7e60a7
@danbraunai-apollo
Permute to identity for plotting mask_vals
3845ca3
@danbraunai-apollo
Remove post_relu_act_recon config arg
3bb654c
@danbraunai-apollo
Remove code from global scope in plotting
ebee911
@danbraunai-apollo
Handle deprecated 'post_relu_act_recon' arg.
0b3f61d
@danbraunai-apollo
Use mps if available
931b6f3
@danbraunai-apollo
Avoid mps as it breaks tms
19d7181
@danbraunai-apollo
Untie gates in TMS
8560f1b
@danbraunai-apollo
Allow for detached inputs to gates and use target_out in random_mask_…
79391e9 
…recon
@danbraunai-apollo
Add GateMLP
cd23609
@danbraunai-apollo
Remove bias_val and train_bias config args
96939c2
danbraunai-apollo and others added 29 commits March 17, 2025 16:00
@danbraunai-apollo
Load env vars when running sweeps too
5afdc92
@danbraunai-apollo
Add layerwise recon (#263)
e80f874 
* Add layerwise recon

* Add layerwise_random_recon_loss

* Protect the eyes of mathematicians
@danbraunai-apollo
Remove transformer-lens dependency
16992e5
@danbraunai-apollo
Use new random masks for layerwise_random_masks
7f6a94b
@danbraunai-apollo
Add jaxtyping to dependencies
5c632f9
@danbraunai-apollo
Add einops dependency
5981df6
@danbraunai-apollo
Use calc_recon_mse in calc_random_masks_mse_loss for consistency
fcff304
@danbraunai-apollo
Set bias to zero in GateMLP mlp_out
7ac2a42
@danbraunai-apollo
WIP: Swap components with Llama nn.Linear modules
037caf1
@danbraunai-apollo
Fix nn.Linear shape and handle masked components
1a1dcaf
@danbraunai-apollo
WIP: Add lm_decomposition script
993da44
@danbraunai-apollo
Fix module paths
fccc189
@danbraunai-apollo
WIP: Add param_match_loss
3fcf593
@danbraunai-apollo
Add layerwise recon losses
aa7cacf
@danbraunai-apollo
Add lp sparsity loss
82b505a
@danbraunai-apollo
Minor comment and config clean
96ae954
@danbraunai-apollo
Make components a submodule of SSModel and update model loading
cb12ed1
@danbraunai-apollo
Add SSModel.from_pretrained()
d3a7c76
@danbraunai-apollo
WIP: Fix download with weights_only=True
1425354
@danbraunai-apollo
Merge branch 'main' into feature/lm
8ba8ca9
@danbraunai-apollo
Calc mask l0 for lms
7a23520
@danbraunai-apollo
Merge branch 'main' into feature/lm
2706112
@danbraunai-apollo
Fix missing GateMLP type references
0103c0c
@danbraunai-apollo
Merge branch 'feature/lm' into feature/lm-temp
bcd3e09
@danbraunai-apollo
Update component_viz for new model format
60fa3cc
@danbraunai-apollo
Plot mean components during apd run
04bcbe1
@danbraunai-apollo
Re-organise wandb logging
c2bdda1
@danbraunai-apollo
Add streamlit dashboard for lm (#2)
072085e 
* WIP: Add dashboard

* Create base_cache_dir if it doesn't exist

* Functional dashboard

* Add simple-stories-train and datasets to pyproject.toml
@danbraunai-apollo
Remove unused set_nested_module_attr function
04a2138
@danbraunai-apollo danbraunai-apollo merged commit f47c5b4 into dev Apr 22, 2025
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@danbraunai-apollo