qat-convert.py

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()