3.2.2 AWQ算法的层实现

3.2.2 AWQ算法的层实现

AWQ算法需要对不同类型的神经网络层进行专门的量化实现。本节将详细介绍线性层和卷积层的AWQ量化实现。

线性层量化实现

class AWQLinearQuantizationWrapper(nn.Module):
    """AWQ线性层量化包装器"""
    
    def __init__(self, linear_layer: nn.Linear, config: AWQConfig,
                 scaling_manager: ScalingFactorManager):
        super().__init__()
        self.config = config
        self.scaling_manager = scaling_manager
        
        # 保存原始层参数
        self.in_features = linear_layer.in_features
        self.out_features = linear_layer.out_features
        self.original_weight = linear_layer.weight.data.clone()
        self.original_bias = linear_layer.bias.data.clone() if linear_layer.bias is not None else None
        
        # 量化权重
        self.quantized_weight = nn.Parameter(self.original_weight.clone())
        self.quantization_error = torch.tensor(0.0)
        
        # 缩放因子参数
        self.weight_scale = nn.Parameter(torch.ones(1))
        self.bias_scale = nn.Parameter(torch.ones(1)) if self.original_bias is not None else None
        
        # 初始化
        self._initialize_quantization()
        
    def _initialize_quantization(self):
        """初始化量化参数"""
        # 计算初始缩放因子
        weight_scale = self.scaling_manager.calculate_adaptive_scaling(
            self.original_weight, "weight"
        )
        
        self.weight_scale.data.fill_(weight_scale)
        
        if self.original_bias is not None:
            bias_scale = self.scaling_manager.calculate_adaptive_scaling(
                self.original_bias, "bias"
            )
            self.bias_scale.data.fill_(bias_scale)
            
        # 应用初始量化
        self._apply_quantization()
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 应用量化
        quantized_weight = self._quantize_tensor(self.quantized_weight, self.weight_scale)
        quantized_bias = self._quantize_tensor(self.original_bias, self.bias_scale) if self.original_bias is not None else None
        
        # 线性变换
        output = torch.nn.functional.linear(x, quantized_weight, quantized_bias)
        
        # 计算量化误差
        self._calculate_quantization_error()
        
        return output
    
    def _quantize_tensor(self, tensor: torch.Tensor, scale: nn.Parameter) -> torch.Tensor:
        """量化张量"""
        if tensor is None:
            return None
            
        # 应用缩放和量化
        scaled_tensor = tensor * scale
        quantized = torch.round(scaled_tensor)
        
        # 量化范围限制
        if self.config.symmetric_quantization:
            max_val = 2**(self.config.quantization_bits - 1) - 1
            quantized = torch.clamp(quantized, -max_val - 1, max_val)
        else:
            max_val = 2**self.config.quantization_bits - 1
            quantized = torch.clamp(quantized, 0, max_val)
        
        # 反量化
        dequantized = quantized / scale
        
        return dequantized
    
    def _apply_quantization(self):
        """应用量化"""
        # 量化权重
        self.quantized_weight.data = self._quantize_tensor(
            self.original_weight, self.weight_scale
        )
        
        # 量化偏置
        if self.original_bias is not None and self.bias_scale is not None:
            quantized_bias = self._quantize_tensor(
                self.original_bias, self.bias_scale
            )
            self.quantized_weight.data += quantized_bias.unsqueeze(1)
    
    def _calculate_quantization_error(self):
        """计算量化误差"""
        # 计算权重量化误差
        weight_error = torch.mean((self.original_weight - self.quantized_weight)**2)
        
        # 更新量化误差
        self.quantization_error = weight_error
        
    def get_quantization_error(self) -> torch.Tensor:
        """获取量化误差"""
        return self.quantization_error
    
    def update_scaling_factors(self, new_weight_scale: float, new_bias_scale: float = None):
        """更新缩放因子"""
        self.weight_scale.data.fill_(new_weight_scale)
        
        if self.bias_scale is not None and new_bias_scale is not None:
            self.bias_scale.data.fill_(new_bias_scale)
            
        # 重新应用量化
        self._apply_quantization()
​

卷积层量化实现

class AWQConvQuantizationWrapper(nn.Module):
    """AWQ卷积层量化包装器"""
    
    def __init__(self, conv_layer: nn.Conv2d, config: AWQConfig,
                 scaling_manager: ScalingFactorManager):
        super().__init__()
        self.config = config
        self.scaling_manager = scaling_manager
        
        # 保存原始层参数
        self.in_channels = conv_layer.in_channels
        self.out_channels = conv_layer.out_channels
        self.kernel_size = conv_layer.kernel_size
        self.stride = conv_layer.stride
        self.padding = conv_layer.padding
        self.dilation = conv_layer.dilation
        self.groups = conv_layer.groups
        
        self.original_weight = conv_layer.weight.data.clone()
        self.original_bias = conv_layer.bias.data.clone() if conv_layer.bias is not None else None
        
        # 量化权重
        self.quantized_weight = nn.Parameter(self.original_weight.clone())
        self.quantization_error = torch.tensor(0.0)
        
        # 缩放因子参数（逐通道）
        self.weight_scales = nn.Parameter(torch.ones(self.out_channels))
        self.bias_scales = nn.Parameter(torch.ones(self.out_channels)) if self.original_bias is not None else None
        
        # 初始化
        self._initialize_quantization()
        
    def _initialize_quantization(self):
        """初始化量化参数"""
        # 计算每个输出通道的缩放因子
        for i in range(self.out_channels):
            channel_weight = self.original_weight[i]
            channel_scale = self.scaling_manager.calculate_adaptive_scaling(
                channel_weight, f"conv_weight_channel_{i}"
            )
            self.weight_scales.data[i] = channel_scale
            
        # 量化偏置
        if self.original_bias is not None and self.bias_scales is not None:
            for i in range(self.out_channels):
                channel_bias = self.original_bias[i]
                channel_scale = self.scaling_manager.calculate_adaptive_scaling(
                    channel_bias, f"conv_bias_channel_{i}"
                )
                self.bias_scales.data[i] = channel_scale
                
        # 应用初始量化
        self._apply_quantization()
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 应用量化
        quantized_weight = self._quantize_weight(self.quantized_weight, self.weight_scales)
        quantized_bias = self._quantize_bias(self.original_bias, self.bias_scales) if self.original_bias is not None else None
        
        # 卷积操作
        output = torch.nn.functional.conv2d(
            x, quantized_weight, quantized_bias,
            stride=self.stride, padding=self.padding,
            dilation=self.dilation, groups=self.groups
        )
        
        # 计算量化误差
        self._calculate_quantization_error()
        
        return output
    
    def _quantize_weight(self, weight: torch.Tensor, scales: nn.Parameter) -> torch.Tensor:
        """量化权重"""
        # 逐通道量化
        quantized_weight = torch.zeros_like(weight)
        
        for i in range(self.out_channels):
            channel_weight = weight[i]
            channel_scale = scales[i]
            
            # 量化单个通道
            scaled_channel = channel_weight * channel_scale
            quantized_channel = torch.round(scaled_channel)
            
            # 量化范围限制
            if self.config.symmetric_quantization:
                max_val = 2**(self.config.quantization_bits - 1) - 1
                quantized_channel = torch.clamp(quantized_channel, -max_val - 1, max_val)
            else:
                max_val = 2**self.config.quantization_bits - 1
                quantized_channel = torch.clamp(quantized_channel, 0, max_val)
            
            # 反量化
            dequantized_channel = quantized_channel / channel_scale
            quantized_weight[i] = dequantized_channel
            
        return quantized_weight
    
    def _quantize_bias(self, bias: torch.Tensor, scales: nn.Parameter) -> torch.Tensor:
        """量化偏置"""
        if bias is None or scales is None:
            return None
            
        quantized_bias = torch.zeros_like(bias)
        
        for i in range(bias.size(0)):
            channel_bias = bias[i]
            channel_scale = scales[i]
            
            # 量化单个通道偏置
            scaled_bias = channel_bias * channel_scale
            quantized_bias[i] = torch.round(scaled_bias) / channel_scale
            
        return quantized_bias
    
    def _apply_quantization(self):
        """应用量化"""
        # 量化权重
        self.quantized_weight.data = self._quantize_weight(
            self.original_weight, self.weight_scales
        )
        
        # 量化偏置
        if self.original_bias is not None and self.bias_scales is not None:
            quantized_bias = self._quantize_bias(
                self.original_bias, self.bias_scales
            )
            # 将偏置加到权重的第一个点上
            for i in range(self.out_channels):
                self.quantized_weight.data[i] += quantized_bias[i].unsqueeze(1).unsqueeze(2).unsqueeze(3)
    
    def _calculate_quantization_error(self):
        """计算量化误差"""
        # 计算权重量化误差
        weight_error = torch.mean((self.original_weight - self.quantized_weight)**2)
        
        # 更新量化误差
        self.quantization_error = weight_error
        
    def get_quantization_error(self) -> torch.Tensor:
        """获取量化误差"""
        return self.quantization_error
​