AWQ 代码深入剖析

# awq_results 是一个字典
awq_results = {
        "scale": [],
        "clip": [],
    }

@torch.no_grad()
def scale_gelu_fc(gelu, fc, scales):
    # 进行网络层类型检查
    assert isinstance(gelu, (nn.GELU, BloomGelu, GELUActivation))
    assert isinstance(fc, nn.Linear)

    # scales.view(1, -1) 将缩放因子重塑为 2D 张量, 以匹配权重矩阵的维度
    # mul_ 是就地乘法操作（节省内存），直接修改 fc.weight 的值
    fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))

    # 遍历全连接层的所有参数，确保没有任何参数变成 NaN 值，这是一个重要的数值稳定性检查。
    for p in fc.parameters():
        assert torch.isnan(p).sum() == 0

@torch.no_grad()
def get_act_scale(x):
    # x.abs().view(-1, x.shape[-1]): 重塑张量维度 -> [b*s, h]
    # x.abs().view(-1, x.shape[-1]).mean(0): 沿着第0维计算平均值 -> [h,]
    return x.abs().view(-1, x.shape[-1]).mean(0)

# _search_module_scale 是 auto_scale_block 内的函数
def _search_module_scale(block, linears2scale: list, x, kwargs={}):
    """
    在给定的模块和输入特征下，搜索最佳的缩放因子。

    参数:
        block (nn.Module): 当前要搜索缩放因子的模块。
        linears2scale (list of nn.Linear): 需要缩放的全连接层列表。
        x (torch.Tensor): 输入特征张量。
        kwargs (dict): 其他关键字参数。

    返回:
        torch.Tensor: 最佳缩放因子，形状为 [out_features]。
    """
    x = x.to(next(block.parameters()).device)  # 将输入特征移动到模块所在设备
    with torch.no_grad():
        org_out = block(x, **kwargs)  # 原始输出
        if isinstance(org_out, tuple):
            org_out = org_out[0]

    x_max = get_act_scale(x)  # 获取激活的缩放因子

    best_error = float("inf")
    best_ratio = -1
    best_scales = None

    n_grid = 20  # 缩放比例的网格大小
    history = []  # 保存每个比例的误差

    # 保存原始状态字典
    org_sd = {k: v.cpu() for k, v in block.state_dict().items()}

    for ratio in range(n_grid):
        ratio = ratio * 1 / n_grid  # 当前比例
        scales = x_max.pow(ratio).clamp(min=1e-4).view(-1)  # s^alpha 计算缩放因子
        scales = scales / (scales.max() * scales.min()).sqrt()  # 归一化缩放因子(论文没有描述这步操作)

        for fc in linears2scale:
            # 权重乘以缩放因子，对显著权重进行放大以降低量化误差, 对应公式(2)中的 w*s
            fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))
            # 使用权重量化函数 w_quantize_func 执行权重量化操作，对应公式(4) W⋅diag(s))(diag(s)^−1
            fc.weight.data = w_quantize_func(fc.weight.data) / (scales.view(1, -1))

        out = block(x, **kwargs)  # 量化后的输出
        if isinstance(out, tuple):
            out = out[0]

        # 计算均方误差
        loss = ( (org_out - out).float().pow(2).mean().item() )
        history.append(loss)

        # 如果当前误差小于最佳误差，更新最佳误差和最佳比例
        is_best = loss < best_error
        if is_best:
            best_error = loss
            best_ratio = ratio
            best_scales = scales

        # 恢复原始状态字典
        block.load_state_dict(org_sd)

    if best_ratio == -1:
        print(history)
        raise Exception("Failed to find a suitable scaling ratio.")

    best_scales = best_scales.view(-1)  # 调整形状

    assert torch.isnan(best_scales).sum() == 0, "缩放因子包含NaN值"
        return best_scales.detach()

scales = x_max.pow(ratio).clamp(min=1e-4).view(-1)  # 计算缩放因子
scales = scales / (scales.max() * scales.min()).sqrt()  # 为防止 scales 过大或过小，需要多一步归一化操作

# n_grid 取值 20
for ratio in range(n_grid):
    ratio = ratio * 1 / n_grid  # 当前比例
    for fc in linears2scale:
        # 权重乘以缩放因子，对显著权重进行放大以降低量化误差, 对应公式(2)中的 w*s
        fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))
        # 使用权重量化函数 w_quantize_func 执行权重量化操作，对应公式(4) W⋅diag(s))(diag(s)^−1
        fc.weight.data = w_quantize_func(fc.weight.data) / (scales.view(1, -1))

# 处理 LlamaDecoderLayer 的层
if isinstance(module, LlamaDecoderLayer):
    
    # 处理 self-attention 的查询、键、值投影层
    scales_list.append(
        # `_auto_get_scale` 是对 `_search_module_scale` 的一层包装，作用是获取某个层组的缩放因子。
        _auto_get_scale(
            prev_op=module.input_layernorm,
            layers=[
                module.self_attn.q_proj,
                module.self_attn.k_proj,
                module.self_attn.v_proj,
            ],
            inp=input_feat["self_attn.q_proj"],
            module2inspect=module.self_attn,
            kwargs=module_kwargs,
        )
    )
    # 处理 self-attention 的输出投影层
    if module.self_attn.v_proj.weight.shape == module.self_attn.o_proj.weight.shape:
        scales_list.append(
            _auto_get_scale(
                prev_op=module.self_attn.v_proj,
                layers=[module.self_attn.o_proj],
                inp=input_feat["self_attn.o_proj"],
            )
        )
    scales_list.append(
        _auto_get_scale(
            prev_op=module.post_attention_layernorm,
            layers=[module.mlp.gate_proj, module.mlp.up_proj],
            inp=input_feat["mlp.gate_proj"],
            module2inspect=module.mlp,
        )
    )
    scales_list.append(
        _auto_get_scale(
            prev_op=module.mlp.up_proj,
            layers=[module.mlp.down_proj],
            inp=input_feat["mlp.down_proj"],
        )
    )

# Weight Quantization: 自动裁剪层权重以适应量化
@torch.no_grad()
def auto_clip_layer(
    w, input_feat, n_bit, q_config, n_grid=20, max_shrink=0.5, n_sample_token=512
):
    """
    自动裁剪层权重以最小化量化误差。

    参数:
        w (torch.Tensor): 权重矩阵，形状为 [out_features, in_features]。
        input_feat (torch.Tensor): 输入特征，形状为 [n_token, in_features]。
        n_bit (int): 量化位数，例如8表示INT8量化。
        q_config (dict): 量化配置参数，例如 {"q_group_size": 128}。
        n_grid (int): 搜索缩放比例的网格大小，默认20。
        max_shrink (float): 最大缩放比例，默认为0.5。
        n_sample_token (int): 采样的token数量，默认512。

    返回:
        torch.Tensor: 最佳的最大裁剪值，形状为 [out_features, n_group]。
    """
    assert w.dim() == 2, "权重矩阵必须是二维的"
    org_w_shape = w.shape  # 原始权重形状

    # 获取分组大小，如果q_group_size > 0则使用该值，否则使用权重的列数
    group_size = (
        q_config["q_group_size"] if q_config["q_group_size"] > 0 else w.shape[1]
    )

    # 调整输入特征的形状
    input_feat = input_feat.view(-1, input_feat.shape[-1])  # [n_token, in_features]
    input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size)  # [1, n_token, n_group, group_size]
    input_feat = input_feat[:, 0 :: input_feat.shape[1] // n_sample_token]  # 采样输入特征
    w = w.reshape(w.shape[0], 1, -1, group_size)  # [out_features, 1, n_group, group_size]

    # 定义每个批次处理的输出通道数，防止显存溢出
    oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64
    assert w.shape[0] % oc_batch_size == 0, "输出通道数必须是 oc_batch_size 的倍数"

    w_all = w  # 保存所有权重
    best_max_val_all = []  # 保存所有批次的最佳最大值

    # 遍历每个输出通道批次
    for i_b in range(w.shape[0] // oc_batch_size):
        # 获取当前批次的权重
        w_batch = w_all[i_b * oc_batch_size : (i_b + 1) * oc_batch_size]

        # 计算每个输出通道的原始最大绝对值
        org_max_val = w_batch.abs().amax(dim=-1, keepdim=True)  # [oc_batch_size, 1, n_group, 1]

        best_max_val = org_max_val.clone()  # 初始化最佳最大值为原始最大值
        min_errs = torch.ones_like(org_max_val) * 1e9  # 初始化最小误差为很大的值
        input_feat = input_feat.to(w_batch.device)  # 将输入特征移动到权重所在设备

        # 计算原始输出
        org_out = (input_feat * w_batch).sum(dim=-1)  # [oc_batch_size, n_token, n_group]

        # 遍历缩放比例网格
        for i_s in range(int(max_shrink * n_grid)):
            # 计算当前缩放比例
            shrink_ratio = 1 - i_s / n_grid
            max_val = org_max_val * shrink_ratio  # 缩放后的最大值
            min_val = -max_val  # 对称最小值

            # 裁剪权重
            cur_w = torch.clamp(w_batch, min_val, max_val)

            # 量化权重
            q_w = pseudo_quantize_tensor(cur_w, n_bit=n_bit, **q_config)

            # 计算量化后的输出
            cur_out = (input_feat * q_w).sum(dim=-1)

            # 计算量化误差（均方误差）
            err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape)  # [oc_batch_size, 1, n_group, 1]

            # 更新最小误差和最佳最大值
            cur_best_idx = err < min_errs
            min_errs[cur_best_idx] = err[cur_best_idx]
            best_max_val[cur_best_idx] = max_val[cur_best_idx]

            # 释放不再需要的张量
            del cur_w
            del cur_out

        # 保存当前批次的最佳最大值
        best_max_val_all.append(best_max_val)

    # 合并所有批次的最佳最大值
    best_max_val = torch.cat(best_max_val_all, dim=0)  # [out_features, 1, n_group, 1]

    # 清理内存
    del input_feat
    del org_out
    gc.collect()
    torch.cuda.empty_cache()

    # 返回最佳最大值，去除多余的维度
    return best_max_val.squeeze(1)  # [out_features, n_group]

@torch.no_grad()
def run_awq():
    # 逐层处理模型
    for i in tqdm.tqdm(range(len(layers)), desc="Running AWQ..."):
        layer = layers[i]
        layer = layer.cuda()
        named_linears = get_named_linears(layer)

        # 定义一个钩子，用于缓存所有线性层的输入特征
        def cache_input_hook(m, x, y, name, feat_dict):
            x = x[0]
            x = x.detach().cpu()
            feat_dict[name].append(x)

        input_feat = defaultdict(list)
        handles = []
        for name in named_linears:
            handles.append(
                named_linears[name].register_forward_hook(
                    functools.partial(cache_input_hook, name=name, feat_dict=input_feat)
                )
            )
        inps = inps.to(next(layer.parameters()).device)  # 确保输入在正确的设备上
        # 获取输出作为下一层的输入
        inps = layer(inps, **layer_kwargs)[0]
        for h in handles:
            h.remove()

        # 将所有输入特征拼接
        input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()}
        ###省略代码####
        scales_list = auto_scale_block(layer, layer_kwargs, w_bit=w_bit,q_config=q_config, input_feat=input_feat,)
        apply_scale(layers[i], scales_list, input_feat_dict=input_feat)
        
        clip_list = auto_clip_block(layer,w_bit=w_bit,q_config=q_config,input_feat=input_feat,)
        apply_clip(layer, clip_list)

def pseudo_quantize_tensor(
    w, n_bit=8, zero_point=True, q_group_size=-1, inplace=False, get_scale_zp=False
):
    org_w_shape = w.shape
    if q_group_size > 0:
        assert org_w_shape[-1] % q_group_size == 0
        # num_groups = weight.shape[1] // group_size
        w = w.reshape(-1, q_group_size) # 权重分组
    assert w.dim() == 2
    if zero_point:
        max_val = w.amax(dim=1, keepdim=True)
        min_val = w.amin(dim=1, keepdim=True)
        max_int = 2**n_bit - 1
        min_int = 0
        scales = (max_val - min_val).clamp(min=1e-5) / max_int
        zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
    else:  # we actually never used this
        assert min_val is None
        max_val = w.abs().amax(dim=1, keepdim=True)
        max_val = max_val.clamp(min=1e-5)
        max_int = 2 ** (n_bit - 1) - 1
        min_int = -(2 ** (n_bit - 1))
        scales = max_val / max_int
        zeros = 0

    assert torch.isnan(scales).sum() == 0
    assert torch.isnan(w).sum() == 0

    if inplace:
        (
            (w.div_(scales).round_().add_(zeros)).clamp_(min_int, max_int).sub_(zeros)
        ).mul_(scales)
    else:
        w = (
            torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros
        ) * scales
    assert torch.isnan(w).sum() == 0

    w = w.reshape(org_w_shape)

    if get_scale_zp:
        return w, scales.view(w.shape[0], -1), zeros.view(w.shape[0], -1)
    else:
        return w

w = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
q_group_size = 2
w = w.reshape(-1, q_group_size)
# 结果：tensor([[1, 2],
#               [3, 4],
#               [5, 6],
#               [7, 8]])

class WQLinear(nn.Module):
    def __init__(self, w_bit, group_size, in_features, out_features, bias, dev, dtype=torch.float16):
        """
        构造函数，用于初始化 4-bit 权重量化的线性层。
        
        参数:
          w_bit: 权重量化的比特数，目前仅支持 4-bit。
          group_size: 分组大小，用于分组量化；如果传入 -1，则默认按 in_features 分组。
          in_features: 输入特征数。
          out_features: 输出特征数。
          bias: 是否需要 bias（偏置）。
          dev: 设备（如 'cuda'）。
          dtype: 存储缩放因子和 scaled_zeros 的数据类型，默认 torch.float16。
        """
        super().__init__()

        # 目前仅支持 4-bit 量化
        if w_bit not in [4]:
            raise NotImplementedError("Only 4-bit are supported for now.")

        # 保存基本参数
        self.in_features = in_features
        self.out_features = out_features
        self.w_bit = w_bit
        # 如果 group_size 为 -1，则使用 in_features 作为分组大小
        self.group_size = group_size if group_size != -1 else in_features
        # 分割迭代次数，固定为8
        self.split_k_iters = 8
        # 内部交织因子，固定为4
        self.interleave = 4

        # 简单检查：保证输入特征数能够整除分组大小
        assert self.in_features % self.group_size == 0
        # 保证输出特征数能够被 (32 / w_bit) 整除
        assert out_features % (32 // self.w_bit) == 0
        # pack_num: 每32位中能存储的量化数目（例如 32//4 = 8）
        pack_num = 32 // self.w_bit
        # int16_pack_num: 每16位中能存储的量化数目（例如 16//4 = 4）
        int16_pack_num = 16 // self.w_bit

        # 检查输出特征数应能被交织因子整除
        assert out_features % (self.interleave) == 0

        # 初始化 qweight 缓冲区，用于存储量化后的权重，数据类型为 int16
        # qweight 的尺寸为: [out_features // interleave, in_features // int16_pack_num * interleave]
        self.register_buffer(
            "qweight",
            torch.zeros(
                (
                    out_features // self.interleave,
                    in_features // int16_pack_num * self.interleave,
                ),
                dtype=torch.int16,
                device=dev,
            ),
        )
        # 初始化 scales 缓冲区，用于存储缩放因子
        # 尺寸为: [calculate_zeros_width(in_features, group_size) * pack_num, out_features]
        self.register_buffer(
            "scales",
            torch.zeros(
                (
                    calculate_zeros_width(in_features, self.group_size) * pack_num,
                    out_features,
                ),
                dtype=dtype,
                device=dev,
            ),
        )
        # 初始化 scaled_zeros 缓冲区，用于存储经过缩放的零点偏移
        self.register_buffer(
            "scaled_zeros",
            torch.zeros(
                (
                    calculate_zeros_width(in_features, self.group_size) * pack_num,
                    out_features,
                ),
                dtype=dtype,
                device=dev,
            ),
        )

        # 如果 bias 为 True，则初始化 bias 缓冲区，尺寸为 [out_features]
        if bias:
            self.register_buffer(
                "bias", torch.zeros((out_features), dtype=dtype, device=dev)
            )
        else:
            self.bias = None

@classmethod
def from_linear(
    cls, linear, w_bit, group_size, init_only=False, scales=None, zeros=None
):
    """
    类方法，用于从一个全精度的线性层构造量化后的 WQLinear 对象。
    
    参数:
        linear: 全精度线性层（如 nn.Linear）的实例。
        w_bit: 权重量化的比特数（目前仅支持4-bit）。
        group_size: 分组大小，用于分组量化。
        init_only: 如果为 True，则仅初始化对象，不进行实际的量化操作。
        scales: 缩放因子矩阵，用于量化时反算权重。
        zeros: 零点矩阵，用于量化时调整偏置。
    
    返回:
        构造好的 WQLinear 对象。
    """
    # 根据 full precision linear 层的参数构造一个 WQLinear 实例
    awq_linear = cls(
        w_bit,
        group_size,
        linear.in_features,
        linear.out_features,
        linear.bias is not None,
        linear.weight.device,
        dtype=linear.weight.data.dtype
    )
    if init_only:  # 如果仅初始化，不进行量化，则直接返回
        return awq_linear

    # 实际量化时需要提供 scales 和 zeros 信息
    assert scales is not None and zeros is not None
    # scale_zeros 为每个分组中零点和缩放因子的乘积，用于权重偏移调整
    scale_zeros = zeros * scales

    dtype = scales.dtype

    pack_num = 32 // awq_linear.w_bit
    # 构造 qscales，用于存放量化时的缩放因子，尺寸为:
    # [scales.shape[0], calculate_zeros_width(in_features, group_size) * pack_num]
    qscales = torch.zeros(
        (
            scales.shape[0],
            calculate_zeros_width(linear.in_features, group_size) * pack_num,
        ),
        dtype=dtype,
        device=scales.device,
    )
    # 将提供的 scales 填充到 qscales 的前面部分
    qscales[:, : scales.shape[1]] = scales
    # 将 qscales 转置后存入 awq_linear.scales（并确保内存连续）
    awq_linear.scales = qscales.transpose(1, 0).contiguous()
    # 如果原始线性层有 bias，则复制 bias 到量化层
    if linear.bias is not None:
        awq_linear.bias = linear.bias.clone().to(dtype)

    # 量化权重：对每个输入特征 idx 进行处理
    intweight = []
    for idx in range(awq_linear.in_features):
        # 对每个输入特征，计算对应的量化权重：
        # 公式为：round((原始权重 + scale_zeros) / qscales)
        # 注意这里 idx // group_size 表示同一组内共享相同的 scales 与 zeros 参数
        intweight.append(
            torch.round(
                (linear.weight.data[:, idx] + scale_zeros[:, idx // group_size])
                / qscales[:, idx // group_size]
            ).to(torch.int)[:, None]
        )
    # 将每个输入特征的量化结果拼接成一个矩阵
    intweight = torch.cat(intweight, dim=1)
    # 转换数据类型为 int32
    intweight = intweight.to(dtype=torch.int32)
    # 使用 pack_intweight 将 intweight 按照 interleave=4 和 kstride=64 进行打包，存入 qweight
    awq_linear.qweight = pack_intweight(
        intweight.contiguous(), interleave=4, kstride=64
    )

    # 对 zeros 进行数据类型转换为 int32
    zeros = zeros.to(dtype=torch.int32)
    # 创建与 qscales 同形状的 scaled_zeros，用于存储经过缩放后的零点
    scaled_zeros = torch.zeros_like(qscales)
    # 计算 scaled_zeros：公式为 - (qscales * zeros)，将 zeros 转换为 float32 后乘以 qscales，再转换回目标 dtype
    scaled_zeros[:, : scales.shape[1]] = -(
        qscales[:, : scales.shape[1]] * (zeros.to(torch.float32))
    ).to(dtype)
    # 存储 scaled_zeros（转置后确保内存连续）
    awq_linear.scaled_zeros = scaled_zeros.transpose(1, 0).contiguous()

    return awq_linear

def test_pack_intweight_edge_case(self):
    """
    测试特殊情况下的打包，如所有权重都是最大值15或最小值0。
    """
    N = 8
    K = 128
    interleave = 4
    kstride = 32

    # 测试所有权重为15的情况
    unpacked_qweight_max = torch.full((N, K), 15, dtype=torch.int64).to(self.device)
    packed_qweight_max = pack_intweight(unpacked_qweight_max, interleave, kstride)
    expected_packed_value_max = 15 | (15 << 4) | (15 << 8) | (15 << 12)
    self.assertTrue(torch.all(packed_qweight_max == expected_packed_value_max),
                    "Packed weight value mismatch when all weights are 15.")

    # 测试所有权重为0的情况
    unpacked_qweight_min = torch.zeros((N, K), dtype=torch.int64).to(self.device)
    packed_qweight_min = pack_intweight(unpacked_qweight_min, interleave, kstride)
    expected_packed_value_min = 0
    self.assertTrue(torch.all(packed_qweight_min == expected_packed_value_min),
                    "Packed weight value mismatch when all weights are 0.")