Reduce memory use when loading model (#327)

* No grad when extracting weights * Remove circular reference
TransformerLensOrg · Jun 21, 2023 · 3cd9436 · 3cd9436
1 parent 4f3e143
commit 3cd9436
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/transformer_lens/hook_points.py b/transformer_lens/hook_points.py
@@ -167,6 +167,8 @@ def setup(self):
         self.mod_dict = {}
         self.hook_dict: Dict[str, HookPoint] = {}
         for name, module in self.named_modules():
+            if name == "":
+                continue
             module.name = name
             self.mod_dict[name] = module
             if "HookPoint" in str(type(module)):

diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py
@@ -940,6 +940,10 @@ def get_pretrained_state_dict(
                 )
 
             # Load model weights, and fold in layer norm weights
+
+        for param in hf_model.parameters():
+            param.requires_grad = False
+
         if cfg.original_architecture == "GPT2LMHeadModel":
             state_dict = convert_gpt2_weights(hf_model, cfg)
         elif cfg.original_architecture == "GPTNeoForCausalLM":