Checkpoint转换：safetensors/GGUF

import jsonimport structimport numpy as npfrom typing import Dict, List, Tuple, Optionalfrom dataclasses import dataclassimport hashlibimport os  def simulate_safetensors_save(tensors, filepath):    header = {}    offset = 0    tensor_data_list = []    for name, tensor in sorted(tensors.items()):        tensor_bytes = tensor.tobytes()        header[name] = {            "dtype": str(tensor.dtype),            "shape": list(tensor.shape),            "data_offsets": [offset, offset + len(tensor_bytes)]        }        tensor_data_list.append(tensor_bytes)        offset += len(tensor_bytes)    header_json = json.dumps(header, ensure_ascii=False)    header_bytes = header_json.encode('utf-8')    padding = (8 - len(header_bytes) % 8) % 8    header_bytes += b' ' * padding    with open(filepath, 'wb') as f:        f.write(struct.pack('<Q', len(header_bytes)))        f.write(header_bytes)        for data in tensor_data_list:            f.write(data)    file_size = os.path.getsize(filepath)    return {"header_size": len(header_bytes), "num_tensors": len(tensors), "file_size": file_size, "format": "safetensors"}  def simulate_safetensors_load(filepath):    with open(filepath, 'rb') as f:        header_len = struct.unpack('<Q', f.read(8))[0]        header_bytes = f.read(header_len)        header = json.loads(header_bytes.decode('utf-8'))        data_start = 8 + header_len        tensors = {}        for name, info in header.items():            if not isinstance(info, dict) or "data_offsets" not in info:                continue            start, end = info["data_offsets"]            dtype = np.dtype(info["dtype"])            shape = tuple(info["shape"])            f.seek(data_start + start)            data = f.read(end - start)            tensor = np.frombuffer(data, dtype=dtype).reshape(shape).copy()            tensors[name] = tensor        return tensors  def simulate_pickle_save(tensors, filepath):    np.savez(filepath, **tensors)    file_size = os.path.getsize(filepath + ".npz")    return {"num_tensors": len(tensors), "file_size": file_size, "format": "pt/npz (模拟)", "security_note": "真实pickle格式存在RCE风险，加载不可信文件需谨慎"}  @dataclassclass QuantizationConfig:    bits: int = 4    group_size: int = 32  def quantize_q4(tensor, config):    flat = tensor.flatten().astype(np.float32)    n_elements = len(flat)    n_groups = (n_elements + config.group_size - 1) // config.group_size    padded = np.pad(flat, (0, n_groups * config.group_size - n_elements))    grouped = padded.reshape(n_groups, config.group_size)    max_vals = np.max(np.abs(grouped), axis=1)    scales = max_vals / 7.0    scales = np.where(scales == 0, 1e-8, scales)    quantized = np.round(grouped / scales[:, None]).astype(np.int8)    quantized = np.clip(quantized, -8, 7)    packed = np.zeros((n_groups, config.group_size // 2), dtype=np.uint8)    for i in range(config.group_size // 2):        low = quantized[:, i*2] & 0x0F        high = (quantized[:, i*2+1] & 0x0F) << 4        packed[:, i] = low | high    return packed, scales.astype(np.float16)  def dequantize_q4(packed, scales, original_shape, config):    n_groups = packed.shape[0]    group_size = config.group_size    quantized = np.zeros((n_groups, group_size), dtype=np.int8)    for i in range(group_size // 2):        byte = packed[:, i]        low = (byte & 0x0F).astype(np.int8)        high = ((byte >> 4) & 0x0F).astype(np.int8)        low = np.where(low > 7, low - 16, low)        high = np.where(high > 7, high - 16, high)        quantized[:, i*2] = low        quantized[:, i*2+1] = high    dequantized = quantized * scales[:, None]    flat = dequantized.flatten()    n_elements = 1    for s in original_shape:        n_elements *= s    return flat[:n_elements].reshape(original_shape)  def tensor_hash(tensor):    return hashlib.md5(tensor.tobytes()).hexdigest()  def compare_formats(original, converted):    results = {}    for name in original:        if name not in converted:            results[name] = "missing"            continue        orig = original[name]        conv = converted[name]        max_diff = np.max(np.abs(orig - conv))        mean_diff = np.mean(np.abs(orig - conv))        results[name] = {"max_diff": float(max_diff), "mean_diff": float(mean_diff), "exact_match": bool(np.array_equal(orig, conv))}    return results  def demo():    print("=" * 60)    print("模型Checkpoint格式转换工具箱演示")    print("=" * 60)    np.random.seed(42)    test_tensors = {        "model.embedding.weight": np.random.randn(1000, 256).astype(np.float32),        "model.layers.0.attention.q_proj.weight": np.random.randn(256, 256).astype(np.float32),        "model.layers.0.attention.k_proj.weight": np.random.randn(256, 256).astype(np.float32),        "model.layers.0.attention.v_proj.weight": np.random.randn(256, 256).astype(np.float32),        "model.layers.0.attention.o_proj.weight": np.random.randn(256, 256).astype(np.float32),        "model.layers.0.mlp.gate_proj.weight": np.random.randn(256, 512).astype(np.float32),        "model.layers.0.mlp.down_proj.weight": np.random.randn(512, 256).astype(np.float32),        "model.lm_head.weight": np.random.randn(1000, 256).astype(np.float32),    }    total_params = sum(t.size for t in test_tensors.values())    print(f"\n模拟模型参数数量: {total_params:,}")    print(f"原始FP32体积: {total_params * 4 / 1024 / 1024:.2f} MB")    print("\n【1】safetensors格式")    print("-" * 50)    st_info = simulate_safetensors_save(test_tensors, "/tmp/test_model.safetensors")    print(f"文件大小: {st_info['file_size'] / 1024 / 1024:.2f} MB")    print(f"张量数量: {st_info['num_tensors']}")    print(f"头部大小: {st_info['header_size']} 字节")    print("特点: 安全（无代码执行）、mmap零拷贝加载、加载速度快")    loaded_st = simulate_safetensors_load("/tmp/test_model.safetensors")    st_diff = compare_formats(test_tensors, loaded_st)    exact_count = sum(1 for v in st_diff.values() if isinstance(v, dict) and v["exact_match"])    print(f"\n加载验证: {exact_count}/{len(test_tensors)} 个张量精确匹配")    print("结论: safetensors转换是无损的")    print("\n【2】PyTorch pt格式（pickle模拟）")    print("-" * 50)    pt_info = simulate_pickle_save(test_tensors, "/tmp/test_model.pt")    print(f"文件大小: {pt_info['file_size'] / 1024 / 1024:.2f} MB")    print(f"张量数量: {pt_info['num_tensors']}")    print(f"安全提示: {pt_info['security_note']}")    print("特点: 兼容性好、灵活、但存在pickle反序列化安全风险")    print("\n【3】GGUF量化格式")    print("-" * 50)    quant_config = QuantizationConfig(bits=4, group_size=32)    packed, scales = quantize_q4(test_tensors["model.layers.0.attention.q_proj.weight"], quant_config)    gguf_size = packed.nbytes + scales.nbytes    orig_size = test_tensors["model.layers.0.attention.q_proj.weight"].nbytes    print(f"量化位数: 4位")    print(f"分组大小: 32")    print(f"测试层原始大小: {orig_size / 1024:.1f} KB")    print(f"测试层量化后大小: {gguf_size / 1024:.1f} KB")    print(f"压缩比: {orig_size / gguf_size:.1f}x")    print("\n【量化精度验证】")    test_name = "model.layers.0.attention.q_proj.weight"    test_tensor = test_tensors[test_name]    dequant = dequantize_q4(packed, scales, test_tensor.shape, quant_config)    max_err = np.max(np.abs(test_tensor - dequant))    mean_err = np.mean(np.abs(test_tensor - dequant))    print(f"测试层: {test_name}")    print(f"最大误差: {max_err:.6f}")    print(f"平均误差: {mean_err:.6f}")    print(f"相对误差: {mean_err / np.mean(np.abs(test_tensor)) * 100:.2f}%")    print("结论: 量化是有损转换，精度损失在可接受范围内")    print("\n演示完成！")  if __name__ == "__main__":    demo()