ディープラーニングモデル量子化完全ガイド：INT8・INT4・GPTQ・AWQ・GGUFをマスターする

import torch
import numpy as np

# 各データ型のメモリサイズを確認
x_fp32 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.float32)
x_fp16 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.float16)
x_bf16 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.bfloat16)

print(f"FP32: {x_fp32.element_size()} bytes per element")  # 4 bytes
print(f"FP16: {x_fp16.element_size()} bytes per element")  # 2 bytes
print(f"BF16: {x_bf16.element_size()} bytes per element")  # 2 bytes

# 7Bパラメータモデルのメモリ計算
params = 7e9
fp32_memory_gb = params * 4 / 1e9
fp16_memory_gb = params * 2 / 1e9
int8_memory_gb = params * 1 / 1e9
int4_memory_gb = params * 0.5 / 1e9

print(f"\n7Bモデルのメモリ要件:")
print(f"FP32: {fp32_memory_gb:.1f} GB")   # 28.0 GB
print(f"FP16: {fp16_memory_gb:.1f} GB")   # 14.0 GB
print(f"INT8: {int8_memory_gb:.1f} GB")   # 7.0 GB
print(f"INT4: {int4_memory_gb:.1f} GB")   # 3.5 GB

q = clamp(round(x / scale) + zero_point, q_min, q_max)

x_approx = scale * (q - zero_point)

import torch
import numpy as np

def symmetric_quantize(x: torch.Tensor, num_bits: int = 8):
    """対称量子化の実装"""
    q_max = 2 ** (num_bits - 1) - 1  # INT8では127
    q_min = -q_max  # -127

    # スケールの計算
    max_abs = x.abs().max()
    scale = max_abs / q_max

    # 量子化
    q = torch.clamp(torch.round(x / scale), q_min, q_max).to(torch.int8)

    return q, scale

def asymmetric_quantize(x: torch.Tensor, num_bits: int = 8):
    """非対称量子化の実装"""
    q_max = 2 ** num_bits - 1  # UINT8では255
    q_min = 0

    # スケールとzero_pointの計算
    min_val = x.min()
    max_val = x.max()
    scale = (max_val - min_val) / (q_max - q_min)
    zero_point = q_min - torch.round(min_val / scale)
    zero_point = torch.clamp(zero_point, q_min, q_max).to(torch.int32)

    # 量子化
    q = torch.clamp(torch.round(x / scale) + zero_point, q_min, q_max).to(torch.uint8)

    return q, scale, zero_point

def dequantize(q: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor = None):
    """逆量子化"""
    if zero_point is None:
        return scale * q.float()
    return scale * (q.float() - zero_point.float())

# テスト
x = torch.randn(100)
print(f"元データの範囲: [{x.min():.4f}, {x.max():.4f}]")

# 対称量子化
q_sym, scale_sym = symmetric_quantize(x)
x_reconstructed_sym = dequantize(q_sym, scale_sym)
error_sym = (x - x_reconstructed_sym).abs().mean()
print(f"対称量子化の平均誤差: {error_sym:.6f}")

# 非対称量子化
q_asym, scale_asym, zp_asym = asymmetric_quantize(x)
x_reconstructed_asym = dequantize(q_asym, scale_asym, zp_asym)
error_asym = (x - x_reconstructed_asym).abs().mean()
print(f"非対称量子化の平均誤差: {error_asym:.6f}")

import torch

def per_group_quantize(weight: torch.Tensor, group_size: int = 128, num_bits: int = 4):
    """Per-Group量子化の実装"""
    rows, cols = weight.shape

    # グループに分割
    weight_grouped = weight.reshape(-1, group_size)

    # グループごとの最大/最小値
    max_vals = weight_grouped.max(dim=1, keepdim=True)[0]
    min_vals = weight_grouped.min(dim=1, keepdim=True)[0]

    q_max = 2 ** num_bits - 1  # INT4では15

    # スケールの計算
    scales = (max_vals - min_vals) / q_max
    zero_points = torch.round(-min_vals / scales)

    # 量子化
    q = torch.clamp(torch.round(weight_grouped / scales) + zero_points, 0, q_max)

    # 逆量子化
    weight_dequant = scales * (q - zero_points)
    weight_dequant = weight_dequant.reshape(rows, cols)

    return q, scales, zero_points, weight_dequant

# 例：Transformer重みの量子化
weight = torch.randn(4096, 4096)  # Llamaスタイルの重み
q, scales, zp, weight_dequant = per_group_quantize(weight, group_size=128, num_bits=4)

error = (weight - weight_dequant).abs().mean()
print(f"Per-Group INT4量子化の平均誤差: {error:.6f}")
print(f"圧縮率: {weight.element_size() * weight.numel() / (q.numel() / 2 + scales.numel() * 4):.2f}x")

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from datasets import load_dataset

def collect_calibration_data(model_name: str, num_samples: int = 128):
    """キャリブレーションデータの収集"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # WikiText-2またはC4データセットがよく使われる
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

    texts = []
    for item in dataset:
        if len(item['text'].strip()) > 100:
            texts.append(item['text'].strip())
        if len(texts) >= num_samples:
            break

    # トークン化
    encoded = [
        tokenizer(text, return_tensors="pt", max_length=2048, truncation=True)
        for text in texts
    ]

    return encoded

class MinMaxCalibrator:
    """Min-Maxキャリブレーター"""

    def __init__(self):
        self.min_val = float("inf")
        self.max_val = float("-inf")

    def update(self, tensor: torch.Tensor):
        self.min_val = min(self.min_val, tensor.min().item())
        self.max_val = max(self.max_val, tensor.max().item())

    def compute_scale_zp(self, num_bits: int = 8, symmetric: bool = True):
        q_max = 2 ** (num_bits - 1) - 1 if symmetric else 2 ** num_bits - 1

        if symmetric:
            max_abs = max(abs(self.min_val), abs(self.max_val))
            scale = max_abs / q_max
            zero_point = 0
        else:
            scale = (self.max_val - self.min_val) / q_max
            zero_point = -round(self.min_val / scale)

        return scale, zero_point

import numpy as np
from scipy import stats

class HistogramCalibrator:
    """ヒストグラムベースのキャリブレーター（KLダイバージェンス最小化）"""

    def __init__(self, num_bins: int = 2048):
        self.num_bins = num_bins
        self.histogram = None
        self.bin_edges = None

    def update(self, tensor: torch.Tensor):
        data = tensor.detach().float().numpy().flatten()

        if self.histogram is None:
            self.histogram, self.bin_edges = np.histogram(data, bins=self.num_bins)
        else:
            new_hist, _ = np.histogram(data, bins=self.bin_edges)
            self.histogram += new_hist

    def compute_optimal_range(self, num_bits: int = 8):
        """KLダイバージェンスを最小化する最適範囲を探索"""
        num_quantized_bins = 2 ** num_bits - 1

        best_kl = float("inf")
        best_threshold = None

        for i in range(num_quantized_bins, len(self.histogram)):
            reference = self.histogram[:i].copy().astype(float)
            reference /= reference.sum()

            quantized = np.zeros(i)
            bin_size = i / num_quantized_bins

            for j in range(num_quantized_bins):
                start = int(j * bin_size)
                end = int((j + 1) * bin_size)
                quantized[start:end] = reference[start:end].sum() / (end - start)

            quantized = np.where(quantized == 0, 1e-10, quantized)
            reference_clipped = np.where(reference == 0, 1e-10, reference)

            kl = stats.entropy(reference_clipped, quantized)

            if kl < best_kl:
                best_kl = kl
                best_threshold = self.bin_edges[i]

        return -best_threshold, best_threshold

import torch
import math
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_perplexity(model, tokenizer, text: str, device: str = "cuda"):
    """パープレキシティの計算"""
    encodings = tokenizer(text, return_tensors="pt")
    input_ids = encodings.input_ids.to(device)

    max_length = 1024
    stride = 512

    nlls = []
    prev_end_loc = 0

    for begin_loc in range(0, input_ids.size(1), stride):
        end_loc = min(begin_loc + max_length, input_ids.size(1))
        trg_len = end_loc - prev_end_loc

        input_ids_chunk = input_ids[:, begin_loc:end_loc]
        target_ids = input_ids_chunk.clone()
        target_ids[:, :-trg_len] = -100

        with torch.no_grad():
            outputs = model(input_ids_chunk, labels=target_ids)
            neg_log_likelihood = outputs.loss

        nlls.append(neg_log_likelihood)
        prev_end_loc = end_loc

        if end_loc == input_ids.size(1):
            break

    ppl = torch.exp(torch.stack(nlls).mean())
    return ppl.item()

# PPL比較例
# FP16:          PPL ~5.68
# INT8:          PPL ~5.71 (~0.5%増)
# INT4 (GPTQ):   PPL ~5.89 (~3.7%増)
# INT4 (naive):  PPL ~6.52 (~14.8%増)

import torch
import torch.nn as nn
import torch.nn.functional as F

class FakeQuantize(nn.Module):
    """Fake量子化モジュール"""

    def __init__(self, num_bits: int = 8, symmetric: bool = True):
        super().__init__()
        self.num_bits = num_bits
        self.symmetric = symmetric

        self.register_buffer('scale', torch.tensor(1.0))
        self.register_buffer('zero_point', torch.tensor(0))
        self.register_buffer('fake_quant_enabled', torch.tensor(1))

        if symmetric:
            self.q_min = -(2 ** (num_bits - 1))
            self.q_max = 2 ** (num_bits - 1) - 1
        else:
            self.q_min = 0
            self.q_max = 2 ** num_bits - 1

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.fake_quant_enabled[0] == 0:
            return x

        # 学習中に指数移動平均でスケールを更新
        if self.training:
            with torch.no_grad():
                if self.symmetric:
                    max_abs = x.abs().max()
                    new_scale = max_abs / self.q_max
                else:
                    new_scale = (x.max() - x.min()) / (self.q_max - self.q_min)

                self.scale.copy_(0.9 * self.scale + 0.1 * new_scale)

        # Fake量子化：量子化して逆量子化
        x_scaled = x / self.scale
        x_clipped = torch.clamp(x_scaled, self.q_min, self.q_max)
        x_rounded = torch.round(x_clipped)
        x_dequant = x_rounded * self.scale

        return x_dequant

class STERound(torch.autograd.Function):
    """round()のためのStraight-Through Estimator"""

    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, grad_output):
        # round()を通じて勾配をそのまま通過（恒等近似）
        return grad_output

class STEClamp(torch.autograd.Function):
    """clamp()のためのStraight-Through Estimator"""

    @staticmethod
    def forward(ctx, x, min_val, max_val):
        ctx.save_for_backward(x)
        ctx.min_val = min_val
        ctx.max_val = max_val
        return torch.clamp(x, min_val, max_val)

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        # clamp範囲内のみ勾配を通過
        grad = grad_output * ((x >= ctx.min_val) & (x <= ctx.max_val)).float()
        return grad, None, None

class QATLinear(nn.Module):
    """QATが適用された線形層"""

    def __init__(self, in_features, out_features, num_bits=8):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.weight_fake_quant = FakeQuantize(num_bits=num_bits)
        self.act_fake_quant = FakeQuantize(num_bits=num_bits, symmetric=False)

    def forward(self, x):
        # 活性化の量子化
        x_q = self.act_fake_quant(x)
        # 重みの量子化
        w_q = self.weight_fake_quant(self.linear.weight)
        # FP32計算（実際の展開ではINT8）
        return F.linear(x_q, w_q, self.linear.bias)

# QAT学習ワークフロー
import torch.optim as optim
from torch.quantization import prepare_qat, convert

def train_qat_model(model, train_loader, num_epochs=10):
    """QAT学習の例"""

    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
    model_prepared = prepare_qat(model.train())

    optimizer = optim.Adam(model_prepared.parameters(), lr=1e-5)

    for epoch in range(num_epochs):
        for batch in train_loader:
            inputs, labels = batch
            outputs = model_prepared(inputs)
            loss = F.cross_entropy(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # INT8モデルに変換
    model_prepared.eval()
    model_quantized = convert(model_prepared)

    return model_quantized

import torch
from torch.ao.quantization import (
    get_default_qconfig,
    get_default_qat_qconfig,
    prepare,
    prepare_qat,
    convert
)

# 静的量子化（PTQ）
def static_quantization_example():
    """静的量子化の例"""
    model = MyModel()
    model.eval()

    # バックエンド設定（fbgemm: x86、qnnpack: ARM）
    model.qconfig = get_default_qconfig('fbgemm')

    # キャリブレーション準備
    model_prepared = prepare(model)

    # キャリブレーションデータから統計を収集
    with torch.no_grad():
        for data in calibration_loader:
            model_prepared(data)

    # INT8モデルに変換
    model_quantized = convert(model_prepared)

    return model_quantized

# 動的量子化（LSTM、Linearに有効）
def dynamic_quantization_example():
    """動的量子化の例"""
    model = MyModel()

    model_quantized = torch.quantization.quantize_dynamic(
        model,
        {nn.Linear, nn.LSTM},  # 量子化する層の型
        dtype=torch.qint8
    )

    return model_quantized

from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization import QConfigMapping

def fx_quantization_example(model, calibration_data):
    """FXグラフモード量子化"""
    model.eval()

    qconfig_mapping = QConfigMapping().set_global(
        get_default_qconfig('fbgemm')
    )

    example_inputs = (torch.randn(1, 3, 224, 224),)

    model_prepared = prepare_fx(
        model,
        qconfig_mapping,
        example_inputs
    )

    with torch.no_grad():
        for batch in calibration_data:
            model_prepared(batch)

    model_quantized = convert_fx(model_prepared)

    return model_quantized

argmin_Q ||WX - QX||_F^2

import torch
import math

def gptq_quantize_weight(weight: torch.Tensor,
                          hessian: torch.Tensor,
                          num_bits: int = 4,
                          group_size: int = 128,
                          damp_percent: float = 0.01):
    """
    GPTQアルゴリズムによる重みの量子化

    Args:
        weight: [out_features, in_features] 重み行列
        hessian: [in_features, in_features] ヘッセ行列 (H = 2 * X @ X.T)
        num_bits: 量子化ビット数
        group_size: グループサイズ
        damp_percent: ヘッセ行列安定化のための減衰比
    """
    W = weight.clone().float()
    n_rows, n_cols = W.shape

    # ヘッセ行列のダンピング（数値安定性）
    H = hessian.clone().float()
    dead_cols = torch.diag(H) == 0
    H[dead_cols, dead_cols] = 1
    W[:, dead_cols] = 0

    damp = damp_percent * H.diag().mean()
    H.diagonal().add_(damp)

    # Cholesky分解による逆ヘッセ行列
    H_inv = torch.linalg.cholesky(H)
    H_inv = torch.cholesky_inverse(H_inv)
    H_inv = torch.linalg.cholesky(H_inv, upper=True)

    Q = torch.zeros_like(W)
    Losses = torch.zeros_like(W)

    q_max = 2 ** (num_bits - 1) - 1

    for col_idx in range(n_cols):
        w_col = W[:, col_idx]
        h_inv_diag = H_inv[col_idx, col_idx]

        # Per-groupスケールの計算
        if group_size != -1 and col_idx % group_size == 0:
            group_end = min(col_idx + group_size, n_cols)
            w_group = W[:, col_idx:group_end]
            max_abs = w_group.abs().max(dim=1)[0].unsqueeze(1)
            scale = max_abs / q_max
            scale = torch.clamp(scale, min=1e-8)

        # 量子化
        q_col = torch.clamp(torch.round(w_col / scale.squeeze()), -q_max, q_max)
        q_col = q_col * scale.squeeze()
        Q[:, col_idx] = q_col

        # 量子化誤差
        err = (w_col - q_col) / h_inv_diag
        Losses[:, col_idx] = err ** 2 / 2

        # 残りの重みへ誤差を伝播（核心ステップ！）
        W[:, col_idx + 1:] -= err.unsqueeze(1) * H_inv[col_idx, col_idx + 1:].unsqueeze(0)

    return Q, Losses

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch

def quantize_with_gptq(
    model_name: str,
    output_dir: str,
    bits: int = 4,
    group_size: int = 128
):
    """AutoGPTQによるモデル量子化"""

    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

    quantize_config = BaseQuantizeConfig(
        bits=bits,
        group_size=group_size,
        damp_percent=0.01,
        desc_act=False,
        sym=True,
        true_sequential=True
    )

    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config=quantize_config
    )

    # キャリブレーションデータの準備
    from datasets import load_dataset
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

    calibration_data = []
    for text in dataset["text"][:128]:
        if len(text.strip()) > 50:
            encoded = tokenizer(
                text.strip(),
                return_tensors="pt",
                max_length=2048,
                truncation=True
            )
            calibration_data.append(encoded["input_ids"].squeeze())

    print(f"GPTQ {bits}bit量子化開始...")
    model.quantize(calibration_data)

    model.save_quantized(output_dir, use_safetensors=True)
    tokenizer.save_pretrained(output_dir)

    print(f"量子化完了: {output_dir}")
    return model, tokenizer


def load_gptq_model(model_dir: str, device: str = "cuda"):
    """GPTQ量子化モデルの読み込み"""

    model = AutoGPTQForCausalLM.from_quantized(
        model_dir,
        device=device,
        use_triton=False,
        disable_exllama=False,
        inject_fused_attention=True,
        inject_fused_mlp=True
    )

    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    return model, tokenizer

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

def quantize_with_awq(
    model_name: str,
    output_dir: str,
    bits: int = 4,
    group_size: int = 128
):
    """AutoAWQによるモデル量子化"""

    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True
    )

    model = AutoAWQForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
        use_cache=False
    )

    quant_config = {
        "zero_point": True,
        "q_group_size": group_size,
        "w_bit": bits,
        "version": "GEMM"
    }

    print(f"AWQ {bits}bit量子化開始...")
    model.quantize(tokenizer, quant_config=quant_config)

    model.save_quantized(output_dir)
    tokenizer.save_pretrained(output_dir)

    print(f"AWQ量子化完了: {output_dir}")
    return model

def load_awq_model(model_dir: str, device: str = "cuda"):
    """AWQ量子化モデルの読み込み"""

    model = AutoAWQForCausalLM.from_quantized(
        model_dir,
        fuse_layers=True,
        trust_remote_code=True,
        safetensors=True
    )

    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    return model, tokenizer

# llama.cppのクローンとビルド
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp

# CUDAサポート付きビルド
cmake -B build -DGGML_CUDA=ON
cmake --build build --config Release -j $(nproc)

# CPUのみのビルド
cmake -B build
cmake --build build --config Release -j $(nproc)

# HuggingFaceモデルをGGUFに変換
python convert_hf_to_gguf.py \
    --model meta-llama/Llama-2-7b-hf \
    --outfile llama2-7b-f16.gguf \
    --outtype f16

# Q4_K_Mに量子化
./build/bin/llama-quantize \
    llama2-7b-f16.gguf \
    llama2-7b-q4_k_m.gguf \
    Q4_K_M

# 推論の実行
./build/bin/llama-cli \
    -m llama2-7b-q4_k_m.gguf \
    -p "The future of AI is" \
    -n 100 \
    --ctx-size 4096 \
    --threads 8 \
    --n-gpu-layers 35

from llama_cpp import Llama

# モデルの読み込み
llm = Llama(
    model_path="./llama2-7b-q4_k_m.gguf",
    n_ctx=4096,
    n_gpu_layers=35,
    n_threads=8,
    verbose=False
)

# テキスト生成
output = llm(
    "Once upon a time",
    max_tokens=200,
    temperature=0.7,
    top_p=0.9,
    stop=["</s>", "\n\n"]
)

print(output["choices"][0]["text"])

# チャット補完フォーマット
response = llm.create_chat_completion(
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is machine learning?"}
    ],
    max_tokens=500,
    temperature=0.7
)

print(response["choices"][0]["message"]["content"])

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# INT8モデルの読み込み
model_8bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    load_in_8bit=True,
    device_map="auto"
)

def print_model_size(model, label):
    """モデルのメモリ使用量を表示"""
    total_params = sum(p.numel() for p in model.parameters())
    total_bytes = sum(
        p.numel() * p.element_size() for p in model.parameters()
    )
    print(f"{label}: {total_params/1e9:.2f}Bパラメータ、{total_bytes/1e9:.2f} GB")

print_model_size(model_8bit, "INT8モデル")
# INT8モデル: 6.74Bパラメータ、約7.0 GB

import bitsandbytes as bnb
from transformers import BitsAndBytesConfig

# NF4量子化設定（QLoRA）
bnb_config_nf4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,  # 二重量子化
)

# FP4量子化設定
bnb_config_fp4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="fp4",
    bnb_4bit_compute_dtype=torch.float16,
)

# モデルの読み込み
model_4bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config_nf4,
    device_map="auto"
)

print_model_size(model_4bit, "NF4モデル")
# NF4モデル: 6.74Bパラメータ、約4.0 GB（二重量子化込み）

# QLoRAファインチューニングの設定
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model_4bit = prepare_model_for_kbit_training(model_4bit)

lora_config = LoraConfig(
    r=64,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model_lora = get_peft_model(model_4bit, lora_config)
model_lora.print_trainable_parameters()
# trainable params: 4,194,304 || all params: 3,504,607,232 || trainable%: 0.1197

Y = (X * diag(s)^(-1)) * (diag(s) * W)
  = X_smooth * W_smooth

def smooth_quantize(
    model,
    calibration_samples,
    alpha: float = 0.5
):
    """
    SmoothQuantを適用する

    Args:
        alpha: 転移強度（0=重みのみ、1=活性化のみ）
               推奨: 0.5（均等分配）
    """

    act_scales = {}

    def collect_scales(name):
        def hook(module, input, output):
            inp = input[0].detach()
            if inp.dim() == 3:
                inp = inp.reshape(-1, inp.size(-1))

            channel_max = inp.abs().max(dim=0)[0]

            if name not in act_scales:
                act_scales[name] = channel_max
            else:
                act_scales[name] = torch.maximum(act_scales[name], channel_max)
        return hook

    handles = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            handles.append(module.register_forward_hook(collect_scales(name)))

    with torch.no_grad():
        for sample in calibration_samples:
            model(**sample)

    for h in handles:
        h.remove()

    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) and name in act_scales:
            act_scale = act_scales[name]
            weight_scale = module.weight.abs().max(dim=0)[0]

            # マイグレーションスケールの計算
            smooth_scale = (act_scale ** alpha) / (weight_scale ** (1 - alpha))
            smooth_scale = torch.clamp(smooth_scale, min=1e-5)

            # 重みにスケールを適用
            module.weight.data = module.weight.data / smooth_scale.unsqueeze(0)

    return model, act_scales

import torch

def spqr_quantize(weight: torch.Tensor,
                   num_bits: int = 3,
                   outlier_threshold_percentile: float = 1.0):
    """
    SpQR量子化（簡略版）

    核心：上位p%の外れ値をFP16で格納し、残りを低ビット量子化
    """

    threshold = torch.quantile(weight.abs(), 1 - outlier_threshold_percentile / 100)
    outlier_mask = weight.abs() > threshold

    # 外れ値の格納（FP16）
    outlier_values = weight.clone()
    outlier_values[~outlier_mask] = 0

    # 残りの量子化
    regular_weight = weight.clone()
    regular_weight[outlier_mask] = 0

    q_max = 2 ** (num_bits - 1) - 1
    group_size = 16

    rows, cols = regular_weight.shape
    regular_grouped = regular_weight.reshape(-1, group_size)

    max_abs = regular_grouped.abs().max(dim=1, keepdim=True)[0]
    scales = max_abs / q_max
    scales = torch.clamp(scales, min=1e-8)

    q = torch.clamp(torch.round(regular_grouped / scales), -q_max, q_max).to(torch.int8)
    regular_dequant = (scales * q.float()).reshape(rows, cols)

    reconstructed = regular_dequant + outlier_values

    error = (weight - reconstructed).abs().mean().item()
    print(f"外れ値の比率: {outlier_mask.float().mean():.2%}")
    print(f"平均再構成誤差: {error:.6f}")

    return q, scales, outlier_values, outlier_mask

import time
import torch
import GPUtil

def benchmark_quantization(model, tokenizer, device="cuda", num_runs=50):
    """量子化モデルのベンチマーク"""

    prompt = "The history of artificial intelligence began"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    if device == "cuda":
        torch.cuda.synchronize()
        gpu = GPUtil.getGPUs()[0]
        memory_used_gb = gpu.memoryUsed / 1024

    # ウォームアップ
    with torch.no_grad():
        for _ in range(5):
            outputs = model.generate(
                **inputs,
                max_new_tokens=50,
                do_sample=False
            )

    # 速度測定
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.time()

    with torch.no_grad():
        for _ in range(num_runs):
            outputs = model.generate(
                **inputs,
                max_new_tokens=50,
                do_sample=False
            )

    if device == "cuda":
        torch.cuda.synchronize()
    elapsed = time.time() - start

    avg_time = elapsed / num_runs
    tokens_per_second = 50 / avg_time

    return {
        "memory_gb": memory_used_gb,
        "avg_time_ms": avg_time * 1000,
        "tokens_per_second": tokens_per_second
    }

# 結果例（A100 80GB、Llama-2-7B）
benchmark_results = {
    "FP16": {"memory_gb": 13.5, "tokens_per_second": 52.3, "ppl": 5.68},
    "INT8 (bitsandbytes)": {"memory_gb": 7.8, "tokens_per_second": 38.1, "ppl": 5.71},
    "INT4 GPTQ": {"memory_gb": 4.5, "tokens_per_second": 65.2, "ppl": 5.89},
    "INT4 AWQ": {"memory_gb": 4.3, "tokens_per_second": 68.7, "ppl": 5.86},
    "Q4_K_M (GGUF, CPU)": {"memory_gb": 4.1, "tokens_per_second": 45.2, "ppl": 5.91},
    "INT4 NF4": {"memory_gb": 4.0, "tokens_per_second": 31.5, "ppl": 5.94},
}

def recommend_quantization(
    task: str,
    model_size_b: float,
    gpu_memory_gb: float,
    cpu_only: bool = False,
    fine_tuning_needed: bool = False
):
    """タスクと環境に基づいて量子化を推奨"""

    recommendations = []

    if cpu_only:
        recommendations.append({
            "method": "GGUF Q4_K_M",
            "reason": "CPU推論に最適化、llama.cppベース",
            "library": "llama-cpp-python"
        })
        return recommendations

    if fine_tuning_needed:
        recommendations.append({
            "method": "bitsandbytes NF4 + QLoRA",
            "reason": "ファインチューニング可能、LoRAアダプター学習で約4GBのオーバーヘッド",
            "library": "bitsandbytes + peft"
        })
        return recommendations

    fp16_memory = model_size_b * 2
    int8_memory = model_size_b * 1
    int4_memory = model_size_b * 0.5

    if fp16_memory <= gpu_memory_gb * 0.8:
        recommendations.append({
            "method": "FP16 (ベースライン)",
            "reason": "メモリ十分、最高品質",
            "memory_gb": fp16_memory
        })

    if int4_memory <= gpu_memory_gb * 0.8:
        recommendations.append({
            "method": "AWQ INT4",
            "reason": "高速推論、優れた品質",
            "library": "autoawq",
            "memory_gb": int4_memory
        })
        recommendations.append({
            "method": "GPTQ INT4",
            "reason": "最高のINT4品質、量子化プロセスは遅め",
            "library": "auto-gptq",
            "memory_gb": int4_memory
        })

    return recommendations

# 使用例
recommendations = recommend_quantization(
    task="chat",
    model_size_b=7.0,
    gpu_memory_gb=16.0,
    fine_tuning_needed=False
)

for rec in recommendations:
    print(f"\n方法: {rec['method']}")
    print(f"理由: {rec['reason']}")
    if 'library' in rec:
        print(f"ライブラリ: {rec['library']}")
    if 'memory_gb' in rec:
        print(f"予想メモリ: {rec['memory_gb']:.1f} GB")