[Minor] Fix gradient backprop trainables with upstream interventions

.github/pull_request_template.md

@@ -1,15 +1,15 @@
-## 📝 Description
+## Description
 
 [Describe Your Changes Here ...]
 
 ## Testing Done
 
 [Describe Your Changes Here ...]
 
-# Checklist:
+## Checklist:
 
 - [ ] My PR title strictly follows the format: `[Your Priority] Your Title`
 - [ ] I have attached the testing log above
 - [ ] I provide enough comments to my code
 - [ ] I have changed documentations
-- [ ] I have added tests for my changes
+- [ ] I have added tests for my changes

pyvene/models/gpt2/modelings_intervenable_gpt2.py

@@ -22,7 +22,7 @@
  "attention_value_output": ("h[%s].attn.c_proj", CONST_INPUT_HOOK),
  "head_attention_value_output": ("h[%s].attn.c_proj", CONST_INPUT_HOOK, (split_head_and_permute, "n_head")),
  "attention_weight": ("h[%s].attn.attn_dropout", CONST_INPUT_HOOK),
- "attention_output": ("h[%s].attn", CONST_OUTPUT_HOOK),
+ "attention_output": ("h[%s].attn.resid_dropout", CONST_OUTPUT_HOOK),
  "attention_input": ("h[%s].attn", CONST_INPUT_HOOK),
  "query_output": ("h[%s].attn.c_attn", CONST_OUTPUT_HOOK, (split_three, 0)),
  "key_output": ("h[%s].attn.c_attn", CONST_OUTPUT_HOOK, (split_three, 1)),

pyvene/models/intervenable_base.py

@@ -555,11 +555,11 @@ def _gather_intervention_output(
  self._intervention_reverse_link[representations_key]
  ]
  else:
- # cold gather
- original_output = output
  # data structure casting
  if isinstance(output, tuple):
- original_output = output[0]
+ original_output = output[0].clone()
+ else:
+ original_output = output.clone()
  # gather subcomponent
  original_output = output_to_subcomponent(
  original_output,
@@ -788,7 +788,7 @@ def hook_callback(model, args, kwargs, output=None):
  output = kwargs[list(kwargs.keys())[0]]
  else:
  output = args
-
+ 
  selected_output = self._gather_intervention_output(
  output, key, unit_locations_base[key_i]
  )
@@ -838,26 +838,15 @@ def hook_callback(model, args, kwargs, output=None):
  self._intervention_reverse_link[key]
  ] = intervened_representation.clone()
 
- # very buggy due to tensor version
- if self.model_has_grad:
- # TODO: figure out how to allow this!
- if isinstance(output, tuple):
- raise ValueError(
- "Model grad is not allowed when "
- "intervening output is tuple type."
- )
- output_c = output.clone()
- # patched in the intervned activations in-place
+ if isinstance(output, tuple):
  _ = self._scatter_intervention_output(
- output_c, intervened_representation, key, unit_locations_base[key_i]
+ output[0], intervened_representation, key, unit_locations_base[key_i]
  )
- output = output_c.clone()
  else:
- # patched in the intervned activations in-place
  _ = self._scatter_intervention_output(
  output, intervened_representation, key, unit_locations_base[key_i]
  )
-
+ 
  self._intervention_state[key].inc_setter_version()
 
  handlers.append(module_hook(hook_callback, with_kwargs=True))
@@ -1073,7 +1062,6 @@ def _wait_for_forward_with_serial_intervention(
  def _broadcast_unit_locations(
  self,
  batch_size,
- intervention_group_size,
  unit_locations
  ):
  if self.mode == "parallel":
@@ -1086,33 +1074,33 @@ def _broadcast_unit_locations(
  k = "sources->base"
  if isinstance(v, int):
  if is_base_only:
- _unit_locations[k] = (None, [[[v]]*batch_size]*intervention_group_size)
+ _unit_locations[k] = (None, [[[v]]*batch_size]*len(self.interventions))
  else:
  _unit_locations[k] = (
- [[[v]]*batch_size]*intervention_group_size, 
- [[[v]]*batch_size]*intervention_group_size
+ [[[v]]*batch_size]*len(self.interventions), 
+ [[[v]]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  elif len(v) == 2 and isinstance(v[0], int) and isinstance(v[1], int):
  _unit_locations[k] = (
- [[[v[0]]]*batch_size]*intervention_group_size, 
- [[[v[1]]]*batch_size]*intervention_group_size
+ [[[v[0]]]*batch_size]*len(self.interventions), 
+ [[[v[1]]]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  elif len(v) == 2 and v[0] == None and isinstance(v[1], int):
- _unit_locations[k] = (None, [[[v[1]]]*batch_size]*intervention_group_size)
+ _unit_locations[k] = (None, [[[v[1]]]*batch_size]*len(self.interventions))
  self.use_fast = True
  elif len(v) == 2 and isinstance(v[0], int) and v[1] == None:
- _unit_locations[k] = ([[[v[0]]]*batch_size]*intervention_group_size, None)
+ _unit_locations[k] = ([[[v[0]]]*batch_size]*len(self.interventions), None)
  self.use_fast = True
  elif isinstance(v, list) and get_list_depth(v) == 1:
  # [0,1,2,3] -> [[[0,1,2,3]]], ...
  if is_base_only:
- _unit_locations[k] = (None, [[v]*batch_size]*intervention_group_size)
+ _unit_locations[k] = (None, [[v]*batch_size]*len(self.interventions))
  else:
  _unit_locations[k] = (
- [[v]*batch_size]*intervention_group_size, 
- [[v]*batch_size]*intervention_group_size
+ [[v]*batch_size]*len(self.interventions), 
+ [[v]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  else:
@@ -1125,27 +1113,27 @@ def _broadcast_unit_locations(
  for k, v in unit_locations.items():
  if isinstance(v, int):
  _unit_locations[k] = (
- [[[v]]*batch_size]*intervention_group_size, 
- [[[v]]*batch_size]*intervention_group_size
+ [[[v]]*batch_size]*len(self.interventions), 
+ [[[v]]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  elif len(v) == 2 and isinstance(v[0], int) and isinstance(v[1], int):
  _unit_locations[k] = (
- [[[v[0]]]*batch_size]*intervention_group_size, 
- [[[v[1]]]*batch_size]*intervention_group_size
+ [[[v[0]]]*batch_size]*len(self.interventions), 
+ [[[v[1]]]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  elif len(v) == 2 and v[0] == None and isinstance(v[1], int):
- _unit_locations[k] = (None, [[[v[1]]]*batch_size]*intervention_group_size)
+ _unit_locations[k] = (None, [[[v[1]]]*batch_size]*len(self.interventions))
  self.use_fast = True
  elif len(v) == 2 and isinstance(v[0], int) and v[1] == None:
- _unit_locations[k] = ([[[v[0]]]*batch_size]*intervention_group_size, None)
+ _unit_locations[k] = ([[[v[0]]]*batch_size]*len(self.interventions), None)
  self.use_fast = True
  elif isinstance(v, list) and get_list_depth(v) == 1:
  # [0,1,2,3] -> [[[0,1,2,3]]], ...
  _unit_locations[k] = (
- [[v]*batch_size]*intervention_group_size, 
- [[v]*batch_size]*intervention_group_size
+ [[v]*batch_size]*len(self.interventions), 
+ [[v]*batch_size]*len(self.interventions)
  )
  self.use_fast = True
  else:
@@ -1191,16 +1179,15 @@ def _broadcast_sources(
  def _broadcast_subspaces(
  self,
  batch_size,
- intervention_group_size,
  subspaces
  ):
  """Broadcast simple subspaces input"""
  _subspaces = subspaces
  if isinstance(subspaces, int):
- _subspaces = [[[subspaces]]*batch_size]*intervention_group_size
+ _subspaces = [[[subspaces]]*batch_size]*len(self.interventions)
 
  elif isinstance(subspaces, list) and isinstance(subspaces[0], int):
- _subspaces = [[subspaces]*batch_size]*intervention_group_size
+ _subspaces = [[subspaces]*batch_size]*len(self.interventions)
  else:
  # TODO: subspaces is easier to add more broadcast majic.
  pass
@@ -1292,13 +1279,11 @@ def forward(
  return self.model(**base), None
 
  # broadcast
- unit_locations = self._broadcast_unit_locations(
- get_batch_size(base), len(self._intervention_group), unit_locations)
+ unit_locations = self._broadcast_unit_locations(get_batch_size(base), unit_locations)
  sources = [None]*len(self._intervention_group) if sources is None else sources
  sources = self._broadcast_sources(sources)
  activations_sources = self._broadcast_source_representations(activations_sources)
- subspaces = self._broadcast_subspaces(
- get_batch_size(base), len(self._intervention_group), subspaces)
+ subspaces = self._broadcast_subspaces(get_batch_size(base), subspaces)
 
  self._input_validation(
  base,
@@ -1412,13 +1397,11 @@ def generate(
  unit_locations = {"base": 0}
 
  # broadcast
- unit_locations = self._broadcast_unit_locations(
- get_batch_size(base), len(self._intervention_group), unit_locations)
+ unit_locations = self._broadcast_unit_locations(get_batch_size(base), unit_locations)
  sources = [None]*len(self._intervention_group) if sources is None else sources
  sources = self._broadcast_sources(sources)
  activations_sources = self._broadcast_source_representations(activations_sources)
- subspaces = self._broadcast_subspaces(
- get_batch_size(base), len(self._intervention_group), subspaces)
+ subspaces = self._broadcast_subspaces(get_batch_size(base), subspaces)
 
  self._input_validation(
  base,

pyvene/models/interventions.py

@@ -13,7 +13,7 @@ class Intervention(torch.nn.Module):
 
  def __init__(self, **kwargs):
  super().__init__()
- self.trainble = False
+ self.trainable = False
  self.is_source_constant = False
 
  self.use_fast = kwargs["use_fast"] if "use_fast" in kwargs else False
@@ -87,7 +87,7 @@ class TrainableIntervention(Intervention):
 
  def __init__(self, **kwargs):
  super().__init__(**kwargs)
- self.trainble = True
+ self.trainable = True
  self.is_source_constant = False
 
  def tie_weight(self, linked_intervention):
@@ -204,7 +204,7 @@ class VanillaIntervention(Intervention, LocalistRepresentationIntervention):
  def __init__(self, **kwargs):
  super().__init__(**kwargs)
 
- def forward(self, base, source, subspaces=None):  
+ def forward(self, base, source, subspaces=None): 
  return _do_intervention_by_swap(
  base,
  source if self.source_representation is None else self.source_representation,
@@ -478,7 +478,7 @@ def __init__(self, **kwargs):
  self.pca_std = torch.nn.Parameter(
  torch.tensor(pca_std, dtype=torch.float32), requires_grad=False
  )
- self.trainble = False
+ self.trainable = False
 
  def forward(self, base, source, subspaces=None):
  base_norm = (base - self.pca_mean) / self.pca_std