Update Python version to 3.11, enhance weight processing by adding layer number…

Update Python version to 3.11, enhance weight processing by adding layer number extraction, and implement output directory cleanup for optimized and origin generation scripts.

.gitignore

@@ -2,8 +2,7 @@ __pycache__
 .DS_Store
 .venv
 .vscode
-outputs
-outputs-qwen
+outputs-*
 weights
 model.safetensors
 001-H-LLM

.python-version

-3.12
+3.11

hllm/config.py

@@ -24,21 +24,30 @@ class CFG:
             "v.pkl",
             "q.pkl",
         ]
+        self.group_numbers_batch = [
+            8,
+            32,
+            32,
+            32,
+            32,
+            32,
+            32,
+        ]
+
         self.run_weights = "down.pkl"  # 用于赋值
-        self.safetensors = "model.safetensors"
+        self.group_number = 32
 
-        self.weights_dir = "001-H-LLM/qwen"
-        self.mapped_weights_dir = "001-H-LLM/qwen/mapped_weights"
+        self.safetensors_path = "001-H-LLM/qwen/model.safetensors"
 
-        self.verify_generate_activation_on_exist = False
+        self.weights_dir = "001-H-LLM/qwen/weights"
+        self.mapped_weights_dir = "001-H-LLM/qwen/mapped_weights"
 
         self.num_workers = 64
         self.value_range = [-6, -4, -3, -2, -1.5, -1, -0.5, 0.5, 1, 1.5, 2, 3, 4, 6]
         self.python_path = sys.executable
 
-        self.group_number = 32
-
-        self.output_dir = "outputs"
+        self.output_dir = "outputs-qwen"
+        self.verify_generate_activation_on_exist = False
 
         os.makedirs(self.weights_dir, exist_ok=True)
         os.makedirs(self.output_dir, exist_ok=True)

hllm/eda/generate_quant_weights.py

 import os
+import shutil
 import numpy as np
 import safetensors
 import pickle
+import re
 from tqdm import tqdm
 from hllm.eda.utils_quant import quant_and_dequant
 from hllm.config import CFG
@@ -26,10 +28,26 @@ def is_substring_in_list(substring, string_list):
     return any(s in substring for s in string_list)
 
 
+def get_layer_number(key):
+    # 使用正则表达式匹配layers后面的数字
+    match = re.search(r"layers\.(\d+)\.", key)
+    if match:
+        return int(match.group(1))
+    return -1  # 对于不包含层号的键（如norm），返回-1
+
+
 def run(config: CFG):
-    file_path = os.path.join(config.weights_dir, config.safetensors)
+    shutil.rmtree(config.weights_dir, ignore_errors=True)
+    os.makedirs(config.weights_dir, exist_ok=True)
+    print("Start generating quant weights")
+    file_path = config.safetensors_path
     with safetensors.safe_open(file_path, framework="pt") as f:
-        for i, key in enumerate(tqdm(f.keys())):
+        # 首先收集所有键并按层号排序
+        all_keys = [(key, get_layer_number(key)) for key in f.keys()]
+        sorted_keys = sorted(all_keys, key=lambda x: x[1])
+
+        for key, _ in tqdm(sorted_keys):
+            # print(key)
             if not is_substring_in_list(key, name_dict.keys()):
                 continue
             tensor = f.get_tensor(key)
@@ -40,6 +58,19 @@ def run(config: CFG):
 
     for k in weights.keys():
         weights[k] = np.array(weights[k])
-        file_path = os.path.join(config.weights_dir, f"{k}.pkl")
-        with open(file_path, "wb") as f:
+        # 保存权重文件
+        pkl_path = os.path.join(config.weights_dir, f"{k}.pkl")
+        with open(pkl_path, "wb") as f:
             pickle.dump(weights[k], f)
+
+        # 保存权重信息到txt文件
+        info_path = os.path.join(config.weights_dir, f"{k}_info.txt")
+        with open(info_path, "w") as f:
+            f.write(f"Weight name: {k}\n")
+            f.write(f"Shape: {weights[k].shape}\n")
+            f.write(f"Data type: {weights[k].dtype}\n")
+            f.write(f"Max value: {np.max(weights[k])}\n")
+            f.write(f"Min value: {np.min(weights[k])}\n")
+            f.write(f"Mean value: {np.mean(weights[k])}\n")
+            f.write(f"Standard deviation: {np.std(weights[k])}\n")
+    print("Quant weights generated at", config.weights_dir)

hllm/eda/mapping_weights.py

@@ -3,6 +3,7 @@ import os
 from hllm.config import CFG
 import pickle
 from tqdm import tqdm
+import shutil
 
 
 def mapping_weights(weights, value_range):
@@ -13,9 +14,10 @@ def mapping_weights(weights, value_range):
 
 
 def run(config: CFG):
+    shutil.rmtree(config.mapped_weights_dir, ignore_errors=True)
+    os.makedirs(config.mapped_weights_dir, exist_ok=True)
     print("Start mapping weights")
-    path_dir = os.path.join(config.mapped_weights_dir)
-    os.makedirs(path_dir, exist_ok=True)
+    path_dir = config.mapped_weights_dir
     value_range = config.value_range
     for file in os.listdir(config.weights_dir):
         if file in config.run_weights_batch:

hllm/optimized/generate_fsm.py

 # %%
+import shutil
 import sys
 import numpy as np
 from pyrilog import (
@@ -535,6 +536,8 @@ def process_task(i, name, weights_file_name, matrix, H, L, VN, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

hllm/optimized/generate_info.py

 import os
 import json
 import pickle
+import shutil
 import numpy as np
 from tqdm import tqdm
 from multiprocessing import Pool
@@ -34,6 +35,8 @@ def process_weight(args):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.mapped_weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file_name}")

hllm/optimized/generate_layer_mux.py

@@ -21,6 +21,7 @@ import os
 from hllm.config import CFG
 from tqdm import tqdm
 import pickle
+import shutil
 from hllm.utils import calculate_WW
 
 
@@ -106,6 +107,8 @@ def process_task(i, name, weights_file_name, matrix, H, L, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

hllm/optimized/generate_mid_wrapper.py

 # %%
 import json
 import sys
+import shutil
 import numpy as np
 from pyrilog import (
     VerilogGenerator,
@@ -68,18 +69,18 @@ def generate_module(
                 sw_ports[f"WT_{j}_out_S"] = f"WT_{j}_out_S[{i}]"
                 sw_ports[f"WT_{j}_out_C"] = f"WT_{j}_out_C[{i}]"
             add_instance(
-                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP*GN+i}",
-                f"Sub_wrapper_{cur_GP*GN+i}",
+                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP * GN + i}",
+                f"Sub_wrapper_{cur_GP * GN + i}",
                 None,
                 sw_ports,
             )
             tcl.add_dependency(
                 f"Sub_wrapper",
-                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"WT_group",
-                f"WT_group_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"WT_group_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"Mid_wrapper",
@@ -87,7 +88,7 @@ def generate_module(
             )
             tcl.add_dependency(
                 f"Layer_mux",
-                f"Layer_mux_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"Layer_mux_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"FSM",
@@ -96,11 +97,11 @@ def generate_module(
             for j in range(len(value_range)):
                 tcl.add_dependency(
                     f"Mux_wrapper",
-                    f"Mux_wrapper_tp_{weights_file_name}_vc_{cur_GP*GN+i}_value_{j}.sv",
+                    f"Mux_wrapper_tp_{weights_file_name}_vc_{cur_GP * GN + i}_value_{j}.sv",
                 )
                 tcl.add_dependency(
                     f"Mux",
-                    f"Mux_tp_{weights_file_name}_vc_{cur_GP*GN+i}_value_{j}.sv",
+                    f"Mux_tp_{weights_file_name}_vc_{cur_GP * GN + i}_value_{j}.sv",
                 )
             for line in ww_list[cur_GP * GN + i]:
                 tcl.add_dependency(
@@ -146,6 +147,8 @@ def process_task(i, name, weights_file_name, ww_list, H, L, VN, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

hllm/optimized/generate_mux.py

@@ -22,6 +22,7 @@ from hllm.config import CFG
 from tqdm import tqdm
 import pickle
 import json
+import shutil
 
 
 # %%
@@ -119,6 +120,8 @@ def process_task(i, name, weights_file_name, matrix, H, L, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.mapped_weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file_name}")

hllm/optimized/generate_mux_wrapper.py

@@ -23,6 +23,7 @@ from hllm.config import CFG
 from tqdm import tqdm
 import pickle
 import json
+import shutil
 
 
 # %%
@@ -105,6 +106,7 @@ def process_task(i, name, weights_file_name, H, L, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
     os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.mapped_weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]

hllm/origin/generate_fsm.py

 # %%
+import shutil
 import sys
 import numpy as np
 from pyrilog import (
@@ -537,6 +538,11 @@ def process_task(i, name, weights_file_name, matrix, H, L, VN, config: CFG):
 def run(name: str, config: CFG):
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
+    output_dir = os.path.join(config.output_dir, name, weights_file_name)
+
+    shutil.rmtree(output_dir, ignore_errors=True)
+    os.makedirs(output_dir, exist_ok=True)
+
     print(f"Processing {weights_file}")
     with open(weights_file, "rb") as f:
         matrixs = pickle.load(f)
@@ -555,5 +561,4 @@ def run(name: str, config: CFG):
                 result = future.result()
             except Exception as e:
                 print(f"Generating {result} failed with an error: {e}")
-    file_dir = os.path.join(config.output_dir, name, weights_file_name)
-    print("Files generated in", file_dir)
+    print("Files generated in", output_dir)

hllm/origin/generate_hn.py

 # %%
 import sys
+import shutil
 import numpy as np
 from pyrilog import ModuleBlock, add_parameter, add_input, add_output, add_assign
 from concurrent.futures import ProcessPoolExecutor, as_completed
@@ -85,6 +86,8 @@ def process_task(i, name, weights_file_name, matrix, H, L, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

hllm/origin/generate_layer_mux.py

 # %%
 import sys
+import shutil
 import numpy as np
 from pyrilog import (
     VerilogGenerator,
@@ -105,6 +106,8 @@ def process_task(i, name, weights_file_name, matrix, H, L, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

hllm/origin/generate_mid_wrapper.py

 # %%
 import sys
+import shutil
 import numpy as np
 from pyrilog import (
     VerilogGenerator,
@@ -66,26 +67,26 @@ def generate_module(
                 sw_ports[f"WT_{j}_out_S"] = f"WT_{j}_out_S[{i}]"
                 sw_ports[f"WT_{j}_out_C"] = f"WT_{j}_out_C[{i}]"
             add_instance(
-                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP*GN+i}",
-                f"Sub_wrapper_{cur_GP*GN+i}",
+                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP * GN + i}",
+                f"Sub_wrapper_{cur_GP * GN + i}",
                 None,
                 sw_ports,
             )
             tcl.add_dependency(
                 f"HN",
-                f"HN_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"HN_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"WT_group",
-                f"WT_group_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"WT_group_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"Layer_mux",
-                f"Layer_mux_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"Layer_mux_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"Sub_wrapper",
-                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP*GN+i}.sv",
+                f"Sub_wrapper_tp_{weights_file_name}_vc_{cur_GP * GN + i}.sv",
             )
             tcl.add_dependency(
                 f"Mid_wrapper",
@@ -138,6 +139,8 @@ def process_task(i, name, weights_file_name, ww_list, H, L, VN, config: CFG):
 
 
 def run(name: str, config: CFG):
+    shutil.rmtree(config.output_dir, ignore_errors=True)
+    os.makedirs(config.output_dir, exist_ok=True)
     weights_file = os.path.join(config.weights_dir, config.run_weights)
     weights_file_name = os.path.splitext(os.path.basename(weights_file))[0]
     print(f"Processing {weights_file}")

main.py

@@ -52,10 +52,9 @@ def run_optimized(config: CFG):
 def run_weights_preprocess(config: CFG):
     import hllm.eda.generate_quant_weights as generate_quant_weights
     import hllm.eda.mapping_weights as generate_mapping_weights
-    
+
     generate_quant_weights.run(config=config)
     generate_mapping_weights.run(config=config)
-    
 
 
 def run_verify(config: CFG):
@@ -68,8 +67,9 @@ def run_verify(config: CFG):
 
 
 def batch_run(config: CFG):
-    for weights in config.run_weights_batch:
-        config.run_weights = weights
+    for i in range(len(config.run_weights_batch)):
+        config.run_weights = config.run_weights_batch[i]
+        config.group_number = config.group_numbers_batch[i]
         run_origin(config)
         run_optimized(config)
         run_verify(config)
@@ -77,8 +77,8 @@ def batch_run(config: CFG):
 
 if __name__ == "__main__":
     config = CFG()
-    # run_weights_preprocess(config)
+    run_weights_preprocess(config)
     # run_origin(config)
     # run_optimized(config)
     # run_verify()
-    batch_run(config)
\ No newline at end of file
+    # batch_run(config)

todo.md

+- 增加生成前删除原有的文件，目前只改了 optimized 的 fsm