`lrp_patching.ipynb` with `GroupedQueryAttention` (e.g. `Qwen/Qwen2-0.5B`)?

[Ja1Zhou]

Hi! I wonder what is the correct way to apply RelP to models with grouped query attention?

I tried Qwen/Qwen2-0.5B which is a model mentioned in the paper:

# model = HookedTransformer.from_pretrained("gpt2-small")
model = HookedTransformer.from_pretrained("Qwen/Qwen2-0.5B")

model.cfg.use_lrp = True
model.cfg.LRP_rules = ['LN-rule', 'Identity-rule', 'Half-rule']

model.set_use_attn_result(True)
model.set_use_hook_mlp_in(True)
model.set_use_attn_in(True)
model.set_use_split_qkv_input(True)

But throws AssertionError: Can't use attn_in with GroupedQueryAttention, please use split_qkv_input instead