Open gsarti opened 3 months ago
Hey @gsarti,
this is really awesome! Thank you for figuring this out, I would like to add this to the documentation. We just submitted the camera-ready version to ICML 2024, so now I have a little bit more time to spend on this repository again.
I ran your code and noticed that the scaling of the relevances is not perfect for plotting the heatmaps in LaTeX (otherwise it is fine!). If I compute the max and min values of inseq_relevance
, at least one value should be 1 or -1. In other words, if you run
# scaling wrong for LaTeX plotting
pdf_heatmap(tokens, inseq_relevance, path='heatmap_inseq.pdf', backend='xelatex')
it looks different to
# scaling correct
inseq_relevance2 = inseq_relevance / inseq_relevance.abs().max()
pdf_heatmap(tokens, inseq_relevance2, path='heatmap_inseq2.pdf', backend='xelatex')
Otherwise, your example is 100% fine! Do you have an inseq specific command to obtain this kind of scaling, or do I have to scale the relevance with inseq_relevance2 = inseq_relevance / inseq_relevance.abs().max()
?
Regarding your question:
The Composite
class is able to change ONLY nn.Module(s). This means, that we have to modify all lines of source code containing e.g. +
or torch.nn.softmax
etc. by replacing them with the LXT drop-in replacements.
The reason why I added AttentionValueMatmul
is a lilttle bit confusing: I just wanted to compare our AttnLRP implementation against the CP-LRP implementation (Ali et. al, 2022).
So, I converted thetoch.matmul
operation into an nn.Module so that I can apply different Composites without having to reload a complete new model source code everytime I change the LRP implementation.
I am actually working on a torch.fx
implementation that is able to replace these operations automatically. In my smaller tests, it is already working quite well! However, there is a big drawback: torch.fx
traced models do not support gradient checkpointing! And this is necessary to run LRP on large LLMs. This is why, I unfortunately still recommend modifying the source code. I plan to add a range of popular models including BERT, GPT and CLIP.
Best (:
Hi @rachtibat, thanks for your answer! It would be awesome to have the example available in the LXT docs, and I also plan to add it to our Inseq docs soon.
Regarding your question about rescaling, can it be that you're using v0.6? rescale=True
only works when installing from main (pip install git+https://github.com/inseq-team/inseq.git@main
) and should perform exactly the operation you mention, but I suspect it is silently ignored as kwargs
in previous versions.
Nice to hear that you're digging deeper on the torch.fx
side! Could you provide a pointer for the issue with torch.fx
and gradient checkpointing? I tried the example here with torch.compile
, and it worked for me. Is it a problem specific to symbolic_trace
?
Hey @gsarti,
Thanks, I didn't try out rescale=True
and thank you for pointing me to the issue. Unfortunately, I could not solve the issue, it is indeed related to symbolic_trace
.
I assume the problem is that symbolic_trace
removes all checkpoint()
operations from the graphs. Here is an issue explaining it but nobody responded.
I set up a small example demonstrating that a traced module does not perform activation checkpointing:
import torch
from torch.utils.checkpoint import checkpoint
import torch.fx as fx
def forward_hook(module, input, output):
if input[0].requires_grad:
print(f"track gradient in submodule")
else:
print(f"do not track gradient in submodule")
class SubModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear2 = torch.nn.Linear(1024, 2048)
self.linear3 = torch.nn.Linear(2048, 1024)
self.linear4 = torch.nn.Linear(2048, 1024)
# apply forward hook
self.linear3.register_forward_hook(forward_hook)
def forward(self, x):
x = self.linear2(x)
x = self.linear3(x) + self.linear4(x)
return x
class myModel(torch.nn.Module):
def __init__(self, grad_checkpoint):
super().__init__()
self.linear1 = torch.nn.Linear(1024, 1024)
self.submodule = SubModule()
self.checkpoint = grad_checkpoint
def forward(self, x):
x = self.linear1(x)
if self.checkpoint:
x = checkpoint(self.submodule, x, use_reentrant=True)
else:
x = self.submodule(x)
x = x.sum()
return x
def run(grad_checkpoint, use_graph_module):
device = "cuda" if torch.cuda.is_available() else "cpu"
model = myModel(grad_checkpoint).to(device)
x = torch.randn((1, 1024), requires_grad=True, device=device)
print("### Start forward pass")
if use_graph_module:
tracer = fx.Tracer()
graph = tracer.trace(model)
gm = fx.GraphModule(model, graph)
gm.recompile()
# gm.print_readable()
out = gm(x)
else:
out = model(x)
print("### Start backward pass")
out.backward()
print("gradient shape", x.grad.shape)
if __name__ == "__main__":
# works as expected
print("---- no checkpointing ----")
run(grad_checkpoint=False, use_graph_module=False)
print("---- no checkpointing & torch.fx ----")
run(grad_checkpoint=False, use_graph_module=True)
print("---- enabled checkpointing ----")
run(grad_checkpoint=True, use_graph_module=False)
# does not perform activation checkpointing!
print("---- enabled checkpointing & torch.fx ----")
run(grad_checkpoint=True, use_graph_module=True)
Yes, I just confirmed it by adding this line at the top
import torch
from torch.utils.checkpoint import checkpoint
import torch.fx as fx
# do not trace inside checkpoint()
fx.wrap('checkpoint')
Then, activation checkpointing works, but torch.fx is not tracing the content of the checkpointed layers, which is defeating the purpose why we use torch.fx in the first place.
If I use torch==2.3
and compile()
like in your example, it works! So, there might be a way! Maybe it is possible to extract a graph from compile()
.
Hey @rachtibat,
Great job with the implementation! I just wanted to report that I managed to reproduce the results of the TinyLLaMA Quickstart demo using Inseq as wrapper for the attribution process:
I think this could be useful since Inseq can easily extract scores for an entire generated sequence without having to perform forward passes step by step. It also allows for easy post-processing and visualization of attribution scores, which could make its usage simpler for LXT users!
I only had a question related to the current need to override model implementations to define custom propagation rules. My understanding is that most operations, including in-place ones (as mentioned here) could be overridden via the
Composite
class. However, I see you still need to perform slight adaptations to the architecture, e.g. defining customProjSiluMultiplication
andAttentionValueMatmul
for LLaMA. How challenging would it be to have a registrableComposite
accounting for all operations requiring overriding without having to subclass the originaltransformers
models? Is there really no way out of this?