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
state_dicts = [{} for _ in range(mp)]
for file_path in tqdm(glob(os.path.join(hf_ckpt_path, "*.safetensors"))):
with safe_open(file_path, framework="pt", device="cpu") as f:
for name in f.keys():
if "model.layers.61" in name:
continue
param: torch.Tensor = f.get_tensor(name)
if name.startswith("model."):
name = name[len("model."):]
name = name.replace("self_attn", "attn")
name = name.replace("mlp", "ffn")
name = name.replace("weight_scale_inv", "scale")
name = name.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:
assert param.size(dim) % mp == 0
shard_size = param.size(dim) // mp
new_param = param.narrow(dim, i * shard_size, shard_size).contiguous()
state_dicts[i][name] = new_param
try:
with safe_open(file_path, framework="pt", device="cpu") as f:
for name in f.keys():
if "model.layers.61" in name:
continue
param: torch.Tensor = f.get_tensor(name)
if name.startswith("model."):
name = name[len("model."):]
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:
assert param.size(dim) % mp == 0
shard_size = param.size(dim) // mp
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)