deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-04-20 02:28:57 -04:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Added Try-Catch and Memory Optimization in Convert.py

Browse Source
This commit is contained in:
AzazAhmedLipu79 2025-01-28 12:15:22 +06:00
parent b5d872ead0
commit 41029d972d
1 changed files with 42 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
inference/convert.py
View File

@ -47,33 +47,49 @@ def main(hf_ckpt_path, save_path, n_experts, mp):
n_local_experts = n_experts // mp n_local_experts = n_experts // mp
state_dicts = [{} for _ in range(mp)] state_dicts = [{} for _ in range(mp)]
for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))): for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))):
with safe_open(file_path, framework="pt", device="cpu") as f: try:
for name in f.keys(): with safe_open(file_path, framework="pt", device="cpu") as f:
if "model.layers.61" in name: for name in f.keys():
continue if "model.layers.61" in name:
param: torch.Tensor = f.get_tensor(name) continue
if name.startswith("model."): param: torch.Tensor = f.get_tensor(name)
name = name[len("model."):] if name.startswith("model."):
name = name.replace("self_attn", "attn") name = name[len("model."):]
name = name.replace("mlp", "ffn") name = name.replace("self_attn", "attn").replace("mlp", "ffn")
name = name.replace("weight_scale_inv", "scale") name = name.replace("weight_scale_inv", "scale").replace("e_score_correction_bias", "bias")
name = name.replace("e_score_correction_bias", "bias")
key = name.split(".")[-2] key = name.split(".")[-2]
assert key in mapping assert key in mapping
new_key, dim = mapping[key] new_key, dim = mapping[key]
name = name.replace(key, new_key) name = name.replace(key, new_key)
for i in range(mp):
new_param = param for i in range(mp):
if "experts" in name and "shared_experts" not in name: new_param = param
idx = int(name.split(".")[-3]) if "experts" in name and "shared_experts" not in name:
if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts: idx = int(name.split(".")[-3])
continue if idx < i * n_local_experts or idx >= (i + 1) * n_local_experts:
elif dim is not None: continue
assert param.size(dim) % mp == 0 elif dim is not None:
shard_size = param.size(dim) // mp assert param.size(dim) % mp == 0
new_param = param.narrow(dim, i * shard_size, shard_size).contiguous() shard_size = param.size(dim) // mp
state_dicts[i][name] = new_param new_param = param.narrow(dim, i * shard_size, shard_size).contiguous()
state_dicts[i][name] = new_param
if len(state_dicts[i]) > 1000:
save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors"))
state_dicts[i].clear() # Free up memory
except Exception as e:
print(f"Error processing {file_path}: {e}")
# Final save step to ensure no data is lost
for i in range(mp):
if state_dicts[i]: # Only save if there's data left
save_file(state_dicts[i], os.path.join(save_path, f"model{i}-mp{mp}.safetensors"))
os.makedirs(save_path, exist_ok=True) os.makedirs(save_path, exist_ok=True)