naklecha / llama3-from-scratch Goto Github PK

import torch
import transformers
import json

model = transformers.AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B')
tokenizer = transformers.AutoTokenizer.from_pretrained('meta-llama/Meta-Llama-3-8B',model_max_length=512,padding_side="right",use_fast=False,)

dim = model.config.hidden_size
n_layers = model.config.num_hidden_layers
n_heads = model.config.num_attention_heads
n_kv_heads = model.config.num_key_value_heads
vocab_size = model.config.vocab_size
#The following two lines of code seem to be useless
multiple_of = model.config.multiple_of if model.config.multiple_of else 1024
ffn_dim_multiplier = model.config.ffn_dim_multiplier if model.config.ffn_dim_multiplier else 1.3

norm_eps = model.config.rms_norm_eps
rope_theta = model.config.rope_theta

prompt = "the answer to the ultimate question of life, the universe, and everything is "
tokens = tokenizer.encode(prompt)
print(tokens)
tokens = torch.tensor(tokens)
prompt_split_as_tokens = [tokenizer.decode([token.item()]) for token in tokens]
print(prompt_split_as_tokens)
embedding_layer = torch.nn.Embedding(vocab_size, dim)
tok_embeddings_weight = model.model.embed_tokens.weight
embedding_layer.weight.data.copy_(tok_embeddings_weight)
token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16)
token_embeddings_unnormalized.shape

def rms_norm(tensor, norm_weights):
    return (tensor * torch.rsqrt(tensor.pow(2).mean(-1, keepdim=True) + norm_eps)) * norm_weights

#RoPE
zero_to_one_split_into_64_parts = torch.tensor(range(64))/64
zero_to_one_split_into_64_parts

freqs = 1.0 / (rope_theta ** zero_to_one_split_into_64_parts)
freqs

freqs_for_each_token = torch.outer(torch.arange(17), freqs)
freqs_cis = torch.polar(torch.ones_like(freqs_for_each_token), freqs_for_each_token)
freqs_cis.shape

final_embedding = token_embeddings_unnormalized
for layer in range(n_layers):
    qkv_attention_store = []
    layer_embedding_norm = rms_norm(final_embedding, model.model.layers[layer].post_attention_layernorm.weight)
    q_layer = model.model.layers[layer].self_attn.q_proj.weight
    q_layer = q_layer.view(n_heads, q_layer.shape[0] // n_heads, dim)
    k_layer = model.model.layers[layer].self_attn.k_proj.weight
    k_layer = k_layer.view(n_kv_heads, k_layer.shape[0] // n_kv_heads, dim)
    v_layer = model.model.layers[layer].self_attn.v_proj.weight
    v_layer = v_layer.view(n_kv_heads, v_layer.shape[0] // n_kv_heads, dim)
    w_layer = model.model.layers[layer].self_attn.o_proj.weight.to(torch.bfloat16)
    for head in range(n_heads):
        q_layer_head = q_layer[head]
        k_layer_head = k_layer[head//4]
        v_layer_head = v_layer[head//4]
        q_per_token = torch.matmul(layer_embedding_norm, q_layer_head.T)
        k_per_token = torch.matmul(layer_embedding_norm, k_layer_head.T)
        v_per_token = torch.matmul(layer_embedding_norm, v_layer_head.T).to(torch.bfloat16)
        q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
        q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
        q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis)
        q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
        k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
        k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
        k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)
        k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
        qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(128)**0.5
        mask = torch.full((len(token_embeddings_unnormalized), len(token_embeddings_unnormalized)), float("-inf"))
        mask = torch.triu(mask, diagonal=1)
        qk_per_token_after_masking = qk_per_token + mask
        qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
        qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
        qkv_attention_store.append(qkv_attention)
    stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
    #w_layer = model[f"layers.{layer}.attention.wo.weight"]
    embedding_delta = torch.matmul(stacked_qkv_attention, w_layer.T)
    embedding_after_edit = final_embedding + embedding_delta
    embedding_after_edit_normalized = rms_norm(embedding_after_edit, model.model.layers[layer].input_layernorm.weight)
    w1 = model.model.layers[layer].mlp.gate_proj.weight
    w2 = model.model.layers[layer].mlp.down_proj.weight
    w3 = model.model.layers[layer].mlp.up_proj.weight
    output_after_feedforward = torch.matmul(torch.functional.F.silu(torch.matmul(embedding_after_edit_normalized, w1.T)) * torch.matmul(embedding_after_edit_normalized, w3.T), w2.T)
    final_embedding = embedding_after_edit+output_after_feedforward
	
final_embedding = rms_norm(final_embedding, model.model.norm.weight)
final_embedding.shape

logits = torch.matmul(final_embedding[-1], model.lm_head.weight.T)
logits.shape

next_token = torch.argmax(logits, dim=-1)
next_token
>>>tensor(8168)
tokenizer.decode([next_token.item()])
>>>';&'
my question：Why is the result not 42?

layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shape