llama2-22b / frankenllama_22b.py

Upload frankenmerge script

2bece07 10 months ago

No virus

6.14 kB

	#!/usr/bin/env python3
	# Charles O. Goddard
	# 7/20/2023
	"""Script used to generate the base frankenmerge. Output will need fine-tuning to be useful."""

	import copy
	import torch
	from torch import Tensor, nn
	import transformers

	from transformers.models.llama.modeling_llama import (
	LlamaForCausalLM,
	LlamaDecoderLayer,
	)
	from transformers import LlamaForCausalLM, LlamaConfig

	import torch
	import transformers
	import numpy as np


	MODEL_NAME_13B = "meta-llama/Llama-2-13b-hf" # primary model
	MODEL_NAME_33B = "huggyllama/llama-30b" # donor
	BLOCK_DIAGONAL = True
	# If BLOCK_DIAGONAL is set to True, each tensor in the resultant model will form a
	# block diagonal matrix, as illustrated below:

	# a a a 0 0
	# a a a 0 0
	# a a a 0 0
	# 0 0 0 b b
	# 0 0 0 b b

	# In this configuration, the states (hidden and intermediate) from the original
	# and donor models are completely decoupled. That is, the hidden states
	# corresponding to the original model remain unchanged, and the new dimensions
	# added from the donor model do not depend on the hidden states of the original model.

	# If BLOCK_DIAGONAL is set to False, the tensors will instead have the following form:

	# a a a 0 0
	# a a a 0 0
	# a a a 0 0
	# b b b b b
	# b b b b b

	# In this case, the output of the newly added attention heads depends on the hidden
	# state values as if they were part of the donor model. Although the original model's
	# hidden states remain unchanged in either case, interaction between the new and old
	# features will result in features of varying usefulness.


	class NoInit:
	def __enter__(self):
	def noop(args, *kwargs):
	pass

	(k, u, n) = (
	torch.nn.init.kaiming_uniform_,
	torch.nn.init.uniform_,
	torch.nn.init.normal_,
	)
	torch.nn.init.kaiming_uniform_ = noop
	torch.nn.init.uniform_ = noop
	torch.nn.init.normal_ = noop

	transformers.modeling_utils._init_weights = False
	self.funcs = (k, u, n)

	def __exit__(self, *args):
	(k, u, n) = self.funcs
	(
	torch.nn.init.kaiming_uniform_,
	torch.nn.init.uniform_,
	torch.nn.init.normal_,
	) = (
	k,
	u,
	n,
	)
	transformers.modeling_utils._init_weights = True


	def format_kmb(n, digits=None):
	n = int(n)
	if n < 1000:
	return str(n)
	elif n < 1000_000:
	return f"{round(n/1000, digits)}k"
	elif n < 1000 * 1000 * 1000:
	return f"{round(n/(1000*1000), digits)}m"
	else:
	return f"{round(n/(100010001000), digits)}b"


	def count_params(model):
	model_parameters = filter(lambda p: p.requires_grad, model.parameters())
	params = sum([np.prod(p.size()) for p in model_parameters])
	return int(params)


	torch.set_default_dtype(torch.float16)

	config_13b: LlamaConfig = LlamaConfig.from_pretrained(MODEL_NAME_13B)
	config_33b: LlamaConfig = LlamaConfig.from_pretrained(MODEL_NAME_33B)
	config_more = copy.deepcopy(config_13b)
	config_more.intermediate_size = config_33b.intermediate_size
	config_more.hidden_size = config_33b.hidden_size
	config_more.num_key_value_heads = config_33b.num_key_value_heads
	config_more.num_attention_heads = config_33b.num_key_value_heads

	print(config_more)

	with NoInit():
	model = LlamaForCausalLM(config_more)

	print(f"{format_kmb(count_params(model), 3)} parameters")


	def merge_tensors_inplace(dest: Tensor, s0: Tensor, s1: Tensor, block_diagonal: bool):
	dest.zero_()
	if block_diagonal:
	dest[s0.shape[0] :, s0.shape[1] :] = s1[
	s0.shape[0] : dest.shape[0],
	s0.shape[1] : dest.shape[1],
	]
	else:
	dest[s0.shape[0] :, :] = s1[
	s0.shape[0] : dest.shape[0],
	: dest.shape[1],
	]
	dest[: s0.shape[0], : s0.shape[1]] = s0


	with NoInit():
	donor_13b = (
	LlamaForCausalLM.from_pretrained(MODEL_NAME_13B).to(torch.float16).eval()
	)
	donor_33b = (
	LlamaForCausalLM.from_pretrained(MODEL_NAME_33B).to(torch.float16).eval()
	)

	with torch.no_grad():
	for layer_idx in range(len(model.model.layers)):
	layer: LlamaDecoderLayer = model.model.layers[layer_idx]
	l13: LlamaDecoderLayer = donor_13b.model.layers[layer_idx]
	l33: LlamaDecoderLayer = donor_33b.model.layers[layer_idx]

	for name in ("q_proj", "k_proj", "v_proj", "o_proj"):
	dest: nn.Linear = getattr(layer.self_attn, name)
	s13: nn.Linear = getattr(l13.self_attn, name)
	s33: nn.Linear = getattr(l33.self_attn, name)
	merge_tensors_inplace(dest.weight, s13.weight, s33.weight, BLOCK_DIAGONAL)

	for name in ("up_proj", "gate_proj", "down_proj"):
	dest: nn.Linear = getattr(layer.mlp, name)
	s13: nn.Linear = getattr(l13.mlp, name)
	s33: nn.Linear = getattr(l33.mlp, name)
	merge_tensors_inplace(dest.weight, s13.weight, s33.weight, BLOCK_DIAGONAL)

	layer.input_layernorm.weight[:] = l33.input_layernorm.weight[
	: layer.input_layernorm.weight.shape[0]
	]
	layer.input_layernorm.weight[
	: l13.input_layernorm.weight.shape[0]
	] = l13.input_layernorm.weight
	layer.post_attention_layernorm.weight[:] = l33.post_attention_layernorm.weight[
	: layer.post_attention_layernorm.weight.shape[0]
	]
	layer.post_attention_layernorm.weight[
	: l13.post_attention_layernorm.weight.shape[0]
	] = l13.post_attention_layernorm.weight

	# have initial output depend on only original llama2-13b features, so model
	# starts unimpaired and can learn to incorporate the new features as well
	model.lm_head.weight.zero_()
	model.lm_head.weight[
	: donor_13b.lm_head.weight.shape[0], : donor_13b.lm_head.weight.shape[1]
	] = donor_13b.lm_head.weight

	merge_tensors_inplace(
	model.model.embed_tokens.weight,
	donor_13b.model.embed_tokens.weight,
	donor_33b.model.embed_tokens.weight,
	BLOCK_DIAGONAL,
	)

	model.save_pretrained("./llama2-22b/", safe_serialization=True)