Repository: TransformerLensOrg/TransformerLens
Issue category: Performance Issue

Issue:
[Bug Report] `refactor_factored_attn_matrices` uses lots of GPU memory
**Describe the bug**
When loading models with `refactor_factored_attn_matrices=True` and `device="cuda"`, they take up a lot of space on GPU. Thanks to @callummcdougall for noticing this. See example code at bottom for snippets and statistics.

**System Info**
vast.ai 2*3090 devbox, on branch main.

**Additional context**
Doing `import gc; gc.collect(); import torch; torch.cuda.empty_cache()` after the refactor stuff would help but if someone could look into further whether we can keep the permanent refactor objects on GPU (or not) that would be great and save some more trouble! If nothing else I can add a warning if we enter the refactor function with device equal to GPU.

### Checklist

- [ :white_check_mark:  ] I have checked that there is no similar [issue](https://github.com/neelnanda-io/TransformerLens/issues) in the repo (**required**)

Example code

On my machine with 3090s, this uses 46% of GPU memory:

```
model = HookedTransformer.from_pretrained(
    "solu-10l", # "NeelNanda/SoLU_10L1280W_C4_Code"
    device="cuda",    
    refactor_factored_attn_matrices=True,
)
```

I check memory with 

```
_r = torch.cuda.memory_reserved(0)
_a = torch.cuda.memory_allocated(0)
_p = 100*(_r-_a)/_r 
print(f"Percentage of reserved memory that is allocated: {_p:.2f}%")
```

After `import gc; gc.collect(); import torch; torch.cuda.empty_cache()` this is down to 6%

However, with this code

```
model = HookedTransformer.from_pretrained(
    "solu-10l", # "NeelNanda/SoLU_10L1280W_C4_Code"
    device="cpu",
    refactor_factored_attn_matrices=True,
)
model.to("cuda")
```

just 0.16% of my memory is used


Hints:
Is `refactor_factored_attn_matrices` really the issue here? As far as I can tell, using `from_pretrained` to load models to the GPU directly always takes up much more GPU memory than first loading them to the CPU and then moving them to the GPU. When I use your snippet, setting `refactor_factored_attn_matrices=True` only increases the total memory usage by a small amount. 

Patch:
diff --git a/transformer_lens/HookedEncoder.py b/transformer_lens/HookedEncoder.py
index e470731d3..254a72a10 100644
--- a/transformer_lens/HookedEncoder.py
+++ b/transformer_lens/HookedEncoder.py
@@ -229,7 +229,7 @@ def from_pretrained(
             official_model_name, cfg, hf_model, **from_pretrained_kwargs
         )
 
-        model = cls(cfg, tokenizer, move_to_device)
+        model = cls(cfg, tokenizer, move_to_device=False)
 
         dtype = from_pretrained_kwargs.get("torch_dtype", None)
         if dtype is not None:
@@ -237,6 +237,9 @@ def from_pretrained(
 
         model.load_state_dict(state_dict, strict=False)
 
+        if move_to_device:
+            model.to(cfg.device)
+
         print(f"Loaded pretrained model {model_name} into HookedTransformer")
 
         return model
diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py
index fb9540853..581068fc3 100644
--- a/transformer_lens/HookedTransformer.py
+++ b/transformer_lens/HookedTransformer.py
@@ -84,10 +84,6 @@ def __init__(
             )
         self.cfg = cfg
 
-        assert (
-            self.cfg.n_devices == 1 or move_to_device
-        ), "If n_devices > 1, must move_to_device"
-
         if tokenizer is not None:
             self.set_tokenizer(tokenizer)
         elif self.cfg.tokenizer_name is not None:
@@ -157,7 +153,7 @@ def __init__(
             # the second gets the next n_layers // n_devices blocks ... the last gets the last n_layers // n_devices
             # blocks, the final
             # normalization layer (if it exists) and the unembed layer
-            HookedTransformer.move_model_modules_to_device(self)
+            self.move_model_modules_to_device()
 
         # Helper variable to store a small (10K-20K) dataset of training data. Empty by default, can be loaded with
         # load_sample_training_dataset
@@ -748,22 +744,21 @@ def cpu(self):
         # Wrapper around cuda that also changes self.cfg.device
         return self.to("cpu")
 
-    @classmethod
-    def move_model_modules_to_device(cls, model: "HookedTransformer"):
-        model.embed.to(devices.get_device_for_block_index(0, model.cfg))
-        model.hook_embed.to(devices.get_device_for_block_index(0, model.cfg))
-        if model.cfg.positional_embedding_type != "rotary":
-            model.pos_embed.to(devices.get_device_for_block_index(0, model.cfg))
-            model.hook_pos_embed.to(devices.get_device_for_block_index(0, model.cfg))
-        if hasattr(model, "ln_final"):
-            model.ln_final.to(
-                devices.get_device_for_block_index(model.cfg.n_layers - 1, model.cfg)
+    def move_model_modules_to_device(self):
+        self.embed.to(devices.get_device_for_block_index(0, self.cfg))
+        self.hook_embed.to(devices.get_device_for_block_index(0, self.cfg))
+        if self.cfg.positional_embedding_type != "rotary":
+            self.pos_embed.to(devices.get_device_for_block_index(0, self.cfg))
+            self.hook_pos_embed.to(devices.get_device_for_block_index(0, self.cfg))
+        if hasattr(self, "ln_final"):
+            self.ln_final.to(
+                devices.get_device_for_block_index(self.cfg.n_layers - 1, self.cfg)
             )
-        model.unembed.to(
-            devices.get_device_for_block_index(model.cfg.n_layers - 1, model.cfg)
+        self.unembed.to(
+            devices.get_device_for_block_index(self.cfg.n_layers - 1, self.cfg)
         )
-        for i, block in enumerate(model.blocks):
-            block.to(devices.get_device_for_block_index(i, model.cfg))
+        for i, block in enumerate(self.blocks):
+            block.to(devices.get_device_for_block_index(i, self.cfg))
 
     @classmethod
     def from_pretrained(
@@ -778,7 +773,6 @@ def from_pretrained(
         hf_model=None,
         device=None,
         n_devices=1,
-        move_state_dict_to_device=True,
         tokenizer=None,
         move_to_device=True,
         **from_pretrained_kwargs,
@@ -824,9 +818,6 @@ def from_pretrained(
                 default will load to CUDA if available, else CPU.
             n_devices (int, optional): The number of devices to split the model
                 across. Defaults to 1. If greater than 1, `device` must be cuda.
-            move_state_dict_to_device (bool): Whether to move the state dict to the
-                relevant device before processing and loading in the weights.
-                Defaults to True.
             tokenizer (*optional): The tokenizer to use for the model. If not
                 provided, it is inferred from cfg.tokenizer_name or initialized to None.
                 If None, then the model cannot be passed strings, and d_vocab must be explicitly set.
@@ -882,7 +873,7 @@ def from_pretrained(
         )
 
         # Create the HookedTransformer object
-        model = cls(cfg, tokenizer, move_to_device)
+        model = cls(cfg, tokenizer, move_to_device=False)
 
         dtype = from_pretrained_kwargs.get("torch_dtype", None)
         if dtype is not None:
@@ -894,9 +885,11 @@ def from_pretrained(
             center_writing_weights=center_writing_weights,
             center_unembed=center_unembed,
             refactor_factored_attn_matrices=refactor_factored_attn_matrices,
-            move_state_dict_to_device=move_state_dict_to_device,
         )
 
+        if move_to_device:
+            model.move_model_modules_to_device()
+
         print(f"Loaded pretrained model {model_name} into HookedTransformer")
 
         return model
@@ -961,7 +954,6 @@ def load_and_process_state_dict(
         center_unembed: bool = True,
         fold_value_biases: bool = True,
         refactor_factored_attn_matrices: bool = False,
-        move_state_dict_to_device: bool = True,
     ):
         """Method to load a state dict into the model, and to apply processing to simplify it. The state dict is assumed
         to be in the HookedTransformer format.
@@ -984,39 +976,10 @@ def load_and_process_state_dict(
                 output bias to the layer, and make it easier to interpret the head's output.
             refactor_factored_attn_matrices (bool, optional): Whether to convert the factored
                 matrices (W_Q & W_K, and W_O & W_V) to be "even". Defaults to False
-            move_state_dict_to_device (bool, optional): Whether to move the state dict to the device of the model.
-                Defaults to True.
             model_name (str, optional): checks the model name for special cases of state dict loading. Only used for
                 Redwood 2L model currently
         """
 
-        assert (
-            self.cfg.n_devices == 1 or move_state_dict_to_device
-        ), "If n_devices > 1, move_state_dict_to_device must be True"
-
-        if move_state_dict_to_device:
-            for k, v in state_dict.items():
-                if k.startswith("embed") or k.startswith("pos_embed"):
-                    state_dict[k] = v.to(
-                        devices.get_device_for_block_index(0, self.cfg)
-                    )
-                elif k.startswith("ln_final") or k.startswith("unembed"):
-                    state_dict[k] = v.to(
-                        devices.get_device_for_block_index(
-                            self.cfg.n_layers - 1, self.cfg
-                        )
-                    )
-                elif k.startswith("blocks"):
-                    state_dict[k] = v.to(
-                        devices.get_device_for_block_index(
-                            int(k.split(".")[1]), self.cfg
-                        )
-                    )
-                else:
-                    raise KeyError(
-                        f"State Dict contains a key not in the HookedTransformer format: {k}"
-                    )
-
         state_dict = self.fill_missing_keys(state_dict)
         if fold_ln:
             if self.cfg.normalization_type not in ["LN", "LNPre"]:
@@ -1356,7 +1319,6 @@ def process_weights_(
         center_writing_weights: bool = True,
         center_unembed: bool = True,
         refactor_factored_attn_matrices: bool = False,
-        move_state_dict_to_device: bool = True,
     ):
         """
         Wrapper around load_and_process_state_dict to allow for in-place processing of the weights. This is useful if
@@ -1380,7 +1342,6 @@ def process_weights_(
             center_writing_weights=center_writing_weights,
             center_unembed=center_unembed,
             refactor_factored_attn_matrices=refactor_factored_attn_matrices,
-            move_state_dict_to_device=move_state_dict_to_device,
         )
 
     @torch.inference_mode()
diff --git a/transformer_lens/utils.py b/transformer_lens/utils.py
index 8c805ccc2..d1d2ffb03 100644
--- a/transformer_lens/utils.py
+++ b/transformer_lens/utils.py
@@ -55,11 +55,8 @@ def download_file_from_hf(
         **select_compatible_kwargs(kwargs, hf_hub_download),
     )
 
-    # Load to the CPU device if CUDA is not available
-    map_location = None if torch.cuda.is_available() else torch.device("cpu")
-
     if file_path.endswith(".pth") or force_is_torch:
-        return torch.load(file_path, map_location=map_location)
+        return torch.load(file_path, map_location="cpu")
     elif file_path.endswith(".json"):
         return json.load(open(file_path, "r"))
     else:


Test patch:
diff --git a/tests/acceptance/test_hooked_transformer.py b/tests/acceptance/test_hooked_transformer.py
index 85d4002b5..f708b0301 100644
--- a/tests/acceptance/test_hooked_transformer.py
+++ b/tests/acceptance/test_hooked_transformer.py
@@ -6,6 +6,7 @@
 from transformers import AutoConfig
 
 from transformer_lens import HookedTransformer
+from transformer_lens.components import LayerNormPre
 from transformer_lens.loading_from_pretrained import OFFICIAL_MODEL_NAMES
 from transformer_lens.utils import clear_huggingface_cache
 
@@ -141,6 +142,15 @@ def test_from_pretrained_dtype():
     assert model.W_K.dtype == torch.bfloat16
 
 
+def test_process_weights_inplace():
+    """Check that process_weights_ works"""
+    model = HookedTransformer.from_pretrained_no_processing("gpt2-small")
+    model.process_weights_()
+    loss = model.forward(text, return_type="loss")
+    assert (loss.item() - loss_store["gpt2-small"]) < 4e-5
+    assert isinstance(model.ln_final, LayerNormPre)
+
+
 def test_from_pretrained_revision():
     """
     Check that the from_pretrained parameter `revision` (= git version) works
diff --git a/tests/unit/test_svd_interpreter.py b/tests/unit/test_svd_interpreter.py
index 6aba89d32..0f64f7ee2 100644
--- a/tests/unit/test_svd_interpreter.py
+++ b/tests/unit/test_svd_interpreter.py
@@ -6,8 +6,7 @@
 
 MODEL = "solu-2l"
 VECTOR_TYPES = ["OV", "w_in", "w_out"]
-ATOL = 1e-4  # Absolute tolerance - how far does a float have to be before we consider it no longer equal?
-# ATOL is set to 1e-4 because the tensors we check on are also to 4 decimal places.
+ATOL = 2e-4  # Absolute tolerance - how far does a float have to be before we consider it no longer equal?
 model = HookedTransformer.from_pretrained(MODEL)
 unfolded_model = HookedTransformer.from_pretrained(MODEL, fold_ln=False)
 second_model = HookedTransformer.from_pretrained("solu-3l")