I've tried modifying the embed_to_distrib function in pyvene/models/basic_utils.py to also support llama models as such:
def embed_to_distrib(model, embed, log=False, logits=False):
"""Convert an embedding to a distribution over the vocabulary"""
if "gpt2" in model.config.architectures[0].lower():
with torch.inference_mode():
vocab = torch.matmul(embed, model.wte.weight.t())
if logits:
return vocab
return lsm(vocab) if log else sm(vocab)
elif "llama" in model.config.architectures[0].lower():
with torch.inference_mode():
vocab = model.lm_head(embed)
if logits:
return vocab
return lsm(vocab) if log else sm(vocab)
It seems to work fine when doing causal tracing (see images below):
Single restored layer
MLP layer
Attention layer
Would this be the correct approach to do so on Llama model, and would it be of interest for pyvene?
Suggestion / Feature Request
I've tried modifying the
embed_to_distrib
function inpyvene/models/basic_utils.py
to also support llama models as such:It seems to work fine when doing causal tracing (see images below):
Single restored layer![Screenshot 2024-07-18 at 10 41 40](https://github.com/user-attachments/assets/2014eca1-ffb2-4eeb-9f50-c145bd52c8a7)
MLP layer![Screenshot 2024-07-18 at 10 42 27](https://github.com/user-attachments/assets/e13aeb7e-4006-4a4d-af06-2eecb63ae81b)
Attention layer![Screenshot 2024-07-18 at 10 42 41](https://github.com/user-attachments/assets/cc7f0c6d-a312-4532-a5e3-cc1651b154c0)
Would this be the correct approach to do so on Llama model, and would it be of interest for pyvene?