first commit

dev/main.py

@@ -0,0 +1,160 @@
+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']
+    )