Refactor training

[jbloomAus]

Description

Mostly refactoring work with no functional changes:

• some cleaning up of previous regressions (not saving as much in our checkpoints)
• moving towards better abstractions / boundaries but not finished (more changes coming, but likely won't touch the high level api again this week).
• updated docs to clean up a little bit and remove stale code. Still work to do.

[codecov]

Codecov Report

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

Project coverage is 66.96%. Comparing base (cf94845) to head (2282f5b).

Files	Patch %	Lines
sae_lens/training/train_sae_on_language_model.py	72.72%	29 Missing and 10 partials ⚠️
sae_lens/training/lm_runner.py	28.57%	35 Missing ⚠️
sae_lens/training/checkpointing.py	20.00%	20 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #158      +/-   ##
==========================================
- Coverage   67.31%   66.96%   -0.36%     
==========================================
  Files          18       19       +1     
  Lines        1802     1704      -98     
  Branches      284      266      -18     
==========================================
- Hits         1213     1141      -72     
+ Misses        523      506      -17     
+ Partials       66       57       -9     

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