chore: import upstream snapshot with attribution

kt-kernel/scripts/check.py

@@ -0,0 +1,277 @@
+import os
+
+# insert the path of the project
+import sys
+
+# sys.path.insert(0, "/home/azure/ktransformers")
+import argparse
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+import re
+from collections import defaultdict
+import itertools
+import os
+import torch
+import numpy as np
+
+tensor_from_amx = [".mlp.experts."]  # todo: add keys in gguf that should be used in the final tensor
+
+
+def safe_open_binary_to_tensor(file_path):
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"文件不存在: {file_path}")
+
+    if not os.access(file_path, os.R_OK):
+        raise PermissionError(f"没有权限读取文件: {file_path}")
+
+    try:
+        with open(file_path, "rb") as f:
+            binary_data = f.read()
+
+        np_array = np.frombuffer(binary_data, dtype=np.int8)
+
+        tensor = torch.from_numpy(np_array)
+
+        return tensor
+
+    except Exception as e:
+        raise IOError(f"file process error: {str(e)}")
+
+
+def read_safetensor_keys_from_folder(folder_path) -> dict:
+    """
+    :param folder_path: folder path
+    :return: key_to_file_map
+    """
+    # check if the folder path is exist
+    if not os.path.exists(folder_path):
+        raise FileNotFoundError(f"GGUF dir not found: {folder_path}")
+    if os.path.isfile(folder_path):
+        folder_path = os.path.dirname(folder_path)
+
+    key_to_file_map = {}
+
+    found_safetensor = False
+    for root, dirs, files in os.walk(folder_path):
+        # sort files
+        files = sorted(files)
+        for file in files:
+            if file.endswith(".safetensors"):
+                found_safetensor = True
+                file_path = os.path.join(root, file)
+                try:
+                    with safe_open(file_path, framework="pt") as f:
+                        for key in f.keys():
+                            if "model.layers.61" in key:
+                                # skip MTP layer
+                                continue
+                            # try:
+                            #     if int(key.split('.')[2]) > 4:
+                            #         continue
+                            # except:
+                            #     pass
+                            key_to_file_map[key] = file_path
+                except Exception as e:
+                    print(f"Error reading Safetensor file {file_path}: {e}")
+
+    if not found_safetensor:
+        raise FileNotFoundError(f"No Safetensor files found in {folder_path}")
+
+    return key_to_file_map
+
+
+def read_amx_tensor_from_folder(folder_path, keys) -> dict:
+    layer_list = [f"_layer_{i}" for i in range(3, 61)]
+    numa_list = ["_numa_0", "_numa_1"]
+
+    down_list = [f"INT4_down_{i}_quant_.kt" for i in range(256)]
+    gate_list = [f"INT4_gate_{i}_quant_.kt" for i in range(256)]
+    up_list = [f"INT4_up_{i}_quant_.kt" for i in range(256)]
+    down_scale_list = [f"INT4_down_{i}_scale_.kt" for i in range(256)]
+    gate_scale_list = [f"INT4_gate_{i}_scale_.kt" for i in range(256)]
+    up_scale_list = [f"INT4_up_{i}_scale_.kt" for i in range(256)]
+    target = ["ffn_up_exps", "ffn_down_exps", "ffn_gate_exps"]
+    tensor_file_map = {}
+    for key in keys:
+        layer = int(key.split(".")[1])
+        if layer < 3:
+            continue
+        layer_path = f"_layer_{layer}"
+        # concatenate the path layer/numa/(down|gate|up)_(0-255)_3670016Byte_quant_.kt
+        # store the path in the tensor_file_map
+        # key = key+'.idx.weight'
+        # scale_key = key+'.idx.scale'
+        for numa_idx, numa in enumerate(numa_list):
+            # TODO: 256 should be a variable
+            for i in range(256):
+                prefix_key = ".".join(key.split(".")[:-1])
+
+                experts_key = prefix_key + f".{i}.numa.{numa_idx}.weight"
+                scale_key = prefix_key + f".{i}.numa.{numa_idx}.scale"
+                if "down" in experts_key:
+                    tensor_file_map[experts_key] = os.path.join(folder_path, layer_path, numa, down_list[i])
+                    tensor_file_map[scale_key] = os.path.join(folder_path, layer_path, numa, down_scale_list[i])
+                elif "gate" in experts_key:
+                    tensor_file_map[experts_key] = os.path.join(folder_path, layer_path, numa, gate_list[i])
+                    tensor_file_map[scale_key] = os.path.join(folder_path, layer_path, numa, gate_scale_list[i])
+                elif "up" in experts_key:
+                    tensor_file_map[experts_key] = os.path.join(folder_path, layer_path, numa, up_list[i])
+                    tensor_file_map[scale_key] = os.path.join(folder_path, layer_path, numa, up_scale_list[i])
+    return tensor_file_map
+
+
+# def translate_name(name:str)->str:
+#     """
+#     :param name: name of the tensor
+#     :return: translated name
+#     """
+#     name = translate_name_to_gguf(name)
+#     name = name.replace(".up_proj.", ".ffn_up_exps.")
+#     name = name.replace(".down_proj.", ".ffn_down_exps.")
+#     name = name.replace(".gate_proj.", ".ffn_gate_exps.")
+#     name = name.replace(".ffn_gate_inp.e_score_correction_bias", ".exp_probs_b.bias")
+#     return name
+
+
+def _clean_keys(keys):
+    keys = list(keys)
+    target = ["ffn_up_exps", "ffn_down_exps", "ffn_gate_exps"]
+    # only keep the keys that contain the target
+    keys = [key for key in keys if any(target_key in key for target_key in target) and "ggml_type" not in key]
+    return keys
+
+
+def combine_tensor_sources(safetensor_path, amx_path):
+    safetensor_tensor_file_map = read_safetensor_keys_from_folder(safetensor_path)
+
+    keys = _clean_keys(safetensor_tensor_file_map.keys())
+
+    amx_tensor_file_map = read_amx_tensor_from_folder(amx_path, keys)
+    target_tensor_map = {}
+    for key in safetensor_tensor_file_map.keys():
+        if "_exps." in key:
+            continue
+
+        target_tensor_map[key] = safetensor_tensor_file_map[key]
+
+    for key in amx_tensor_file_map.keys():
+        target_tensor_map[key] = amx_tensor_file_map[key]
+
+    return target_tensor_map
+
+
+def write_combined_tensor(target_tensor_map: dict, output_path: str):
+    # Ensure output directory exists
+    os.makedirs(output_path, exist_ok=True)
+
+    # Cache for safetensor file handles and GGUF loaders
+    safetensors_cache = {}
+    amx_cache = {}
+
+    # Group tensors by layer
+    layer_groups = defaultdict(list)
+    non_layer_keys = []
+    layer_pattern = re.compile(r"blk\.(\d+)\.")
+
+    for key in target_tensor_map:
+        match = layer_pattern.search(key)
+        if match:
+            layer_groups[int(match.group(1))].append(key)
+        else:
+            non_layer_keys.append(key)
+
+    # Calculate the number of shards
+    total_shards = len(layer_groups) + (1 if non_layer_keys else 0) - 1
+
+    shard_idx = 0
+    # Save non-layer tensors to the first shard if they exist
+    if non_layer_keys:
+        tensors = {}
+        for key in non_layer_keys:
+            file_path = target_tensor_map[key]
+            tensor = None
+            ggml_type = None
+            if file_path.endswith(".safetensors"):
+                if file_path not in safetensors_cache:
+                    safetensors_cache[file_path] = safe_open(file_path, framework="pt")
+                f = safetensors_cache[file_path]
+                tensor = f.get_tensor(key)
+            elif file_path.endswith(".kt"):
+                tensor = safe_open_binary_to_tensor(file_path)
+            else:
+                raise ValueError(f"Unsupported file format: {file_path}")
+            tensors[key] = tensor
+
+        output_file = os.path.join(output_path, f"model-{shard_idx:05}-of-{total_shards:05}.safetensors")
+        print(f"Saving non-layer tensors to {output_file}")
+        save_file(tensors, output_file)
+        shard_idx += 1
+
+    # Save each layer's tensors to subsequent shards
+    for layer_num in sorted(layer_groups.keys()):
+        layer_keys = layer_groups[layer_num]
+        tensors = {}
+        for key in layer_keys:
+            file_path = target_tensor_map[key]
+            tensor = None
+            ggml_type = None
+            if file_path.endswith(".safetensors"):
+                if file_path not in safetensors_cache:
+                    safetensors_cache[file_path] = safe_open(file_path, framework="pt")
+                f = safetensors_cache[file_path]
+                tensor = f.get_tensor(key)
+                tensor_info = tensor.shape
+            elif file_path.endswith(".kt"):
+                tensor = safe_open_binary_to_tensor(file_path)
+            else:
+                raise ValueError(f"Unsupported file format: {file_path}")
+            tensors[key] = tensor
+
+        output_file = os.path.join(output_path, f"model-{shard_idx:05}-of-{total_shards:05}.safetensors")
+        print(f"Saving layer {layer_num} to {output_file}")
+        save_file(tensors, output_file)
+        shard_idx += 1
+    return
+
+
+def main():
+    # 输入已经处理过的混合模型路径，提前处理好的amx路径，输出路径
+    parser = argparse.ArgumentParser(description="Read parameters from Safetensor and GGUF files")
+    parser.add_argument(
+        "--safetensor_path",
+        type=str,
+        help="Path to the Safetensor file",
+        default="/mnt/data/models/DeepSeek-R1-GGML-FP8-Hybrid/DeepSeek-R1-IQ1S-FP8",
+    )
+    parser.add_argument(
+        "--amx_path", type=str, help="Path to the GGUF file", default="/mnt/data/models/DeepSeek-R1-INT4"
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="Path to the output file",
+        default="/mnt/data/models/DeepSeek-R1-GGML-FP8-Hybrid/DeepSeek-R1-AMXQ4-FP8",
+    )
+
+    # print all the arguments
+    print("All the arguments:")
+    print(parser.parse_args())
+
+    # 解析命令行参数
+    args = parser.parse_args()
+
+    safetensor_path = args.safetensor_path
+    amx_path = args.amx_path
+    output_path = args.output_path
+
+    target_tensor_map = combine_tensor_sources(safetensor_path, amx_path)
+    for key, value in target_tensor_map.items():
+        print(f"{key}: {value}")
+    write_combined_tensor(target_tensor_map, output_path)
+
+    return
+
+
+if __name__ == "__main__":
+    main()