How we train saes replication

[jbloomAus]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

You have tested formatting, typing and unit tests (acceptance tests not currently in use)

• I have run make check-ci to check format and linting. (you can run make format to format code if needed.)

Performance Check.

If you have implemented a training change, please indicate precisely how performance changes with respect to the following metrics:

• L0
• CE Loss
• MSE Loss

Please links to wandb dashboards with a control and test group: https://api.wandb.ai/links/jbloom/syh38mi2

[codecov]

Codecov Report

Attention: Patch coverage is 84.39306% with 27 lines in your changes are missing coverage. Please review.

Project coverage is 64.40%. Comparing base (4d12e7a) to head (eb35e67).
Report is 1 commits behind head on main.

Files	Patch %	Lines
sae_lens/training/cache_activations_runner.py	80.76%	7 Missing and 3 partials ⚠️
sae_lens/training/activations_store.py	71.87%	8 Missing and 1 partial ⚠️
sae_lens/training/train_sae_on_language_model.py	75.00%	1 Missing and 2 partials ⚠️
sae_lens/training/config.py	88.23%	1 Missing and 1 partial ⚠️
sae_lens/training/evals.py	0.00%	2 Missing ⚠️
sae_lens/training/optim.py	95.65%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #123      +/-   ##
==========================================
+ Coverage   59.91%   64.40%   +4.48%     
==========================================
  Files          17       17              
  Lines        1654     1753      +99     
  Branches      277      289      +12     
==========================================
+ Hits          991     1129     +138     
+ Misses        602      560      -42     
- Partials       61       64       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.