Create convert_for_g2_draft.py

convert_for_g2_draft.py

@@ -0,0 +1,78 @@
+import os
+from safetensors import safe_open
+from safetensors.torch import save_file
+import torch
+from typing import Dict, Set
+
+
+def get_tensors(file_path: str) -> Dict[str, torch.Tensor]:
+    tensors = {}
+    with safe_open(file_path, framework="pt", device="cpu") as f:
+        for k in f.keys():
+            tensors[k] = f.get_tensor(k)
+    return tensors
+
+
+def get_quantized_modules(tensor_keys, keyword: str = "scale") -> Set[str]:
+    """
+    Extract module prefixes like model.layers.0.mlp.down_proj from keys like:
+    model.layers.0.mlp.down_proj.weight_scale
+    """
+    quantized_modules = set()
+    for key in tensor_keys:
+        if keyword in key:
+            parts = key.split(".")
+            if len(parts) >= 2:
+                mod_prefix = ".".join(parts[:-1])
+                quantized_modules.add(mod_prefix)
+    return quantized_modules
+
+
+def modify_quantized_tensors(tensors: Dict[str, torch.Tensor], quantized_modules: Set[str]) -> Dict[str, torch.Tensor]:
+    """
+    For each quantized module:
+    - weight        --> divide by 2
+    - weight_scale  --> multiply by 2
+    - input_scale   --> multiply by 2
+    """
+    modified = {}
+    for key, tensor in tensors.items():
+        modified_tensor = tensor
+        for mod in quantized_modules:
+            if key == f"{mod}.weight":
+                modified_tensor = (tensor.to(torch.float32) / 2).to(torch.float8_e4m3fn)
+            elif key == f"{mod}.weight_scale":
+                modified_tensor = tensor * 2
+            elif key == f"{mod}.input_scale":
+                modified_tensor = tensor * 2
+        modified[key] = modified_tensor
+    return modified
+
+
+def process_folder(folder_path: str, output_folder: str):
+    os.makedirs(output_folder, exist_ok=True)
+
+    for file in os.listdir(folder_path):
+        if not file.endswith(".safetensors"):
+            continue
+
+        file_path = os.path.join(folder_path, file)
+        print(f"Processing: {file_path}")
+        try:
+            tensors = get_tensors(file_path)
+            quantized_modules = get_quantized_modules(tensors.keys())
+            modified_tensors = modify_quantized_tensors(tensors, quantized_modules)
+
+            output_path = os.path.join(output_folder, file)
+            save_file(modified_tensors, output_path)
+            print(f"Saved modified tensors to: {output_path}")
+
+        except Exception as e:
+            print(f"Failed to process {file_path}: {e}")
+
+
+if __name__ == "__main__":
+    input_folder = "/mnt/disk5/tencent/Hunyuan-7B-Instruct-FP8"
+    output_folder = "/mnt/disk5/tencent/Hunyuan-7B-Instruct-FP8-modified"
+
+    process_folder(input_folder, output_folder)