SINQによるgpt-ossの再量子化

Huaweiが開発したSINQという手法は、従来のHuggingFaceモデルを任意に再量子化できるという特徴を持っている。

非常に手軽に量子化できるため、有用性がかなり高い方法だと思うが、実際に量子化する上でいくつか問題になった落とし穴をメモしておく。

題材として使ったのは、gpt-oss-20bで、これはすでに4ビット量子化された状態なので、これを3ビットに再量子化するプロセスを見ながら落とし穴について調べていく。

https://github.com/huawei-csl/SINQ

まず、SINQをインストールする。この手順は公式と同じである。

# 1. Clone the repository
git clone https://github.com/huawei-csl/SINQ.git
cd SINQ

# 2. Install dependencies
pip install -r req.txt

# 3. Install SINQ
pip install .

これで、まず量子化だけ先に行う。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sinq.patch_model import AutoSINQHFModel
from sinq.sinqlinear import BaseQuantizeConfig

model_name ="openai/gpt-oss-20b"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16,device_map="cuda:0")
tokenizer = AutoTokenizer.from_pretrained(model_name)

quant_cfg = BaseQuantizeConfig(
    nbits=3,            # quantization bit-width
    group_size=64,     # group size
    tiling_mode="1D",   # tiling strategy
    method="sinq"       # quantization method ("asinq" for the calibrated version)
)

AutoSINQHFModel.quantize_model(
    model,
    tokenizer=tokenizer,
    quant_config=quant_cfg,
    compute_dtype=torch.bfloat16,
    device="cuda:0"
)

from sinq.patch_model import AutoSINQHFModel

save_dir = "openai-gpt-oss-20b-sinq-3bit"  # any path
AutoSINQHFModel.save_quantized(model, save_dir, verbose=True) # model is an already sinq-quantized model

from sinq.patch_model import AutoSINQHFModel
import torch

qmodel = AutoSINQHFModel.from_quantized(
    save_dir,
    device="cuda:0",
    compute_dtype=torch.bfloat16, 
)

# (optional) quick smoke test
prompt = "Explain neural network quantization in one sentence."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
with torch.inference_mode():
    out_ids = qmodel.generate(**inputs, max_new_tokens=32, do_sample=False)
print(tokenizer.decode(out_ids[0], skip_special_tokens=True))

このプログラムを実行すると、「openai-gpt-oss-20b-sinq-3bit」というディレクトリに量子化後のモデルが得られる。

ここでハマったのは、group_sizeでブロックを分割しなければならない時、ブロックサイズは複数あるため、最大公約数を取らなければならないということ。gpt-oss-20bの場合は、ブロックサイズは2880と4096だったので、最大公約数は64となった。

このモデルを推論しようとすると、また別のハマりどころが出現する。

gpt-ossの内部でテンソルを生成する部分があり、これがSINQで量子化されたテンソルと異なるデバイスに配置されることがあるのでそこでエラーが頻発した。

これを回避するパッチを当てる。

from transformers import AutoTokenizer                                                                                     
model_name ="openai/gpt-oss-20b"                                                                                           
                                                                                                                           
tokenizer = AutoTokenizer.from_pretrained(model_name)                                                                      
from sinq.patch_model import AutoSINQHFModel                                                                               
save_dir = "openai-gpt-oss-20b-sinq-3bit"  # any path                                                                      
                                                                                                                           
from sinq.patch_model import AutoSINQHFModel                                                                               
import torch                                                                                                               
qmodel = AutoSINQHFModel.from_quantized(                                                                                   
     save_dir,                                                                                                             
     device="cuda:0",
     compute_dtype=torch.bfloat16, 
 )

def fix_attention_sinks(model, device="cuda:0"):
    fixed = 0
    for mod in model.modules():
        if mod.__class__.__name__ == "GptOssAttention":
            if hasattr(mod, "sinks"):
                sinks = getattr(mod, "sinks")
                if isinstance(sinks, torch.Tensor) and sinks.device.type == "meta":
                    # まず消す
                    delattr(mod, "sinks")
                    # buffer として再登録（ゼロ埋め実体化）
                    real = torch.zeros_like(sinks, device=device)
                    mod.register_buffer("sinks", real, persistent=True)
                    print(f"[fix_attention_sinks] re-registered sinks as buffer on {device}")
                    fixed += 1 
    return fixed

fixed = fix_attention_sinks(qmodel, device="cuda:0")
print("sinks fixed:", fixed)

def _fix_sinks_pre_hook(mod, *args, **kwargs):
    # args[0] が hidden_states のはずなので、その device に合わせる
    target_dev = None
    for a in args:
        if torch.is_tensor(a): 
            target_dev = a.device
            break
    if target_dev is None:
        return
    if hasattr(mod, "sinks") and isinstance(mod.sinks, torch.Tensor):
        if mod.sinks.device != target_dev:
            mod.register_buffer("sinks", mod.sinks.to(target_dev), persistent=True)

# SelfAttention モジュールにだけフックを仕込む
for m in qmodel.modules():
    if m.__class__.__name__.lower().endswith("attention"):
        m.register_forward_pre_hook(_fix_sinks_pre_hook)

prompt = "Q:日本で一番高い山は?\nA:"
inputs = tokenizer(prompt, return_tensors="pt")


try:
    embed_dev = qmodel.model.embed_tokens.weight.device
except Exception:
    embed_dev = next(qmodel.parameters()).device  # フォールバック

for k, v in inputs.items():
    if torch.is_tensor(v):
        inputs[k] = v.to(embed_dev, non_blocking=True)

import os
os.environ["TORCH_COMPILE_DISABLE"] = "1"  # torch.compile 経路を一旦無効化

with torch.inference_mode():
     out_ids = qmodel.generate(**inputs, max_new_tokens=32, do_sample=False, use_cache=False)
print(tokenizer.decode(out_ids[0], skip_special_tokens=True)) 

厳密には、gpt-ossをもとに推論する場合、チャットテンプレートに載せなければならないが、本稿ではその部分は省略してある。

無事、日本語であっても推論できることが確認できた。

SINQを使えば、gpt-oss-20bを2ビットに再量子化することもできる。つまり、現状の半分以下のVRAMでも動かすことができるということになる。

SINQは、量子化と同時に高速化、省VRAM化も実現できるので、使い方によっては強力な道具になるだろう。