Device aware cache (#250)

* Add hook example

* Fix docstrings

* Convert device to torch.device

* Respect provided cfg device

* Keep cached tensors on device

* Test cache device

* Handle sequence on GPU

* Keep tensor on device

* Test plain cuda

* Update test_multi_gpu.py

fix spelling of cache

---------

Co-authored-by: Joseph Bloom <[email protected]>

tests/acceptance/test_multi_gpu.py

@@ -90,11 +90,17 @@ def test_device_separation_and_cache(gpt2_medium_on_1_device, n_devices):
  gpt2_logits_n_devices, gpt2_cache_n_devices = model_n_devices.run_with_cache(
  gpt2_tokens, remove_batch_dim=True)
 
+ # Make sure the tensors in cache remain on their respective devices
+ for i in range(model_n_devices.cfg.n_layers):
+ expected_device = get_device_for_block_index(i, cfg=model_n_devices.cfg)
+ cache_device = gpt2_cache_n_devices[f"blocks.{i}.mlp.hook_post"].device
+ assert cache_device == expected_device
+
  assert torch.allclose(gpt2_logits_1_device.to("cpu"),
  gpt2_logits_n_devices.to("cpu"))
  for key in gpt2_cache_1_device.keys():
- assert torch.allclose(gpt2_cache_1_device[key],
- gpt2_cache_n_devices[key])
+ assert torch.allclose(gpt2_cache_1_device[key].to("cpu"),
+ gpt2_cache_n_devices[key].to("cpu"))
 
  cuda_devices = set()
  n_params_on_device = {}
@@ -114,3 +120,27 @@ def test_device_separation_and_cache(gpt2_medium_on_1_device, n_devices):
  print(
  f"Number of devices: {n_devices}, Model loss (1 device): {loss_1_device}, Model loss ({n_devices} devices): {loss_n_devices}, Time taken (1 device): {elapsed_time_1_device:.4f} seconds, Time taken ({n_devices} devices): {elapsed_time_n_devices:.4f} seconds"
  )
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 CUDA devices")
+def test_cache_device():
+ model = HookedTransformer.from_pretrained("gpt2-small", device="cuda:1")
+
+ logits, cache = model.run_with_cache("Hello there")
+ assert norm_device(cache["blocks.0.mlp.hook_post"].device) == norm_device(torch.device("cuda:1"))
+
+ logits, cache = model.run_with_cache("Hello there", device="cpu")
+ assert norm_device(cache["blocks.0.mlp.hook_post"].device) == norm_device(torch.device("cpu"))
+
+ model.to("cuda")
+ logits, cache = model.run_with_cache("Hello there")
+ assert norm_device(cache["blocks.0.mlp.hook_post"].device) == norm_device(logits.device)
+
+
+def norm_device(device):
+ """
+ Convenience function to normalize device strings for comparison.
+ """
+ device_str = str(device)
+ if device_str.startswith("cuda") and ':' not in device_str:
+ device_str += ':0'
+ return device_str

transformer_lens/head_detector.py

@@ -122,9 +122,10 @@ def get_duplicate_token_head_detection_pattern(model: HookedTransformer, sequenc
  model: Model being used.
  sequence: String being fed to the model."""
 
- sequence = model.to_tokens(sequence)
- token_pattern = [np.array(sequence) for i in range(sequence.shape[-1])]
- token_pattern = np.concatenate(token_pattern, axis=0)
+ sequence = model.to_tokens(sequence).detach().cpu()
+
+ # Repeat sequence to create a square matrix.
+ token_pattern = sequence.repeat(sequence.shape[-1], 1).numpy()
 
  # If token_pattern[i][j] matches its transpose, then token j and token i are duplicates.
  eq_mask = np.equal(token_pattern, token_pattern.T).astype(int)

transformer_lens/hook_points.py

@@ -280,15 +280,13 @@ def add_caching_hooks(
  Args:
  names_filter (NamesFilter, optional): Which activations to cache. Can be a list of strings (hook names) or a filter function mapping hook names to booleans. Defaults to lambda name: True.
  incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
- device (_type_, optional): The device to store on. Defaults to CUDA if available else CPU.
+ device (_type_, optional): The device to store on. Defaults to same device as model.
  remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). Defaults to False.
  cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
 
  Returns:
  cache (dict): The cache where activations will be stored.
  """
- if device is None:
- device = "cuda" if torch.cuda.is_available() else "cpu"
  if cache is None:
  cache = {}
 
@@ -340,8 +338,7 @@ def run_with_cache(
  list of str, or a function that takes a string and returns a bool. Defaults to None, which
  means cache everything.
  device (str or torch.Device, optional): The device to cache activations on. Defaults to the
- model device. Note that this must be set if the model does not have a model.cfg.device
- attribute. WARNING: Setting a different device than the one used by the model leads to
+ model device. WARNING: Setting a different device than the one used by the model leads to
  significant performance degradation.
  remove_batch_dim (bool, optional): If True, removes the batch dimension when caching. Only
  makes sense with batch_size=1 inputs. Defaults to False.
@@ -382,7 +379,7 @@ def get_caching_hooks(
  Args:
  names_filter (NamesFilter, optional): Which activations to cache. Can be a list of strings (hook names) or a filter function mapping hook names to booleans. Defaults to lambda name: True.
  incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
- device (_type_, optional): The device to store on. Defaults to CUDA if available else CPU.
+ device (_type_, optional): The device to store on. Keeps on the same device as the layer if None.
  remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). Defaults to False.
  cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
 
@@ -391,8 +388,6 @@ def get_caching_hooks(
  fwd_hooks (list): The forward hooks.
  bwd_hooks (list): The backward hooks. Empty if incl_bwd is False.
  """
- if device is None:
- device = "cuda" if torch.cuda.is_available() else "cpu"
  if cache is None:
  cache = {}
 

transformer_lens/utilities/devices.py

@@ -9,10 +9,25 @@ def get_device_for_block_index(
  cfg: HookedTransformerConfig,
  device: Optional[Union[torch.device, str]] = None,
 ):
+ """
+ Determine the device for a given layer index based on the model configuration.
+ 
+ This function assists in distributing model layers across multiple devices. The distribution
+ is based on the configuration's number of layers (cfg.n_layers) and devices (cfg.n_devices).
+
+ Args:
+ index (int): Model layer index.
+ cfg (HookedTransformerConfig): Model and device configuration.
+ device (Optional[Union[torch.device, str]], optional): Initial device used for determining the target device.
+ If not provided, the function uses the device specified in the configuration (cfg.device).
+
+ Returns:
+ torch.device: The device for the specified layer index.
+ """
  assert cfg.device is not None
  layers_per_device = cfg.n_layers // cfg.n_devices
  if device is None:
  device = cfg.device
- if isinstance(device, torch.device):
-  device = device.type
- return torch.device(device, index // layers_per_device)
+ device = torch.device(device)
+ device_index = (device.index or 0) + (index // layers_per_device)
+ return torch.device(device.type, device_index)