See this notebook for basic usage of the Edit module.
Standardization
The Alterations module contains simple bespoke name maps for different models.
You can pass unified=True to any model that inherits from NNsight model. This will load a dictionary attr_map into the model which is passed into ._envoy.
The Envoy maps standardized names onto bespoke model attributes, and represents custom names in a defined __repr__ and __str__ method.
model = LanguageModel("openai-community/gpt2", device_map="cuda:0", unified=True)
Module Edits
The Edit module is designed for dynamically editing models during runtime. Under the hood, it uses TorchDynamo to recompile PyTorch code into Torch FX graphs.
Edit Module
The Edit module contains three classes: Edit, Compiler, and Editor.
Edit
The Edit object defines some edit a user wishes to make on the compiled FX graph.
Compiler
The compiler is responsible for grouping and optimizing edits. A couple important methods:
group_by_hierarchy takes all the edits passed into the model. It groups them into batches of edits on the same modules. These batches are ordered so that lower level edits (...h.0.attn) are made before higher level edits (...transformer).
get_backend takes a batch of edits and creates a custom backend for that specific module. It updates edit_tracker to keep track of which target instance is being edited.
Editor
The Editor is a simple context manager which is entered if edits are passed into a LanguageModel.
Utils and Relevant Changes
A fold method was added to envoy which folds in _orig_mod for easier attribute access.
The Edit module contains various utils for printing and inspecting TorchDynamo compiled modules.
TODO
[ ] Figure out another name besides "unified"
[ ] Add docstrings to methods
[ ] Add color coded legend to print_gm
[ ] Add options to print more detail when compiling graphs for users (?)
[ ] Could use organization/style suggestions from @JadenFiotto-Kaufman
Module Naming Standards and Module Edits
See this notebook for basic usage of the Edit module.
Standardization
The
Alterationsmodule contains simple bespoke name maps for different models.You can pass
unified=Trueto any model that inherits from NNsight model. This will load a dictionaryattr_mapinto the model which is passed into._envoy.The
Envoymaps standardized names onto bespoke model attributes, and represents custom names in a defined__repr__and__str__method.Module Edits
The Edit module is designed for dynamically editing models during runtime. Under the hood, it uses TorchDynamo to recompile PyTorch code into Torch FX graphs.
Edit Module
The Edit module contains three classes:
Edit,Compiler, andEditor.EditTheEditobject defines some edit a user wishes to make on the compiled FX graph.CompilerThe compiler is responsible for grouping and optimizing edits. A couple important methods:group_by_hierarchytakes all the edits passed into the model. It groups them into batches of edits on the same modules. These batches are ordered so that lower level edits (...h.0.attn) are made before higher level edits (...transformer).get_backendtakes a batch of edits and creates a custom backend for that specific module. It updatesedit_trackerto keep track of which target instance is being edited.EditorTheEditoris a simple context manager which is entered if edits are passed into a LanguageModel.Utils and Relevant Changes
A
foldmethod was added to envoy which folds in_orig_modfor easier attribute access.The Edit module contains various utils for printing and inspecting TorchDynamo compiled modules.
TODO
print_gm