自定义量化Linear层scale参数grad为None的原因及解决方法

• 直接修改self.fc.weight.data绕过了PyTorch计算图追踪：在forward函数中，你将dequant_weight直接赋值给self.fc.weight.data，这属于对张量底层数据的硬修改，不会被记录到计算图中。后续调用self.fc(x)时，使用的是修改后的权重，但这个权重和scale、zero_point之间没有计算图连接，反向传播时梯度自然无法传递到这两个参数上。
• 原Linear层的权重参数self.fc.weight本身可训练，但你在forward中直接覆盖了它的data，导致原权重的梯度和自定义量化参数的梯度都无法正确传递。

核心思路是不直接修改原Linear层的权重数据，而是基于量化反量化后的权重手动完成线性计算，让计算图完整追踪scale和zero_point的梯度。修改后的代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 实现直通估计器的round操作
def round_ste(x):
    return (x.round() - x).detach() + x

class QLinear(nn.Module):
    def __init__(self, input_dim, out_dim, bits=8):
        super(QLinear, self).__init__()
        self.up = 2 ** bits - 1
        self.down = 0
        self.fc = nn.Linear(input_dim, out_dim)
        weight = self.fc.weight.data
        # 初始化可训练的scale和zero_point
        self.scale = nn.Parameter(torch.tensor((torch.max(weight) - torch.min(weight)) / (self.up - self.down)), requires_grad=True)
        self.zero_point = nn.Parameter(torch.tensor(self.down - (torch.min(weight) / self.scale).round()), requires_grad=True)

    def forward(self, x):
        weight = self.fc.weight
        # 量化-反量化流程，保留计算图连接
        quant_weight = round_ste(weight / self.scale) + self.zero_point
        quant_weight = torch.clamp(quant_weight, self.down, self.up)
        dequant_weight = (quant_weight - self.zero_point) * self.scale
        
        # 手动用反量化后的权重执行线性变换，不修改原fc的权重数据
        bias = self.fc.bias if self.fc.bias is not None else None
        return F.linear(x, dequant_weight, bias)

class QNet(nn.Module):
    def __init__(self):
        super(QNet, self).__init__()
        self.fc1 = QLinear(28 * 28, 100)
        self.fc2 = QLinear(100, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        x = F.softmax(x, dim=1)  # 指定dim避免多批次输入时的计算错误
        return x

• 替换self.fc(x)为F.linear(x, dequant_weight, bias)：直接使用反量化后的权重调用函数式API，让计算图完整记录scale、zero_point到输出的路径，确保反向传播时梯度能正常传递。
• 修复softmax维度问题：原代码未指定dim，多批次输入时会导致计算逻辑错误，添加dim=1符合分类任务的常规逻辑。
• 补充round_ste的正确实现：直通估计器让梯度绕过round操作直接传递，保证量化操作不阻断梯度流。

内容的提问来源于stack exchange，提问作者zz z