143 lines
5.8 KiB
Python
143 lines
5.8 KiB
Python
"""
|
|
torchrun --nproc_per_node=1 convert_hf_to_picotron.py --save_path smollm.pth
|
|
"""
|
|
import os
|
|
import argparse
|
|
from tqdm import tqdm
|
|
import torch, torch.distributed as dist
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from utils import set_all_seed
|
|
import lovely_tensors as lt; lt.monkey_patch()
|
|
|
|
from model import Llama
|
|
from distributed.process_group_manager import setup_process_group_manager
|
|
|
|
def sanity_check_weights(model, model_hf, picotron_to_hf):
|
|
|
|
total, fail = 0, 0
|
|
|
|
state_dict = model.state_dict()
|
|
state_dict_hf = model_hf.state_dict()
|
|
|
|
for name, name_hf in picotron_to_hf.items():
|
|
|
|
param_hf = state_dict_hf[name_hf]
|
|
param = state_dict[name]
|
|
|
|
total += 1
|
|
try:
|
|
torch.testing.assert_close(param_hf, param, rtol=1e-10, atol=1e-10)
|
|
except AssertionError as e:
|
|
print(f"{name_hf} and {name} are not equal")
|
|
fail += 1
|
|
|
|
if fail == 0:
|
|
print("All parameters are equal")
|
|
else:
|
|
AssertionError(f"{fail}/{total} parameters are not equal")
|
|
|
|
def sanity_check_generation(model, model_hf, model_name, prompt, max_new_tokens):
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
|
input_ids_hf = tokenizer.encode(prompt, return_tensors="pt").to(device=model_hf.device)
|
|
input_ids = input_ids_hf.clone().to(device=model_hf.device)
|
|
|
|
for _ in range(max_new_tokens):
|
|
# picotron model
|
|
seq_len = input_ids.shape[1]
|
|
position_index = torch.arange(seq_len).view(1, -1).to(device=model_hf.device)
|
|
|
|
logits = model(input_ids=input_ids, position_ids=position_index)
|
|
next_token = torch.argmax(logits, dim=-1)
|
|
input_ids = torch.cat([input_ids, next_token[:, -1].unsqueeze(-1)], dim=-1)
|
|
|
|
# HF model
|
|
logits_hf = model_hf(input_ids_hf).logits
|
|
next_token_hf = torch.argmax(logits_hf[:, -1, :], dim=-1)
|
|
input_ids_hf = torch.cat([input_ids_hf, next_token_hf.unsqueeze(0)], dim=-1)
|
|
|
|
# Assert logits are equal
|
|
torch.testing.assert_close(logits, logits_hf, atol=1e-4, rtol=1e-4)
|
|
|
|
print("Input prompt:\n", prompt)
|
|
print("Reference model output:\n", tokenizer.decode(input_ids_hf[0], skip_special_tokens=True))
|
|
print("picotron model output:\n", tokenizer.decode(input_ids[0], skip_special_tokens=True))
|
|
|
|
def get_weights_mapping(model_hf, to_hf):
|
|
|
|
hf_to_picotron = {}
|
|
|
|
hf_to_picotron["model.embed_tokens.weight"] = "embedding.weight"
|
|
hf_to_picotron["model.norm.weight"] = "final_norm.weight"
|
|
hf_to_picotron["lm_head.weight"] = "final_proj.weight"
|
|
|
|
for i in range(model_hf.config.num_hidden_layers):
|
|
# Attention
|
|
hf_to_picotron[f"model.layers.{i}.self_attn.q_proj.weight"] = f"decoder_layers.{i}.attention.q_proj.weight"
|
|
hf_to_picotron[f"model.layers.{i}.self_attn.k_proj.weight"] = f"decoder_layers.{i}.attention.k_proj.weight"
|
|
hf_to_picotron[f"model.layers.{i}.self_attn.v_proj.weight"] = f"decoder_layers.{i}.attention.v_proj.weight"
|
|
hf_to_picotron[f"model.layers.{i}.self_attn.o_proj.weight"] = f"decoder_layers.{i}.attention.o_proj.weight"
|
|
# MLP
|
|
hf_to_picotron[f"model.layers.{i}.mlp.gate_proj.weight"] = f"decoder_layers.{i}.mlp.gate_proj.weight"
|
|
hf_to_picotron[f"model.layers.{i}.mlp.up_proj.weight"] = f"decoder_layers.{i}.mlp.up_proj.weight"
|
|
hf_to_picotron[f"model.layers.{i}.mlp.down_proj.weight"] = f"decoder_layers.{i}.mlp.down_proj.weight"
|
|
|
|
hf_to_picotron[f"model.layers.{i}.input_layernorm.weight"] = f"decoder_layers.{i}.norm_attn.weight"
|
|
hf_to_picotron[f"model.layers.{i}.post_attention_layernorm.weight"] = f"decoder_layers.{i}.norm_mlp.weight"
|
|
|
|
# check if we have takens all keys from the reference model
|
|
for key in hf_to_picotron:
|
|
assert key in model_hf.state_dict(), f"{key} not found in reference model"
|
|
|
|
if to_hf:
|
|
# Mapping from picotron to hf
|
|
picotron_to_hf = {v: k for k, v in hf_to_picotron.items()}
|
|
return picotron_to_hf
|
|
|
|
return hf_to_picotron
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Convert HF llama weights to picotron")
|
|
parser.add_argument("--save_path", type=str, default="smollm.pth")
|
|
parser.add_argument("--model_name", type=str, default="HuggingFaceTB/SmolLM-360M-Instruct")
|
|
parser.add_argument("--prompt", type=str, default="My name is")
|
|
parser.add_argument("--max_new_tokens", type=int, default=50)
|
|
|
|
args = parser.parse_args()
|
|
|
|
local_rank, world_size = int(os.environ["LOCAL_RANK"]), int(os.environ["WORLD_SIZE"])
|
|
|
|
#TODO: add gloo backend for generation
|
|
dist.init_process_group(backend="nccl")
|
|
torch.cuda.set_device(local_rank)
|
|
device = torch.device("cuda", local_rank)
|
|
setup_process_group_manager(tp_size=1, pp_size=1, dp_size=1, cp_size=1)
|
|
set_all_seed(seed=42)
|
|
|
|
model_hf = AutoModelForCausalLM.from_pretrained(args.model_name).to(device)
|
|
|
|
model = Llama(config=model_hf.config, device=device)
|
|
picotron_to_hf = get_weights_mapping(model_hf, to_hf=True)
|
|
|
|
ref_state_dict = model_hf.state_dict()
|
|
|
|
for name, param in tqdm(
|
|
model.named_parameters(),
|
|
total=len(list(model.named_parameters())),
|
|
desc="Converting",
|
|
):
|
|
if name in picotron_to_hf:
|
|
ref_name = picotron_to_hf[name]
|
|
ref_param = ref_state_dict[ref_name]
|
|
param.data.copy_(ref_param)
|
|
|
|
torch.save(model.state_dict(), args.save_path)
|
|
|
|
new_model = Llama(config=model_hf.config, device=device)
|
|
new_model.load_state_dict(torch.load(args.save_path))
|
|
|
|
print("Sanity check weight ...")
|
|
sanity_check_weights(new_model, model_hf, picotron_to_hf)
|
|
print("Sanity check generation ...")
|
|
sanity_check_generation(new_model, model_hf, args.model_name, args.prompt, args.max_new_tokens)
|
|
print("Conversion successful") |