Upload qat-convert.py with huggingface_hub

qat-convert.py

@@ -0,0 +1,176 @@
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+import os
+import safetensors
+
+class SteTernaryQuantizer(nn.Module):
+    def __init__(self, group_size):
+        super().__init__()
+        self.group_size = group_size
+
+    def forward(self, x):
+        org_w_shape = x.shape
+        if self.group_size > 0:
+            assert x.shape[-1] % self.group_size == 0
+            x = x.reshape(-1, self.group_size)
+        elif self.group_size == -1:
+            x = x.reshape(-1, x.shape[-1])
+        assert x.dim() == 2
+        scales = 1.0 / (x.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5))
+        x_q = (torch.clamp(torch.round(x * scales),-1,1) / scales)
+        assert torch.isnan(x_q).sum() == 0
+        x = x.reshape(org_w_shape)
+        x_q = x_q.reshape(org_w_shape)
+        return x_q
+
+class SteIntQuantizer(nn.Module):
+    def __init__(self, bit, group_size):
+        super().__init__()
+        self.bit = bit
+        self.group_size = group_size
+
+    def forward(self, x):
+        org_w_shape = x.shape
+        if self.group_size > 0:
+            assert org_w_shape[-1] % self.group_size == 0
+            x = x.reshape(-1, self.group_size)
+        elif self.group_size == -1:
+            x = x.reshape(-1, x.shape[-1])
+        
+        assert x.dim() == 2
+
+        abs_max_val = x.abs().amax(dim=1, keepdim=True)
+        max_int = 2 ** (self.bit - 1) - 1
+        min_int = - (2 ** (self.bit - 1))
+        scales = abs_max_val.clamp(min=1e-5) / max_int
+
+        assert torch.isnan(scales).sum() == 0
+
+        x_q = (torch.clamp(torch.round(x / scales), min_int, max_int)) * scales
+        
+        assert torch.isnan(x_q).sum() == 0
+
+        x = x.reshape(org_w_shape)
+        x_q = x_q.reshape(org_w_shape)
+
+        return x_q
+
+class SteInt2Quantizer(nn.Module):
+    def __init__(self, group_size):
+        super().__init__()
+        self.group_size = group_size
+
+    def forward(self, x):
+        org_w_shape = x.shape
+        if self.group_size > 0:
+            assert x.shape[-1] % self.group_size == 0
+            x = x.reshape(-1, self.group_size)
+        elif self.group_size == -1:
+            x = x.reshape(-1, x.shape[-1])
+
+        assert x.dim() == 2
+
+        scales = 1.0 / (x.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5) * 1)
+        x_q = (torch.clamp(torch.round(x * scales),-2,1) / scales)
+
+        assert torch.isnan(x_q).sum() == 0
+
+        x = x.reshape(org_w_shape)
+        x_q = x_q.reshape(org_w_shape)
+
+        return x_q
+
+def quantize_model_bin(input_bin_path, output_bin_path, quant_type="ternary", bit=2, group_size=128, device="cuda" if torch.cuda.is_available() else "cpu"):
+    """
+    直接对PyTorch模型bin文件进行量化。
+
+    Args:
+        input_bin_path: 输入模型bin文件路径
+        output_bin_path: 输出量化后的模型bin文件路径
+        quant_type: 量化类型 ("ternary" 或 "int")
+        bit: 整数量化的位数 (仅在 quant_type="int" 时使用)
+        group_size: 量化分组大小
+        device: 运行设备
+    """
+    print(f"加载模型文件: {input_bin_path}...")
+    if input_bin_path.endswith(".bin"):
+        state_dict = torch.load(input_bin_path, map_location=device)
+    elif input_bin_path.endswith(".safetensors"):
+        state_dict = safetensors.load_file(input_bin_path)
+    elif os.path.isdir(input_bin_path) and os.path.exists(os.path.join(input_bin_path, "pytorch_model.bin")):
+        state_dict = torch.load(os.path.join(input_bin_path, "pytorch_model.bin"), map_location=device)
+    elif os.path.isdir(input_bin_path) and os.path.exists(os.path.join(input_bin_path, "model.safetensors")):
+        state_dict = safetensors.load_file(os.path.join(input_bin_path, "model.safetensors"))
+    else:
+        raise ValueError(f"不支持的模型文件类型: {input_bin_path}")
+    
+    print(f"应用 {quant_type} 量化...")
+    if quant_type == "ternary":
+        quantizer = SteTernaryQuantizer(group_size=group_size)
+    elif quant_type == "int":
+        quantizer = SteIntQuantizer(bit=bit, group_size=group_size)
+    elif quant_type == "int2":
+        quantizer = SteInt2Quantizer(group_size=group_size)
+    else:
+        raise ValueError(f"不支持的量化类型: {quant_type}")
+
+    # 统计需要量化的参数数量
+    total_params = sum(1 for k, v in state_dict.items() if ("weight" in k and "layer" in k) or ("fc" in k))
+    
+    # 应用量化
+    with torch.no_grad():
+        for name, param in tqdm(state_dict.items(), total=total_params, desc="量化中"):
+            if (("weight" in name and "layer" in name and param.dim() == 2) or ("fc" in name and param.dim() == 2)):
+                # 对权重进行量化
+                original_weight = param.data.clone()
+                quantized_weight = quantizer(original_weight)
+                state_dict[name] = quantized_weight
+
+                # 打印前几个层的统计信息
+                if total_params > 0:
+                    total_params -= 1
+                    if total_params > total_params - 5:
+                        print(f"层: {name}")
+                        print(f"  原始范围: {original_weight.min():.4f} 到 {original_weight.max():.4f}")
+                        print(f"  量化后范围: {quantized_weight.min():.4f} 到 {quantized_weight.max():.4f}")
+                        print(f"  均方误差: {((original_weight - quantized_weight)**2).mean():.8f}")
+
+    # 保存量化后的模型
+    print(f"保存量化后的模型到: {output_bin_path}...")
+    if output_bin_path.endswith(".bin"):
+        torch.save(state_dict, output_bin_path)
+    elif output_bin_path.endswith(".safetensors"):
+        safetensors.save_file(state_dict, output_bin_path)
+    else:
+        os.makedirs(os.path.dirname(output_bin_path), exist_ok=True)
+        output_bin_path = os.path.join(output_bin_path, "pytorch_model.bin")
+        torch.save(state_dict, output_bin_path)
+    print("完成!")
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="量化PyTorch模型bin文件")
+    parser.add_argument("--input_bin", type=str, required=True, help="输入模型bin文件路径")
+    parser.add_argument("--output", type=str, required=True, help="输出量化后的模型bin文件路径")
+    parser.add_argument("--quant_type", type=str, default="ternary", choices=["ternary", "int", "int2"], help="量化类型")
+    parser.add_argument("--bit", type=int, default=2, help="整数量化的位数")
+    parser.add_argument("--group_size", type=int, default=-1, help="量化分组大小")
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="运行设备")
+    parser.add_argument("--config", type=str, default="", help="model config file")
+
+    args = parser.parse_args()
+    os.makedirs(args.output, exist_ok=True)
+    quantize_model_bin(
+        input_bin_path=args.input_bin,
+        output_bin_path=os.path.join(args.output, "pytorch_model.bin"),
+        quant_type=args.quant_type,
+        bit=args.bit,
+        group_size=args.group_size,
+        device=args.device
+    )
+    if args.config:
+        os.system(f"cp {args.config}/* {args.output}")
+        print(f"复制{args.config}文件到{args.output}")
+if __name__ == "__main__":
+    main()