BitCPM-CANN-3B-unquantized源码解析：modeling_llama.py中的量化实现

【免费下载链接】BitCPM-CANN-3B-unquantized BitCPM-CANN-3B-unquantized 是 BitCPM-CANN-3B 的未量化量化感知训练（QAT）检查点，专为持续预训练和微调而设计。它保留了全精度的潜在权重，并通过 modeling.py 中定义的三元伪量化器（权重 → {-1, 0, 1}，带分组缩放，通过 STE 训练），使模型能够在量化约束下继续学习。 项目地址: https://ai.gitcode.com/OpenBMB/BitCPM-CANN-3B-unquantized

BitCPM-CANN-3B-unquantized是OpenBMB开源社区推出的未量化量化感知训练（QAT）检查点，专为持续预训练和微调设计。它保留全精度潜在权重，通过modeling_llama.py中定义的三元伪量化器实现量化约束下的持续学习，是AI模型部署优化的实用方案。

量化实现核心架构概览

modeling_llama.py文件实现了完整的LLaMA模型架构与量化功能，核心量化组件包括：

• 量化器选择机制：通过get_quantizer函数实现不同量化策略的灵活切换
• 三元伪量化器：核心类SteTernaryQuantizer实现权重{-1,0,1}的三元量化
• 量化线性层：LinearQuantizer类将量化逻辑嵌入标准线性层
• 模型集成：在注意力机制和MLP模块中应用量化线性层

量化器类型与选择逻辑

文件第60-68行定义了量化器选择函数，支持三种量化类型：

def get_quantizer(quant_type="none", bit=4, group_size=128):
    if quant_type == "intsym":
        return SteIntSymQuantizerGPTQ(bit, group_size)
    elif quant_type == "ternary":
        return SteTernaryQuantizer(group_size)
    elif quant_type == "none":
        return NoQuantizer()
    else:
        raise ValueError(f"Unsupported quantization type: {quant_type}")

项目默认使用三元量化（ternary），这是一种高效的模型压缩技术，通过将权重限制为-1、0和1三个离散值，显著降低计算复杂度和内存占用。

三元伪量化器深度解析

核心原理与实现

SteTernaryQuantizer类（109-132行）实现了带分组缩放的三元量化，采用直通估计器（STE）进行梯度计算：

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:
            x = x.reshape(-1, self.group_size)
        elif self.group_size == -1:
            x = x.reshape(-1, x.shape[-1])

        # 计算缩放因子
        scales = 1.0 / (x.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5))
        # 量化为-1,0,1并反缩放
        x_q = (torch.clamp(torch.round(x * scales),-1,1) / scales)

        x_q = x_q.reshape(org_w_shape)
        # STE梯度估计
        return x + (x_q - x).detach()

关键技术点解析

1. 分组量化：通过group_size参数将权重分成多个小组独立量化，平衡量化精度和计算效率

2. 缩放因子计算：使用每组权重的绝对值均值作为缩放基准：

   scales = 1.0 / (x.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5))
   

3. 三元量化过程：

   • 权重乘以缩放因子后四舍五入
   • 钳位到{-1, 0, 1}三个离散值
   • 除以缩放因子恢复量级

4. STE训练技巧：通过(x_q - x).detach()实现直通估计，前向传播使用量化值，反向传播使用原始权重梯度，解决量化操作不可导问题

量化线性层的实现与应用

LinearQuantizer类设计

LinearQuantizer类（141-151行）继承自nn.Linear，将量化逻辑嵌入线性层：

class LinearQuantizer(nn.Linear):
    def __init__(self, in_features, out_features, bias=False, quant_type="ternary", bit=4, group_size=-1):
        super().__init__(in_features, out_features, bias)
        self.quantizer = get_quantizer(quant_type, bit, group_size)

    def forward(self, x):
        weight_tensor = self.quantizer(self.weight)
        x = torch.nn.functional.linear(x, weight_tensor)
        if self.bias is not None:
            x = x + self.bias
        return x

在模型中的应用

量化线性层被广泛应用于模型关键组件：

1. 注意力机制（392-395行）：

   self.q_proj = LinearQuantizer(...)
   self.k_proj = LinearQuantizer(...)
   self.v_proj = LinearQuantizer(...)
   self.o_proj = LinearQuantizer(...)
   

2. MLP模块（328-330行）：

   self.gate_proj = LinearQuantizer(...)
   self.up_proj = LinearQuantizer(...)
   self.down_proj = LinearQuantizer(...)
   

这种设计使量化操作与模型结构紧密结合，在几乎不改变原有代码结构的情况下实现量化功能。

量化效果验证与可视化

项目提供了丰富的量化训练效果对比图，展示了在GPU和NPU上的预训练与SFT（监督微调）损失曲线：

GPU与NPU量化训练对比

BitCPM-CANN-3B-unquantized在GPU上的预训练损失曲线，展示了量化训练的稳定性

BitCPM-CANN-3B-unquantized在NPU上的预训练损失曲线，验证了跨平台量化一致性

SFT微调效果

GPU环境下量化模型监督微调的损失变化

NPU环境下量化模型监督微调的损失变化，与GPU结果高度一致

从损失曲线可以看出，量化模型在训练过程中表现稳定，与全精度模型的收敛趋势基本一致，验证了三元量化方案的有效性。

快速开始使用量化模型

要开始使用BitCPM-CANN-3B-unquantized的量化功能，可按以下步骤操作：

1. 克隆仓库：

   git clone https://gitcode.com/OpenBMB/BitCPM-CANN-3B-unquantized
   

2. 安装依赖：

   cd BitCPM-CANN-3B-unquantized/example
   pip install -r requirements.txt
   

3. 运行量化训练：

   # 预训练
   bash run.sh
   # 或SFT微调
   bash run_sft.sh
   

配置文件example/ds_config.json和example/ds_config_z2.json提供了不同量化配置选项，可根据需求调整。

总结与扩展

BitCPM-CANN-3B-unquantized通过modeling_llama.py中的三元伪量化器实现了高效的模型量化，主要优势包括：

• 精度与效率平衡：采用分组缩放和STE技术，在大幅降低计算成本的同时保持模型性能
• 灵活集成：通过LinearQuantizer类实现即插即用的量化功能
• 跨平台支持：在GPU和NPU上均表现稳定的量化训练效果

未来可探索的方向包括：不同量化粒度的对比实验、量化参数的自适应调整、更高效的量化感知训练策略等。通过这些优化，BitCPM-CANN-3B-unquantized有望在边缘设备部署和大规模推理场景中发挥更大价值。

通过深入理解modeling_llama.py中的量化实现，开发者可以更好地掌握模型压缩技术，为AI模型的高效部署提供有力支持。

【免费下载链接】BitCPM-CANN-3B-unquantized BitCPM-CANN-3B-unquantized 是 BitCPM-CANN-3B 的未量化量化感知训练（QAT）检查点，专为持续预训练和微调而设计。它保留了全精度的潜在权重，并通过 modeling.py 中定义的三元伪量化器（权重 → {-1, 0, 1}，带分组缩放，通过 STE 训练），使模型能够在量化约束下继续学习。 项目地址: https://ai.gitcode.com/OpenBMB/BitCPM-CANN-3B-unquantized