Ablation Hook I created seems to affect the calculations of the gradients on later layers

Realmbird: <p>So far I have created view hooks to view the outputs and ablation hooks to ablate a specific head from attn_outputs.</p>
<ol>
<li>
<p>So far I have confirmed with my view hooks after the ablation hooks that impact of the ablation hooks makes the attn_output values for the ablated hooks to be 0</p>
</li>
<li>
<p>However, the ablation hooks for some reason does not have a impact on the down stream calculations on the later hooks as looking at non ablated hooks or layers with my view hooks before and after ablation or with or without the ablation hooks has 0 change</p>
</li>
</ol>
<p>I would like to ask how could I fix the ablation hooks so that an ablation hook in layer 1 will have the changes be seen in another hooks gradients in layer 2</p>
<p>ablation hook<br>
hooks = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<p>for layer_idx, head_num in reciever_heads_sorted:</p>
<pre><code class="lang-auto">print(f"    Adding ablation hook for head ({layer_idx}, {head_num})")



\# Define ablation hook for this specific head

def create_ablation_hook(target_head_num):

    def ablation_hook(module, input, output):

        attention_output = output\[0\]

        batch_size_tensor, seq_len, hidden_dim = attention_output.shape

        num_heads = 32

        head_dim = hidden_dim // num_heads

        

        \# Zero out the specific head

        reshaped = attention_output.view(batch_size_tensor, seq_len, num_heads, head_dim)

        reshaped\[:, :, target_head_num, :\] = 0

        modified = reshaped.view(batch_size_tensor, seq_len, hidden_dim)

        

        return (modified,) + output\[1:\]

    return ablation_hook



\# Register hook for this layer

attention_layer = model.model.layers\[layer_idx\].self_attn

hook = attention_layer.register_forward_hook(create_ablation_hook(head_num))

hooks.append((hook, layer_idx, head_num))
</code></pre>
<p>view hook</p>
<p>import torch.nn as nn</p>
<p>view_hooks_after = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<p>hidden_states_dict_after = {}</p>
<p>dropout = 0.0  # Set to your model’s dropout if needed</p>
<p>for layer_idx, head_num in all_heads_to_check:</p>
<pre><code class="lang-auto">print(f"    Adding view hook for head ({layer_idx}, {head_num})")



\# Define view hook for this specific head

def create_view_hook(target_head_num, layer_idx):

    

    def view_hook(module, input, output):

        attn_output, attn_weights = output



        batch_size, seq_len, hidden_dim = attn_output.shape

        num_heads = 32

        head_dim = hidden_dim // num_heads

        \# Extract only the output for the target head

        attn_output_reshaped = attn_output.view(batch_size, seq_len, num_heads, head_dim)

        target_head_output = attn_output_reshaped\[:, :, target_head_num, :\]  # (batch, seq, head_dim)

        \# Extract only the attention probs for the target head

        target_head_probs = attn_weights\[:, target_head_num, :, :\]  # (batch, seq, seq)




        hidden_states_dict_after\[(layer_idx, target_head_num)\] = {

            "attn_output": target_head_output.detach().cpu().to(torch.float32),

            "attn_probs": target_head_probs.detach().cpu().to(torch.float32)

    }

        return output



    return view_hook



\# Register hook for this layer

attention_layer = model.model.layers\[layer_idx\].self_attn

hook = attention_layer.register_forward_hook(create_view_hook(head_num, layer_idx))

view_hooks_after.append((hook, layer_idx, head_num))
</code></pre>

Realmbird: <p>Other weird behavior with my hooks was the input gradients were () missing not sure why</p>

Realmbird: <p>Model info</p>
<pre><code class="lang-auto">from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig, AutoModelForCausalLM, pipeline

import torch


model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B" # Or any other suitable model

mname = model_name

tokenizer = AutoTokenizer.from_pretrained(model_name)

# Important: Add a pad token if the tokenizer doesn't have one, especially for decoder models.

if tokenizer.pad_token is None:

    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
</code></pre>

Realmbird: <p>Code to run it</p>
<pre><code class="lang-auto">def get_attention_all_heads(model, tokenizer, text, max_length=1024):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=max_length)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
    outputs = model(**inputs, use_cache=False, output_attentions=True)
    attention = torch.stack([a[0].cpu().to(torch.float32) for a in outputs.attentions])
    del outputs
    torch.cuda.empty_cache()
    return attention

# Usage:
attention_all_heads_no_ablation = get_attention_all_heads(model, tokenizer, generated_text_no_ablation[0])
</code></pre>

valerian.rey: <p>Could you tell us exactly which variable has a different value from what you would expect?</p>
<p>Also, could you try to make a single executable code block to reproduce your issue?</p>

Realmbird: <ul>
<li>Issues I had so far is that when looking at my pytorch hooks such as the ablation and view hooks the input_gradient is missing and is equal to () for the view hooks.</li>
<li>However, the output gradients were still present.</li>
<li>Using the output gradients I ablated the specific layer and head in the model.</li>
<li>Ablating from the output gradients of the model gave the intended affect of having 0s be along the attention head for the attn_output</li>
<li>However, I am not sure why despite being able to ablate the attention output the ablation hook in layer 1 for example is not affecting heads in later layers. If I understand it correctly the ablation hook run in order and the activation in the later layer tensors should be affected by autograd. So far I am not sure why changing the output_attn in a earlier layer did not affect a later hook that looks into a later layer in the model. Is this something to do with the fact my input grad was () or empty?</li>
<li>I am not sure what you mean by singular executable code block but, do you mean copying all the code needed to test the error?</li>
</ul>