Add torch compile

[tomMcGrath]

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 (see W&B dash)
• 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:

This change accelerates training by enabling torch.compile (and providing compilation options) for LM and SAE as config flags. This substantially increases throughput over uncompiled models.

W&B dashboard for performance comparison: https://wandb.ai/tmcgrath/compilation-checking/workspace

[tomMcGrath]

ah apologies I inherited a couple of commits from the earlier CL - will pull these out after that PR is resolved

[codecov]

Codecov Report

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

Project coverage is 64.43%. Comparing base (1a0619c) to head (8bf01f0).
Report is 1 commits behind head on main.

Files	Patch %	Lines
sae_lens/training/lm_runner.py	12.50%	7 Missing ⚠️
sae_lens/training/evals.py	0.00%	2 Missing ⚠️
sae_lens/training/activations_store.py	66.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #129      +/-   ##
==========================================
- Coverage   64.58%   64.43%   -0.16%     
==========================================
  Files          17       17              
  Lines        1762     1777      +15     
  Branches      291      295       +4     
==========================================
+ Hits         1138     1145       +7     
- Misses        560      567       +7     
- Partials       64       65       +1     

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