A minor problem regarding the skip-connection visualization

[YunhaoZhang-Mars]

Thanks for the awesome repository! After going through it step-by-step, I have a better understanding of Llama3 techniques, such as rotary position embedding, grouped key and value, etc.

I found that there might be a minor mistake regarding the skip-connection visualization: the corresponding code is in the section "WE FINALLY HAVE NEW EDITED EMBEDDINGS FOR EACH TOKEN AFTER THE FIRST LAYER" :

layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shape

As embedding_after_edit instead of embedding_after_edit_normalized is used, the visualization should be

[wangdsh]

I found the same problem.

[naklecha]

oh yes you are actually right!!!