初始化项目，由ModelHub XC社区提供模型
Model: PKU-DS-LAB/Fairy2i-W2 Source: Original Platform
This commit is contained in:
ModelHub XC
308
quantization.py
Normal file
308
quantization.py
Normal file
@@ -0,0 +1,308 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_complex_tensor(w_real: torch.Tensor, w_imag: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply PhaseQuant logic to complex weight tensors"""
|
||||
phase = torch.angle(w_real + 1j * w_imag)
|
||||
|
||||
real_pos = (phase >= -math.pi / 4) & (phase < math.pi / 4)
|
||||
real_neg = (phase >= 3 * math.pi / 4) | (phase < -3 * math.pi / 4)
|
||||
imag_pos = (phase >= math.pi / 4) & (phase < 3 * math.pi / 4)
|
||||
imag_neg = (phase >= -3 * math.pi / 4) & (phase < -math.pi / 4)
|
||||
|
||||
mask_real = real_pos | real_neg
|
||||
mask_imag = imag_pos | imag_neg
|
||||
|
||||
s_re = w_real[mask_real].abs().mean() if mask_real.any() else torch.tensor(0.0, device=w_real.device)
|
||||
s_im = w_imag[mask_imag].abs().mean() if mask_imag.any() else torch.tensor(0.0, device=w_imag.device)
|
||||
|
||||
s_re = torch.clamp(s_re, min=1e-6)
|
||||
s_im = torch.clamp(s_im, min=1e-6)
|
||||
if torch.isnan(s_re) or torch.isinf(s_re): s_re = torch.tensor(1e-6, device=w_real.device)
|
||||
if torch.isnan(s_im) or torch.isinf(s_im): s_im = torch.tensor(1e-6, device=w_imag.device)
|
||||
|
||||
qw_real = torch.zeros_like(w_real)
|
||||
qw_imag = torch.zeros_like(w_imag)
|
||||
|
||||
qw_real[real_pos] = 1.0
|
||||
qw_real[real_neg] = -1.0
|
||||
qw_imag[imag_pos] = 1.0
|
||||
qw_imag[imag_neg] = -1.0
|
||||
|
||||
qw_real_scaled = qw_real * s_re
|
||||
qw_imag_scaled = qw_imag * s_im
|
||||
return qw_real_scaled.to(w_real.dtype), qw_imag_scaled.to(w_imag.dtype)
|
||||
|
||||
def apply_complex_inspired_quantization(model: nn.Module):
|
||||
"""Apply complex-inspired quantization to real-valued model"""
|
||||
print("Applying complex-inspired quantization (PhaseQuant-based)...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
A = module.weight.data
|
||||
if A.shape[0] % 2 != 0 or A.shape[1] % 2 != 0:
|
||||
print(f" -> Skipping layer (non-even dimensions): {A.shape}")
|
||||
return
|
||||
|
||||
n, m = A.shape[0] // 2, A.shape[1] // 2
|
||||
A11, A12 = A[:n, :m], A[:n, m:]
|
||||
A21, A22 = A[n:, :m], A[n:, m:]
|
||||
|
||||
U_re = 0.5 * (A11 + A22)
|
||||
U_im = 0.5 * (A21 - A12)
|
||||
W_re = 0.5 * (A11 - A22)
|
||||
W_im = 0.5 * (A12 + A21)
|
||||
|
||||
U_re_q, U_im_q = quantize_complex_tensor(U_re, U_im)
|
||||
W_re_q, W_im_q = quantize_complex_tensor(W_re, W_im)
|
||||
|
||||
A11_q = W_re_q + U_re_q
|
||||
A12_q = W_im_q - U_im_q
|
||||
A21_q = W_im_q + U_im_q
|
||||
A22_q = -W_re_q + U_re_q
|
||||
|
||||
A_quant_top = torch.cat([A11_q, A12_q], dim=1)
|
||||
A_quant_bottom = torch.cat([A21_q, A22_q], dim=1)
|
||||
A_quant = torch.cat([A_quant_top, A_quant_bottom], dim=0)
|
||||
|
||||
module.weight.data = A_quant.to(A.dtype)
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
print("Complex-inspired quantization completed.")
|
||||
return model
|
||||
|
||||
def apply_bitnet_quantization(model: nn.Module):
|
||||
"""Apply BitNet 1-bit quantization to real-valued model"""
|
||||
print("Applying BitNet (true 1-bit, affine) quantization to real-valued model...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
scale = module.weight.data.abs().mean()
|
||||
alpha = module.weight.data.mean()
|
||||
centered_weights = module.weight.data - alpha
|
||||
binarized_weights = torch.where(centered_weights > 0, 1.0, -1.0)
|
||||
module.weight.data = binarized_weights.to(module.weight.data.dtype) * scale
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
print("BitNet quantization completed.")
|
||||
return model
|
||||
|
||||
def apply_bitnet_1_58bit_quantization_standard(model: nn.Module):
|
||||
"""Apply BitNet 1.58-bit quantization to real-valued model (quantize to {-1, 0, +1})"""
|
||||
print("Applying BitNet 1.58-bit (absmean threshold) quantization to real-valued model...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
W = module.weight.data
|
||||
gamma = W.abs().mean()
|
||||
W_normalized = W / (gamma + 1e-5)
|
||||
W_quantized = torch.clamp(torch.round(W_normalized), -1.0, 1.0)
|
||||
module.weight.data = W_quantized.to(W.dtype) * gamma
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
print("BitNet 1.58-bit (absmean threshold) quantization completed.")
|
||||
return model
|
||||
|
||||
def apply_bitnet_1_58bit_quantization_variant(model: nn.Module, threshold: float = 0.5):
|
||||
"""Apply BitNet 1.58-bit quantization to real-valued model (quantize to {-1, 0, +1})"""
|
||||
print("Applying BitNet 1.58-bit (ternary) quantization to real-valued model...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
gamma = module.weight.data.abs().mean()
|
||||
normalized_weights = module.weight.data / (gamma + 1e-5)
|
||||
adaptive_threshold = threshold
|
||||
ternary_weights = torch.zeros_like(normalized_weights)
|
||||
ternary_weights[normalized_weights > adaptive_threshold] = 1.0
|
||||
ternary_weights[normalized_weights < -adaptive_threshold] = -1.0
|
||||
module.weight.data = ternary_weights.to(module.weight.data.dtype) * gamma
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
print("BitNet 1.58-bit quantization completed.")
|
||||
return model
|
||||
|
||||
def minmax_1bit_quantize_dequantize(w: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply 1-bit Min-Max quantization and dequantization to weight tensor"""
|
||||
min_val = w.min()
|
||||
max_val = w.max()
|
||||
scale = (max_val - min_val) / 1.0
|
||||
zero_point = min_val
|
||||
|
||||
if abs(scale) < 1e-9:
|
||||
return w
|
||||
|
||||
quantized_w = torch.round((w - zero_point) / scale)
|
||||
dequantized_w = quantized_w * scale + zero_point
|
||||
|
||||
return dequantized_w.to(w.dtype)
|
||||
|
||||
def apply_minmax_1bit_quantization(model: nn.Module):
|
||||
"""Apply Min-Max 1-bit quantization to real-valued model"""
|
||||
print("Applying Min-Max (1-bit) quantization to real-valued model...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
module.weight.data = minmax_1bit_quantize_dequantize(module.weight.data)
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
|
||||
print("Min-Max 1-bit quantization completed.")
|
||||
return model
|
||||
|
||||
def symmetric_minmax_1bit_quantize_dequantize(w: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply symmetric 1-bit Min-Max quantization to weight tensor (quantize to {-1, 1})"""
|
||||
max_abs = w.abs().max()
|
||||
scale = max_abs
|
||||
|
||||
if scale < 1e-9:
|
||||
return w
|
||||
|
||||
quantized_w = (w / scale).sign()
|
||||
dequantized_w = quantized_w * scale
|
||||
|
||||
return dequantized_w.to(w.dtype)
|
||||
|
||||
def apply_symmetric_minmax_1bit_quantization(model: nn.Module):
|
||||
"""Apply symmetric Min-Max 1-bit quantization to real-valued model"""
|
||||
print("Applying symmetric Min-Max (1-bit, to {-1, 1}) quantization to real-valued model...")
|
||||
|
||||
@torch.no_grad()
|
||||
def quantize_linear_layer(module: nn.Linear):
|
||||
module.weight.data = symmetric_minmax_1bit_quantize_dequantize(module.weight.data)
|
||||
|
||||
model.apply(lambda module: quantize_linear_layer(module) if isinstance(module, nn.Linear) else None)
|
||||
|
||||
print("Symmetric Min-Max 1-bit quantization completed.")
|
||||
return model
|
||||
|
||||
class BitNetQuantSTE(torch.autograd.Function):
|
||||
"""BitNet STE: quantize in forward, pass gradients in backward"""
|
||||
@staticmethod
|
||||
def forward(ctx, w):
|
||||
scale = w.abs().mean()
|
||||
alpha = w.mean()
|
||||
centered_w = w - alpha
|
||||
binarized_w = torch.where(centered_w > 0, 1.0, -1.0).to(w.dtype)
|
||||
quantized_w = binarized_w * scale
|
||||
return quantized_w
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output
|
||||
|
||||
|
||||
class BitNet1_58QuantSTE(torch.autograd.Function):
|
||||
"""BitNet 1.58-bit STE: quantize to {-1, 0, +1}, pass gradients in backward"""
|
||||
@staticmethod
|
||||
def forward(ctx, w):
|
||||
gamma = w.abs().mean()
|
||||
w_normalized = w / (gamma + 1e-5)
|
||||
w_quantized = torch.clamp(torch.round(w_normalized), -1.0, 1.0)
|
||||
quantized_w = (w_quantized * gamma).to(w.dtype)
|
||||
return quantized_w
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output
|
||||
|
||||
|
||||
class PhaseQuantSTE(torch.autograd.Function):
|
||||
"""Complex-Phase STE: quantize in forward, pass gradients in backward"""
|
||||
@staticmethod
|
||||
def forward(ctx, w_real, w_imag):
|
||||
phase = torch.angle(w_real + 1j * w_imag)
|
||||
|
||||
real_pos = (phase >= -math.pi / 4) & (phase < math.pi / 4)
|
||||
real_neg = (phase >= 3 * math.pi / 4) | (phase < -3 * math.pi / 4)
|
||||
imag_pos = (phase >= math.pi / 4) & (phase < 3 * math.pi / 4)
|
||||
imag_neg = (phase >= -3 * math.pi / 4) & (phase < -math.pi / 4)
|
||||
|
||||
mask_real = real_pos | real_neg
|
||||
mask_imag = imag_pos | imag_neg
|
||||
|
||||
s_re = w_real[mask_real].abs().mean() if mask_real.any() else torch.tensor(0.0, device=w_real.device)
|
||||
s_im = w_imag[mask_imag].abs().mean() if mask_imag.any() else torch.tensor(0.0, device=w_imag.device)
|
||||
|
||||
s_re = torch.clamp(s_re, min=1e-6)
|
||||
s_im = torch.clamp(s_im, min=1e-6)
|
||||
|
||||
qw_real = torch.zeros_like(w_real)
|
||||
qw_imag = torch.zeros_like(w_imag)
|
||||
|
||||
qw_real[real_pos] = 1.0
|
||||
qw_real[real_neg] = -1.0
|
||||
qw_imag[imag_pos] = 1.0
|
||||
qw_imag[imag_neg] = -1.0
|
||||
|
||||
qw_real_scaled = qw_real * s_re
|
||||
qw_imag_scaled = qw_imag * s_im
|
||||
|
||||
return qw_real_scaled.to(w_real.dtype), qw_imag_scaled.to(w_imag.dtype)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_w_real, grad_w_imag):
|
||||
return grad_w_real, grad_w_imag
|
||||
|
||||
class PhaseQuantSTE_V2(torch.autograd.Function):
|
||||
"""Two-step residual quantization"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, w_real: torch.Tensor, w_imag: torch.Tensor):
|
||||
qw_real_o1, qw_imag_o1 = PhaseQuantSTE.apply(w_real, w_imag)
|
||||
error_real = w_real - qw_real_o1
|
||||
error_imag = w_imag - qw_imag_o1
|
||||
qw_real_o2, qw_imag_o2 = PhaseQuantSTE.apply(error_real, error_imag)
|
||||
qw_real = qw_real_o1 + qw_real_o2
|
||||
qw_imag = qw_imag_o1 + qw_imag_o2
|
||||
return qw_real, qw_imag
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_real, grad_imag):
|
||||
return grad_real, grad_imag
|
||||
|
||||
|
||||
class PhaseQuantSTE_V3(torch.autograd.Function):
|
||||
"""Three-step residual quantization"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, w_real: torch.Tensor, w_imag: torch.Tensor):
|
||||
qw_real_o1, qw_imag_o1 = PhaseQuantSTE.apply(w_real, w_imag)
|
||||
error_real_1 = w_real - qw_real_o1
|
||||
error_imag_1 = w_imag - qw_imag_o1
|
||||
qw_real_o2, qw_imag_o2 = PhaseQuantSTE.apply(error_real_1, error_imag_1)
|
||||
error_real_2 = error_real_1 - qw_real_o2
|
||||
error_imag_2 = error_imag_1 - qw_imag_o2
|
||||
qw_real_o3, qw_imag_o3 = PhaseQuantSTE.apply(error_real_2, error_imag_2)
|
||||
qw_real = qw_real_o1 + qw_real_o2 + qw_real_o3
|
||||
qw_imag = qw_imag_o1 + qw_imag_o2 + qw_imag_o3
|
||||
return qw_real, qw_imag
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_real, grad_imag):
|
||||
return grad_real, grad_imag
|
||||
|
||||
|
||||
class PhaseQuantSTE_V4(torch.autograd.Function):
|
||||
"""Four-step residual quantization"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, w_real: torch.Tensor, w_imag: torch.Tensor):
|
||||
qw_real_o1, qw_imag_o1 = PhaseQuantSTE.apply(w_real, w_imag)
|
||||
error_real_1 = w_real - qw_real_o1
|
||||
error_imag_1 = w_imag - qw_imag_o1
|
||||
qw_real_o2, qw_imag_o2 = PhaseQuantSTE.apply(error_real_1, error_imag_1)
|
||||
error_real_2 = error_real_1 - qw_real_o2
|
||||
error_imag_2 = error_imag_1 - qw_imag_o2
|
||||
qw_real_o3, qw_imag_o3 = PhaseQuantSTE.apply(error_real_2, error_imag_2)
|
||||
error_real_3 = error_real_2 - qw_real_o3
|
||||
error_imag_3 = error_imag_2 - qw_imag_o3
|
||||
qw_real_o4, qw_imag_o4 = PhaseQuantSTE.apply(error_real_3, error_imag_3)
|
||||
qw_real = qw_real_o1 + qw_real_o2 + qw_real_o3 + qw_real_o4
|
||||
qw_imag = qw_imag_o1 + qw_imag_o2 + qw_imag_o3 + qw_imag_o4
|
||||
return qw_real, qw_imag
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_real, grad_imag):
|
||||
return grad_real, grad_imag
|
||||
Reference in New Issue
Block a user