first commit

2026-10-23 15:00:59 +08:00
commit 4627899899
156 changed files with 398139 additions and 0 deletions
--- a/dev/worker.py
+++ b/dev/worker.py
@@ -0,0 +1,117 @@
+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()