1. 量子化された重み W_q (INT8/INT4) をメモリから読み込み
2. 逆量子化して W_dequant (FP16) を得る
3. 活性化 X (FP16) と行列積: Y = X @ W_dequant

import torch
import torch.nn as nn


def symmetric_quantize(x: torch.Tensor, bits: int = 8) -> tuple:
    """
    対称量子化

    Args:
        x: 入力テンソル
        bits: 量子化ビット数

    Returns:
        x_q: 量子化されたテンソル (int)
        scale: スケール
    """
    qmax = 2 ** (bits - 1) - 1
    qmin = -2 ** (bits - 1)

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

    # 量子化
    x_q = torch.round(x / scale).clamp(qmin, qmax).to(torch.int8)

    return x_q, scale


def symmetric_dequantize(x_q: torch.Tensor, scale: float) -> torch.Tensor:
    """
    逆量子化

    Args:
        x_q: 量子化されたテンソル
        scale: スケール

    Returns:
        x_dequant: 逆量子化されたテンソル
    """
    return x_q.float() * scale


# 使用例
torch.manual_seed(42)
W = torch.randn(256, 512)  # 元の重み

# 量子化
W_q, scale = symmetric_quantize(W, bits=8)
print(f"元の重み: dtype={W.dtype}, メモリ={W.numel() * 4} bytes")
print(f"量子化後: dtype={W_q.dtype}, メモリ={W_q.numel() * 1} bytes")

# 逆量子化
W_dequant = symmetric_dequantize(W_q, scale)

# 量子化誤差
error = (W - W_dequant).abs().mean()
print(f"平均絶対誤差: {error:.6f}")

def group_quantize(x: torch.Tensor, group_size: int = 128, bits: int = 4) -> tuple:
    """
    グループごとの量子化

    Args:
        x: 入力テンソル (in_features, out_features)
        group_size: グループサイズ
        bits: 量子化ビット数

    Returns:
        x_q: 量子化されたテンソル
        scales: グループごとのスケール
    """
    in_features, out_features = x.shape
    assert in_features % group_size == 0, "in_features must be divisible by group_size"

    qmax = 2 ** (bits - 1) - 1
    qmin = -2 ** (bits - 1)

    # グループに分割
    x_grouped = x.view(-1, group_size, out_features)
    n_groups = x_grouped.shape[0]

    # グループごとにスケール計算
    scales = x_grouped.abs().max(dim=1, keepdim=True).values / qmax
    scales = scales.clamp(min=1e-10)  # ゼロ除算防止

    # 量子化
    x_q = torch.round(x_grouped / scales).clamp(qmin, qmax)
    x_q = x_q.view(in_features, out_features).to(torch.int8)

    scales = scales.squeeze(1)  # (n_groups, out_features)

    return x_q, scales


def group_dequantize(
    x_q: torch.Tensor,
    scales: torch.Tensor,
    group_size: int = 128,
) -> torch.Tensor:
    """
    グループごとの逆量子化
    """
    in_features, out_features = x_q.shape

    x_q_grouped = x_q.view(-1, group_size, out_features).float()
    scales_expanded = scales.unsqueeze(1)  # (n_groups, 1, out_features)

    x_dequant = x_q_grouped * scales_expanded
    return x_dequant.view(in_features, out_features)


# 使用例
W = torch.randn(512, 1024)

# INT4 グループ量子化
W_q, scales = group_quantize(W, group_size=128, bits=4)
W_dequant = group_dequantize(W_q, scales, group_size=128)

error = (W - W_dequant).abs().mean()
print(f"INT4 グループ量子化の平均絶対誤差: {error:.6f}")

# スケールのメモリ
print(f"元のメモリ: {W.numel() * 4} bytes")
print(f"量子化後: 重み {W_q.numel() * 0.5} bytes + スケール {scales.numel() * 2} bytes")

class QuantizedLinear(nn.Module):
    """量子化されたLinear層"""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bits: int = 8,
        group_size: int = None,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.bits = bits
        self.group_size = group_size

        # 量子化された重みを登録（学習は想定しない）
        self.register_buffer('weight_q', torch.zeros(in_features, out_features, dtype=torch.int8))
        self.register_buffer('scales', torch.zeros(in_features // (group_size or in_features), out_features))
        self.register_buffer('bias', None)

    @classmethod
    def from_float(cls, linear: nn.Linear, bits: int = 8, group_size: int = None):
        """float Linear層から量子化版を作成"""
        quant_linear = cls(
            linear.in_features,
            linear.out_features,
            bits=bits,
            group_size=group_size,
        )

        W = linear.weight.data.T  # (in, out)

        if group_size:
            W_q, scales = group_quantize(W, group_size, bits)
        else:
            W_q, scale = symmetric_quantize(W, bits)
            scales = scale.expand(1, linear.out_features)

        quant_linear.weight_q.copy_(W_q)
        quant_linear.scales.copy_(scales)

        if linear.bias is not None:
            quant_linear.bias = linear.bias.data.clone()

        return quant_linear

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 逆量子化
        if self.group_size:
            W = group_dequantize(self.weight_q, self.scales, self.group_size)
        else:
            W = self.weight_q.float() * self.scales

        # 行列積
        out = x @ W

        if self.bias is not None:
            out = out + self.bias

        return out


# 使用例
linear = nn.Linear(512, 1024)
quant_linear = QuantizedLinear.from_float(linear, bits=4, group_size=128)

x = torch.randn(2, 10, 512)
y_original = linear(x)
y_quantized = quant_linear(x)

print(f"出力の差（平均絶対誤差）: {(y_original - y_quantized).abs().mean():.6f}")

import matplotlib.pyplot as plt
import numpy as np

# 重みの分布をシミュレート
np.random.seed(42)
weights = np.random.randn(10000) * 0.1

# INT8量子化
qmax = 127
scale = np.max(np.abs(weights)) / qmax
weights_q = np.round(weights / scale).clip(-128, 127)
weights_dequant = weights_q * scale

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# 元の分布
ax = axes[0]
ax.hist(weights, bins=50, color='steelblue', alpha=0.7, edgecolor='black')
ax.set_title('Original Weights (FP32)', fontsize=12)
ax.set_xlabel('Value')
ax.set_ylabel('Count')

# 量子化後の分布
ax = axes[1]
ax.hist(weights_q, bins=50, color='coral', alpha=0.7, edgecolor='black')
ax.set_title('Quantized Weights (INT8)', fontsize=12)
ax.set_xlabel('Quantized Value')
ax.set_ylabel('Count')

# 量子化誤差
ax = axes[2]
error = weights - weights_dequant
ax.hist(error, bins=50, color='seagreen', alpha=0.7, edgecolor='black')
ax.set_title('Quantization Error', fontsize=12)
ax.set_xlabel('Error')
ax.set_ylabel('Count')
ax.axvline(x=0, color='red', linestyle='--')

plt.tight_layout()
plt.show()