Add run_with_cache_until and offload_params_after method to transformer_lens.hook_point.HookedRootModule to do early stopping & remove unnecessary parameters.
Replace run_with_cache in activation caching with run_with_cache_until to improve speed of activation generation.
This PR includes:
run_with_cache_untilandoffload_params_aftermethod totransformer_lens.hook_point.HookedRootModuleto do early stopping & remove unnecessary parameters.run_with_cachein activation caching withrun_with_cache_untilto improve speed of activation generation.