初始化项目，由ModelHub XC社区提供模型
Model: chargoddard/llama2-22b Source: Original Platform
This commit is contained in:
ModelHub XC
188
frankenllama_22b.py
Normal file
188
frankenllama_22b.py
Normal file
@@ -0,0 +1,188 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user