isolate_submodel

isolate_submodel#

Utility for capturing and analyzing a small subcomputation in a larger model.

This utility is designed to enable quickly exploring the behavior of small parts of a large model, and allows reproducing the in-context behavior of the subpart without actually having to run the full model. It works by capturing the intermediate activations immediately before and after the selected submodel, and also isolating any state or shared parameters used by it. This makes it possible to call the submodel in isolation on the original intermediate activation before it, and compare its output to the original intermediate activation after it.

Classes

IsolatedSubmodel

An isolated part of a submodel, with its saved inputs and outputs.

IsolationCapturer

Helper object that captures information necessary to isolate the submodel.

Functions

call_and_extract_submodel(...)

Calls a model with an argument, and captures the selected submodel.