mirror of
https://github.com/deepseek-ai/DeepSeek-V3.git
synced 2025-02-23 14:18:57 -05:00
128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
import os
|
|
import shutil
|
|
import mmap
|
|
import threading
|
|
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
from tqdm import tqdm, trange
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import torch
|
|
from safetensors.torch import safe_open, save_file
|
|
from collections import defaultdict
|
|
|
|
mapping = {
|
|
"embed_tokens": ("embed", 0),
|
|
"input_layernorm": ("attn_norm", None),
|
|
"post_attention_layernorm": ("ffn_norm", None),
|
|
"q_proj": ("wq", 0),
|
|
"q_a_proj": ("wq_a", None),
|
|
"q_a_layernorm": ("q_norm", None),
|
|
"q_b_proj": ("wq_b", 0),
|
|
"kv_a_proj_with_mqa": ("wkv_a", None),
|
|
"kv_a_layernorm": ("kv_norm", None),
|
|
"kv_b_proj": ("wkv_b", 0),
|
|
"o_proj": ("wo", 1),
|
|
"gate": ("gate", None),
|
|
"gate_proj": ("w1", 0),
|
|
"down_proj": ("w2", 1),
|
|
"up_proj": ("w3", 0),
|
|
"norm": ("norm", None),
|
|
"lm_head": ("head", 0),
|
|
"scale": ("scale", None),
|
|
}
|
|
|
|
# Thread Lock for Safe Dictionary Access
|
|
state_lock = threading.Lock()
|
|
|
|
def fast_copy(src: Path, dst: Path):
|
|
"""Efficiently copies large files using shutil for optimal memory usage"""
|
|
if dst.exists():
|
|
dst.unlink() # Remove file if it already exists
|
|
if src.stat().st_size < 10 * 1024 * 1024: # If file < 10MB, use shutil
|
|
shutil.copyfile(src, dst)
|
|
else:
|
|
with open(src, "rb") as f_src, open(dst, "wb") as f_dst:
|
|
shutil.copyfileobj(f_src, f_dst, length=16*1024*1024)
|
|
|
|
def copy_token_file(file_path, save_path):
|
|
"""Helper function for parallel copying of token files"""
|
|
fast_copy(file_path, Path(save_path) / file_path.name)
|
|
|
|
def inner_safe_open(name: str, f, mp, state_dicts, n_local_experts):
|
|
"""Processes tensor files and maps keys correctly"""
|
|
with torch.no_grad():
|
|
param: torch.Tensor = f.get_tensor(name)
|
|
name = name[len("model."):] if name.startswith("model.") else name
|
|
name = name.replace("self_attn", "attn").replace("mlp", "ffn")
|
|
name = name.replace("weight_scale_inv", "scale").replace("e_score_correction_bias", "bias")
|
|
key = name.split(".")[-2]
|
|
assert key in mapping
|
|
new_key, dim = mapping[key]
|
|
name = name.replace(key, new_key)
|
|
|
|
for i in range(mp):
|
|
new_param = param
|
|
if "experts" in name and "shared_experts" not in name:
|
|
idx = int(name.split(".")[-3])
|
|
if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts:
|
|
continue
|
|
elif dim is not None:
|
|
shard_size = param.size(dim) // mp
|
|
new_param = param.narrow(dim, i * shard_size, shard_size).contiguous()
|
|
|
|
# Lock to avoid race conditions
|
|
with state_lock:
|
|
state_dicts[i][name] = new_param
|
|
|
|
def process_file(file_path, mp, state_dicts, n_local_experts):
|
|
"""Processes a single safetensor file"""
|
|
with safe_open(file_path, framework="pt", device="cpu") as f:
|
|
for name in f.keys():
|
|
if "model.layers.61" not in name:
|
|
inner_safe_open(name, f, mp, state_dicts, n_local_experts)
|
|
|
|
def main(hf_ckpt_path, save_path, n_experts, mp):
|
|
"""Converts and saves model checkpoint files into a specified format."""
|
|
n_local_experts = n_experts // mp
|
|
|
|
# Use defaultdict to prevent key errors in multi-threading
|
|
state_dicts = [defaultdict(dict) for _ in range(mp)]
|
|
|
|
file_list = list(Path(hf_ckpt_path).glob("*.safetensors"))
|
|
token_files = list(Path(hf_ckpt_path).glob("*token*"))
|
|
|
|
Path(save_path).mkdir(parents=True, exist_ok=True)
|
|
|
|
# Parallel Processing with ThreadPoolExecutor
|
|
with ThreadPoolExecutor() as executor:
|
|
futures = {
|
|
executor.submit(process_file, file, mp, state_dicts, n_local_experts): file
|
|
for file in file_list
|
|
}
|
|
for future in tqdm(as_completed(futures), desc="Processing safetensors", total=len(file_list)):
|
|
future.result() # Ensure exceptions are raised
|
|
|
|
# Save processed model shards
|
|
for i in trange(mp, desc="Saving model shards"):
|
|
save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors"))
|
|
|
|
# Parallel Token File Copying
|
|
with ThreadPoolExecutor() as executor:
|
|
futures = {
|
|
executor.submit(copy_token_file, file, save_path): file
|
|
for file in token_files
|
|
}
|
|
for future in tqdm(as_completed(futures), desc="Copying token files", total=len(token_files)):
|
|
future.result() # Ensure exceptions are raised
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--hf-ckpt-path", type=str, required=True)
|
|
parser.add_argument("--save-path", type=str, required=True)
|
|
parser.add_argument("--n-experts", type=int, required=True)
|
|
parser.add_argument("--model-parallel", type=int, required=True)
|
|
args = parser.parse_args()
|
|
|
|
assert args.n_experts % args.model_parallel == 0, "n_experts must be divisible by model_parallel"
|
|
main(args.hf_ckpt_path, args.save_path, args.n_experts, args.model_parallel)
|