chore: more flexible training losses

[chanind]

Description

This PR refactors the TrainStepOutput in TrainingSAE.training_forward_pass() to allow passing losses as a dict rather than a hardcoded set of specific losses. The current implementation requires always providing values for ghost grad loss and auxiliary loss, even when it doesn't make sense for the training method. As SAELens supports more architectures, there will likely be more losses like this that only make sense for certain architectures, so it makes sense to let the train forward step define the losses that make sense for the circumstances.

This should also make it easier for researchers to hack on the SAELens SAEs since it means they can define any losses they want to track.

This PR also allows the losses to remain as tensors rather than being floats. The current implementation requires calling .item() on each loss, which forces a synchronization with the GPU and likely slows down training slightly. We only ever use the losses for logging, which happens by default once every 100 steps, so this synchronization is unnecessary.

Type of change

Please delete options that are not relevant.

• Refactor
• 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.)

[codecov]

Codecov Report

Attention: Patch coverage is 90.90909% with 2 lines in your changes missing coverage. Please review.

Project coverage is 66.29%. Comparing base (0b56d03) to head (97b03f5).

Files with missing lines	Patch %	Lines
sae_lens/training/sae_trainer.py	80.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #357      +/-   ##
==========================================
+ Coverage   66.17%   66.29%   +0.11%     
==========================================
  Files          25       25              
  Lines        3367     3370       +3     
  Branches      429      429              
==========================================
+ Hits         2228     2234       +6     
+ Misses       1020     1018       -2     
+ Partials      119      118       -1     

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