160 lines
4.4 KiB
Python
160 lines
4.4 KiB
Python
import os
|
|
import torch
|
|
import json
|
|
import multiprocessing as mp
|
|
|
|
from glob import glob
|
|
try:
|
|
from natsort import natsorted
|
|
except ModuleNotFoundError:
|
|
# Model shards use zero-padded numeric names, so lexical order is equivalent.
|
|
natsorted = sorted
|
|
|
|
from worker import worker
|
|
from utils import (
|
|
load_quant_config,
|
|
copy_files,
|
|
get_ignore,
|
|
process_and_save_info,
|
|
)
|
|
|
|
|
|
def main(
|
|
input_dir,
|
|
output_dir,
|
|
model_type,
|
|
mode,
|
|
num_workers,
|
|
custom_dequant_config,
|
|
custom_quant_config,
|
|
retry_files
|
|
):
|
|
print(f"Make output_dir: {output_dir}")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
print(f"Copy files from {input_dir} to {output_dir}")
|
|
copy_files(input_dir, output_dir)
|
|
|
|
safetensor_files = natsorted(glob(os.path.join(input_dir, "*.safetensors")))
|
|
if retry_files is not None:
|
|
safetensor_files_ = []
|
|
for sf in safetensor_files:
|
|
for rf in retry_files:
|
|
if sf.endswith(rf):
|
|
safetensor_files_.append(sf)
|
|
safetensor_files = safetensor_files_
|
|
assert len(safetensor_files) > 0
|
|
|
|
num_gpus = torch.cuda.device_count()
|
|
num_workers = min(num_workers, min(num_gpus, len(safetensor_files)))
|
|
assert num_workers > 0, "No CUDA devices found or num_workers incorrect"
|
|
|
|
print(f"{num_workers} GPUs will be used")
|
|
print(f"Found {len(safetensor_files)} safetensor files")
|
|
|
|
# ============================
|
|
# 关键：统一使用 spawn context
|
|
# ============================
|
|
ctx = mp.get_context("spawn")
|
|
|
|
task_queue = ctx.Queue()
|
|
manager = ctx.Manager()
|
|
global_lock = ctx.Lock()
|
|
|
|
weight_map_shared = manager.dict()
|
|
quant_map_shared = manager.dict()
|
|
|
|
for sf in safetensor_files:
|
|
# TODO: add a filter for specified files
|
|
task_queue.put(sf)
|
|
try:
|
|
quant_config = load_quant_config(input_dir)
|
|
except Exception:
|
|
quant_config = None
|
|
|
|
processes = []
|
|
|
|
index_file = os.path.join(input_dir, "model.safetensors.index.json")
|
|
tensor_index = None
|
|
if os.path.exists(index_file):
|
|
print("Loading tensor index")
|
|
with open(index_file) as f:
|
|
tensor_index = json.load(f)["weight_map"]
|
|
|
|
for gpu_id in range(num_workers):
|
|
p = ctx.Process(
|
|
target=worker,
|
|
args=(
|
|
global_lock,
|
|
task_queue,
|
|
gpu_id,
|
|
weight_map_shared,
|
|
quant_map_shared,
|
|
output_dir,
|
|
mode,
|
|
model_type,
|
|
quant_config,
|
|
custom_dequant_config,
|
|
custom_quant_config,
|
|
tensor_index
|
|
),
|
|
)
|
|
p.start()
|
|
processes.append(p)
|
|
|
|
for p in processes:
|
|
p.join()
|
|
|
|
failed_workers = [(p.pid, p.exitcode) for p in processes if p.exitcode != 0]
|
|
missing_shards = [
|
|
os.path.basename(sf) for sf in safetensor_files
|
|
if not os.path.isfile(os.path.join(output_dir, os.path.basename(sf)))
|
|
]
|
|
if failed_workers or missing_shards:
|
|
raise RuntimeError(
|
|
f"Quantization incomplete: workers={failed_workers}, missing_shards={missing_shards}"
|
|
)
|
|
|
|
output_model = output_dir.split(os.sep)[-1]
|
|
qms_output_dir = os.path.join("online_quant_configs", output_model)
|
|
os.makedirs(qms_output_dir, exist_ok=True)
|
|
|
|
process_and_save_info(
|
|
input_dir,
|
|
output_dir,
|
|
qms_output_dir,
|
|
dict(weight_map_shared),
|
|
mode in ("quant", "dequant+quant"),
|
|
custom_quant_config.get("symmetric"),
|
|
ignore=None if mode == "dequant" else get_ignore(custom_quant_config),
|
|
group_size=custom_quant_config.get("group_size", -1),
|
|
)
|
|
|
|
if not("dequant" in mode and "+quant" not in mode):
|
|
with open(os.path.join(qms_output_dir, "quant_map.json"), "w") as f:
|
|
json.dump(dict(quant_map_shared), f, indent=2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import yaml
|
|
from argparse import ArgumentParser
|
|
|
|
mp.set_start_method("spawn", force=True)
|
|
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--config", type=str)
|
|
args = parser.parse_args()
|
|
|
|
with open(args.config, "r") as f:
|
|
cfg = yaml.safe_load(f)
|
|
|
|
main(
|
|
cfg['paths']['checkpoint_dir'],
|
|
cfg['paths']['output_dir'],
|
|
cfg['model']['model_type'],
|
|
cfg['runtime']['mode'],
|
|
cfg['runtime']['num_workers'],
|
|
cfg['dequant'],
|
|
cfg['quant'],
|
|
cfg['retry_files']
|
|
) |