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

我们采用的量化框架是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文件路径。至此,全部步骤完成,可以顺利运行。
