118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import multiprocessing as mp
|
|
import traceback
|
|
|
|
from safetensors.torch import save_file
|
|
import random
|
|
import time
|
|
import fcntl
|
|
import shutil
|
|
|
|
|
|
def save_weights(save_path, output):
|
|
tmp_path = save_path + ".tmp"
|
|
save_file(output, tmp_path)
|
|
|
|
with open(tmp_path, "rb") as f:
|
|
os.fsync(f.fileno())
|
|
os.replace(tmp_path, save_path)
|
|
|
|
|
|
# TODO: 分开加载，不覆盖
|
|
|
|
def worker(
|
|
global_lock,
|
|
task_queue: mp.Queue,
|
|
gpu_id: int,
|
|
weight_map_shared,
|
|
quant_map_shared,
|
|
output_dir: str,
|
|
mode: str, # dequant, quant, dequant+quant
|
|
model_type: str | None = None,
|
|
quant_config: dict | None = None, # from original config.json
|
|
custom_dequant_config: dict | None = None, # from .yml
|
|
custom_quant_config: dict | None = None, # from .yml
|
|
tensor_index: dict | None = None,
|
|
input_dir: str | None = None
|
|
):
|
|
try:
|
|
torch.cuda.set_device(gpu_id)
|
|
device = torch.device(f"cuda:{gpu_id}")
|
|
print(f"[Worker {gpu_id}] Started on {device}")
|
|
|
|
from functions.dequant import dequant
|
|
from functions.quant import quant
|
|
|
|
while True:
|
|
try:
|
|
sf = task_queue.get(timeout=5)
|
|
except Exception:
|
|
print(f"[Worker {gpu_id}] Task queue empty, exiting")
|
|
break
|
|
|
|
try:
|
|
print(f"[Worker {gpu_id}] Processing {sf}")
|
|
|
|
output = None
|
|
|
|
if "dequant" in mode:
|
|
output = dequant(
|
|
file_name=sf,
|
|
quant_config=quant_config,
|
|
custom_dequant_config=custom_dequant_config,
|
|
device=device,
|
|
model_type=model_type,
|
|
tensor_index=tensor_index,
|
|
input_dir=os.path.dirname(sf)
|
|
)
|
|
|
|
if "+quant" in mode:
|
|
output, quant_map = quant(
|
|
state_dict=output,
|
|
custom_quant_config=custom_quant_config,
|
|
device=device,
|
|
model_type=model_type
|
|
)
|
|
else:
|
|
if "quant" in mode:
|
|
output, quant_map = quant(
|
|
file_name=sf,
|
|
custom_quant_config=custom_quant_config,
|
|
device=device,
|
|
model_type=model_type
|
|
)
|
|
else:
|
|
raise ValueError(f"only support quant, dequant, and dequant+quant, get {mode}")
|
|
|
|
file_name = os.path.basename(sf)
|
|
|
|
for wname, _ in output.items():
|
|
weight_map_shared[wname] = file_name
|
|
|
|
try:
|
|
for wname, dtype in quant_map.items():
|
|
quant_map_shared[wname] = dtype
|
|
except Exception as e:
|
|
pass
|
|
|
|
save_path = os.path.join(output_dir, file_name)
|
|
sleep_time = random.uniform(0, 2)
|
|
time.sleep(sleep_time)
|
|
|
|
with global_lock:
|
|
print(f"[Worker {gpu_id}] Saving {file_name}")
|
|
save_weights(save_path, output)
|
|
|
|
# 主动释放显存
|
|
del output
|
|
torch.cuda.empty_cache()
|
|
|
|
except Exception as e:
|
|
print(f"[Worker {gpu_id}] Error processing {sf}")
|
|
traceback.print_exc()
|
|
|
|
except Exception:
|
|
traceback.print_exc()
|