llama2-22b / frankenllama_22b.py
chargoddard's picture
Upload frankenmerge script
2bece07
raw
history blame
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/(1000*1000*1000), 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)