Proposal

We currently have a TODO in load_model.py#L98 to support the stop_at_layer param in HookedProxyLM for huggingface models. Adding support for this will save compute when we're just trying to extract LLM activations at a specific layer, since there's no need to extract activations at later layers.

An example of how to do this is in baukit/nethook.py . We just need to throw an exception after we process the layer we care about to stop processing, and then catch that exception before returning to the user.

Checklist

• I have checked that there is no similar issue in the repo (required)