[Python] Calculate the gradient for attention but also the FFN for pretrained LLMs

I would like to return the gradient with respect to specific layers and the FFN layer in the Transformer architecture of pre-trained LLMs from the hugging-face model. Is that even possible?

I am working with the code of this repo which is the following:

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-medium-4k-instruct")
model = AutoModelForCausalLM.from_pretrained(
"microsoft/Phi-3-medium-4k-instruct", # note: check spelling if you get error
device_map="auto",
torch_dtype=torch.float16, # or torch.float32 if preferred
trust_remote_code=True
)

# Create a pipeline
generator = pipeline(
"text-generation",
model = model,
tokenizer = tokenizer,
return_full_text= False,
max_new_tokens = 100,
do_sample = False
)

# Prepare a prompt
prompt = "Whats is the co-capital of Greece according to the country's public opinion?"
inputs = tokenizer(prompt, return_tensors="pt")
inputs = inputs.to("cuda:0") # send inputs to cuda

# Run the model with attention outputs enabled
# Make sure to pass output_attentions=True
outputs = model(input_ids=inputs.input_ids, output_attentions=True)

# outputs.attentions is a tuple with one element per layer
# Each element is a tensor of shape (batch_size, num_heads, seq_len, seq_len)
attentions = outputs.attentions

# Generate output
output = generator(prompt)
print(output[0]["generated_text"])


How to return the gradient concerning the attention layer. Then what to do with other components of the same network?

Continue reading...