File size: 5,519 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | import os
import re
from argparse import ArgumentParser
from dataclasses import dataclass
from glob import glob
from typing import Generator
import torch
from safetensors.torch import safe_open
from tqdm import tqdm
from transformers import AutoConfig
from veomni.models import build_tokenizer, save_model_weights
@dataclass
class StateDictIterator:
filepath: str
def __iter__(self) -> Generator[tuple[str, torch.Tensor], None, None]:
if self.filepath.endswith(".safetensors"):
with safe_open(self.filepath, framework="pt", device="cpu") as f:
for key in f.keys():
yield key, f.get_tensor(key)
else:
state_dict = torch.load(self.filepath, map_location="cpu", weights_only=True, mmap=True)
for key in state_dict.keys():
yield key, state_dict[key]
def moe_merge(state_dict: dict[str, torch.Tensor], config) -> dict[str, torch.Tensor]:
new_state_dict: dict[str, torch.Tensor] = dict()
processed_keys: set[str] = set()
num_layers = config.num_hidden_layers
num_experts = config.num_experts
first_k_dense_replace = config.first_k_dense_replace
print(f"Merging {num_layers} layers with {num_experts} experts each")
proj_types = ["gate_proj", "up_proj", "down_proj"]
for layer_id in range(first_k_dense_replace, num_layers):
for proj_type in proj_types:
expert_weights = []
current_expert_keys = []
for expert_id in range(num_experts):
expert_key = f"model.layers.{layer_id}.mlp.experts.{expert_id}.{proj_type}.weight"
assert expert_key in state_dict, f"Missing key: {expert_key}"
expert_weights.append(state_dict[expert_key])
current_expert_keys.append(expert_key)
assert len(expert_weights) == num_experts
merged_weight = torch.stack(expert_weights, dim=0)
new_key = f"model.layers.{layer_id}.mlp.experts.{proj_type}"
new_state_dict[new_key] = merged_weight
processed_keys.update(current_expert_keys)
for key in current_expert_keys:
del state_dict[key]
print(f"✓ Layer {layer_id}.{proj_type}: {expert_weights[0].shape} -> {merged_weight.shape}")
del expert_weights
for key, tensor in state_dict.items():
if key not in processed_keys:
new_state_dict[key] = tensor
return new_state_dict
def split_moe_experts(
merged_state_dict: dict[str, torch.Tensor], config
) -> dict[str, torch.Tensor]:
split_state_dict: dict[str, torch.Tensor] = dict()
num_experts = getattr(config, "num_experts", None)
if num_experts is None:
raise ValueError(
"Could not find the 'num_experts' attribute in the configuration. "
"Please ensure the correct model configuration is passed."
)
proj_types = ["gate_proj", "up_proj", "down_proj"]
merged_key_pattern = re.compile(
r"model\.layers\.(\d+)\.mlp\.experts\.(" + "|".join(proj_types) + r")$"
)
for key, merged_tensor in merged_state_dict.items():
match = merged_key_pattern.match(key)
if match:
layer_id = match.group(1)
proj_type = match.group(2)
if not (merged_tensor.dim() > 1 and merged_tensor.shape[0] == num_experts):
raise ValueError(
f"Tensor '{key}' has an unexpected shape {merged_tensor.shape}. "
f"Its first dimension should be equal to the number of experts ({num_experts})."
)
for expert_id in range(num_experts):
expert_tensor = merged_tensor[expert_id]
original_key = f"model.layers.{layer_id}.mlp.experts.{expert_id}.{proj_type}.weight"
split_state_dict[original_key] = expert_tensor
print(f"Unmerged {key} -> {num_experts} individual expert weights")
else:
split_state_dict[key] = merged_tensor
return split_state_dict
def main(input_path, output_path):
torch.set_default_dtype(torch.bfloat16)
os.makedirs(output_path, exist_ok=True)
config = AutoConfig.from_pretrained(input_path, trust_remote_code=True)
tokenizer = build_tokenizer(input_path)
safetensor_files = list(glob(os.path.join(input_path, "*.safetensors")))
safetensor_files.sort()
state_dict_iterators = [StateDictIterator(shard_file) for shard_file in safetensor_files]
state_dict = {}
for state_dict_iterator in tqdm(state_dict_iterators, desc="Loading checkpoint shards"):
for name, tensor in state_dict_iterator:
state_dict[name] = tensor.cpu()
if args.mode == "merge":
new_state_dict = moe_merge(state_dict, config)
elif args.mode == "split":
new_state_dict = split_moe_experts(state_dict, config)
else:
raise ValueError("unsupport mode")
state_dict.clear()
model_assets = [config, tokenizer]
save_model_weights(output_path, new_state_dict, model_assets=model_assets)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("-i", "--input-path", type=str, required=True)
parser.add_argument("-o", "--output-path", type=str, required=True)
parser.add_argument("-m", "--mode", type=str, default="merge", choices=["merge", "split"])
args = parser.parse_args()
main(args.input_path, args.output_path)
|