Views
No views yet
1import torch
2import gc
3
4from transformers import LlamaConfig, LlamaForCausalLM, AutoModelForCausalLM, AutoTokenizer
5from tqdm import tqdm
6
7def unload_model(model):
8 del model
9 gc.collect()
10 if torch.cuda.is_available():
11 torch.cuda.empty_cache()
12
13def create_llama_config(exaone_config):
14 return LlamaConfig(
15 vocab_size=exaone_config.vocab_size,
16 hidden_size=exaone_config.hidden_size,
17 intermediate_size=exaone_config.intermediate_size,
18 num_hidden_layers=exaone_config.num_layers,
19 num_attention_heads=exaone_config.num_attention_heads,
20 max_position_embeddings=exaone_config.max_position_embeddings,
21 rms_norm_eps=exaone_config.layer_norm_epsilon,
22 num_key_value_heads=exaone_config.num_key_value_heads,
23 rope_theta=exaone_config.rope_theta,
24 bos_token_id=exaone_config.bos_token_id,
25 eos_token_id=exaone_config.eos_token_id,
26 pad_token_id=exaone_config.pad_token_id,
27 attention_bias=False,
28 )
29
30def copy_embedding_weights(llama_model, exaone_model):
31 llama_model.model.embed_tokens.weight.data = exaone_model.transformer.wte.weight.data.to(torch.float16)
32
33def copy_layer_weights(llama_layer, exaone_layer):
34 # Self-attention
35 llama_layer.self_attn.q_proj.weight.data = exaone_layer.attn.attention.q_proj.weight.data.to(torch.float16)
36 llama_layer.self_attn.k_proj.weight.data = exaone_layer.attn.attention.k_proj.weight.data.to(torch.float16)
37 llama_layer.self_attn.v_proj.weight.data = exaone_layer.attn.attention.v_proj.weight.data.to(torch.float16)
38 llama_layer.self_attn.o_proj.weight.data = exaone_layer.attn.attention.out_proj.weight.data.to(torch.float16)
39 # MLP
40 llama_layer.mlp.gate_proj.weight.data = exaone_layer.mlp.c_fc_0.weight.data.to(torch.float16)
41 llama_layer.mlp.up_proj.weight.data = exaone_layer.mlp.c_fc_1.weight.data.to(torch.float16)
42 llama_layer.mlp.down_proj.weight.data = exaone_layer.mlp.c_proj.weight.data.to(torch.float16)
43 # Layer Norms
44 llama_layer.input_layernorm.weight.data = exaone_layer.ln_1.weight.data.to(torch.float16)
45 llama_layer.post_attention_layernorm.weight.data = exaone_layer.ln_2.weight.data.to(torch.float16)
46
47def copy_final_weights(llama_model, exaone_model):
48 llama_model.model.norm.weight.data = exaone_model.transformer.ln_f.weight.data.to(torch.float16)
49 llama_model.lm_head.weight.data = exaone_model.lm_head.weight.data.to(torch.float16)
50
51def port_exaone_to_llama(exaone_model_path, llama_model_path):
52 print("Loading EXAONE model and tokenizer...")
53 exaone_model = AutoModelForCausalLM.from_pretrained(exaone_model_path, torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True)
54 exaone_tokenizer = AutoTokenizer.from_pretrained(exaone_model_path, trust_remote_code=True)
55 exaone_config = exaone_model.config
56
57 print("Creating Llama configuration...")
58 llama_config = create_llama_config(exaone_config)
59
60 print("Initializing Llama model...")
61 llama_model = LlamaForCausalLM(llama_config)
62 llama_model.to(torch.float16)
63 llama_model.to('cpu')
64
65 print("Copying weights...")
66 with torch.no_grad():
67 copy_embedding_weights(llama_model, exaone_model)
68
69 for i in tqdm(range(exaone_config.num_layers), desc="Copying layers"):
70 copy_layer_weights(llama_model.model.layers[i], exaone_model.transformer.h[i])
71 if i % 10 == 0: # Garbage collection every 10 layers
72 gc.collect()
73 if torch.cuda.is_available():
74 torch.cuda.empty_cache()
75
76 copy_final_weights(llama_model, exaone_model)
77
78 print("Unloading EXAONE model to free memory...")
79 unload_model(exaone_model)
80
81 print(f"Saving ported Llama model and tokenizer to {llama_model_path}")
82 llama_model.save_pretrained(llama_model_path, safe_serialization=True, max_shard_size="1GB")
83 exaone_tokenizer.save_pretrained(llama_model_path)
84
85 print("Unloading Llama model...")
86 unload_model(llama_model)
87
88 print(f"EXAONE model successfully ported to Llama format and saved at {llama_model_path}")
89
90if __name__ == "__main__":
91 exaone_model_path = "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct"
92 llama_model_path = "./exa_llamafied"
93 port_exaone_to_llama(exaone_model_path, llama_model_path)