perf: support early-stopping in HF models #430
Conversation
| stop_hook = None | ||
| if stop_at_layer is not None: | ||
| if return_type != "logits": | ||
| raise NotImplementedError( |
There was a problem hiding this comment.
The fact that it's a NotImplementedError instead of eg ValueError or RuntimeError implies that it can be supported in the future -- is that true?
If it's true what is it -- returning (None, None) instead?
There was a problem hiding this comment.
Yeah there's nothing keeping this from being implemented in the future - it's just that SAELens currently doesn't use more of TLens API, so there's no point in implementing support for more options here beyond what SAELens actually uses. This class is just a wrapper around Huggingface models to make them look like a TLens model for SAELens to work wtih.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #430 +/- ##
==========================================
+ Coverage 74.44% 74.55% +0.10%
==========================================
Files 19 19
Lines 3146 3183 +37
Branches 456 464 +8
==========================================
+ Hits 2342 2373 +31
- Misses 648 650 +2
- Partials 156 160 +4 ☔ View full report in Codecov by Sentry. |
| @@ -72,6 +75,7 @@ def __init__(self, model: torch.nn.Module, tokenizer: PreTrainedTokenizerBase): | |||
| super().__init__() | |||
| self.model = model | |||
| self.tokenizer = tokenizer | |||
| self.decoder_block_matcher = guess_decoder_block_matcher(model) | |||
There was a problem hiding this comment.
Is effectively getting the layer names from _guess_block_matcher_from_layers() the best we can do? I'm afraid it's error-prone and implicit.
Description
This PR supports early stopping in the huggingface wrapper. This is a perf improvement, so we don't need to run the whole model when all we want to do is extract some activations from an intermediate layer.
Fixes #429
Type of change
Please delete options that are not relevant.
Checklist:
You have tested formatting, typing and tests
make check-cito check format and linting. (you can runmake formatto format code if needed.)