Spaces:

iblfe
/

test

Runtime error

App Files Files Community

test / src /export_hf_checkpoint.py

iblfe

Upload folder using huggingface_hub

b585c7f verified 5 months ago

raw history blame contribute delete

No virus

9.92 kB

	import os
	import json
	import shutil
	import subprocess

	import torch
	from accelerate import infer_auto_device_map, dispatch_model
	from accelerate.utils import get_balanced_memory
	from peft import PeftModel
	from transformers import PreTrainedModel


	def do_export():
	BASE_MODEL = 'h2oai/h2ogpt-4096-llama2-13b-chat'
	LORA_WEIGHTS = 'Llama-2-13b-chat-hf.h2oaiopenassistant_oasst1_h2ogpt_llama2_chat.1_epochs.b2aed9250804d815c258976c98ce968bacd88389.7'
	OUTPUT_NAME = "h2ogpt-oasst1-4096-llama2-13b"

	BASE_MODEL = 'meta-llama/Llama-2-7b-chat-hf'
	LORA_WEIGHTS = 'Llama-2-7b-chat-hf.h2oaiopenassistant_oasst1_h2ogpt_llama2_chat.1_epochs.0c6b906f73b5639fd1d53c74fecbc9cf64f0f225.8'
	OUTPUT_NAME = "h2ogpt-oasst1-4096-llama2-7b"

	BASE_MODEL = 'meta-llama/Llama-2-70b-chat-hf'
	LORA_WEIGHTS = 'Llama-2-70b-chat-hf.h2oaiopenassistant_oasst1_h2ogpt_llama2_chat.1_epochs.0c6b906f73b5639fd1d53c74fecbc9cf64f0f225.6'
	OUTPUT_NAME = "h2ogpt-oasst1-4096-llama2-70b"

	base_model = os.getenv('BASE_MODEL')
	output = os.getenv('MODEL')
	# for testing
	if base_model and output:
	BASE_MODEL = base_model
	LORA_WEIGHTS = output + ".lora"
	OUTPUT_NAME = output

	llama_type = "llama" in BASE_MODEL
	as_pytorch = False # False -> HF

	from loaders import get_loaders
	model_loader, tokenizer_loader, conditional_type = (
	get_loaders(model_name=BASE_MODEL, reward_type=False, llama_type=llama_type))

	tokenizer = tokenizer_loader.from_pretrained(
	BASE_MODEL,
	local_files_only=False,
	resume_download=True,
	)
	tokenizer.save_pretrained(OUTPUT_NAME)

	base_model = model_loader(
	BASE_MODEL,
	load_in_8bit=False,
	trust_remote_code=True,
	torch_dtype=torch.float16,
	device_map={"": "cpu"},
	)

	print(base_model)
	if llama_type:
	layers = base_model.model.layers
	first_weight = layers[0].self_attn.q_proj.weight
	else:
	if any([x in BASE_MODEL.lower() for x in ["pythia", "h2ogpt", "gpt-neox"]]):
	layers = base_model.gpt_neox.base_model.layers
	first_weight = layers[0].attention.query_key_value.weight
	elif any([x in BASE_MODEL.lower() for x in ["falcon"]]):
	first_weight = base_model.transformer.h._modules['0'].self_attention.query_key_value.weight
	else:
	layers = base_model.transformer.base_model.h
	first_weight = layers[0].attn.q_proj.weight
	first_weight_old = first_weight.clone()

	lora_model = PeftModel.from_pretrained(
	base_model,
	LORA_WEIGHTS,
	device_map={"": "cpu"},
	torch_dtype=torch.float16,
	)

	assert torch.allclose(first_weight_old, first_weight)

	# merge weights TODO: include all lora_target_modules, not just default ones
	if llama_type:
	merged_model = lora_model.merge_and_unload()
	# for layer in lora_model.base_model.model.model.layers:
	# layer.self_attn.q_proj.merge_weights = True
	# layer.self_attn.k_proj.merge_weights = True
	# layer.self_attn.v_proj.merge_weights = True
	# layer.self_attn.o_proj.merge_weights = True
	else:
	if any([x in BASE_MODEL.lower() for x in ["pythia", "gpt-neox"]]):
	for layer in lora_model.base_model.gpt_neox.base_model.layers:
	layer.attention.query_key_value.merge_weights = True
	merged_model = lora_model
	else:
	merged_model = lora_model.merge_and_unload()
	# for layer in lora_model.base_model.transformer.base_model.h:
	# layer.attn.q_proj.merge_weights = True
	# layer.attn.v_proj.merge_weights = True

	# max_memory = get_balanced_memory(merged_model)
	# device_map = infer_auto_device_map(merged_model, max_memory=max_memory)
	# merged_model = dispatch_model(
	# merged_model,
	# device_map=device_map,
	# )
	merged_model.eval()
	print(merged_model)

	# did we do anything?
	assert not torch.allclose(first_weight_old, first_weight)

	merged_model_sd = merged_model.state_dict()

	if as_pytorch:
	# FIXME - might not be generic enough still
	params = {
	"dim": base_model.config.hidden_size,
	"n_heads": base_model.config.num_attention_heads,
	"n_layers": base_model.config.num_hidden_layers,
	"norm_eps": base_model.config.layer_norm_eps,
	"vocab_size": base_model.config.vocab_size,
	}
	n_layers = params["n_layers"]
	n_heads = params["n_heads"]
	dim = params["dim"]
	dims_per_head = dim // n_heads
	base = 10000.0
	inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))

	def permute(w):
	return (
	w.view(n_heads, dim // n_heads // 2, 2, dim).transpose(1, 2).reshape(dim, dim)
	)


	def unpermute(w):
	return (
	w.view(n_heads, 2, dim // n_heads // 2, dim).transpose(1, 2).reshape(dim, dim)
	)


	def translate_state_dict_key(k):
	if "gpt-neoxt" in BASE_MODEL.lower():
	k = k.replace("gpt_neox.model.", "")
	else:
	k = k.replace("base_model.model.", "")
	if k == "model.embed_tokens.weight":
	return "tok_embeddings.weight"
	elif k == "model.norm.weight":
	return "norm.weight"
	elif k == "lm_head.weight":
	return "output.weight"
	elif k.startswith("model.layers."):
	layer = k.split(".")[2]
	if k.endswith(".self_attn.q_proj.weight"):
	return f"layers.{layer}.attention.wq.weight"
	elif k.endswith(".self_attn.k_proj.weight"):
	return f"layers.{layer}.attention.wk.weight"
	elif k.endswith(".self_attn.v_proj.weight"):
	return f"layers.{layer}.attention.wv.weight"
	elif k.endswith(".self_attn.o_proj.weight"):
	return f"layers.{layer}.attention.wo.weight"
	elif k.endswith(".mlp.gate_proj.weight"):
	return f"layers.{layer}.feed_forward.w1.weight"
	elif k.endswith(".mlp.down_proj.weight"):
	return f"layers.{layer}.feed_forward.w2.weight"
	elif k.endswith(".mlp.up_proj.weight"):
	return f"layers.{layer}.feed_forward.w3.weight"
	elif k.endswith(".input_layernorm.weight"):
	return f"layers.{layer}.attention_norm.weight"
	elif k.endswith(".post_attention_layernorm.weight"):
	return f"layers.{layer}.ffn_norm.weight"
	elif k.endswith("rotary_emb.inv_freq") or "lora" in k:
	return None
	else:
	print(layer, k)
	raise NotImplementedError
	else:
	print(k)
	raise NotImplementedError


	new_state_dict = {}
	for k, v in merged_model_sd.items():
	new_k = translate_state_dict_key(k)
	if new_k is not None:
	if "wq" in new_k or "wk" in new_k:
	new_state_dict[new_k] = unpermute(v)
	else:
	new_state_dict[new_k] = v

	os.makedirs("./ckpt", exist_ok=True)

	torch.save(new_state_dict, "./ckpt/consolidated.00.pth")

	with open("./ckpt/params.json", "w") as f:
	json.dump(params, f)
	else:
	# deloreanized_sd = {
	# k.replace("base_model.model.", ""): v
	# for k, v in merged_model_sd.items()
	# if "lora" not in k
	# }
	merged_model.config.custom_pipelines = {
	"text-generation": {
	"impl": "h2oai_pipeline.H2OTextGenerationPipeline",
	"pt": "AutoModelForCausalLM"
	}
	}
	PreTrainedModel.save_pretrained(
	merged_model,
	OUTPUT_NAME,
	# state_dict=deloreanized_sd,
	# max_shard_size="5GB",
	)

	do_copy(OUTPUT_NAME)
	test_copy()


	def do_copy(OUTPUT_NAME):
	dest_file = os.path.join(OUTPUT_NAME, "h2oai_pipeline.py")
	shutil.copyfile("src/h2oai_pipeline.py", dest_file)
	os.system("""sed -i 's/from enums.*//g' %s""" % dest_file)
	os.system("""sed -i 's/from stopping.*//g' %s""" % dest_file)
	os.system("""sed -i 's/from prompter.*//g' %s""" % dest_file)
	os.system("""cat %s\|grep -v "from enums import PromptType" >> %s""" % ('src/enums.py', dest_file))
	os.system("""cat %s\|grep -v "from enums import PromptType" >> %s""" % ('src/prompter.py', dest_file))
	os.system("""cat %s\|grep -v "from enums import PromptType" >> %s""" % ('src/stopping.py', dest_file))


	TEST_OUTPUT_NAME = "test_output"


	def test_copy():
	if os.path.isdir(TEST_OUTPUT_NAME):
	shutil.rmtree(TEST_OUTPUT_NAME)
	os.makedirs(TEST_OUTPUT_NAME, exist_ok=False)
	do_copy(TEST_OUTPUT_NAME)
	shutil.copy('src/export_hf_checkpoint.py', TEST_OUTPUT_NAME)
	os.environ['DO_COPY_TEST'] = '1'
	os.chdir(TEST_OUTPUT_NAME)
	output = subprocess.check_output(['python', 'export_hf_checkpoint.py'])
	print(output)


	def inner_test_copy():
	"""
	pytest -s -v export_hf_checkpoint.py::test_copy
	:return:
	"""
	# test imports
	# below supposed to look bad in pycharm, don't fix!
	from h2oai_pipeline import get_stopping, get_prompt, H2OTextGenerationPipeline
	assert get_stopping
	assert get_prompt
	assert H2OTextGenerationPipeline


	if __name__ == '__main__':
	if os.getenv('DO_COPY_TEST'):
	inner_test_copy()
	else:
	do_export()
	# uncomment for raw isolated test, but test is done every time for each export now
	# test_copy()