Update Llamafy code

README.md

@@ -40,46 +40,47 @@ def create_llama_config(exaone_config):
 
 def copy_embedding_weights(llama_model, exaone_model):
     """Copy embedding weights from EXAONE to LLaMA model."""
-    llama_model.model.embed_tokens.weight.data = exaone_model.transformer.wte.weight.data.to(llama_model.device)
+    llama_model.model.embed_tokens.weight.data = exaone_model.transformer.wte.weight.data
 
-def copy_layer_weights(llama_layer, exaone_layer, device):
+def copy_layer_weights(llama_layer, exaone_layer):
     """Copy weights for a single layer from EXAONE to LLaMA model."""
     # Self-attention
-    llama_layer.self_attn.q_proj.weight.data = exaone_layer.attn.attention.q_proj.weight.data.to(device)
-    llama_layer.self_attn.k_proj.weight.data = exaone_layer.attn.attention.k_proj.weight.data.to(device)
-    llama_layer.self_attn.v_proj.weight.data = exaone_layer.attn.attention.v_proj.weight.data.to(device)
-    llama_layer.self_attn.o_proj.weight.data = exaone_layer.attn.attention.out_proj.weight.data.to(device)
+    llama_layer.self_attn.q_proj.weight.data = exaone_layer.attn.attention.q_proj.weight.data
+    llama_layer.self_attn.k_proj.weight.data = exaone_layer.attn.attention.k_proj.weight.data
+    llama_layer.self_attn.v_proj.weight.data = exaone_layer.attn.attention.v_proj.weight.data
+    llama_layer.self_attn.o_proj.weight.data = exaone_layer.attn.attention.out_proj.weight.data
     # MLP
-    llama_layer.mlp.gate_proj.weight.data = exaone_layer.mlp.c_fc_0.weight.data.to(device)
-    llama_layer.mlp.up_proj.weight.data = exaone_layer.mlp.c_fc_1.weight.data.to(device)
-    llama_layer.mlp.down_proj.weight.data = exaone_layer.mlp.c_proj.weight.data.to(device)
+    llama_layer.mlp.gate_proj.weight.data = exaone_layer.mlp.c_fc_0.weight.data
+    llama_layer.mlp.up_proj.weight.data = exaone_layer.mlp.c_fc_1.weight.data
+    llama_layer.mlp.down_proj.weight.data = exaone_layer.mlp.c_proj.weight.data
     # Layer Norms
-    llama_layer.input_layernorm.weight.data = exaone_layer.ln_1.weight.data.to(device)
-    llama_layer.post_attention_layernorm.weight.data = exaone_layer.ln_2.weight.data.to(device)
+    llama_layer.input_layernorm.weight.data = exaone_layer.ln_1.weight.data
+    llama_layer.post_attention_layernorm.weight.data = exaone_layer.ln_2.weight.data
 
 def copy_final_weights(llama_model, exaone_model):
     """Copy final layer norm and LM head weights from EXAONE to LLaMA model."""
-    llama_model.model.norm.weight.data = exaone_model.transformer.ln_f.weight.data.to(llama_model.device)
-    llama_model.lm_head.weight.data = exaone_model.lm_head.weight.data.to(llama_model.device)
+    llama_model.model.norm.weight.data = exaone_model.transformer.ln_f.weight.data
+    llama_model.lm_head.weight.data = exaone_model.lm_head.weight.data
 
 def port_exaone_to_llama(exaone_model_path, llama_model_path):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     print("Loading EXAONE model...")
-    exaone_model = load_model(exaone_model_path, ExaoneForCausalLM)
+    exaone_model = load_model(exaone_model_path, ExaoneForCausalLM).to(device)
     exaone_config = exaone_model.config
 
     print("Creating LLaMA configuration...")
     llama_config = create_llama_config(exaone_config)
 
     print("Initializing LLaMA model...")
-    llama_model = LlamaForCausalLM(llama_config)
-    llama_model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    llama_model = LlamaForCausalLM(llama_config).to(device)
 
     print("Copying weights...")
     copy_embedding_weights(llama_model, exaone_model)
 
     for i in range(exaone_config.num_layers):
         print(f"Copying weights for layer {i+1}/{exaone_config.num_layers}")
-        copy_layer_weights(llama_model.model.layers[i], exaone_model.transformer.h[i], llama_model.device)
+        copy_layer_weights(llama_model.model.layers[i], exaone_model.transformer.h[i])
 
     copy_final_weights(llama_model, exaone_model)