游乐游手机版
首页/AI热点日报/热点详情

Quarot量化Qwen3模型实现在线推理教程

类型:热点整理2026-07-19
本文深入探索Qwen3-8B模型的W8A8在线旋转量化技术,并完整演示如何将其集成到transformers框架中实现推理运行。 首先明确几个核心要点:要掌握Quarot旋转量化,关键在于实现在线旋转——这需要修改模型的forward计算图。本文将带您完整走一遍流程,涵盖配置编写、模型修改、权重存储

本文深入探索Qwen3-8B模型的W8A8在线旋转量化技术,并完整演示如何将其集成到transformers框架中实现推理运行。

首先明确几个核心要点:要掌握Quarot旋转量化,关键在于实现在线旋转——这需要修改模型的forward计算图。本文将带您完整走一遍流程,涵盖配置编写、模型修改、权重存储,以及最终在transformers环境中正常推理的步骤。

使用quarot量化qwen3并实现在线推理

我们采用的量化框架是llmc,推理框架则使用transformers(当然,你也可以替换为vllm或sglang等框架)。

配置详解:在线旋转量化YAML文件编写

首先按照llmc要求配置好环境,然后在configs/quantization目录下新建一个mine子目录,放置在线旋转的配置文件,例如命名为quarot_w_a.yml

llmc原生仅支持opt和llama的在线旋转,因此model的type需设置为Llama——好在Qwen3与Llama结构相似,可以直接复用。

base:
  seed: &seed42
model:
  type: Llama
  path: /data/gongoubo/checkpoints/Qwen/Qwen3-8B
  torch_dtype: auto
quant:
  method: Quarot
  weight:
    bit: 8
    symmetric: True
    granularity: per_channel
    group_size: -1
    calib_algo: minmax
  act:
    bit: 8
    symmetric: True
    granularity: per_token
  special:
    rotate_mode: hadamard
    fp32_had: True
    online_rotate: True
sa ve:
  sa ve_trans: True
  sa ve_fake: True
  sa ve_vllm: True
  sa ve_path: /data/gongoubo/checkpoints/Qwen/llmc/Qwen3-8B-w8a8-online

这里我们采用W8A8量化方案:权重使用per_channel粒度,激活使用per_token粒度。几个保存选项的含义需要明确:

  • save_trans:仅保存旋转后的权重(未量化)。
  • save_fake:保存旋转+量化后再反量化恢复的权重(模拟量化效果)。
  • save_vllm:保存旋转后真正量化了的权重。

由于transformers本身不支持W8A8的整数推理,后续加载推理时,我们将使用save_fake保存的版本。

模型修改核心:Attention和MLP的forward函数

在线旋转的关键在于修改模型的两个模块:Attention的v、o投影,以及MLP的up、down投影。下面直接展示代码修改——

MLP部分:在__init__中注册在线旋转器,在forward中对激活值执行Hadamard变换。

class Qwen3MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

        had_K_tensor, K_tensor = get_hadK(self.intermediate_size)
        self.rotater = Rotater(
            online_full_had=True,   # MLP使用全Hadamard变换
            online_partial_had=False,
            fp32_had=True,
            K=K_tensor,
            had_K=had_K_tensor,
            had_dim=None,
        )
        print(f'enable online rotate for Qwen2MLP')

    def forward(self, x):
        act = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
        act = self.rotater.rotate(act)
        down_proj = self.down_proj(act)
        return down_proj

Attention部分:对Q、K进行RoPE编码后,在计算注意力之前对value执行部分Hadamard变换;注意力输出后,再对output进行一次旋转。注意,Q和K已通过RoPE处理,因此旋转操作需放在注意力计算之前。

class Qwen3Attention(nn.Module):
    def __init__(self, config, layer_idx):
        # ... 省略部分初始化代码,与原始一致 ...
        had_K_tensor, K_tensor = get_hadK(self.num_head)
        self.rotater = Rotater(
            online_full_had=False,
            online_partial_had=True,
            fp32_had=True,
            K=K_tensor,
            had_K=had_K_tensor,
            had_dim=self.head_dim,
        )
        print("enable Qwen3Attention")

    def forward(self, hidden_states, position_embeddings, attention_mask, ...):
        # 常规QKV投影 + QK做RoPE
        query_states = self.q_norm(self.q_proj(hidden_states).view(...)).transpose(1,2)
        key_states = self.k_norm(self.k_proj(hidden_states).view(...)).transpose(1,2)
        value_states = self.v_proj(hidden_states).view(...).transpose(1,2)
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        # 在线旋转 value(以及 Q/K 的部分,如果需要)
        # 这里旋转是放到 attention 计算之前对 value 做?注意原代码中旋转是放在后面,我们保留原样
        # 为了清晰,省略中间部分(原代码中旋转在注意力输出后)
        attn_output, attn_weights = some_attention_forward(...)
        attn_output = attn_output.reshape(-1, self.num_head * self.head_dim)
        attn_output = self.rotater.rotate(attn_output)
        attn_output = attn_output.reshape(...).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights

推理环节:加载量化模型并直接运行

完成所有修改后,推理过程变得非常简单——正常加载fake_quant_model目录下的权重和配置文件即可。注意需要导入修改过的modeling_qwen3_online_llmc模块(即您刚才修改好结构的文件)。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "7"
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from modeling_qwen3_online_llmc import Qwen3ForCausalLM

model_name = "/data/gongoubo/checkpoints/Qwen/llmc/Qwen3-8B-w8a8-online/fake_quant_model/"
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)

# 如果词嵌入与LM头共享权重,需要先解绑,以避免加载时报维度不一致
if config.tie_word_embeddings:
    config.tie_word_embeddings = False

model = Qwen3ForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, config=config).to("cuda:0")
model.eval()

# 构造输入
message = "<|im_start|>user\n你是谁?<|im_end|>\n<|im_start|>assistant\n"
input_ids = tokenizer.encode(message, return_tensors="pt").to(model.device)
output = model.generate(input_ids, max_new_tokens=256, do_sample=False)
print(tokenizer.decode(output[0]))

整个过程就是这么简单。请记得替换为您自己修改好的modeling_qwen3文件路径。至此,全部步骤完成,可以顺利运行。

来源:https://www.53ai.com/news/finetuning/2025073001593.html

相关热点

继续查看同栏目近期热点。

延伸阅读

补充最近整理过的热点入口。